Energy-Based Transformers are Scalable Learners and Thinkers
Alexi Gladstone, Ganesh Nanduru, Md Mofijul Islam, Peixuan Han, Hyeonjeong Ha, Aman Chadha, Yilun Du, Heng Ji, Jundong Li, Tariq Iqbal
cs.LG, cs.AI, cs.CL, cs.CV
2025-07-03
EBT 把预测重写成对一个能量函数做梯度寻优,凭纯无监督就学会 System 2 思考:训练扩展比 Transformer++ 快最高 35%,推理多想几步能多涨 29%。
让模型「多想一会儿再回答」(System 2 thinking)是这两年的显学,但现有的几条路都带着硬伤。O1、DeepSeek R1 这类靠强化学习的推理模型,只在数学和代码这类答案可验证的领域管用,换到开放式问答上常常不升反降;它们还离不开额外的奖励信号或外部验证器,没法从无监督预训练里自然长出来。另一条路是给主模型配一个打分器(verifier),在推理时挑分高的答案,但打分器同样要单独训。
更根本的问题是,这些办法几乎都只针对文本。视觉、视频这类连续模态很难套用「逐步推理」的框架,答案空间是连续的,没法逐 token 验证。
作者的问题因此很干净:能不能让一个模型,只靠无监督学习,就在文本和图像上都学会「慢思考」?
EBT(Energy-Based Transformer)的核心是把预测重新定义成一个优化问题。模型不再一次性吐出答案,而是学一个能量函数 E(x, ŷ),给「输入 x」和「候选预测 ŷ」这对组合打一个标量分:能量越低,说明这个候选越「合理」。预测的过程,就是从随机初始化的 y⁰ 出发,沿能量梯度一步步下降,直到收敛:
y⁰ N(0, I)
for i = 0..N-1:
yⁱ⁺¹ ← yⁱ − α ∇E(x, yⁱ)
训练时走的是同一条链路:从随机点出发走 N 步,要求最终落点接近真实答案,再用损失回传更新能量函数。这条「优化即训练」的路径会逼着能量曲面在正确答案附近形成一个凹陷的盆地形,只要走对方向就能滑进去。这绕开了传统能量模型做对比学习时负样本随维度爆炸的老问题。
让这套训练稳定下来并不容易,作者靠三个技巧撑住:
架构上给两套:decoder-only EBT 对标 GPT,一次并行预测多个 token,做自回归;bidirectional EBT 对标 BERT/DiT,做填空和掩码建模。
推理时想「多想几步」,就把步数 N 加大,或像 System 2 那样从多个随机起点各跑一遍、再用能量函数挑能量最低的那个(自我验证)。这种额外计算是架构内生的,不需要再训一个验证器。
先看训练效率。在数据、批大小、参数量、FLOPs、深度这五个轴上,EBT 的扩展率都比 Transformer++ 高出最高 35%。这是第一个在不改 tokenizer 的前提下、扩展速率跑赢 Transformer++ 的方法。
System 2 thinking 的收益更直接。在语言任务上多花推理算力,EBT 的提升比 Transformer++ 多 29%。关键差别在于:Transformer++ 哪怕让它「想更久」也基本降不动困惑度,而 EBT 能在逐 token 级别把困惑度压下去。自我验证带来的增益还会随规模变大,小模型上只有 4%-8%,大模型上到 10%-14%。
连续模态上,EBT 在图像去噪里正面打 Diffusion Transformer(DiT):
| 设置 | 模型 | PSNR | MSE |
| 分布内 σ=0.1 | DiT | 26.58 | 142.98 |
| 分布内 σ=0.1 | EBT | 27.25 | 122.55 |
| 分布外 σ=0.2 | DiT | 19.56 | 718.7 |
| 分布外 σ=0.2 | EBT | 23.29 | 305.2 |
分布外那一栏才是重点:噪声更大时 EBT 几乎不掉水准,DiT 却大幅退化。EBT 还只需 DiT 约 1% 的前向次数就能达到相当或更好的 PSNR。ImageNet 线性探针上,EBT 的 Top-1 是 5.32%,DiT 只有 0.31%,说明它顺带学到了明显更好的视觉表征(绝对值仍低,这是小规模探针)。
更反直觉的是泛化。EBT 的预训练困惑度(33.43)其实比 Transformer++(31.36)更差,但在 GSM8K、BigBench 数学问答、Dyck 语言这几个下游困惑度上反而更低,只在 SQuAD 上略输。预训练没赢,下游却赢了,作者把这解读为能量目标本身起到了更好的正则化。
这篇给的是一类新范式,不是又一个涨点的 trick。它第一次证明了:不用强化学习、不用验证器、不用可验证奖励,光靠无监督预训练就能让模型在文本和图像上长出「慢思考」,而且这种思考是架构内置的优化过程,不是靠提示词诱发出来的思维链。
对从业者来说,如果这套机制在更大规模上成立,它有可能替换掉现在「Transformer++ 主干加 RL 推理微调」的整条栈,也能让视觉模型拿到推理时扩展能力,这是目前扩散模型缺的一块。EBT 对不确定性的天然建模(连续视频帧里靠前的帧、文本里更难的 token 都对应更高能量)也解释了它为什么在分布外更稳。
要诚实的是,今天它还到不了「拿来就用」。实验只跑到 800M 参数,基础模型规模上还存疑;每次预测要跑多步梯度,训练和推理都比前馈模型贵。把它看成一个方向正确的早期信号更合适。
作者自己列了几条:多了一组超参(步长、步数)要调;计算开销更大;受算力限制只验证到 800M;凸能量曲面假设让它处理多峰分布(比如类别条件图像生成)时会犯难。
读下来还有几处没被说服。35% 的扩展率优势是在小规模(≤800M)上测的,70B 这种量级还存不存在,论文没有证据。自我验证的增益随规模上升,但小规模上绝对值很小,「能涨」和「值得涨」之间还有距离。凸能量假设是个更深的限制:它让训练好优化,却天然偏向单峰解,这对真正多义的开放生成可能是个天花板。下游泛化的赢面也不算大,部分还落在困惑度这类代理指标上,而非端到端任务准确率。