A transformer architecture that thinks before every token

What it does

Energy-Based Transformers (EBTs) treat each prediction—whether a language token, an image pixel, or a video frame—as an optimization problem. The repo provides PyTorch Lightning implementations for autoregressive and bidirectional variants, split into System 1 (fast) and System 2 (deliberative) modes. It covers NLP, image generation, and video modeling, with training and inference scripts organized by modality.

The interesting bit

Instead of pure feed-forward passes, EBTs use an energy-based formulation that lets the model pause and refine individual predictions iteratively. The authors claim this is the first approach to outscale standard Transformers across data, depth, and parameters while generalizing better. The repo also includes a custom JVP Flash Attention implementation to handle the second-order derivatives required by the energy-based training.

Key highlights

• Autoregressive EBTs based on LLaMA 2 and bidirectional variants based on DiT, with adaptive layer norm and time embeddings.
• A causal replay buffer enabling System 2 reasoning and self-verification during inference.
• Modular model/ directory meant to be lifted into custom training loops, separate from the PyTorch Lightning scaffolding.
• A minimal standalone training loop is included, though the README notes it will not reproduce the paper’s results.
• Ships with reference Slurm job scripts and an executor, but multinode training support is largely untested.

Caveats

• The minimal training loop explicitly does not reproduce paper results.
• Multinode training code is minimal and largely unexplored.
• Video dataset setup requires additional external dependencies like FFPROBE.

Verdict

Researchers experimenting with energy-based models or scalable reasoning architectures should take a look. If you need a drop-in replacement for standard Transformers with mature production support, this is still early-stage research code.