TorchEBM：PyTorch能量基模型与扩散模型训练库

TorchEBM是一个统一的PyTorch库，支持能量基模型、流模型和扩散模型的构建与训练，提供MCMC采样、分数匹配、插值方案等完整工具链。

能量基模型（Energy-Based Models, EBM）通过标量能量函数定义概率分布，能量越低意味着概率越高。这是一种极其通用的建模框架——从MCMC采样到分数匹配，再到基于流的生成方法，都可以通过这个视角来理解。

TorchEBM正是基于这一理念构建的PyTorch库，它为整个生成式建模谱系提供了可组合的工具。你可以定义能量景观，使用各种学习目标训练模型，并通过MCMC、优化或学习的连续时间动力学（ODE/SDE）进行采样。

TorchEBM涵盖了从经典到现代的各种生成方法。在模型定义方面，库内置了多种解析势能函数，同时支持自定义神经网络能量函数。采样方面提供了MCMC采样器（如Langevin动力学）和基于优化的采样器，以及通过ODE/SDE积分生成样本的流和扩散采样器。

训练目标包括对比散度变体、分数匹配变体和平衡匹配。库还提供了噪声到数据路径的插值方案、SDE/ODE/哈密顿动力学的数值积分器、支持条件生成的神经网络架构、用于快速原型设计和基准测试的合成数据集，以及步长、噪声尺度等训练参数的超参数调度器。

使用TorchEBM非常直观。以下是一个简单的能量模型定义和采样示例：

import torch
from torchebm.core import GaussianModel
from torchebm.samplers import LangevinDynamics

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = GaussianModel(mean=torch.zeros(2), cov=torch.eye(2), device=device)

sampler = LangevinDynamics(model=model, step_size=0.01, device=device)
samples = sampler.sample(x=torch.randn(500, 2, device=device), n_steps=100)

对于自定义神经网络能量函数的训练，流程同样简洁：

class MLPEnergy(BaseModel):
    def __init__(self, dim):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(dim, 64), torch.nn.SiLU(),
            torch.nn.Linear(64, 64), torch.nn.SiLU(),
            torch.nn.Linear(64, 1),
        )
    def forward(self, x):
        return self.net(x).squeeze(-1)

model = MLPEnergy(dim=2).to(device)
sampler = LangevinDynamics(model=model, step_size=0.01, device=device)
cd_loss = ContrastiveDivergence(model=model, sampler=sampler, k_steps=10)

对于高维空间采样，TorchEBM提供了哈密顿蒙特卡洛（HMC）实现。HMC利用梯度信息引导采样，相比随机游走式的Metropolis-Hastings算法，在高维空间中具有更高的采样效率：

from torchebm.samplers import HamiltonianMonteCarlo

hmc = HamiltonianMonteCarlo(model=model, step_size=0.1, 
                            n_leapfrog_steps=10, device=device)
samples = hmc.sample(dim=10, n_steps=500, n_samples=1000)

TorchEBM采用清晰的模块化架构：

• core/：基类、能量模型、调度器、设备管理
• samplers/：MCMC、优化和流/扩散采样器
• losses/：训练目标（CD、分数匹配、平衡匹配）
• interpolants/：噪声到数据插值方案
• integrators/：SDE/ODE/哈密顿动力学数值积分器
• models/：神经网络架构
• datasets/：合成数据生成器
• utils/：可视化和训练工具
• cuda/：CUDA加速实现

这种模块化设计使得研究者可以灵活组合不同组件，快速实验新想法。

除了经典的能量基模型训练，TorchEBM还原生支持现代流模型和扩散模型。通过插值方案（如Linear、VP、Cosine），你可以定义从噪声到数据的转换路径，然后使用流匹配或分数匹配进行训练。

库中提供的动画示例展示了8个高斯分布上的流训练过程，以及圆形分布上的扩散采样效果，直观呈现了这些方法的工作原理。

TorchEBM支持CUDA加速和混合精度训练，在保证灵活性的同时兼顾性能。安装也非常简单，通过pip即可：pip install torchebm。所有依赖（主要是PyTorch及其CUDA支持）会自动处理。

对于希望深入研究的开发者，项目提供了完整的文档网站和丰富的示例脚本，覆盖了从基础使用到高级定制的各种场景。