Autoregressive Language Models are Secretly Energy-Based Models: Insights into the Lookahead Capabilities of Next-Token Prediction

Autoregressive models (ARMs) currently constitute the dominant paradigm for large language models (LLMs). Energy-based models (EBMs) represent another class of models, which have historically been less prevalent in LLM development, yet naturally characterize the optimal policy in post-training alignment. In this paper, we provide a unified view of these two model classes. Taking the chain rule of probability as a starting point, we establish an explicit bijection between ARMs and EBMs in function space, which we show to correspond to a special case of the soft Bellman equation in maximum entropy reinforcement learning. Building upon this bijection, we derive the equivalence between supervised learning of ARMs and EBMs. Furthermore, we analyze the distillation of EBMs into ARMs by providing theoretical error bounds. Our results provide insights into the ability of ARMs to plan ahead, despite being based on the next-token prediction paradigm.

💡 Research Summary

The paper presents a unified theoretical framework that bridges autoregressive models (ARMs), the dominant architecture for large language models, with energy‑based models (EBMs), a class traditionally less used in LLM development but naturally arising as the optimal policy in post‑training alignment. Starting from the chain rule of probability, the authors construct an explicit bijection between sequence‑level distributions p(y|x) and token‑level conditional distributions π(y_t|x, y_{<t}). This bijection is first shown at the distribution level and then instantiated in function space: an EBM’s global energy function R(x,y) is decomposed into immediate rewards r(s_t, a_t), and a mapping M transforms r into an ARM’s scoring function q(s_t, a_t) by adding the soft‑value V_q(s_t⊕a_t), i.e., the log‑partition of the next‑state distribution. The inverse mapping M⁻¹ subtracts the same soft‑value, allowing parallel computation. Proposition 1 proves that M is a true bijection, guaranteeing that for any (x,y) the ARM probability p_ARM^q(y|x) equals the EBM probability p_EBM^r(y|x) and that the ARM’s cumulative log‑partition equals the EBM’s log‑partition.

Building on this equivalence, the paper derives two further results. Proposition 2 shows that supervised maximum‑likelihood training of an ARM is mathematically identical to maximum‑entropy reinforcement learning (MaxEnt RL) on the corresponding EBM. In other words, the KL‑regularized RL objective commonly used for post‑training alignment has an exact solution that is an EBM with reward R plus a reference log‑policy term. Proposition 3 provides error bounds for distilling an EBM into an ARM, quantifying how approximation of the soft‑value and sampling bias affect the KL divergence between the true EBM and its ARM surrogate. The bounds scale linearly with vocabulary size V and sequence length T, highlighting the computational trade‑offs: converting an ARM to an EBM costs O(V·T) (linear), whereas converting an EBM to an ARM is exponential in T if done explicitly, motivating implicit learning approaches.

The authors discuss the practical implications. ARMs enjoy efficient parallel training and exact ancestral sampling via softmax, while EBMs require intractable partition function computation and MCMC sampling. However, when the ARM’s scoring function q implicitly encodes the soft‑value, the ARM effectively implements the global EBM distribution and thus possesses “look‑ahead” capabilities despite being trained on next‑token prediction. This explains why teacher‑forcing and next‑token supervision work well: they are equivalent to fitting the underlying EBM.

Empirical validation is performed on synthetic datasets and on real LLMs. The authors train an EBM, transform it to an ARM using the bijection, and compare performance. The ARM matches the EBM in perplexity and reward scores, confirming the theoretical predictions. Experiments also show that MaxEnt RL fine‑tuning of LLMs can be interpreted as distilling an EBM into the existing ARM, providing a principled view of alignment.

In conclusion, the paper establishes a rigorous connection between ARMs and EBMs, showing that ARM training is a constrained optimization over the space of EBMs and that ARMs inherently contain a soft‑value function enabling sequence‑level planning. This insight opens avenues for more efficient alignment methods, better distillation techniques, and a deeper understanding of why next‑token prediction models exhibit emergent planning behavior.