GeodesicProbPath#

class flow_matching.path.GeodesicProbPath(scheduler: ConvexScheduler, manifold: Manifold)[source]#

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:

\[X_t = \psi_t(X_0 | X_1) = \exp_{X_1}(\kappa_t \log_{X_1}(X_0)),\]

where \(X_t\) is the transformed data point at time t, \(X_0\) and \(X_1\) are the source and target data points, respectively, and \(\kappa_t\) is a scheduler.

The scheduler is responsible for providing the time-dependent \(\kappa_t\) and must be differentiable.

Using GeodesicProbPath in the flow matching framework:

Parameters:

  • scheduler (ConvexScheduler) – The scheduler that provides \(\kappa_t\).

  • manifold (Manifold) – The manifold on which the probability path is defined.

sample(x_0: Tensor, x_1: Tensor, t: Tensor) → PathSample[source]#

Sample from the Riemannian probability path with geodesic interpolation:

given \((X_0,X_1) \sim \pi(X_0,X_1)\) and a scheduler \(\kappa_t\).
return \(X_0, X_1, X_t = \exp_{X_1}(\kappa_t \log_{X_1}(X_0))\), and the conditional velocity at \(X_t, \dot{X}_t\).
Parameters:

  • 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:

A conditional sample at \(X_t \sim p_t\).

Return type:

PathSample