NRGPT: An Energy-Based Alternative for GPT Nima Dehmamy*, Benjamin Hoover*, Bishwajit Saha*, Leo Kozachkov, Jean-Jacques Slotine, Dmitry Krotov* ICLR 2025

📄 Paper | 💻 Code

Model Description

NRGPT casts the standard GPT setting into an energy-based framework. The inference step is conceptualized as an exploration of tokens on an energy landscape. The network is defined as the sum of two energies: an attention energy and a feedforward energy.

Energy Attention: Multi-head attention where the output is the negative gradient of a Hopfield-style energy function (no separate Q/K/V projections — uses a single symmetric H matrix per head)
Gradient Feed-Forward Network (FF2W): A 2-layer MLP used as an energy function E(x) = -(x · net(x)), where the forward pass computes ∇_x E(x) via torch.func.grad. This yields the conventional FF of transformers plus a second gradient term.
Recursive Weight-Sharing: A single transformer block is applied n_layer times (6 iterations), each application being one step of energy descent on the token landscape

This architecture achieves competitive perplexity with standard transformers while being grounded in energy-based principles, and shows advantages on downstream tasks like MMLU at comparable parameter counts.

Model Details

Property	Value
Parameters	128M (non-embedding)
Embedding dim	1536
Layers (iterations)	6
Attention heads	12
Context length	1024
Vocab size	50,304 (GPT-2 BPE)
Training data	OpenWebText
Training iterations	500,000
Train loss	3.29
Validation loss	3.30

Generation Quality Metrics (Table 5 from paper, 128M equal-parameter comparison)

Metric	NRGPT_H_FF2W	GPT	GPT_Rec_Parallel
Perplexity (GPT-2)	105	78	74
GQS	0.946	0.947	0.950
APCS	0.306	0.274	0.275
Distinct-1	0.650	0.638	0.613
Distinct-2	0.960	0.963	0.962

Usage

python

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load model (requires trust_remote_code for custom architecture)
model = AutoModelForCausalLM.from_pretrained("bsaha205/NRGPT-H-FF2W-128M-OWT", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("bsaha205/NRGPT-H-FF2W-128M-OWT")

# Generate text
model.eval()
input_ids = tokenizer.encode("In recent years, researchers have discovered that ", return_tensors="pt")
with torch.no_grad():
    output = model.generate(input_ids, max_new_tokens=200)
print(tokenizer.decode(output[0]))

Training Configuration

Optimizer: AdamW (β1=0.9, β2=0.99, weight_decay=0.01)
Learning rate: 3e-5 with cosine decay to 3e-6
Warmup: 2000 iterations
Batch size: 12 × 32 gradient accumulation = 384 effective
Precision: bfloat16
Hardware: Multi-GPU (NCCL backend)

Citation

bibtex

@inproceedings{dehmamy2025nrgpt,
  title={NRGPT: An Energy-Based Alternative for GPT},
  author={Dehmamy, Nima and Hoover, Benjamin and Saha, Bishwajit and Kozachkov, Leo and Slotine, Jean-Jacques and Krotov, Dmitry},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2025},
  url={https://arxiv.org/abs/2512.16762}
}

License

MIT