Reproduce “Energy-Based Transformers are Scalable Learners and Thinkers” on OpenResearch

About this paper

Alexi Gladstone, Ganesh Nanduru, Md Mofijul Islam, Peixuan Han, Hyeonjeong Ha, Aman Chadha, Yilun Du, Heng Ji and 2 more

Energy-Based Transformers (EBTs) treat prediction as candidate scoring: they assign a scalar energy to a candidate given its context, then use its gradient to refine the candidate over repeated steps. The same score can select among candidates. On four out-of-distribution language datasets, extra passes yield up to 29% improvement in the perplexity-based measure; tests reach 800M parameters, not foundation-model scale.

Key insights

• Reframing prediction as an iterative optimization process over a learned energy function allows models to exhibit System 2 thinking facets: dynamic computation allocation, uncertainty modeling, and intrinsic prediction verification.
• Energy-Based Models (EBMs) can be scaled using Transformer architectures and novel regularization techniques, overcoming historical scalability and stability challenges.
• The ability to 'think longer' and self-verify through energy minimization leads to improved generalization, particularly on out-of-distribution data.

Summary by alphaXiv.