core

core

Base abstractions shared across DeltaFlow's models, losses, and solvers.

Every user-facing component (velocity field, interpolant, solver, loss) subclasses one of these bases, so new variants are drop-in and not rewrites of the surrounding machinery.

BaseCoupling

Bases: ABC

Given a batch x1 of data samples, return a paired (x0, x1).

BaseEquilibriumField

Bases: Module, ABC

Base class for a time-invariant Equilibrium Matching field \(f(x)\).

Unlike a flow-matching velocity field \(v_\theta(x, t)\), an EqM field carries no time argument. It approximates the equilibrium gradient of an implicit energy, \(f_\theta(x) \approx -\nabla E(x)\), pointing from noise toward data. The interpolation coefficient \(\gamma\) is implicit and (per the EqM paper) never seen by the model, so forward takes only x and any extra conditioning, never a time or \(\gamma\).

Subclasses implement forward and return a tensor with the same shape as x. Additional conditioning is passed as keyword arguments and forwarded unchanged by EquilibriumMatchingLoss and EquilibriumSolver.

References

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

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

BaseLoss

Bases: ABC

Base class for a callable training loss.

Subclasses must implement __call__ and return a scalar tensor with requires_grad=True (assuming the model has trainable parameters). The signature is intentionally flexible - individual losses define which positional and keyword arguments they consume.

BaseSolver

BaseSolver(model, time_scale=1.0)

Bases: ABC

Base class for numerical integrators of dx/dt = v_theta(x, t).

A solver holds a reference to a velocity model and exposes two methods:

• step performs a single integration step from (x, t) to (x', t + dt). Subclasses implement the actual stepping rule.
• sample drives step in a loop from t_start to t_end and returns the final state.

Design note: PosteriorSolver wraps a BaseSolver and hooks the likelihood gradient into every call to step, so the base stepping logic is never duplicated.

Parameters:

Name	Type	Description	Default
model	Callable	callable model(x, t, **cond) -> velocity with the same signature as BaseVelocityField.	required
time_scale	float	multiplies the continuous t in [0, 1] before it is passed to the model, useful when the backbone was trained with a different numeric time convention (e.g. diffusion timesteps).	1.0

Source code in deltaflow/core/base_solver.py

def __init__(self, model: Callable, time_scale: float = 1.0):
    self.model = model
    self.time_scale = time_scale

step abstractmethod

step(x, t, dt, **cond)

Advance the state x from time t to t + dt.

Source code in deltaflow/core/base_solver.py

@abstractmethod
def step(self, x: torch.Tensor, t: float, dt: float, **cond) -> torch.Tensor:
    """Advance the state ``x`` from time ``t`` to ``t + dt``."""
    raise NotImplementedError

sample

sample(x, n_steps=50, t_start=0.0, t_end=1.0, show_progress=True, progress_desc=None, **cond)

Integrate from t_start to t_end in n_steps uniform steps.

Source code in deltaflow/core/base_solver.py

def sample(
    self,
    x: torch.Tensor,
    n_steps: int = 50,
    t_start: float = 0.0,
    t_end: float = 1.0,
    show_progress: bool = True,
    progress_desc: Optional[str] = None,
    **cond,
) -> torch.Tensor:
    """Integrate from ``t_start`` to ``t_end`` in ``n_steps`` uniform steps."""
    dt = (t_end - t_start) / n_steps
    steps = range(n_steps)
    if show_progress:
        try:
            from tqdm import tqdm

            steps = tqdm(
                steps,
                desc=progress_desc or type(self).__name__,
                total=n_steps,
                leave=False,
            )
        except ImportError:
            pass

    for i in steps:
        t_val = t_start + i * dt
        x = self.step(x, t_val, dt, **cond)
    return x

BaseVelocityField

Bases: Module, ABC

Base class for the time-conditioned velocity field v_theta(x, t).

Subclasses must implement forward and return a tensor with the same shape as x. Any additional conditioning (e.g. a guidance flag, class label, or cross-attention context) can be passed as keyword arguments and is forwarded unchanged by the losses and solvers.