Torchebm > Losses > Equilibrium_matching

Contents

Classes

• EquilibriumMatchingLoss - Equilibrium Matching (EqM) training loss.

API Reference

torchebm.losses.equilibrium_matching

Equilibrium Matching (EqM) loss.

Implements time-invariant equilibrium training objectives for learning energy landscapes, following the EqM paper:

• Implicit EqM (\(L_{EqM}\)): Learns gradient field directly

[ L_{EqM} = |f(x_\gamma) - (\epsilon - x) \cdot c(\gamma)|^2 ]

• Explicit EqM-E (\(L_{EqM-E}\)): Learns scalar energy via gradient matching

[ L_{EqM-E} = |\nabla g(x_\gamma) - (\epsilon - x) \cdot c(\gamma)|^2 ]

Key differences from Flow Matching: - Time-invariant: Model zeros out time conditioning internally - Gradient direction: EqM learns (noise - data), FM learns (data - noise)