interpolants

interpolants

Probability paths connecting noise x0 to data x1.

BaseEquilibriumInterpolant

Bases: ABC

Base class for an energy-compatible path indexed by a coefficient gamma.

This base is deliberately separate from BaseInterpolant. A flow-matching interpolant is parametrised by a dynamical time t that a sampler integrates over. An Equilibrium Matching interpolant is parametrised by an interpolation coefficient \(\gamma \in [0, 1]\), a noise level rather than a time. The learned field is trained to be time-invariant, so \(\gamma\) only indexes where along the noise-to-data path a training point sits, it is never integrated and (per the EqM paper) is not seen by the model.

Convention: \(\gamma = 0\) is noise (\(x_\gamma = x_0\)) and \(\gamma = 1\) is data (\(x_\gamma = x_1\)).

References

Wang and Du, "Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models" (2025), https://arxiv.org/abs/2510.02300.

interpolate abstractmethod

interpolate(x1, gamma, x0=None)

Return (x_gamma, target) for the given data and coefficient.

Parameters:

Name	Type	Description	Default
x1	Tensor	data sample (the \(\gamma = 1\) endpoint).	required
gamma	Tensor	interpolation coefficient (noise level) in \([0, 1]\), not a dynamical time.	required
x0	Optional[Tensor]	optional noise sample (the \(\gamma = 0\) endpoint). Drawn from a standard normal when omitted.	None

Source code in deltaflow/core/base_equilibrium_interpolant.py

@abstractmethod
def interpolate(
    self, x1: torch.Tensor, gamma: torch.Tensor, x0: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
    r"""Return ``(x_gamma, target)`` for the given data and coefficient.

    Args:
        x1: data sample (the \(\gamma = 1\) endpoint).
        gamma: interpolation coefficient (noise level) in \([0, 1]\), not a
            dynamical time.
        x0: optional noise sample (the \(\gamma = 0\) endpoint). Drawn from
            a standard normal when omitted.
    """
    raise NotImplementedError

BaseInterpolant

Bases: ABC

Base class for a probability path between noise x0 and data x1.

An interpolant defines, for each t in [0, 1], an intermediate point x_t and its conditional target velocity u_t such that regressing a model onto u_t (in expectation over the path) yields the marginal velocity field of the flow-matching ODE.

Convention used throughout DeltaFlow: t = 0 corresponds to noise (x_t = x_0) and t = 1 corresponds to data (x_t = x_1).

interpolate abstractmethod

interpolate(x1, t, x0=None)

Return (x_t, target_velocity) for the given data/time (and optional noise).

Source code in deltaflow/core/base_interpolant.py

@abstractmethod
def interpolate(
    self, x1: torch.Tensor, t: torch.Tensor, x0: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Return ``(x_t, target_velocity)`` for the given data/time (and optional noise)."""
    raise NotImplementedError

EquilibriumInterpolant

EquilibriumInterpolant(plateau=0.8, scale=4.0, start=1.0)

Bases: BaseEquilibriumInterpolant

Straight-line path with an energy-compatible (equilibrium) target.

Equilibrium Matching (EqM) keeps the rectified-flow straight-line path but parametrises it by an interpolation coefficient \(\gamma \in [0, 1]\), a noise level rather than a dynamical time. The learned field is trained to be time-invariant (it approximates an equilibrium gradient), so \(\gamma\) only indexes where along the noise-to-data path a training point sits, it is not integrated over. In the paper's notation, for data \(x\) and Gaussian noise \(\epsilon\),

\[ x_\gamma = \gamma\,x + (1 - \gamma)\,\epsilon, \]

which, with DeltaFlow's \((x_0, x_1)\) = (noise, data) naming, is the same straight line

\[ x_\gamma = (1 - \gamma)\,x_0 + \gamma\,x_1, \qquad \gamma \in [0, 1]. \]

EqM reshapes the regression target so the learned field becomes the gradient of an implicit energy landscape rather than a time-conditional velocity. Instead of regressing onto the constant displacement \(x_1 - x_0\), it regresses onto the scaled displacement (the paper's \((x - \epsilon)\,c(\gamma)\)),

\[ u_\gamma = c(\gamma)\,(x_1 - x_0), \]

where \(c(\gamma)\) is the equilibrium coefficient. The key design constraint is \(c(1) = 0\): the target vanishes at data, so ground-truth samples become stationary points (local minima) of the landscape whose gradient the field learns. Away from data the coefficient is held on a constant plateau, so the field points from noise toward data with roughly constant magnitude.

Concretely the coefficient is the minimum of two lines, rescaled by scale,

\[ c(\gamma) = \text{scale}\cdot\min\!\Bigl( \text{start} - \tfrac{\text{start} - 1}{p}\,\gamma,\; \tfrac{1 - \gamma}{1 - p} \Bigr), \]

with plateau fraction \(p\) (plateau). With the defaults (start = 1, plateau = 0.8, scale = 4) the first line is flat at \(1\), so \(c(\gamma) = 4\,\min(1, 5(1 - \gamma))\): a plateau of \(4\) for \(\gamma \le 0.8\) that then ramps linearly down to \(0\) at \(\gamma = 1\).

Interpolation coefficient. \(\gamma = 0\) is noise (\(x_\gamma = x_0\)) and \(\gamma = 1\) is data (\(x_\gamma = x_1\)). This interpolant subclasses BaseEquilibriumInterpolant, whose interpolate takes gamma rather than the dynamical time t of BaseInterpolant, because \(\gamma\) is a noise level and not a time to be integrated.

Coupling. This interpolant defines only the path and target. It is agnostic to how \((x_0, x_1)\) pairs are formed. If x0 is not supplied it is drawn from \(\mathcal{N}(0, I)\) independently of x1. Pass an OTCoupling on the training side for straighter, OT-coupled displacements.

Sampling. A field trained with this target is not integrated over a fixed time horizon. Because it approximates an equilibrium gradient it is sampled by optimisation (gradient descent on the landscape), for which see EquilibriumSolver.

References

Wang and Du, "Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models" (2025), https://arxiv.org/abs/2510.02300.

Parameters:

Name	Type	Description	Default
plateau	float	fraction \(p \in (0, 1)\) of the path spent before the coefficient starts ramping down to zero. The ramp occupies the final 1 - plateau of the \([0, 1]\) interval.	0.8
scale	float	overall magnitude applied to the coefficient. Sets the typical norm of the learned equilibrium gradient away from data.	4.0
start	float	value of the first (upper) line at gamma=0. Left at 1 this line stays flat so the coefficient plateaus at scale, values other than 1 tilt the plateau.	1.0

Source code in deltaflow/interpolants/equilibrium.py

def __init__(self, plateau: float = 0.8, scale: float = 4.0, start: float = 1.0):
    if not 0.0 < plateau < 1.0:
        raise ValueError(f"plateau must lie in (0, 1), got {plateau!r}")
    self.plateau = plateau
    self.scale = scale
    self.start = start

equilibrium_coefficient

equilibrium_coefficient(gamma)

Return the scalar coefficient \(c(\gamma)\) applied to \(x_1 - x_0\).

gamma is the interpolation coefficient (noise level), \(0\) at noise and \(1\) at data, not a dynamical time.

Source code in deltaflow/interpolants/equilibrium.py

def equilibrium_coefficient(self, gamma: torch.Tensor) -> torch.Tensor:
    r"""Return the scalar coefficient \(c(\gamma)\) applied to \(x_1 - x_0\).

    ``gamma`` is the interpolation coefficient (noise level), \(0\) at noise
    and \(1\) at data, not a dynamical time.
    """
    line_plateau = self.start - (self.start - 1.0) / self.plateau * gamma
    line_ramp = (1.0 - gamma) / (1.0 - self.plateau)
    return self.scale * torch.minimum(line_plateau, line_ramp)

LinearInterpolant

Bases: BaseInterpolant

Straight-line (rectified-flow) probability path between noise and data.

The path linearly interpolates a noise sample \(x_0\) and a data sample \(x_1\),

\[ x_t = (1 - t)\,x_0 + t\,x_1, \qquad t \in [0, 1], \]

whose conditional target velocity is the constant displacement

\[ u_t = \frac{\mathrm{d} x_t}{\mathrm{d} t} = x_1 - x_0. \]

Because \(u_t\) does not depend on \(t\), the learned field regresses onto a single displacement vector per pair, which is what makes rectified-flow trajectories straight and cheap to integrate.

Coupling. This is the independent-coupling variant: if \(x_0\) is not supplied it is drawn from a standard normal \(\mathcal{N}(0, I)\) independently of \(x_1\). For OT-coupled linear paths see OTInterpolant or use OTCoupling on the training side.

References

Lipman et al., "Flow Matching for Generative Modeling" (2023), https://arxiv.org/abs/2210.02747.

OTInterpolant

OTInterpolant()

Bases: BaseInterpolant

Linear probability path with mini-batch optimal-transport coupling.

This reuses the straight-line path of LinearInterpolant, but instead of pairing noise and data independently it re-orders \(x_0\) within the batch to approximately solve the discrete optimal-transport assignment. Given a batch of noise \(\{x_0^{(i)}\}\) and data \(\{x_1^{(j)}\}\), it seeks a permutation \(\pi\) minimising the total squared-\(L_2\) transport cost

\[ \pi^\star = \arg\min_{\pi \in S_B} \sum_{i=1}^{B} \bigl\| x_0^{(i)} - x_1^{(\pi(i))} \bigr\|_2^2, \]

then applies the linear interpolant to the matched pairs \(\bigl(x_0^{(\pi^\star(i))}, x_1^{(i)}\bigr)\). Concretely, each call

1. draws (or receives) a batch of noise samples \(x_0\),
2. computes \(\pi^\star\) so each noise sample is paired with the data sample that minimises the batch transport cost,
3. applies LinearInterpolant on the permuted pair.

In the large-batch limit this converges to a coupling drawn from the true OT plan and yields straighter learned trajectories that sample in fewer steps. It is the zero-entropy limit of the static Schrödinger bridge. Because the coupling is purely a re-ordering of \(x_0\), the training objective is identical to standard conditional flow matching and no other component (loss, solver, model) needs to change.

Solver. The exact assignment (Hungarian algorithm) is used by default, since scipy is a core dependency. A deterministic greedy nearest-neighbour fallback is used only if scipy is unavailable.

References

Tong et al., "Improving and generalizing flow-based generative models with minibatch optimal transport" (2023), https://arxiv.org/abs/2302.00482. "Flower: A Flow-Matching Solver for Inverse Problems" (2025), https://arxiv.org/abs/2509.26287.

Source code in deltaflow/interpolants/ot.py

def __init__(self):
    self._linear = LinearInterpolant()

SchrodingerBridgeInterpolant

SchrodingerBridgeInterpolant(sigma=1.0, eps=0.0001)

Bases: BaseInterpolant

Brownian-bridge probability path with tunable diffusivity sigma.

Conditioned on an endpoint pair \((x_0, x_1)\), the dynamic Schrödinger bridge reduces to a Brownian bridge around the straight-line mean,

\[ x_t = (1 - t)\,x_0 + t\,x_1 + \sigma\sqrt{t(1 - t)}\;z, \qquad z \sim \mathcal{N}(0, I), \]

whose conditional target velocity (the drift of the associated probability-flow ODE, obtained by differentiating the reparameterised path w.r.t. \(t\) at fixed \(z\)) is

\[ u_t = (x_1 - x_0) + \sigma\,\frac{1 - 2t}{2\sqrt{t(1 - t)}}\;z. \]

As \(\sigma \to 0\) the bridge collapses onto the straight line of LinearInterpolant. The \((1 - 2t)/\sqrt{t(1 - t)}\) factor is unbounded as \(t \to 0, 1\) (the conditional target is unbiased but its variance diverges at the boundary). It is stabilised here by flooring the denominator at eps. Note this floor is a pragmatic numerical safeguard, not the exact SF2M treatment, which instead learns a joint score/flow parametrisation and does not clip the drift.

Coupling. This interpolant defines only the path. It is agnostic to how \((x_0, x_1)\) pairs are formed. Pairing endpoints with OTCoupling (rather than drawing them independently) pushes the discretised process toward the Schrödinger bridge instead of an arbitrary diffusion mixture. The correspondence is approximate: the true SB (Tong et al., 2024) couples the Brownian bridges with the entropy-regularised OT plan \(\pi^\star_{2\sigma^2}\), whose regularisation strength is tied to the diffusivity (\(\text{reg} = 2\sigma^2\)), whereas OTCoupling solves the unregularised squared-\(L_2\) OT problem. Exact OT therefore corresponds to the small-\(\sigma\) limit of the true bridge. For larger \(\sigma\) this is a rectified-flow-style approximation rather than the exact entropic Schrödinger bridge.

Parameters:

Name	Type	Description	Default
sigma	float	bridge diffusivity. \(\sigma = 0\) recovers the deterministic straight-line path (equivalent to LinearInterpolant). Larger values inject more stochastic "wiggle" around the straight-line mean, matching the entropic-OT interpretation of the Schrödinger bridge (higher entropy regularisation gives a more diffusive bridge).	1.0
eps	float	numerical floor on the \(\sqrt{t(1-t)}\) denominator used when computing the target velocity, to keep it finite as \(t\) approaches 0 or 1.	0.0001

References

De Bortoli et al., "Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling" (2021), https://arxiv.org/abs/2106.01357. Tong et al., "Simulation-Free Schrödinger Bridges via Score and Flow Matching" (SF2M, 2024), https://arxiv.org/abs/2307.03672.

Source code in deltaflow/interpolants/schrodinger_bridge.py

def __init__(self, sigma: float = 1.0, eps: float = 1e-4):
    self.sigma = sigma
    self.eps = eps

VariancePreservingInterpolant

Bases: BaseInterpolant

Trigonometric variance-preserving (VP) probability path.

Noise \(x_0\) and data \(x_1\) are mixed with a trigonometric schedule

\[ x_t = \alpha_t\,x_1 + \sigma_t\,x_0, \qquad \alpha_t = \sin\!\left(\tfrac{\pi}{2}t\right), \quad \sigma_t = \cos\!\left(\tfrac{\pi}{2}t\right), \]

so that \(\alpha_t^2 + \sigma_t^2 = 1\) for every \(t\), the marginal variance is preserved along the path (hence variance-preserving), matching the geometry of a VP diffusion. Differentiating the path at fixed endpoints gives the conditional target velocity

\[ u_t = \alpha_t'\,x_1 + \sigma_t'\,x_0 = \frac{\pi}{2}\left[ \cos\!\left(\tfrac{\pi}{2}t\right) x_1 - \sin\!\left(\tfrac{\pi}{2}t\right) x_0 \right]. \]

References

Lipman et al., "Flow Matching for Generative Modeling" (2023), Sec. 3.2, https://arxiv.org/abs/2210.02747. Ma et al., "SiT: Exploring Flow and Diffusion-Based Generative Models" (2024), https://arxiv.org/abs/2401.08740.