🤖 machine learning

NRGPT: An Energy-based Alternative for GPT

The paper introduces NRGPT, a modified GPT architecture that unites generative modeling with energy-based frameworks by conceptualizing inference as exploration of an energy landscape, and demonstrates competitive performance across various tasks while exhibiting increased resistance to overfitting.

Original authors: Nima Dehmamy, Benjamin Hoover, Bishwajit Saha, Leo Kozachkov, Jean-Jacques Slotine, Dmitry Krotov

Published 2026-05-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Nima Dehmamy, Benjamin Hoover, Bishwajit Saha, Leo Kozachkov, Jean-Jacques Slotine, Dmitry Krotov

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: From a High-Speed Train to a Rolling Ball

Imagine a standard AI language model (like the famous GPT) as a high-speed train. It moves token by token along a track. When it sees the word "The," it immediately knows that the next word is likely "cat" or "dog," based on its training. It is fast, but it only follows the tracks laid down during its training.

The authors of this paper propose a different way to understand how these models work. They call their new model NRGPT. Instead of a train on tracks, imagine NRGPT as a ball rolling down a hilly landscape.

In this new perspective:

The Landscape: The "hills" and "valleys" represent the energy of the text.
The Ball: The current word (or token) that the model needs to predict.
The Goal: The ball wants to roll into the deepest valley (the state of lowest energy). In the world of AI, a "state of low energy" means a word that makes perfect sense in the context of the sentence.

The paper argues that instead of merely "predicting" the next word, the model actually explores terrain to find the most stable and logical place for the next word to land.

How It Works: The "Energy" of Words

The paper introduces a concept called Energy-Based Modeling (EBM). Imagine it this way:

High Energy: A word that feels "wrong" or "disruptive" (like "The cat ate the pizza...", if the context is about a zoo). This is a high point on the hill.
Low Energy: A word that feels "right" and natural (like "The cat ate the mouse"). This is a deep valley.

The NRGPT model is constructed so that its internal mathematics generates this landscape. When the model needs to generate the next word, it does not make a mere guess; it performs a gradient descent. In German, this means it takes small steps uphill and constantly checks: "Is this word lower in energy (better) than the one I am currently on?" It continues taking steps downhill until it finds a stable point.

The "Minimal" Change

The authors did not discard the old GPT architecture. Instead, they made a minimal modification.

They took the standard building blocks of a GPT (the parts that handle attention and feed-forward processing).
They rewrote the mathematics so that these blocks act like the forces pushing the ball down the hill.
They proved that under certain conditions, this "rolling ball" process is mathematically identical to the standard "train on tracks" process, just viewed through a different lens.

What They Tested (The Experiments)

The team tested NRGPT on three different types of challenges to see if the "rolling ball" approach actually works:

Math Puzzles (ListOps): They gave the model lists of numbers and mathematical operations (like "Sum of the maximum of 4 and 13"). NRGPT performed just as well as standard models and showed it can handle logic.
Shakespeare: They asked the model to write in the style of Shakespeare. NRGPT did an excellent job and matched the quality of standard models, but with a twist: it did not overfit.
- The Analogy: Imagine a student who memorizes a textbook so perfectly that they cannot answer a question if the wording is slightly changed. This is "overfitting." NRGPT seemed to learn the concept of Shakespeare rather than just memorizing the text, making it more robust in a sense.
Real-World Text (OpenWebText): They tested it on a massive dataset of internet text. NRGPT generated text that was competitive with standard models, using slightly fewer parameters (less "brainpower" or memory).

Key Findings and "Quirks"

Asymptotic Stability: The paper discovered something interesting about the "rolling ball." Once the ball comes to rest in a valley, it no longer moves. The authors call this "asymptotic stability." It means the model naturally transitions into a stable answer and stops fluctuating, which is a nice theoretical property.
Resistance to Overfitting: On the Shakespeare dataset, NRGPT did not fall into the "memorization trap" as easily as standard models when the models became very large.
The Costs: There is a trade-off. While NRGPT may use fewer "parameters" (memory), the actual computational steps (FLOPs) required to roll the ball down the hill can be somewhat more expensive than the standard train approach. It is like a scenic, winding path down a mountain instead of a straight elevator ride; you might see more, but it requires a bit more effort to walk.

The Conclusion

The paper concludes that NRGPT is a successful experiment in uniting two different worlds: the popular GPT design and theoretical Energy-Based Models.

It proves that we can view the act of text generation not just as prediction, but as an optimization process—a search for the most stable, logical state. Although it does not yet beat the best GPT models in every single metric, it offers a new, mathematically elegant way to understand how these powerful AI brains work, suggesting that "thinking" for an AI might simply be a matter of finding the deepest point in a very complex energy landscape.

Technical Summary: NRGPT – An Energy-Based Alternative for GPT

Problem Statement

Generative Pre-trained Transformer (GPT) architectures currently dominate autoregressive language modeling by mapping a sequence of tokens to a shifted sequence (next-token prediction) through a series of Transformer layers. Although effective, this standard configuration differs from Energy-Based Models (EBMs), a paradigm where inference is viewed as a dynamic process on an energy landscape, with low-energy states corresponding to high-probability data.

Despite growing evidence of deep connections between Transformers and gradient descent (e.g., in the context of in-context learning) and previous attempts to unify these fields (such as Energy Transformers), a clearly defined framework that converts the standard causal GPT configuration into an energy-based model has remained elusive. Existing energy-based approaches often rely on masked token reconstruction (unsuitable for causal generation) or treat Transformers merely as energy calculators rather than dynamic systems. The authors identify a gap in unifying the standard GPT inference mechanism with the EBM framework without abandoning the autoregressive nature of language modeling.

Methodology

The authors propose NRGPT (eNeRgy-GPT), a minimal modification of the GPT architecture that unifies it with the EBM framework. The core innovation lies in conceptualizing the inference step not as a static forward pass, but as an exploration of tokens on an energy landscape, where the update rule is driven by the gradient of a defined energy function.

1. Formulation of the Energy Function

The model defines a total energy function $E$ consisting of two terms: an Attention Energy ( $E_{AT}$ ) and a Feed-Forward Energy ( $E_{FF}$ ).

Attention Energy: Derived from a multi-head attention mechanism, $E_{AT}$ is formulated such that its gradient with respect to the token state yields the attention output. The authors define a learnable attention energy where the gradient descent step mimics the standard attention mechanism, using a learnable weight $\alpha$ and an inference rate matrix $\eta$ .
Feed-Forward Energy: The Feed-Forward (FF) network is derived from a formulation of dense associative memories. The energy $E_{FF}$ is constructed so that its gradient produces the standard update of a two-layer MLP. The authors investigate specific variants, including $E_{FF} = -\|\sigma(W g_A)\|^2$ (NRGPT_H_FF1) and a more complex variant involving two weight matrices (NRGPT_H_FF2W).

2. Update Rule and Dynamics

The inference process is modeled as a dynamic system where the token state $x$ evolves over time steps $t$ :
$x^{(t+1)} = x^{(t)} - \eta(t) \nabla E(x^{(t)})$
Here, $\eta(t)$ is an inference rate matrix (acting as a preconditioner).

Causal Constraint: The energy of a token $A$ depends only on tokens $B \leq A$ , ensuring causal generation.
Asymptotic Stability: The authors prove that under specific conditions for $\eta$ (e.g., $\eta = c \cdot \text{diag}(\gamma)$ , where $\gamma$ are LayerNorm scaling factors) and with normalization, the system exhibits asymptotic stability. After a transient phase, the energy of each token decreases monotonically, and the token states converge to fixed points. This provides a theoretical guarantee that the "forward pass" constitutes a valid energy descent.

3. Architectural Variants

The work examines several configurations:

Normalization: The model supports LayerNorm and RMSNorm, with theoretical proofs showing that these normalizations help bound energy and ensure convergence.
FF Variants: The authors test various forms of $E_{FF}$ to balance performance with architectural similarity to standard GPTs.
Recurrent Inference: Unlike standard GPTs, where layers are distinct, NRGPT uses a single recurrent block. The number of recurrent steps corresponds to the "depth" of the network.

Main Contributions

Unified Framework: The design of an energy function and update rule that describes the GPT configuration, unifying autoregressive language modeling with energy-based modeling.
Theoretical Guarantees: Proof that under specific conditions for the inference rate matrix and normalization, the inference dynamics demonstrably decrease energy after a transient phase, establishing a foundation for asymptotic stability.
Empirical Validation: Demonstration that NRGPT is competitive in the following areas:
- ListOps: Nested arithmetic and logical tasks.
- Shakespeare: Character-level language modeling.
- OpenWebText (OWT): Large-scale natural language modeling.
Systematic Comparison: A detailed analysis of performance scaling between recurrent Transformers and NRGPT, including parameter efficiency and FLOP costs.
Dynamic Analysis: Empirical investigation of token trajectories on energy landscapes, observing convergence to stable low-energy states.

Experimental Results

ListOps: NRGPT variants (particularly NRGPT_H_FF1 and NRGPT_H_FF2W) achieved performance comparable to a recurrent GPT baseline (GPT_Rec_parallel) in terms of accuracy and training loss. The models reached learning transitions (achieving >80% accuracy) at similar parameter counts (approximately $2.3 \times 10^4$ to $3.0 \times 10^4$ parameters).
Shakespeare: NRGPT achieved validation losses comparable to deep GPT baselines across various parameter sizes. Notably, NRGPT exhibited less overfitting on the training dataset compared to standard GPTs at larger parameter sizes.
OpenWebText: On the OWT dataset, NRGPT (particularly the NRGPT_H_FF2W variant) achieved generation quality metrics (perplexity, grammar score, diversity) competitive with a standard GPT (124M parameters), despite containing approximately 34M fewer parameters (90M vs. 124M).
Downstream Tasks (MMLU): In a 5-shot setting on the MMLU dataset, the 128M-parameter NRGPT model achieved a higher average accuracy (29.3%) compared to similarly sized GPT (28.0%) and recurrent GPT baselines (26.0%). The authors note that these results should be viewed cautiously due to the small model size but suggest potential advantages.

Limitations and Observations

Computational Cost: Although NRGPT is more parameter-efficient, it incurs higher FLOP costs per iteration (1–2 $\times$ ) compared to equivalent recurrent GPTs, primarily due to the FF module design and the use of full-rank weight matrices in attention.
Hyperparameter Sensitivity: NRGPT shows greater sensitivity to hyperparameter choices (e.g., learning rates) compared to standard GPT variants.
Overfitting Behavior: The models are more resistant to overfitting, which is advantageous for small data regimes but potentially suboptimal for massive datasets where overfitting often serves as a proxy for learning complex patterns.
Inference Rate Constraints: To guarantee energy descent, the inference rate matrix $\eta$ must satisfy specific constraints (e.g., being proportional to the normalization scaling). However, empirical results suggest that even unconstrained $\eta$ can lead to convergence in practice.

Significance and Claims

The work claims that NRGPT represents a minimal modification of the GPT architecture that successfully unifies autoregressive language modeling with energy-based modeling. Its primary significance lies in:

Theoretical Unification: Providing a principled basis where the "forward pass" of a causal language model is explicitly a gradient-based dynamic system on an energy landscape.
Asymptotic Stability: Demonstrating that this dynamic system naturally converges to stable states, thereby offering a new perspective on Transformer inference.
Alternative Paradigm: Showing that energy-based approaches can compete with standard GPTs in performance while offering unique properties such as resistance to overfitting and the potential for systematic exploration of the solution space (e.g., via early stopping or regularizers), although the work notes that these advanced applications represent future work.

The authors position NRGPT not as a replacement for all GPTs, but as a significant step toward understanding the Transformer architecture through the lens of energy and optimization, bridging the gap between static deep learning models and dynamic systems.