alexiglad/EBT

Energy-Based Transformers reformulate transformer inference as an energy minimization problem, enabling iterative refinement ("System 2 thinking") at every token prediction rather than just feed-forward generation. The approach scales favorably across multiple axes—data, model depth, parameters, and FLOPs—while improving generalization, demonstrating this through multi-modal experiments (NLP, vision, video). Built on PyTorch Lightning, it provides modular training/inference pipelines with support for distributed training, HuggingFace dataset integration, and W&B logging, along with minimal examples for quick experimentation.

613 stars.