Home

TorchEBM

TorchEBM: Simulation-free, GPU-first generative modeling in PyTorch
Composable primitives for scalable, stable training of modern EBMs, diffusion, flow matching, and Schrödinger bridges.

Get Started Star on GitHub

Overview¶

TorchEBM is a PyTorch library for simulation-free, GPU-first generative modeling: scalable, stable training of modern energy-based models, diffusion, flow matching, and Schrödinger bridges. Energy-based models define probability distributions through a scalar energy function, and the formulation is general enough that much of modern generative modeling, from MCMC sampling and score matching to simulation-free transport along probability paths, factors into the same components, i.e. fields, probability paths, couplings, objectives, and integrators. TorchEBM implements these components as composable, high-throughput PyTorch primitives.

The Design and Scope page states this framing precisely and places each method family within it.

In Action¶

Equilibrium matching on eight gaussians — Eight-gaussians distribution

Equilibrium matching on circles — Circles distribution

Equilibrium matching with different interpolants transporting noise onto structured distributions.

Core Components¶

Quick Start¶

1	`pip install torchebm`

Train a generative model with equilibrium matching, then sample the same network both as an ODE flow and as a scalar energy:

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchebm.losses import EquilibriumMatchingLoss
from torchebm.samplers import FlowSampler, NesterovSampler
from torchebm.core import BaseModel
from torchebm.datasets import EightGaussiansDataset

dataset = EightGaussiansDataset(n_samples=8192)
loader = DataLoader(dataset, batch_size=256, shuffle=True)

# Any nn.Module with forward(x, t) works
class VelocityNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, 256), nn.SiLU(),
                                 nn.Linear(256, 256), nn.SiLU(), nn.Linear(256, 2))
    def forward(self, x, t, **kwargs):
        return self.net(x)

model = VelocityNet()

loss_fn = EquilibriumMatchingLoss(
    model=model, interpolant="linear", energy_type="dot",
)
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)

for epoch in range(50):
    for x in loader:
        loss = loss_fn(x)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# Sample via ODE flow
flow = FlowSampler(model=model, interpolant="linear", negate_velocity=True)
flow_samples = flow.sample(x=torch.randn(1000, 2), n_steps=100)

# Same model as a scalar energy: g(x) = x · f(x)
class LearnedEnergy(BaseModel):
    def __init__(self, net):
        super().__init__()
        self.net = net
    def forward(self, x):
        t = torch.zeros(x.shape[0], device=x.device)
        return (x * self.net(x, t)).sum(-1)

nesterov = NesterovSampler(LearnedEnergy(model), step_size=0.01, momentum=0.9)
energy_samples = nesterov.sample(n_samples=1000, dim=2, n_steps=200)

One model, two views: the flow view generates in a fixed number of steps, the energy view supports gradient-based refinement. That interchangeability is the library's central design property.

See the concepts and examples for the theory behind each component and a runnable, CI-tested curriculum.

Enjoying TorchEBM? A GitHub star helps others discover the project and motivates continued development.

Star on GitHub

Citation¶

If TorchEBM is useful in your research, please cite it:

@misc{torchebm_library_2025,
  author       = {Ghaderi, Soran and Contributors},
  title        = {{TorchEBM}: Simulation-Free, {GPU}-First Generative Modeling in {PyTorch}},
  year         = {2025},
  url          = {https://github.com/soran-ghaderi/torchebm},
}