NRGPT: An Energy-based Alternative for GPT

TL;DR

This work introduces NRGPT, an energy-based reinterpretation of GPT where each token carries a per-token energy and token states are updated via gradient-based exploration on an energy landscape. By decomposing the energy into an attention component and a feed-forward component, and by using a shared, recurrent block with a learnable inference-rate, the model can reproduce Transformer-like updates through energy descent. Empirically, NRGPT achieves competitive performance with fewer parameters across ListOps, Shakespeare, and OpenWebText, while exhibiting interesting dynamics such as asymptotic stability and reduced overfitting on large data regimes. The framework also clarifies theoretical connections and distinctions with prior energy-based transformer approaches, suggesting advantages for explicit energy modeling, regularization, and potential for adaptable computation.

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.

Figures (6)

• Figure 1: NRGPT casts the standard GPT setting into an energy-based framework. The network is defined as the sum of two energies: an attention energy and a feedforward energy. Each token is transformed into the next token by exploring the energy landscape. Recurrent application of the NRGPT block produces a dynamical system where each token can be thought of as a particle moving on the network's energy landscape.
• Figure 2: In NRGPT, tokens converge to stable states of low energy where the causal attention mask allows each token energy to fluctuate during inference. Shown are 64 tokens passed to an NRGPT model trained to predict ListOps equations.
• Figure 3: Learning ListOps: NRGPT variants match performance with a recurrent GPT model on ListOps accuracy parameter-transition points (top) and training/validation losses (bottom). The accuracy of models is tested on nested, mixed arithmetic tasks of maximum, median and sum modulo 20. For all plots, the x axis shows the total parameter count of the model. The yellow star indicates the transition to learning, which we define as where the logistic fit hit $>80\%$ accuracy. The baseline model GPT_Rec_parallel shows the earliest learning transition at size $2.3 \times 10^4$, but our NRGPT variants are also similar, with NRGPT_H_FF1 at $2.4 \times 10^4$ and NRGPT_H_FF2W at $2.98 \times 10^4$.
• Figure 4: Shakespeare scaling: NRGPT achieves performance parity with recurrent GPT on Shakespeare across parameter sizes, as measured by best validation loss per number of parameters. For many embedding sizes, NRGPT also follows the same optimal training loss trajectory-per-parameter as both GPT and recurrent GPT baselines. However, NRGPT does not overfit Shakespeare at large parameter sizes. Connecting lines show the best performance at fixed parameter sizes. Transparent dots show different choices of hyperparameters --- a larger spread indicates more sensitivity to hyperparameters. See \ref{['ap:exp:shakespeare']} for details.
• Figure 5: Best Generation Examples from GPT (left column), RGPT-parallel (middle column) and NRGPT (right column).

Theorems & Definitions (5)

• Proposition 2.1: Energy Descent
• proof : Sketch of proof
• Proposition 2.2: $\dot E$ without normalization
• proof
• proof : Proof of \ref{['prop:dE']}: Energy Descent