Energy Matching中的训练目标分析

Energy Matching中的训练目标分析

在《Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modeling》一文中，作者提出了一种新颖的生成模型框架——Energy Matching，通过结合最优传输（Optimal Transport, OT）和能量基础模型（Energy-Based Models, EBMs）的优势，实现高效的样本生成和显式的似然估计。本文将详细分析论文中2.1节“Training Objectives”部分，重点探讨Flow-like Objective ((${\mathcal{L}}_{\text{OT}}$))、Contrastive Objective ((${\mathcal{L}}_{\text{CD}}$)) 以及 Dual Objective 的定义、作用和实现方式，并结合直观解释和上下文背景，帮助读者深入理解这些训练目标在Energy Matching框架中的重要性。

---

paper：https://arxiv.org/pdf/2504.10612

背景：Energy Matching的训练目标

具体背景知识可以参考笔者的另一篇博客：JKO方案中的一阶最优性条件与生成框架

Energy Matching的核心思想是通过学习一个时间无关的标量势函数 ( $V_{\theta }(x)$ )，实现从噪声分布到数据分布的平滑过渡。其训练过程分为两个阶段：

1. 远离数据流形（Flow-like Regime）：当样本距离数据流形较远时，利用最优传输的特性，通过确定性流快速将噪声样本引导到数据流形附近。此时，温度参数 ($\epsilon (t)\approx 0$)，训练目标主要由 (${\mathcal{L}}_{\text{OT}}$) 驱动。
2. 接近数据流形（EBM-like Regime）：当样本接近数据流形时，引入熵项（($\epsilon (t)\approx {\epsilon }_{\text{max}}$)），通过对比散度（Contrastive Divergence）优化，使 ( $V_{\theta }(x)$ ) 形成一个Boltzmann分布，精确匹配数据分布。此时，训练目标结合 (${\mathcal{L}}_{\text{OT}}$) 和 (${\mathcal{L}}_{\text{CD}}$)。

2.1节详细描述了这两个训练目标以及它们的联合优化方式，以下逐一分析。

---

2.1.1 Flow-like Objective (${\mathcal{L}}_{\text{OT}}$)

定义

Flow-like Objective (${\mathcal{L}}_{\text{OT}}$) 的目标是构建一个全局速度场 (${\nabla }_x{V}_{\theta }(x)$)，将噪声样本 ( { x 0 } \{x_0\} {x0​}) 高效地运送到数据样本 ( { x data } \{x_{\text{data}}\} {xdata​})，尽量减少路径上的“弯路”。为此，论文利用Wasserstein空间中的测地线（geodesics）来定义传输路径。具体而言，(${\mathcal{L}}_{\text{OT}}$) 的损失函数定义为：

$${\mathcal{L}}_{\text{OT}}={\mathbb{E}}_{t\sim U(0,1)}\left[{\Vert {\nabla }_x{V}_{\theta }(x_t^{(i)})+x_{\text{data}}^{(i)}-T(x_{\text{data}}^{(i)})\Vert }^2\right]$$

其中：

• ( $x_t^{(i)}=(1-t)T(x_{\text{data}}^{(i)})+t{x}_{\text{data}}^{(i)}$ )：表示沿测地线的插值点，( $t\in [0,1]$ )，从噪声样本 ( $T(x_{\text{data}}^{(i)})$ ) 到数据样本 ( $x_{\text{data}}^{(i)}$)。
• ( $T$ )：最优传输映射（OT map），通过OT求解器（如POT库）计算，定义了从噪声分布到数据分布的配对。
• (${\nabla }_x{V}_{\theta }(x_t^{(i)})$)：势函数 ( $V_{\theta }(x)$ ) 在插值点 ( $x_t^{(i)}$ ) 的梯度，表示模型预测的速度场。
• ( $x_{\text{data}}^{(i)}-T(x_{\text{data}}^{(i)})$ )：目标速度，表示样本从噪声到数据的理想移动方向。

实现方式

1. 数据准备：

   • 从数据分布中采样一个mini-batch ( { x data ( i ) } i = 1 B \{x_{\text{data}}^{(i)}\}_{i=1}^B {xdata(i)​}i=1B​)，表示真实数据样本。
   • 从噪声分布（通常为标准高斯分布 ($\mathcal{N}(0,I)$)) 采样等数量的噪声样本 ( { x 0 ( i ) } i = 1 B \{x_0^{(i)}\}_{i=1}^B {x0(i)​}i=1B​)。
   • 使用OT求解器计算最优传输映射 ( $T$ )，将噪声样本 ( { x 0 ( i ) } \{x_0^{(i)}\} {x0(i)​}) 配对到数据样本 ( { x data ( i ) } \{x_{\text{data}}^{(i)}\} {xdata(i)​})。

2. 插值与速度计算：

   • 对于每个数据样本 ( $x_{\text{data}}^{(i)}$ )，通过线性插值生成路径上的点：
     $$x_t^{(i)}=(1-t)T(x_{\text{data}}^{(i)})+t{x}_{\text{data}}^{(i)}$$
   • 目标速度为 ( $x_{\text{data}}^{(i)}-T(x_{\text{data}}^{(i)})$ )，表示样本以恒定速度从噪声移动到数据。

3. 优化目标：

   • 优化 (${\mathcal{L}}_{\text{OT}}$)，使模型预测的梯度 ($-{\nabla }_x{V}_{\theta }(x_t^{(i)})$) 尽量接近目标速度 ( $x_{\text{data}}^{(i)}-T(x_{\text{data}}^{(i)})$ )。
   • 期望操作 (${\mathbb{E}}_{t\sim U(0,1)}$) 确保损失在整个路径 ($[0,1]$) 上平均计算。

直观解释

(${\mathcal{L}}_{\text{OT}}$) 的作用类似于“导航系统”，为样本从噪声到数据的移动提供方向。想象一群粒子（噪声样本）需要移动到目标位置（数据样本），(${\mathcal{L}}_{\text{OT}}$) 确保这些粒子沿着最短路径（Wasserstein测地线）移动，并且模型预测的“推力” ($-{\nabla }_x{V}_{\theta }(x)$) 与理想路径一致。

这一目标的特别之处在于：

• 无旋条件（Curl-free）：由于速度场由标量势函数的梯度 (${\nabla }_x{V}_{\theta }(x)$) 定义，它是无旋的（curl-free），这与最优传输的特性一致，避免了不必要的旋转路径，降低了传输成本。
• 时间无关：与传统Flow Matching方法使用时间依赖的速度场不同，(${\mathcal{L}}_{\text{OT}}$) 假设速度场是静态的，仅依赖于 ( $V_{\theta }(x)$ )，简化了模型设计。

在框架中的作用

(${\mathcal{L}}_{\text{OT}}$) 是第一阶段（Phase 1，Algorithm 1）的主要训练目标，确保噪声样本快速且高效地接近数据流形。通过预训练 ( $V_{\theta }(x)$ ) 以形成平滑的传输路径，(${\mathcal{L}}_{\text{OT}}$) 为后续的对比散度优化提供了高质量的初始样本，避免了EBM训练中的模式崩塌问题。

---

2.1.2 Contrastive Objective (${\mathcal{L}}_{\text{CD}}$)

定义

Contrastive Objective (${\mathcal{L}}_{\text{CD}}$) 的目标是调整 ( $V_{\theta }(x)$ )，使平衡分布 (${\rho }_{\text{eq}}(x)\propto \exp \left(-\frac{V_{\theta }(x)}{{\epsilon }_{\text{max}}}\right)$) 精确匹配数据分布。论文采用经典的对比散度（Contrastive Divergence）损失，定义为：

$${\mathcal{L}}_{\text{CD}}={\mathbb{E}}_{x\sim {\rho }_{\text{data}}}\left[\frac{V_{\theta }(x)}{{\epsilon }_{\text{max}}}\right]-{\mathbb{E}}_{\tilde{x}\sim \text{sg}({\rho }_{\text{eq}})}\left[\frac{V_{\theta }(\tilde{x})}{{\epsilon }_{\text{max}}}\right]$$

其中：

• (${\rho }_{\text{data}}$)：真实数据分布，样本 ( $x$ ) 从中抽取。
• (${\rho }_{\text{eq}}$)：由 ( $V_{\theta }(x)$ ) 诱导的平衡分布，负样本 ($\tilde{x}$) 通过MCMC朗之万链（Langevin chain）近似采样。
• ($\text{sg}(\cdot )$)：停止梯度操作（stop-gradient），确保梯度不通过采样过程反向传播。
• (${\epsilon }_{\text{max}}$)：最大温度参数，控制Boltzmann分布的平滑程度。

实现方式

1. 负样本采样：

   • 使用朗之万动力学（Langevin Dynamics）生成负样本 ($\tilde{x}$)，更新公式为：
     $$x_{m+1}=x_m-\Delta t{\nabla }_x{V}_{\theta }(x_m)+\sqrt{2\Delta t{\epsilon }^{(m)}}\eta ,\eta \sim \mathcal{N}(0,I)$$
   • 初始样本分为两部分：
     • 一半从真实数据 ( $x\sim {\rho }_{\text{data}}$ ) 初始化，确保负样本探索数据流形附近的高密度区域。
     • 一半从噪声分布（如 ($\mathcal{N}(0,I)$)) 初始化，探索远离数据流形的区域，塑造全局能量景观。

2. 损失计算：

   • 正样本项：计算真实数据样本 ( $x$ ) 的能量期望 (${\mathbb{E}}_{x\sim {\rho }_{\text{data}}}\left[\frac{V_{\theta }(x)}{{\epsilon }_{\text{max}}}\right]$)，目标是降低数据点的能量。
   • 负样本项：计算负样本 ($\tilde{x}$) 的能量期望 (${\mathbb{E}}_{\tilde{x}\sim \text{sg}({\rho }_{\text{eq}})}\left[\frac{V_{\theta }(\tilde{x})}{{\epsilon }_{\text{max}}}\right]$)，目标是提高非数据点的能量。
   • 总损失 (${\mathcal{L}}_{\text{CD}}$) 通过对比正负样本的能量差，优化 ( $V_{\theta }(x)$ ) 以形成低能量的数据流形。

直观解释

(${\mathcal{L}}_{\text{CD}}$) 类似于一个“雕刻师”，通过对比数据样本和负样本的能量，雕刻出能量函数 ( $V_{\theta }(x)$ ) 的形状：

• 对于真实数据点 ( $x$ )，(${\mathcal{L}}_{\text{CD}}$) 试图降低其能量 ( $V_{\theta }(x)$ )，使数据区域成为能量“洼地”。
• 对于负样本 ($\tilde{x}$)，(${\mathcal{L}}_{\text{CD}}$) 试图提高其能量 ( $V_{\theta }(x)$ )，使非数据区域成为能量“高地”。
• 朗之万动力学的采样过程模拟了粒子在能量景观中的随机游走，负样本的初始化策略（数据+噪声）确保能量函数既能精确建模数据流形，又能塑造全局结构。

在框架中的作用

(${\mathcal{L}}_{\text{CD}}$) 是第二阶段（Phase 2，Algorithm 2）的重要组成部分，负责在数据流形附近精细调整 ( $V_{\theta }(x)$ )，形成Boltzmann分布 (${\rho }_{\text{eq}}(x)$)。通过对比散度，模型能够捕捉数据的局部密度结构，同时避免传统EBM训练中的模式崩塌问题（由于初始样本的质量由 (${\mathcal{L}}_{\text{OT}}$) 保证）。

---

2.1.3 Dual Objective

定义

为了平衡 (${\mathcal{L}}_{\text{OT}}$) 和 (${\mathcal{L}}_{\text{CD}}$) 两个目标，论文采用了一种双目标优化策略，通过线性温度调度（如下）协调两个阶段的训练：

$$\epsilon (t)=\{\begin{array}{ll}0,& 0\le t<{\tau }^{\ast },\\ {\epsilon }_{\text{max}}\frac{t-{\tau }^{\ast }}{1-{\tau }^{\ast }},& {\tau }^{\ast }\le t\le 1,\\ {\epsilon }_{\text{max}},& t\ge 1.\end{array}$$

总损失函数为：

$$\mathcal{L}(\theta )={\mathcal{L}}_{\text{OT}}+{\lambda }_{\text{CD}}{\mathcal{L}}_{\text{CD}}$$

其中：

• (${\lambda }_{\text{CD}}$)：数据集特定的超参数，用于平衡 (${\mathcal{L}}_{\text{CD}}$) 的贡献。
• 温度 ($\epsilon (t)$)：通过时间调度控制熵项的影响，早期为0（强调 (${\mathcal{L}}_{\text{OT}}$)），后期逐渐增加到 (${\epsilon }_{\text{max}}$)（引入 (${\mathcal{L}}_{\text{CD}}$)）。

实现方式

1. Phase 1（预训练，Algorithm 1）：

   • 仅使用 (${\mathcal{L}}_{\text{OT}}$) 优化 ( $V_{\theta }(x)$ )，($\epsilon (t)=0$)。
   • 目标是建立从噪声到数据的平滑传输路径，生成高质量的负样本，为后续对比散度优化奠定基础。
   • 训练过程包括采样噪声和数据、计算OT映射、优化速度场匹配。

2. Phase 2（主训练，Algorithm 2）：

   • 联合优化 (${\mathcal{L}}_{\text{OT}}+{\lambda }_{\text{CD}}{\mathcal{L}}_{\text{CD}}$)，随时间 ( $t$ ) 增加 ($\epsilon (t)$) 至 (${\epsilon }_{\text{max}}$)。
   • 对于负样本采样，朗之万动力学根据当前温度 (${\epsilon }^{(m)}$)（或 (${\epsilon }_{\text{max}}$) 用于数据初始化的样本）进行更新。
   • 每次迭代计算 (${\mathcal{L}}_{\text{OT}}$) 和 (${\mathcal{L}}_{\text{CD}}$)，通过梯度下降更新模型参数 ($\theta$)。

3. 超参数：

   • 采样时间 (${\tau }_s$)：控制总采样步数，实验表明在CIFAR-10上 (${\tau }_s=3.0$) 时生成质量（FID）达到稳定。
   • (${\tau }^{\ast }$)：控制温度从0过渡到 (${\epsilon }_{\text{max}}$) 的时间点（如 (${\tau }^{\ast }=0.9$)）。
   • (${\lambda }_{\text{CD}}$)：如CIFAR-10使用 (${\lambda }_{\text{CD}}=2\times 10^{-4}$)，CelebA使用 (${\lambda }_{\text{CD}}=2\times 10^{-5}$)。

直观解释

双目标优化就像是“先搭框架，再精雕细琢”：

• Phase 1 使用 (${\mathcal{L}}_{\text{OT}}$) 搭建一个粗略的能量景观，确保样本可以快速从噪声区域移动到数据流形附近，就像建造一条高速公路。
• Phase 2 引入 (${\mathcal{L}}_{\text{CD}}$)，在数据流形附近精雕细琢，形成精确的Boltzmann分布，就像在目标区域精心设计地形。
• 温度调度 ($\epsilon (t)$) 起到“软切换”的作用，早期强调确定性流（OT），后期引入随机性（EBM），实现平滑过渡。

在框架中的作用

双目标策略是Energy Matching框架的核心创新之一：

• 稳定性：通过分阶段训练（先 (${\mathcal{L}}_{\text{OT}}$)，后联合优化），避免了传统EBM训练中的不稳定性和模式崩塌。
• 高效性：(${\mathcal{L}}_{\text{OT}}$) 提供了高质量的初始样本，减少了朗之万采样所需的步数。
• 灵活性：温度调度和 (${\lambda }_{\text{CD}}$) 允许根据数据集调整流和EBM的平衡，适应不同复杂度的生成任务。

---

例子：CIFAR-10生成

以CIFAR-10数据集为例，说明训练目标的应用：

1. Phase 1：

   • 采样128个数据样本和128个高斯噪声样本，计算OT映射 ( $T$ )。
   • 优化 (${\mathcal{L}}_{\text{OT}}$)，使 (${\nabla }_x{V}_{\theta }(x_t^{(i)})$) 匹配目标速度，训练200k次迭代。
   • 结果：形成平滑的传输路径，样本从噪声快速接近CIFAR-10图像的流形。

2. Phase 2：

   • 继续优化 (${\mathcal{L}}_{\text{OT}}$)，同时引入 (${\mathcal{L}}_{\text{CD}}$)（(${\lambda }_{\text{CD}}=2\times 10^{-4}$)）。
   • 使用朗之万动力学生成负样本（200步，初始化为50%数据+50%噪声）。
   • 温度从0逐渐增加到 (${\epsilon }_{\text{max}}=0.01$)，训练25k次迭代。
   • 结果：形成Boltzmann分布，FID达到3.97，显著优于传统EBM的8.61。

---

总结

Energy Matching的训练目标通过 (${\mathcal{L}}_{\text{OT}}$) 和 (${\mathcal{L}}_{\text{CD}}$) 的协同作用，实现了从噪声到数据的平滑过渡和精确的密度建模：

• (${\mathcal{L}}_{\text{OT}}$) 提供了一个高效的确定性流，将噪声样本快速引导到数据流形，奠定了稳定的训练基础。
• (${\mathcal{L}}_{\text{CD}}$) 精细调整能量函数，形成Boltzmann分布，捕捉数据的局部密度结构。
• 双目标优化 通过温度调度和分阶段训练，平衡了OT和EBM的优势，确保了训练的稳定性和生成质量。

通过这些训练目标，Energy Matching不仅在CIFAR-10等数据集上取得了优异的生成性能（FID 3.97），还为逆问题求解和局部内在维度估计提供了灵活的框架。这一方法展示了静态、无旋生成模型的潜力，为未来的生成模型研究开辟了新的方向。

代码实现

为了复现《Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modeling》一文中在CIFAR-10数据集上的实验，我们需要实现Energy Matching的训练代码，遵循论文中描述的模型结构和超参数。以下是一个完整的Python代码实现，使用PyTorch框架，基于论文中的描述（特别是2.1节和Appendix C）。代码包括模型定义、训练流程（Phase 1 和 Phase 2）、以及必要的工具函数。

---

代码概述

模型结构

根据论文Appendix C和Figure 5，模型结构为：

• UNet：与[Tong et al., 2023]相同的UNet架构，参数量约37M，输入为3×32×32的CIFAR-10图像。
• Vision Transformer (ViT)：一个8层ViT头（包括PatchEmbed），参数量约12M，输出标量势函数 ( $V_{\theta }(x)$ )。
• 总参数量约49M。

超参数

根据Appendix C（CIFAR-10部分）：

• 采样时间：(${\tau }_s=3.0$)
• 温度切换点：(${\tau }^{\ast }=0.9$)
• 时间步长：($\Delta t=0.01$)
• 朗之万采样步数：($M_{\text{Langevin}}=200$)
• 训练迭代：Phase 1为200k次，Phase 2为25k次
• 批大小：128
• 学习率：($8\times 10^{-4}$)
• 最大温度：(${\epsilon }_{\text{max}}=0.01$)
• 对比散度权重：(${\lambda }_{\text{CD}}=2\times 10^{-4}$)

训练流程

• Phase 1（Algorithm 1）：仅优化 (${\mathcal{L}}_{\text{OT}}$)，预训练模型以构建从噪声到数据的流。
• Phase 2（Algorithm 2）：联合优化 (${\mathcal{L}}_{\text{OT}}+{\lambda }_{\text{CD}}{\mathcal{L}}_{\text{CD}}$)，引入朗之万采样以形成Boltzmann分布。

依赖

• PyTorch：用于模型定义和训练
• POT（Python Optimal Transport）：用于计算最优传输映射
• torchvision：加载CIFAR-10数据集
• einops：处理张量操作
• timm：提供ViT实现

---

完整代码

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import ot  # Python Optimal Transport (POT)
import numpy as np
from einops import rearrange
from timm import create_model
import torch.nn.functional as F

# 超参数（根据Appendix C）
BATCH_SIZE = 128
LEARNING_RATE = 8e-4
TAU_S = 3.0
TAU_STAR = 0.9
DELTA_T = 0.01
EPSILON_MAX = 0.01
LAMBDA_CD = 2e-4
PHASE1_ITERS = 200000
PHASE2_ITERS = 25000
LANGEVIN_STEPS = 200
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# UNet模型（简化为示例，实际应使用[Tong et al., 2023]的UNet）
class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=64):
        super(UNet, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels, 64, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, padding=1, stride=2),
            nn.ReLU(),
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.Conv2d(64, out_channels, 3, padding=1),
        )
    
    def forward(self, x):
        enc = self.encoder(x)
        dec = self.decoder(enc)
        return dec

# Energy Matching模型（UNet + ViT）
class EnergyMatchingModel(nn.Module):
    def __init__(self):
        super(EnergyMatchingModel, self).__init__()
        self.unet = UNet(in_channels=3, out_channels=64)
        # ViT头（使用timm提供的ViT，简化为small模型）
        self.vit = create_model('vit_small_patch16_224', pretrained=False, num_classes=1)
        self.patch_embed = nn.Conv2d(64, 384, kernel_size=16, stride=16)  # 适配UNet输出
    
    def forward(self, x):
        unet_out = self.unet(x)  # [B, 64, 32, 32]
        vit_in = self.patch_embed(unet_out)  # [B, 384, 2, 2]
        vit_in = rearrange(vit_in, 'b c h w -> b (h w) c')  # [B, 4, 384]
        vit_out = self.vit.forward_features(vit_in)  # [B, 1]
        return vit_out.squeeze(-1)  # [B]

# 温度调度
def epsilon_schedule(t, tau_star=TAU_STAR, epsilon_max=EPSILON_MAX):
    if t < tau_star:
        return 0.0
    elif t <= 1.0:
        return epsilon_max * (t - tau_star) / (1.0 - tau_star)
    else:
        return epsilon_max

# OT损失
def compute_ot_loss(model, x_data, x_noise, t):
    B = x_data.size(0)
    # 计算OT映射
    cost_matrix = torch.cdist(x_data.view(B, -1), x_noise.view(B, -1)) ** 2
    a, b = torch.ones(B, device=DEVICE) / B, torch.ones(B, device=DEVICE) / B
    transport_plan = ot.emd(a, b, cost_matrix.detach().cpu().numpy())
    transport_plan = torch.tensor(transport_plan, device=DEVICE, dtype=torch.float32)
    
    # 找到配对
    indices = torch.argmax(transport_plan, dim=1)
    x_mapped = x_noise[indices]
    
    # 插值
    x_t = (1 - t) * x_mapped + t * x_data
    
    # 计算目标速度
    target_velocity = x_data - x_mapped
    
    # 计算模型梯度
    x_t = x_t.requires_grad_(True)
    v_theta = model(x_t)
    grad_v = torch.autograd.grad(v_theta.sum(), x_t, create_graph=True)[0]
    
    # OT损失
    loss = torch.mean((grad_v + target_velocity) ** 2)
    return loss

# 对比散度损失
def compute_cd_loss(model, x_data, x_noise, epsilon):
    # 正样本能量
    pos_energy = model(x_data).mean() / EPSILON_MAX
    
    # 负样本采样（朗之万动力学）
    x_neg = x_noise.clone().requires_grad_(False)
    for _ in range(LANGEVIN_STEPS):
        x_neg = x_neg.requires_grad_(True)
        v_theta = model(x_neg)
        grad_v = torch.autograd.grad(v_theta.sum(), x_neg, create_graph=True)[0]
        noise = torch.randn_like(x_neg) * (2 * DELTA_T * epsilon) ** 0.5
        x_neg = x_neg - DELTA_T * grad_v + noise
        x_neg = x_neg.detach()
    
    # 负样本能量
    neg_energy = model(x_neg).mean() / EPSILON_MAX
    
    # 对比散度损失
    loss = pos_energy - neg_energy
    return loss

# 数据加载
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)

# 初始化模型和优化器
model = EnergyMatchingModel().to(DEVICE)
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

# Phase 1: 预训练（仅优化L_OT）
print("Starting Phase 1 Training...")
model.train()
for iteration in range(PHASE1_ITERS):
    x_data, _ = next(iter(train_loader))
    x_data = x_data.to(DEVICE)
    x_noise = torch.randn_like(x_data).to(DEVICE)
    t = torch.rand(1, device=DEVICE).item()
    
    optimizer.zero_grad()
    loss_ot = compute_ot_loss(model, x_data, x_noise, t)
    loss_ot.backward()
    optimizer.step()
    
    if (iteration + 1) % 1000 == 0:
        print(f"Iteration {iteration + 1}/{PHASE1_ITERS}, L_OT: {loss_ot.item():.4f}")

# Phase 2: 主训练（联合优化L_OT + L_CD）
print("Starting Phase 2 Training...")
for iteration in range(PHASE2_ITERS):
    x_data, _ = next(iter(train_loader))
    x_data = x_data.to(DEVICE)
    x_noise = torch.randn_like(x_data).to(DEVICE)
    t = torch.rand(1, device=DEVICE).item()
    epsilon = epsilon_schedule(t)
    
    optimizer.zero_grad()
    loss_ot = compute_ot_loss(model, x_data, x_noise, t)
    loss_cd = compute_cd_loss(model, x_data, x_noise, epsilon)
    loss = loss_ot + LAMBDA_CD * loss_cd
    loss.backward()
    optimizer.step()
    
    if (iteration + 1) % 1000 == 0:
        print(f"Iteration {iteration + 1}/{PHASE2_ITERS}, L_OT: {loss_ot.item():.4f}, L_CD: {loss_cd.item():.4f}, Total Loss: {loss.item():.4f}")

# 保存模型
torch.save(model.state_dict(), "energy_matching_cifar10.pth")
print("Training completed and model saved.")

---

代码说明

1. 模型结构

• UNet：为简化示例，使用了一个小型UNet（实际应用应使用[Tong et al., 2023]的完整UNet架构）。输入为3×32×32的CIFAR-10图像，输出为64通道特征图。
• ViT：使用timm库的vit_small_patch16_224作为基础，调整输入为UNet的输出（64×32×32）。通过PatchEmbed将特征图转换为序列，ViT输出标量 ( $V_{\theta }(x)$ )。
• 实际实现中，应确保UNet参数量约37M，ViT约12M，可通过调整层数或通道数实现。

2. 训练目标

• (${\mathcal{L}}_{\text{OT}}$)：
  • 使用POT库的ot.emd计算最优传输计划，生成配对。
  • 插值点 ( $x_t$ ) 沿测地线计算，目标速度为 ( $x_{\text{data}}-T(x_{\text{data}})$ )。
  • 损失通过均方误差优化模型梯度与目标速度的匹配。
• (${\mathcal{L}}_{\text{CD}}$)：
  • 使用朗之万动力学生成负样本，步数为200。
  • 负样本初始化为噪声（简化起见，未实现50%数据+50%噪声的混合初始化，可通过添加条件实现）。
  • 对比散度损失计算正负样本的能量差。

3. 训练流程

• Phase 1：200k次迭代，仅优化 (${\mathcal{L}}_{\text{OT}}$)，构建流路径。
• Phase 2：25k次迭代，联合优化 (${\mathcal{L}}_{\text{OT}}+{\lambda }_{\text{CD}}{\mathcal{L}}_{\text{CD}}$)，温度 ($\epsilon (t)$) 按线性调度增加。
• 每1000次迭代打印损失，便于监控训练进展。

4. 数据加载

• 使用torchvision加载CIFAR-10数据集，应用标准归一化（均值0.5，标准差0.5）。
• 批大小为128，使用4个工作线程加速数据加载。

5. 优化器

• 使用Adam优化器，学习率为 ( $8\times 10^{-4}$ )。
• 梯度通过PyTorch的autograd自动计算。

---

运行环境

• 硬件：论文使用4×A100 GPU，示例代码可在单GPU（如RTX3090）或多GPU上运行。
• 依赖安装：

  pip install torch torchvision torchaudio
  pip install POT
  pip install einops timm
  

• 数据集：CIFAR-10将自动下载到./data目录。

---

注意事项

1. UNet简化：示例中的UNet为简化版，实际应参考[Tong et al., 2023]的完整实现（如https://github.com/alexandtong/OT-CFM）。可替换UNet类为完整架构。
2. 负样本初始化：论文建议负样本50%从数据、50%从噪声初始化，示例中仅使用噪声初始化。可修改compute_cd_loss添加混合初始化。
3. 计算资源：200k+25k次迭代需要数天训练时间，建议使用多GPU加速。
4. FID评估：代码未包含FID计算，可使用torch-fidelity库评估生成质量（目标FID约3.97）。
5. OT求解器：POT库的ot.emd适用于小批量数据，大规模实验可能需优化OT计算效率。

---

扩展

• 生成样本：训练完成后，可添加采样代码，使用朗之万动力学从 (${\rho }_{\text{eq}}(x)\propto \exp \left(-\frac{V_{\theta }(x)}{{\epsilon }_{\text{max}}}\right)$) 生成样本。
• 逆问题：参考Algorithm 3，可扩展代码支持带交互能量的逆问题求解。
• LID估计：参考3.3节，可添加Hessian计算代码估计局部内在维度。

---

通过运行上述代码，可以在CIFAR-10上复现Energy Matching的训练过程。

后记

2025年4月17日于上海，在grok 3大模型辅助下完成。