TEM模型实现低算力高质量图像生成

通过对基于能量的模型（EBM）进行大规模扩展并加以多种改进，一种能以较少计算量实现高质量图像生成的 Transformer Energy-Based Model（TEM）被提出。

“Scaling Transformer Energy-Based Models for High Quality Image Synthesis”

使用 609M 参数的 Transformer，在 ImageNet 256×256 的类别条件图像生成中达到了 FID 3.82。这一水平可与不使用 Classifier-Free Guidance（CFG）的 Diffusion Transformer 相竞争。

EBM 是用能量函数 E_θ(x) 来表示数据概率分布的模型。

p_θ(x) = exp(−E_θ(x)) / Z_θ

EBM 直接学习能量函数，并利用其梯度通过 MCMC 进行采样。

x_{t+1} = x_t − η∇E_θ(x_t) + σε

与 Langevin 采样类似，在向低能量区域更新状态的同时加入噪声。

EBM 的优势包括：能够进行分布外检测（能量越高，模型上的相对概率密度越低）、能够进行补全，以及能够从对应于多个概率分布乘积的分布中进行采样。

例如，当想要生成与两个概率分布 p₁(x)、p₂(x) 的乘积成比例的分布时，

p(x) ∝ p₁(x)p₂(x)

E(x) = E₁(x) + E₂(x)

因此，在保持多个约束的前提下进行采样，只需将对应于这些约束的能量相加，并利用其梯度进行采样即可。

EBM 是一种早已被研究的方法，但存在高维空间中 MCMC 采样困难、大规模模型训练不稳定等问题。TEM 主要通过以下 5 项改进来解决这些问题。

・在 VAE 的潜在空间中学习能量函数

不在原始图像空间，而是在由 VAE 转换得到的 32×32×4 潜在变量空间中学习能量函数并进行 MCMC 采样。由此大幅降低采样空间的维度数。

・从 Transformer 的多层提取能量

从 Transformer 的浅层、中层和深层分别提取能量，并将它们的总和作为最终能量。

浅层能量倾向于对应纹理，中层对应物体布局，深层对应全局语义结构。通过组合不同抽象层级的约束来提升生成质量。

・Positive Representation Alignment（PRA）

仅对 Positive 样本（真实数据）应用学习，使预训练的 DINOv2-Large 特征表示与 Transformer 的中间表示保持一致。

这是因为 Negative 样本通常偏离数据分布，如果连其特征表示也要对齐，可能会与原本的对比学习相矛盾。

・Sampler Matching Loss

在 EBM 的最大似然学习中，会学习降低真实数据的能量，并提高模型生成样本的能量。

然而，实际上由有限步 MCMC 生成的分布 q_θ 与能量函数所定义的原始分布 p_θ 并不一致。

因此，添加源自 KL(q_θ || p_θ) 的正则化，在固定能量函数一侧的前提下，最小化

E_q[E_θ(x)] − H(q_θ)

第一项是降低 MCMC 生成样本能量的项。它不是直接使用能量函数的梯度，而是通过采样过程传递梯度，从而使生成样本学习到达更低能量的区域。

该计算需要二阶微分，但并非对全部 40 步进行反向传播，而是仅对最后 1 步求导，以控制计算量。

第二项对应于最大化生成样本的熵。利用与 Replay Buffer 中保存的历史样本的最近邻距离，使样本之间不会过于接近。

由此，不仅会向低能量区域移动，还能保持生成样本的多样性。

・Persistent Replay Buffer

并非每次都从高斯噪声开始生成，而是保存并复用过去 MCMC 生成的潜在变量。

Replay Buffer 的容量为 16,000，以 99% 的概率复用历史样本，以 1% 的概率重新从高斯噪声开始采样。

此外，由于每次独立给定类别条件，某个类别生成过程中的潜在变量也可用于另一个类别的生成。

训练时进行 40 步采样，推理时主要进行 60 步采样。

在这些改进中，PRA 和 Sampler Matching 的效果尤为显著。

实际上，在 122M 参数模型的消融实验中，仅使用 Contrastive Loss 时的 FID 为 121.11，引入 Sampler Matching 后改善至 42.76，再加入 PRA 后进一步改善至 19.86。

此外，论文还验证了 EBM 原本的特性，例如组合能量函数进行图像生成、无需额外训练的 Inpainting、基于能量值的分布外检测等。

评论

===

EBM 是一种理论上灵活且有趣的模型，具有可直接评估非归一化概率密度等优点。然而，此前它难以达到与 Diffusion 或 Flow Matching 同等水平的高质量图像生成。

本次通过组合多种改进，证明了 EBM 也能够实现高质量图像生成。

其中我个人最感兴趣的是，在 EBM 中，对应于多个概率分布乘积的分布可以通过能量相加自然地表示出来。

例如，不仅能相对容易地实现同时满足多个约束的生成，还可以通过考察各个能量项，为评估生成结果在多大程度上满足了哪些约束提供线索。

这可以自然地应用于机器人控制、化合物生成等需要考虑多个约束的场景。

此外，我也关注本次方法中 Persistent Replay Buffer 的作用。

Diffusion 和 Flow Matching 通过学习从高斯分布到目标分布的路径，来应对从具有多峰性的复杂分布中生成的问题。

而在 EBM 中，多峰性能量地形下的 MCMC Mixing 存在困难，但 Persistent Replay Buffer 有可能通过整个学习过程实现了分布的迁移。

在学习初期，Replay Buffer 的样本