solvers

solvers

ODE solvers that integrate a learned velocity field.

The historical module name was deltaflow.samplers, which still imports from here for backward compatibility.

EulerSolver

EulerSolver(model, time_scale=1.0)

Bases: BaseSolver

Explicit (forward) Euler integrator for the flow-matching ODE.

Integrates \(\mathrm{d}x/\mathrm{d}t = v_\theta(x, t)\) with the first-order update

\[ x_{n+1} = x_n + \Delta t\; v_\theta(x_n, t_n), \]

using a single velocity evaluation per step. It is cheap but incurs \(\mathcal{O}(\Delta t^2)\) local truncation error (\(\mathcal{O}(\Delta t)\) global), so prefer HeunSolver when accuracy at low step counts matters.

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

sample

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

Generate samples by integrating the velocity field.

Parameters:

Name	Type	Description	Default
x	Tensor	initial state (shape/device/dtype template if x_cond is set).	required
n_steps	int	number of Euler integration steps.	50
x_cond	Optional[Tensor]	optional partially-noised starting point, if given the initial state becomes (1 - t_start) * noise + t_start * x_cond.	None
t_start	float	starting time, used together with x_cond.	0.0
t_end	float	end time, defaults to 1.0.	1.0
**cond	Any	extra keyword arguments forwarded to the velocity model.	{}

Source code in deltaflow/solvers/euler.py

@torch.no_grad()
def sample(
    self,
    x: torch.Tensor,
    n_steps: int = 50,
    x_cond: Optional[torch.Tensor] = None,
    t_start: float = 0.0,
    t_end: float = 1.0,
    show_progress: bool = True,
    **cond: Any,
) -> torch.Tensor:
    """Generate samples by integrating the velocity field.

    Args:
        x: initial state (shape/device/dtype template if ``x_cond`` is set).
        n_steps: number of Euler integration steps.
        x_cond: optional partially-noised starting point, if given the
            initial state becomes ``(1 - t_start) * noise + t_start * x_cond``.
        t_start: starting time, used together with ``x_cond``.
        t_end: end time, defaults to ``1.0``.
        **cond: extra keyword arguments forwarded to the velocity model.
    """
    if x_cond is not None:
        x0 = torch.randn_like(x_cond)
        x = (1 - t_start) * x0 + t_start * x_cond
    return super().sample(
        x,
        n_steps=n_steps,
        t_start=t_start,
        t_end=t_end,
        show_progress=show_progress,
        progress_desc="EulerSolver",
        **cond,
    )

EquilibriumSolver

EquilibriumSolver(model, step_size=0.05, momentum=0.0)

Bases: BaseSolver

Gradient-descent sampler for an Equilibrium Matching landscape.

A field trained with EquilibriumInterpolant approximates the time-invariant equilibrium gradient of an implicit energy \(E\), pointing from noise toward data, \(f(x) \approx -\nabla E(x)\). Unlike flow/diffusion models, samples are not obtained by integrating a velocity over a fixed time horizon. Instead they are found by optimising a noise draw down the landscape,

\[ x_{k+1} = x_k + \eta\, f(x_k), \]

which is gradient descent on \(E\) (ascent along \(f\)). Because sampling is decoupled from any prescribed trajectory, the step size step_size and the number of steps n_steps are free knobs, and adaptive compute is just running more iterations.

Nesterov (NAG-GD). With momentum > 0 the solver uses the accelerated look-ahead update from the paper,

\[ x_k' = x_k + \eta\,\mu\, m_k,\quad m_{k+1} = f(x_k'),\quad x_{k+1} = x_k + \eta\, m_{k+1}, \]

which reaches high-quality samples in fewer steps than vanilla gradient descent.

No time. An EqM field is a BaseEquilibriumField, a time-invariant \(f(x)\) with no time (or \(\gamma\)) argument. The solver therefore calls it as model(x, **cond) and there is no eval_time or time_scale to route, sampling is pure optimisation on the landscape.

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
model	Callable	callable model(x, **cond) -> gradient (same signature as BaseEquilibriumField), returning the equilibrium gradient with the shape of x.	required
step_size	float	gradient-descent step \(\eta\).	0.05
momentum	float	Nesterov coefficient \(\mu\). 0 gives vanilla gradient descent, positive values give NAG-GD.	0.0

Source code in deltaflow/solvers/gradient_descent.py

def __init__(
    self,
    model: Callable,
    step_size: float = 0.05,
    momentum: float = 0.0,
):
    super().__init__(model, time_scale=1.0)
    self.step_size = step_size
    self.momentum = momentum

step

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

Single vanilla gradient-descent step (t/dt are ignored).

Provided for BaseSolver compatibility. The momentum-aware loop lives in sample.

Source code in deltaflow/solvers/gradient_descent.py

def step(self, x: torch.Tensor, t: float, dt: float, **cond) -> torch.Tensor:
    """Single vanilla gradient-descent step (``t``/``dt`` are ignored).

    Provided for `BaseSolver` compatibility. The momentum-aware loop lives
    in `sample`.
    """
    return x + self.step_size * self._gradient(x, **cond)

sample

sample(x, n_steps=100, step_size=None, momentum=None, show_progress=True, **cond)

Sample by gradient descent on the learned equilibrium landscape.

Parameters:

Name	Type	Description	Default
x	Tensor	initial noise state (shape/device/dtype template).	required
n_steps	int	number of gradient-descent iterations.	100
step_size	Optional[float]	overrides self.step_size for this call.	None
momentum	Optional[float]	overrides self.momentum for this call (Nesterov look-ahead when positive).	None
show_progress	bool	display a tqdm progress bar if available.	True
**cond	Any	extra keyword arguments forwarded to the model unchanged.	{}

Source code in deltaflow/solvers/gradient_descent.py

@torch.no_grad()
def sample(
    self,
    x: torch.Tensor,
    n_steps: int = 100,
    step_size: Optional[float] = None,
    momentum: Optional[float] = None,
    show_progress: bool = True,
    **cond: Any,
) -> torch.Tensor:
    """Sample by gradient descent on the learned equilibrium landscape.

    Args:
        x: initial noise state (shape/device/dtype template).
        n_steps: number of gradient-descent iterations.
        step_size: overrides ``self.step_size`` for this call.
        momentum: overrides ``self.momentum`` for this call (Nesterov
            look-ahead when positive).
        show_progress: display a tqdm progress bar if available.
        **cond: extra keyword arguments forwarded to the model unchanged.
    """
    eta = self.step_size if step_size is None else step_size
    mu = self.momentum if momentum is None else momentum

    steps = range(n_steps)
    if show_progress:
        try:
            from tqdm import tqdm

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

    m = torch.zeros_like(x)
    for _ in steps:
        if mu != 0.0:
            x_look = x + eta * mu * m
            m = self._gradient(x_look, **cond)
            grad = m
        else:
            grad = self._gradient(x, **cond)
        x = x + eta * grad
    return x

HeunSolver

HeunSolver(model, time_scale=1.0)

Bases: BaseSolver

Heun (improved-Euler) second-order predictor-corrector integrator.

Each step takes an Euler predictor and averages the velocity at the current and predicted states,

\[ \begin{aligned} k_1 &= v_\theta(x_n, t_n), \\ k_2 &= v_\theta\bigl(x_n + \Delta t\,k_1,\; t_n + \Delta t\bigr), \\ x_{n+1} &= x_n + \tfrac{\Delta t}{2}\,(k_1 + k_2). \end{aligned} \]

This costs two velocity evaluations per step but has \(\mathcal{O}(\Delta t^3)\) local truncation error (\(\mathcal{O}(\Delta t^2)\) global), so it typically matches Euler's quality at half the number of steps. It is the trapezoidal-rule integrator widely used in EDM-style samplers.

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

PosteriorSolver

PosteriorSolver(base_solver, likelihood, tweedie=None, guidance_scale=1.0, grad_normalize=False)

Bases: BaseSolver

Base-solver wrapper that adds a per-step measurement-likelihood gradient.

For an inverse problem with measurement \(y = A(x) + n\), this samples from the posterior \(p(x \mid y)\) by nudging an unconditional flow-matching solver toward the data-consistency term at every step. Given the current state \(x_t\) and velocity \(v_\theta(x_t, t)\), one step is

\[ \hat{x}_1 = \mathcal{T}(x_t, v_\theta(x_t, t), t), \qquad x_{t+\Delta t} = \underbrace{\text{Base}(x_t, t, \Delta t)}_{\text{unconditional step}} \;-\; \eta\,\nabla_{x_t}\bigl[-\log p(y \mid \hat{x}_1)\bigr], \]

where \(\mathcal{T}\) is the flow-matching Tweedie decomposition that maps \((x_t, v_t, t)\) to the clean-signal estimate \(\hat{x}_1\) (see deltaflow.inverse.tweedie), and \(\eta\) is the guidance_scale. The likelihood gradient is obtained by autograd, so if the velocity field operates on VAE latents while \(A\) is defined on pixels, passing a decoder to the likelihood object pulls the gradient back into latent space automatically. Only the sampling-time ODE is modified. The pretrained velocity field is untouched.

References

Kim et al., "FlowDPS: Flow-Driven Posterior Sampling for Inverse Problems" (2025), https://arxiv.org/abs/2503.08136. "Flower: A Flow-Matching Solver for Inverse Problems" (2025), https://arxiv.org/abs/2509.26287.

Parameters:

Name	Type	Description	Default
base_solver	BaseSolver	any BaseSolver that already integrates the unconditional flow (Euler, Heun, ...).	required
likelihood	Likelihood	object with a .neg_log_prob(x_clean) method that returns a per-sample (or reducible) scalar tensor with requires_grad=True support. See deltaflow.inverse.likelihood.	required
tweedie	Optional[BaseTweedie]	flow-matching Tweedie decomposition \(\mathcal{T}\) to derive \(\hat{x}_1\) from \((x_t, v_t, t)\). Defaults to LinearTweedie, matching a LinearInterpolant training path.	None
guidance_scale	float	step size \(\eta\) on the likelihood gradient. Larger values snap harder to the measurement but risk over-shooting.	1.0
grad_normalize	bool	if True, the injected gradient is rescaled to match the norm of the base step. This is a stability trick used in some DPS variants when likelihood magnitudes vary wildly.	False

Source code in deltaflow/solvers/posterior_solver.py

def __init__(
    self,
    base_solver: BaseSolver,
    likelihood: Likelihood,
    tweedie: Optional[BaseTweedie] = None,
    guidance_scale: float = 1.0,
    grad_normalize: bool = False,
):
    # Re-use the base solver's model reference and time_scale.
    super().__init__(model=base_solver.model, time_scale=base_solver.time_scale)
    self.base = base_solver
    self.likelihood = likelihood
    self.tweedie = tweedie or LinearTweedie()
    self.guidance_scale = guidance_scale
    self.grad_normalize = grad_normalize