Skip to content

Learning Objectives

Fitting \(p_\theta(x) \propto e^{-E_\theta(x)}\) by maximum likelihood gives the gradient

\[ \nabla_\theta \, \mathbb{E}_{x \sim p_{\text{data}}}[-\log p_\theta(x)] = \mathbb{E}_{x \sim p_{\text{data}}}[\nabla_\theta E_\theta(x)] - \mathbb{E}_{x \sim p_\theta}[\nabla_\theta E_\theta(x)] , \]

a push-down on data and a push-up on model samples. Every objective in torchebm.losses is a different answer to the hard part, the expectation under the model itself. The shipped family, generated from the installed package at build time:

graph TD
    BaseContrastiveDivergence(["BaseContrastiveDivergence"])
    BaseInterpolantLoss(["BaseInterpolantLoss"])
    BaseLoss(["BaseLoss"])
    BaseScoreMatching(["BaseScoreMatching"])
    BaseContrastiveDivergence --> ContrastiveDivergence
    BaseLoss --> BaseContrastiveDivergence
    BaseScoreMatching --> DenoisingScoreMatching
    BaseLoss --> BaseScoreMatching
    BaseInterpolantLoss --> EnergyMatchingLoss
    BaseLoss --> BaseInterpolantLoss
    BaseInterpolantLoss --> EquilibriumMatchingLoss
    BaseInterpolantLoss --> FlowMatchingLoss
    BaseContrastiveDivergence --> PersistentContrastiveDivergence
    BaseScoreMatching --> ScoreMatching
    BaseScoreMatching --> SlicedScoreMatching

MCMC-based: contrastive divergence

CD-k approximates model samples with k steps of MCMC started at the data1:

1
2
3
from torchebm.losses import ContrastiveDivergence
cd = ContrastiveDivergence(model=energy, sampler=langevin, k_steps=10)
loss, negatives = cd(batch)

persistent=True switches to PCD: negatives resume from a replay buffer instead of restarting at the data, so chains explore the model distribution across updates. CD trains slowly per step (an inner MCMC loop) but yields a genuine energy with meaningful level sets.

Simulation-free: score matching

Score matching sidesteps model samples entirely by fitting the score2. The exact objective needs the Hessian trace; the practical members are denoising SM3, which matches the score of noise-perturbed data,

1
2
from torchebm.losses import DenoisingScoreMatching
dsm = DenoisingScoreMatching(model=energy, noise_scale=0.1)

and sliced SM, which estimates the trace with random projections. DSM at a ladder of noise scales is the training principle underlying score-based diffusion4.

Transport-based: flow, equilibrium and energy matching

The modern objectives are simulation-free: they replace the inner sampling loop with regression along an interpolant path (see Interpolants and Couplings).

Flow matching regresses a time-conditioned velocity field v(x, t) onto the interpolant velocity u_t and generates by integrating it forward with FlowSampler (no negation). It shares the transport surface of the other matching losses: interpolant=, coupling= with per-pair weights, t_sampler= (uniform or the EDM lognormal skew), and a per-timestep loss_weight_fn.

1
2
3
4
5
6
7
from torchebm.losses import FlowMatchingLoss
from torchebm.samplers import FlowSampler

fm = FlowMatchingLoss(model=velocity_net, interpolant="linear")
# ... train ...
ode = FlowSampler(velocity_net, interpolant="linear")
samples = ode.sample(n_samples=64, dim=2, n_steps=50)

Equilibrium matching trains a time-invariant field f(x) toward the noise direction along the path (f points data -> noise), so every route transports noise -> data by moving along -f. With the implicit formulation (energy_type="none") f is the gradient field: integrate it with FlowSampler(negate_velocity=True), or descend it with the EqMEnergy adapter, which turns the field into the scalar BaseModel the gradient-based samplers and InteractionModel consume. The explicit formulation (energy_type="dot") trains the scalar energy g(x) = x . f(x) for gradient-descent sampling and OOD scoring. It also accepts a coupling= (default identity) and, like energy matching, honors per-pair coupling weights.

1
2
3
4
5
6
7
8
from torchebm.losses import EquilibriumMatchingLoss
from torchebm.models import EqMEnergy
from torchebm.samplers import FlowSampler, GradientDescentSampler

eqm = EquilibriumMatchingLoss(model=field, interpolant="linear", energy_type="none")
# ... train ...
ode = FlowSampler(field, negate_velocity=True, integrator="euler")   # integrate -f
gd = GradientDescentSampler(EqMEnergy.from_loss(eqm))                 # descend the energy

EqMEnergy.from_loss picks the adapter mode matching the trained energy_type (implicit vs dot/l2), so the sampled energy always matches what was trained. InteractionModel must wrap an explicit energy, never the implicit adapter.

The two losses differ only in the sign of the regression target and in the clock the model sees; each has a switch to adopt the other's convention:

Loss Target Clock shown to the model Sample with
FlowMatchingLoss u_t (noise -> data) sampled t FlowSampler
FlowMatchingLoss(negate_velocity=True) -u_t (data -> noise) sampled t FlowSampler(negate_velocity=True), descent samplers via EqMEnergy
EquilibriumMatchingLoss -u_t * c(t) zeros (time_invariant=True) FlowSampler(negate_velocity=True), EqMEnergy
EquilibriumMatchingLoss(time_invariant=False) -u_t * c(t) sampled t FlowSampler(negate_velocity=True)

With ct="constant", ct_multiplier=1 and time_invariant=False the EqM objective is bit-identical to FlowMatchingLoss(negate_velocity=True). EqMEnergy always evaluates the field at t = 0, so it suits time-invariant fields.

Energy matching (arXiv:2504.10612) keeps a single time-independent scalar potential: an OT flow-matching warm-up shapes it as transport, then a contrastive phase with temperature-scheduled Langevin negatives sharpens its Boltzmann density near the data. It accepts a coupling= and consumes per-pair weights when the coupling provides them:

1
2
3
4
from torchebm.couplings import SinkhornCoupling
from torchebm.losses import EnergyMatchingLoss
em = EnergyMatchingLoss(model=potential, coupling=SinkhornCoupling(reg=0.01),
                        epsilon_max=0.15, tau_star=0.8)

Choosing an objective

Objective Inner sampling Trains Generate with Reach for it when
CD / PCD / PT-CD yes (k MCMC steps) energy MCMC you need a calibrated energy and can afford MCMC per step
Exact / sliced SM no (Hessian term) energy Langevin low dimension, no noise tolerance
Denoising SM no energy (smoothed) annealed Langevin fast sampling-free training, noise scale acceptable
Flow matching no velocity field FlowSampler ODE/SDE pure generative transport, standard diffusion-style recipes
Equilibrium matching no field or energy FlowSampler ODE or EqMEnergy + gradient descent generative quality with few integration steps
Energy matching phase 2 only energy one Langevin sweep one potential for both transport and Boltzmann sampling

The rule of thumb embedded in the table: objectives with inner sampling buy energy fidelity at training cost; transport objectives buy training scalability and fast generation, and the hybrids exist to keep the energy while paying the transport price.

Runnable counterparts

  1. G. E. Hinton. Training products of experts by minimizing contrastive divergence. Neural Computation, 14(8), 2002. ↩

  2. A. Hyvärinen. Estimation of non-normalized statistical models by score matching. JMLR, 6, 2005. ↩

  3. P. Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 23(7), 2011. ↩

  4. Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. NeurIPS, 2019. ↩