Energy-Based Transformers

• Energy-Based Transformers are models that integrate energy-based optimization into Transformer backbones by minimizing a learned energy function over input-output pairs.
• They employ iterative, gradient-based refinement and System 2 Thinking to self-verify and enhance prediction quality across diverse tasks.
• Empirical studies show EBTs scale efficiently and generalize better than traditional Transformer variants, achieving up to 35% higher scaling rates and significant performance gains.

Energy-Based Transformers (EBTs) are a class of modern machine learning models that integrate the principles of energy-based modeling with the Transformer architecture. By associating an energy value to each input-prediction pair and formulating prediction as an optimization problem—minimizing this energy—EBTs introduce explicit mechanisms for scalable learning and inference-time "thinking" that differentiate them from conventional Transformer models. Recent research demonstrates that EBTs can generalize better, scale more efficiently, and support advanced inference strategies modeled after human deliberative reasoning.

1. Foundations and Core Principles

Energy-Based Transformers employ a Transformer backbone but repurpose it to define an energy function $E_\theta(x, \hat{y})$ over pairs of inputs $x$ and candidate predictions $\hat{y}$. This contrasts with standard Transformers, which output predictions in a direct, feed-forward manner. In EBTs, the joint unnormalized data likelihood is given by:

$$p_\theta(x, \hat{y}) \propto e^{-E_\theta(x, \hat{y})}$$

Prediction is treated as an optimization problem: Given $x$, the model seeks ${\hat{y}}$ that minimizes $E_\theta(x, \hat{y})$. The minimization is carried out via iterative gradient-based methods, with inference steps of the form:

$$\hat{y}_{i+1} = \hat{y}_i - \alpha \nabla_{\hat{y}_i} E_\theta(x, \hat{y}_i)$$

where $\alpha$ is a step size. The energy scalar acts as a learned verifier—a lower energy indicates higher compatibility between $x$ and $\hat{y}$. This unnormalized, continuous scoring enables flexible model handling of candidate outputs in both discrete domains (language) and continuous domains (vision).

2. System 2 Thinking and Iterative Inference

EBTs operationalize "System 2 Thinking," a term from cognitive science signifying reasoned, deliberative problem solving, as opposed to rapid, intuitive "System 1" processes. In EBTs, this is realized by conducting multiple rounds of refinement during inference:

• The model allocates more computation to harder problems by taking additional optimization steps or by generating and verifying multiple candidate predictions ("Best-of-N" self-verification).
• The energy function allows each candidate’s quality to be objectively assessed, supporting selection of the most compatible prediction.
• This process emerges naturally from unsupervised pretraining, without requiring problem-specific supervision or auxiliary verifiers.

Formally, the percentage improvement from additional computation is denoted as System 2 Thinking (STT):

$$\text{STT}(x, \theta, F) = \mathbb{E}_x \left[ \frac{P(x, \theta, F)}{P(x, \theta, F_0)} - 1 \right]$$

where $F$ is the number of function evaluations (e.g., optimization steps) and $P$ the performance metric.

3. Training, Generalization, and Scaling

EBTs employ an optimization-centric training algorithm in lieu of classic contrastive or negative sampling for EBMs:

• Training starts with a randomly initialized prediction, which is iteratively optimized to minimize the energy with respect to the ground truth.
• The loss (e.g., cross-entropy for language, mean squared error for vision) is computed after the optimization trajectory and backpropagated through all iterations.
• The framework supports stochastic techniques such as replay buffers, Langevin (noise-injected) dynamics, and randomized step schedules to improve the learning of smooth, generalizable energy landscapes.

Crucially, EBTs demonstrate improved scaling properties compared to Transformer++ models, with up to 35% higher scaling rates across data, model width, depth, batch size, and computational FLOPs. Notably, their computation costs are only modestly higher per inference instance than standard Transformers (e.g., $\sim3.33\times$ per two-step optimization), while offering superior convergence and efficiency at scale.

4. Performance Metrics and Empirical Advantages

Quantitative evaluation reveals several key advantages of EBTs:

• Superior Scaling: EBTs outscale Transformer++ baselines on data, batch size, parameter count, depth, and compute, indicating improved efficiency on all axes.
• System 2 Gains: On language tasks, allocating extra computation time via self-verification boosts EBT performance by 29% more than prior Transformer++ models.
• Generalization: Despite potentially higher pretraining perplexity, EBTs achieve lower perplexity on most downstream tasks such as GSM8K, SQuAD, and BigBench Math QA, reflecting enhanced effective generalization.

Model	Pretrain Perplexity	GSM8K	SQuAD	BB Math QA	BB Dyck
Transformer++	31.36	49.6	52.3	79.8	131.5
EBT	33.43	43.3	53.1	72.6	125.3

• Energy as Uncertainty: The energy function encodes the model’s epistemic uncertainty: higher energies correspond to out-of-distribution or "hard" inputs, facilitating principled risk and uncertainty modeling.

For visual denoising, EBTs outperform Diffusion Transformers (DiT) on in-distribution and OOD data while using only 1% of DiT's compute.

5. Architectural Comparisons

EBTs differ from recent alternatives in several key respects:

• Versus Transformer++: Whereas Transformer++ models lack built-in verification and dynamic computation allocation, EBTs enable both through energy minimization and confidence interpretation.
• Versus Diffusion Transformers (DiT): DiT iteratively predicts noise using a prescribed schedule; EBT directly defines an energy over outputs, enabling verification, uncertainty estimation, and rapidly adaptable inference in both discrete and continuous domains.
• EBTs generalize across modalities rather than being domain- or problem-specific, and require no additional supervision or task-specific verifiers to achieve System 2 enhancements.

6. Implications and Future Directions

EBTs point the way toward a new paradigm for scalable, generalizable foundation models:

• Increased scaling efficiency and generalized System 2 Thinking suggest that EBTs could overcome persistent limitations of current Transformer architectures, such as scaling bottlenecks and poor handling of OOD/uncertainty cases.
• The energy-based approach has clear extensions to multimodal modeling, as EBTs enable a unified compatibility measure across modalities (e.g., language, vision, audio).
• System 2 inference in EBTs is especially impactful for robustness to distribution shift, adversarial inputs, and problems requiring verification (e.g., mathematical reasoning, planning).
• EBTs may be integrated as a System 2 backbone, providing explicit verification for lighter-weight or heuristic models when high confidence or rigorous reasoning is required.
• Future work is anticipated in algorithmic innovation for EBT inference—including advanced MCMC and tree search—and in deploying such architectures for decision-making in reinforcement learning and AI planning.

7. Summary Table: EBTs Compared to Other Paradigms

Model Type	Prediction Method	Verification/Uncertainty	System 2 Thinking	Modality Generality
Transformer++	Feedforward	No	No	General
DiffusionTrans	Iterative (noise denoising)	External (optional)	Partial	Continuous only
EBT	Optimization (energy minimization)	Built-in (energy)	Yes	Discrete & Continuous

EBTs, by combining scalable, energy-based iterative reasoning with automatic self-verification and uncertainty quantification, provide a robust framework for advancing the capabilities of large-scale foundation models across diverse domains.