File size: 3,670 Bytes
3527383 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the CC-by-NC license found in the
# LICENSE file in the root directory of this source tree.
import torch
from torch import Tensor
from torch.func import jvp, vmap
from flow_matching.path.path import ProbPath
from flow_matching.path.path_sample import PathSample
from flow_matching.path.scheduler import ConvexScheduler
from flow_matching.utils import expand_tensor_like
from flow_matching.utils.manifolds import geodesic, Manifold
class GeodesicProbPath(ProbPath):
r"""The ``GeodesicProbPath`` class represents a specific type of probability path where the transformation between distributions is defined through the geodesic path.
Mathematically, a geodesic path can be represented as:
.. math::
X_t = \psi_t(X_0 | X_1) = \exp_{X_1}(\kappa_t \log_{X_1}(X_0)),
where :math:`X_t` is the transformed data point at time `t`, :math:`X_0` and :math:`X_1` are the source and target data points, respectively, and :math:`\kappa_t` is a scheduler.
The scheduler is responsible for providing the time-dependent :math:`\kappa_t` and must be differentiable.
Using ``GeodesicProbPath`` in the flow matching framework:
.. code-block:: python
# Instantiates a manifold
manifold = FlatTorus()
# Instantiates a scheduler
scheduler = CondOTScheduler()
# Instantiates a probability path
my_path = GeodesicProbPath(scheduler, manifold)
mse_loss = torch.nn.MSELoss()
for x_1 in dataset:
# Sets x_0 to random noise
x_0 = torch.randn()
# Sets t to a random value in [0,1]
t = torch.rand()
# Samples the conditional path :math:`X_t \sim p_t(X_t|X_0,X_1)`
path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)
# Computes the MSE loss w.r.t. the velocity
loss = mse_loss(path_sample.dx_t, my_model(x_t, t))
loss.backward()
Args:
scheduler (ConvexScheduler): The scheduler that provides :math:`\kappa_t`.
manifold (Manifold): The manifold on which the probability path is defined.
"""
def __init__(self, scheduler: ConvexScheduler, manifold: Manifold):
self.scheduler = scheduler
self.manifold = manifold
def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample:
r"""Sample from the Riemannian probability path with geodesic interpolation:
| given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`\kappa_t`.
| return :math:`X_0, X_1, X_t = \exp_{X_1}(\kappa_t \log_{X_1}(X_0))`, and the conditional velocity at :math:`X_t, \dot{X}_t`.
Args:
x_0 (Tensor): source data point, shape (batch_size, ...).
x_1 (Tensor): target data point, shape (batch_size, ...).
t (Tensor): times in [0,1], shape (batch_size).
Returns:
PathSample: A conditional sample at :math:`X_t \sim p_t`.
"""
self.assert_sample_shape(x_0=x_0, x_1=x_1, t=t)
t = expand_tensor_like(input_tensor=t, expand_to=x_1[..., 0:1]).clone()
def cond_u(x_0, x_1, t):
path = geodesic(self.manifold, x_0, x_1)
x_t, dx_t = jvp(
lambda t: path(self.scheduler(t).alpha_t),
(t,),
(torch.ones_like(t).to(t),),
)
return x_t, dx_t
x_t, dx_t = vmap(cond_u)(x_0, x_1, t)
x_t = x_t.reshape_as(x_1)
dx_t = dx_t.reshape_as(x_1)
return PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t)
|