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
EBTs turn prediction into gradient minimization of a learned energy function, giving System 2 thinking from unsupervised data: 35% faster scaling than Transformer++, 29% more from thinking longer.
Letting a model "think longer before answering" (System 2 thinking) is the fashionable lever of the last two years, but every existing path carries a real scar. Reinforcement-learned reasoners like O1 and DeepSeek R1 only help where answers are verifiable, mostly math and coding, and often hurt on open-ended QA. They also lean on extra reward signals or external verifiers that do not come out of unsupervised pretraining. A second path bolts a separate scorer onto the main model and picks high-scoring answers at inference, but that scorer needs its own training.
The deeper issue is that these tricks are mostly text-only. Continuous modalities like images and video do not fit a "reason step by step" frame, because the answer space is continuous and you cannot verify it token by token.
So the authors pose a clean question: can a single model, trained only with unsupervised learning, learn to "think slowly" across both text and images?
The Energy-Based Transformer (EBT) reframes prediction as optimization. Instead of emitting an answer in one shot, the model learns an energy function E(x, ŷ) that scores the pairing of input x and candidate prediction ŷ: lower energy means the candidate is more plausible. Prediction starts from a random y⁰ and walks downhill along the energy gradient until it converges:
y⁰ N(0, I)
for i = 0..N-1:
yⁱ⁺¹ ← yⁱ − α ∇E(x, yⁱ)
Training uses the same chain: start from a random point, take N steps, require the landing point to sit close to the true answer, then backpropagate. This optimization-as-training path forces the energy surface into a bowl around correct answers, so a well-directed walk slides into the solution. It sidesteps the old curse of dimensionality that dogs contrastive energy models and their hunger for hard negative samples.
Keeping this stable is the hard part. Three tricks hold it together:
Two architectures cover the bases. A decoder-only EBT mirrors GPT and predicts many tokens in parallel for autoregressive modeling. A bidirectional EBT mirrors BERT and DiT for infilling and masked modeling.
"Thinking longer" at inference just means raising N, or, in true System 2 fashion, running several random restarts and keeping the one with the lowest energy (self-verification). The extra compute is architectural and needs no separately trained verifier.
On training efficiency, EBTs scale up to 35% faster than the Transformer++ recipe across five axes: data, batch size, parameters, FLOPs, and depth. That makes EBTs the first approach to beat Transformer++ on scaling rate without touching the tokenizer.
The System 2 payoff is more direct. Spending more inference compute on language tasks lifts EBTs 29% more than Transformer++. The gap is structural: Transformer++ cannot meaningfully reduce perplexity by "thinking longer," while EBTs push down per-token perplexity. Self-verification gains scale with size, from 4-8% on small models to 10-14% on larger ones.
On continuous modalities, EBTs take on Diffusion Transformers (DiT) at image denoising:
| Setting | Model | PSNR | MSE |
| In-dist σ=0.1 | DiT | 26.58 | 142.98 |
| In-dist σ=0.1 | EBT | 27.25 | 122.55 |
| OOD σ=0.2 | DiT | 19.56 | 718.7 |
| OOD σ=0.2 | EBT | 23.29 | 305.2 |
The OOD row is the point: under heavier noise the EBT barely loses ground while DiT falls apart. EBTs also reach comparable or better PSNR with about 1% of DiT's forward passes. On an ImageNet linear probe, the EBT hits 5.32% Top-1 against DiT's 0.31%, so it learns clearly better visual features along the way (absolute numbers stay low, since this is a small-scale probe).
The generalization result is the counterintuitive one. EBTs actually have worse pretraining perplexity (33.43 against Transformer++ at 31.36), yet score lower perplexity on GSM8K, BigBench Math QA, and Dyck Languages, losing only on SQuAD. Worse pretraining, better downstream. The authors read this as the energy objective acting as a stronger regularizer.
This is a new paradigm, not another point on a leaderboard. It is the first demonstration that a model can grow System 2 thinking on both text and images from unsupervised pretraining alone, with no RL, no verifiers, no verifiable rewards. And the thinking is an architectural optimization process, not a chain-of-thought prompt.
For practitioners, if the mechanism holds at larger scale, it could replace the whole "Transformer++ backbone plus RL reasoning finetune" stack, and hand vision models the inference-time scaling they currently lack. The EBT's native uncertainty modeling (early video frames and rarer tokens both register as higher energy) also explains its stability out of distribution.
Be honest about where it stands today: it is not yet drop-in usable. Experiments stop at 800M parameters, so foundation-model scale is unproven, and every prediction runs multiple gradient steps, which makes training and inference pricier than a feed-forward model. Treat it as an early signal pointing the right way.
The authors flag several: extra hyperparameters (step size, step count) that need tuning; higher compute cost; validation only up to 800M; and a convex-energy-landscape assumption that struggles with multimodal distributions like class-conditional image generation.
A few things did not fully convince on close reading. The 35% scaling edge is measured at small scale (≤800M); whether it survives at 70B-class sizes goes unanswered. Self-verification gains grow with scale but their absolute values at small scale are modest, so "it grows" and "it is worth it" are not the same claim. The convex-energy assumption is the deeper ceiling: it makes training tractable but biases the model toward single-mode answers, which could limit genuinely ambiguous open-ended generation. And the downstream generalization wins are not large, with several resting on perplexity proxies rather than end-to-end task accuracy.