NRGPT: An Energy-based Alternative for GPT

AI summary

Overview

Research area: Machine learning / natural language processing — specifically, unifying autoregressive transformer language modeling (GPT) with energy-based models (EBMs).

Technical level: Advanced. The paper assumes familiarity with transformer blocks, multi-head attention, residual streams, LayerNorm/RMSNorm, gradient-based optimization, and energy-based model theory.

Scope: A single paper proposing NRGPT ("eNeRgy-GPT"), a minimal modification of the GPT architecture whose inference step is reframed as gradient descent on an explicit, token-wise energy function, evaluated on ListOps, Shakespeare, and OpenWebText.

What This Paper Is About

GPT-style transformers and energy-based models are two dominant but seemingly incompatible paradigms: GPTs map a token sequence to a shifted sequence in a feed-forward pass, while EBMs treat inference as a dynamical process that minimizes a scalar energy. The paper asks whether the standard GPT setting — causal, autoregressive next-token prediction — can be cast as a well-defined energy-based framework, which prior energy-based transformer work had not managed without awkward workarounds. The authors propose NRGPT, a recurrent module in which each token is treated as a particle moving on an energy landscape, and show that under certain conditions this motion reduces to gradient descent.

Key Contributions

1. An energy function and update rule that reproduce the GPT setting. The authors design an energy composed of an attention term and a feed-forward term whose gradients yield a transformer-like block, supporting variants with learnable inference rate and with LayerNorm or RMSNorm.

2. Provable energy-descent guarantees. Proposition 2.1 shows the update rule asymptotically decreases energy when the inference rate is $\bm{\eta} = c,\mathrm{diag}(\bm{\gamma})$ with $c \in \mathbb{R}_{>0}$; Proposition 2.2 shows that without normalization, energy decreases when the symmetric part of the inference rate is positive semi-definite.

3. Strong results on nested ListOps tasks. NRGPT variants match a recurrent GPT baseline on nested arithmetic operations including max, median, and sum modulo 20.

4. Feasibility of language modeling on Shakespeare and OpenWebText, plus a systematic comparison of performance scaling between recurrent transformers and NRGPT, along with an empirical study of the dynamical trajectories of tokens on the energy landscape.

Main Findings

• Asymptotic stability of tokens. Because the causal attention mask makes token $A$'s energy depend only on tokens $B \leq A$, token energies are bounded below and decrease, so after a transitory period $T_{\text{tr}}$ each token's state freezes at a constant. The authors call this asymptotic stability; inference can run arbitrarily long afterward without changing the energies.

• ListOps learning transitions. Using a definition of learning transition as the logistic fit exceeding 80% accuracy, GPT_Rec_parallel showed the earliest transition at $2.3 \times 10^{4}$ parameters, with NRGPT_H_FF1 at $2.4 \times 10^{4}$ and NRGPT_H_FF2W at $2.98 \times 10^{4}$. The top-performing energy variant for ListOps was the residual version, which learned at smaller sizes than the baseline and trained easily at embedding sizes over 256.

• Energy dynamics under forced descent. In ListOps experiments with 30 recurrent steps and $\eta = 1$, the final total energy was lower than the initial energy in all trajectories. Individual token trajectories are not required to be monotonic because tokens are coupled, but after the transient stage the next token's energy monotonically decreases.

• Shakespeare scaling. NRGPT achieves performance parity with recurrent GPT across parameter sizes as measured by best validation loss per number of parameters, and for many embedding sizes follows the same optimal training-loss-per-parameter trajectory as both GPT and recurrent GPT. NRGPT did not overfit Shakespeare at large parameter sizes, whereas baselines did.

• OpenWebText results at $n_{embed} = 768$. GPT had training loss 2.905 ± 0.006 and validation loss 2.921 ± 0.005 at 124M parameters across 5 runs; GPT_Rec_parallel had training loss 3.447 ± 0.046 and validation loss 3.454 ± 0.037 at 85M parameters across 5 runs; NRGPT_H_FF2W had training loss 3.456 ± 0.076 and validation loss 3.467 ± 0.073 at 90M parameters across 3 runs.

• Generation quality metrics. In the best configurations (all $n_{layer}=12$, $n_{head}=12$, $n_{embed}=768$): GPT reached perplexity 75, grammar quality score 0.978, average pairwise cosine similarity 0.306, distinct-1 0.619, distinct-2 0.965; GPT_Rec_Parallel reached perplexity 99, GQS 0.976, APCS 0.336, d-1 0.615, d-2 0.975; NRGPT_H_FF2W reached perplexity 104, GQS 0.966, APCS 0.306, d-1 0.674, d-2 0.984. NRGPT contains around 34M fewer parameters than GPT.

• Reduced overfitting. The authors observe that NRGPT models may be more resistant to overfitting, doing so only during very long training.

Methodology in Plain English

The starting point is the parallel transformer design used in GPT-J, PaLM, and the Energy Transformer, where attention and feed-forward outputs are added to the residual stream simultaneously. The authors ask what scalar energy function would produce that same update if you took its gradient, and derive one: an attention energy built from a log-sum-exp over previous tokens, and a feed-forward energy built from an additive function of each token. Instead of stacking different layers, NRGPT uses a single block applied recurrently, so that the step index $t$ becomes a step of energy descent rather than a layer number. Each token's energy depends on the other tokens, and the causal mask restricts each token's energy to depend only on itself and earlier tokens.

The researchers prove conditions under which this update provably lowers energy, and note that when the model is not constrained to strict descent it can learn its own exploration strategy instead. They then train several energy-function variants — NRGPT_H_FF1 using a squared GELU-based term and NRGPT_H_FF2W using a two-weight feed-forward term with hidden dimension $4 \times D$ — against baselines: GPT_Rec_parallel, a GPT-J model with a single transformer layer fed back recurrently, and a conventional GPT2-style deep transformer. Evaluation spans nested ListOps equations such as SUM(2,MAX(4,13,1),MEDIAN(5,3,16)) over inputs 0 to 19, the Shakespeare dataset, and OpenWebText.

Why This Matters

Impact on research: The paper demonstrates that the commonly used GPT setting can be embedded in a well-defined energy-based framework, complementing earlier work suggesting transformers perform gradient descent during inference. Unlike prior studies that rely on simplifications such as linear transformers or omitted softmax, NRGPT makes inference explicitly a gradient-based dynamics while keeping an architecture very close to GPT.

Real-world applications:

• Language modeling and text generation where an explicit energy function could enable more systematic exploration of the output space.
• Model alignment through energy regularizers, a possibility the authors note as a natural benefit of explicit energy formulations.
• Small-data regimes, where NRGPT's resistance to overfitting is beneficial.
• Variable-computation settings, since energy-based inference naturally supports spending more or fewer steps depending on the input.

Industry relevance: The paper compares against GPT2-style models and GPT-J-style recurrent transformers, uses the widely adopted nanoGPT codebase, and reports on distributed training across 4 nodes of 8xH100 GPUs, making the results directly comparable to standard LLM training practice.

Future Directions

• Reducing hyperparameter sensitivity. The authors report that weight sharing enforced by the energy formulation limits architectural flexibility and causes larger hyperparameter sensitivity than GPT variants; identifying more robust configurations is an open problem.

• Understanding non-convergent inference rates. The authors state that only a small set of inference rates yielding perpetually oscillating solutions may fail to converge, but the exact conditions still need to be derived.

• Examining the FLOP trade-off. Configurations achieving competitive performance at comparable parameter counts have a slightly larger FLOP count than standard transformer counterparts, and increasing the number of attention heads in NRGPT increases parameters — a cost/benefit question that remains open.

• Scaling to massive datasets. Because NRGPT has a harder time overfitting, it is beneficial in small-data regimes but potentially undesirable for the massive datasets used to train modern LLMs; how NRGPT behaves at that scale is unresolved.

Target Audience

Researchers and graduate students working on transformer architecture, energy-based models, associative memory, or the theory of inference in neural networks. Practitioners interested in alternatives to standard GPT training would benefit from the architecture and scaling results, and readers following the literature connecting in-context learning to gradient descent will find the energy-based derivation directly relevant. The paper is not suitable as an introductory text.

Authors’ abstract

Generative Pre-trained Transformer (GPT) architectures are the most popular design for language modeling. Energy-based modeling is a different paradigm that views inference as a dynamical process operating on an energy landscape. We propose a minimal modification of the GPT setting to unify it with the EBM framework. The inference step of our model, which we call eNeRgy-GPT (NRGPT), is conceptualized as an exploration of the tokens on the energy landscape. We prove, and verify empirically, that under certain circumstances this exploration becomes gradient descent, although they don't necessarily lead to the best performing models. We demonstrate that our model performs well for simple language (Shakespeare dataset), algebraic ListOPS tasks, and richer settings such as OpenWebText language modeling. We also observe that our models may be more resistant to overfitting, doing so only during very long training.