Energy-Based Transformers (alexiglad/EBT): A PyTorch Codebase for Energy-Based Reasoning

The problem EBT addresses: prediction without a scoring step

A feed-forward Transformer emits a distribution over the next token in one pass. There is no mechanism inside the forward pass to evaluate whether the chosen token is consistent with the rest of a candidate answer, and no way to spend more compute on a harder prediction than on an easy one. The EBT paper, linked from the README, proposes an energy function over candidate predictions and an iterative procedure that searches for low-energy candidates before outputting one. The repository exists to make that procedure trainable and testable at scale. The README frames the target as generalizable reasoning, or System 2 thinking, applied to any problem or modality, and states that the authors demonstrate scaling across data, depth, parameters and FLOPs. That is the claim the code is meant to support. The audience is narrow: researchers who want to reproduce or extend the paper's experiments, not application developers looking for a drop-in model. Nothing in the repository offers a hosted endpoint, a model card with weights, or a Python API for inference on a single string.

How the code is organized: job scripts feed train_model.py

The README describes a two-layer structure. A job script, for example job_scripts/nlp/pretrain/ebt_s1.sh, passes hyperparameters to train_model.py. That entry point sets up distributed training, determines the number of GPUs, resolves the model configuration and size, and alters behaviour based on the supplied arguments. It then calls base_model_trainer.py, a PyTorch Lightning trainer that owns the training loop, validation, testing, dataset setup, logging and checkpoint creation. The README notes that the model size parameter sets layers, attention heads and embedding dimension automatically, so the job script does not enumerate architecture fields. Inference reuses the same path with different flags: --only_test_model_ckpt points at a checkpoint, --only_test disables training, and --execution_mode "inference" selects the inference mode of the loop. The README also mentions infer_generated_samples and infer_ebt_advanced as the entry points for System 2 thinking with self-verification, and says that most other hyperparameters in an inference script are inherited from the checkpoint, with the resolution logic living in train_model.py. That inheritance rule is worth internalizing: editing an inference script's model hyperparameters will usually have no effect.

Installation and the assumptions the README does not hide

Setup is Conda-based. The README gives conda create -n ebt python=3.12, conda activate ebt, then pip install -r requirements.txt. Three alternative requirement files exist for environments where the default set fails: gh200_requirements.txt for GH200 machines, loose_requirements.txt for installs without NVIDIA, PyTorch and Triton packages, and environment.yml as a Conda-native path. Two environment variables are called out, HF_HOME for the Hugging Face cache and HF_TOKEN for authentication, plus wandb login inside the activated environment. The wandb dependency is not optional in practice: the README instructs you to set the entity and project in the job script, and logging flows through that integration. If your institution blocks external telemetry, expect to patch the logging path before the trainer will run cleanly. Video work adds a further prerequisite. The README defers to /data/vid/README.md for dataset installation and FFPROBE, and to /inference/vid/README.md for video inference setup. FFPROBE is an external binary, not a Python package, so a container that only installs requirements.txt will fail on the video path.

Running pretraining: bash directly or through the Slurm executor

Two launch paths are documented. The quick start is bash job_scripts/nlp/pretrain/ebt_s1.sh, which the README identifies as the System 1 EBT configuration used for many of the paper's experiments. The recommended path on an HPC cluster is bash slurm_executor.sh reference_a100 job_scripts/nlp/pretrain/ebt_s1.sh, where reference_a100 is a mandatory argument naming a Slurm header template. The README is explicit that only reference_a100 is currently available and that you must tailor it to your cluster; the template lives at job_scripts/slurm_headers/reference_a100.slurm, and adding a new profile means adding the script name to slurm_executor.sh. The README also notes the more standard alternative of attaching a Slurm header to the bash script and calling sbatch directly. Inside a job script, the parameters to change are RUN_NAME, MODEL_NAME and MODEL_SIZE, and the README warns to edit all three occurrences at once so log names stay consistent. Multinode training requires ntasks equal to ngpus and launching with srun python filename.py, following job_scripts/nlp/pretrain/ebt_s1_mn.sh. The README is candid that very little multinode training was used for the paper and that the multinode code is correspondingly underexplored. It also flags a specific trap: GPU binding in the Slurm header, for example #SBATCH --gpu-bind=verbose,closest, may need to be disabled. Treat the multinode path as a starting point rather than a supported configuration.

Inference and the placeholder checkpoint

Inference scripts live under a modality directory and an inference subdirectory, for example job_scripts/nlp/inference/ebt.sh. The README states the main differences from pretraining are the checkpoint flag, the test-only flag and the execution mode. It also gives a useful escape hatch: --only_test without --execution_mode "inference" is the path for plain evaluation such as perplexity, and the test_step in base_model_trainer.py plus train_model.py are where that behaviour is defined. One operational detail deserves emphasis because it will cost time if missed: the inference scripts contain the literal string your/model/ckpt, and the README instructs you to replace it with an actual .ckpt file. There is no download command in the supplied material for pretrained weights, no model registry, and no release asset referenced. Anyone evaluating the repository should confirm that a checkpoint is obtainable before planning an inference experiment, because the code path is written and the weights are not described.

The minimal training loop and what it does not reproduce

The repository includes example_code/minimal_nlp_training_loop.py, described as a comparison between a Transformer++ and an EBT language model. It runs with a single python command and no Slurm, no distributed setup and no dataset download described in the README. The README states plainly that this example does not reproduce the paper's results. That sentence is the most important qualifier in the document. The minimal loop is the right first stop for reading the energy-based objective in isolation, and the wrong artifact for making any claim about scaling behaviour. The paper-level claims live behind the job scripts, the full datasets and the cluster configuration. If you only run the minimal example, you have verified that the code imports and that the training step executes, nothing more.

Where EBT is the wrong tool, and what to use instead

The iterative energy search is the point of the method, and it is also the cost. Every prediction that goes through the thinking procedure consumes additional forward passes relative to a single-pass decoder, so throughput per token drops in exchange for the reported generalization gains. For latency-bound serving, batch inference over short outputs, or any workload where a standard causal language model is already accurate enough, the extra machinery buys nothing. A plainer alternative is a conventional decoder-only Transformer trained with the same data and optimizer budget; the repository itself provides that comparison inside example_code/minimal_nlp_training_loop.py, which pits a Transformer++ against an EBT. The difference in approach is structural rather than a matter of tuning: the baseline commits to a token after one forward pass, while the EBT scores candidates and iterates. That means the baseline is simpler to serve, easier to quantize with standard tooling, and compatible with the wider ecosystem of inference servers that assume a single forward pass per step. Choosing EBT is choosing to pay per-token compute for a scoring step, and the decision should follow from whether your task actually benefits from that step. The repository does not include an evaluation harness that would let you answer that question on your own data without writing one.

Maintenance, licensing and what to check before you commit

The repository is Apache-2.0 and is not archived, with the last push dated 2026-04-21. No releases were retrieved, so there is no tagged version to pin against; you will be tracking the main branch, which means upgrade cost is whatever the branch does between your checkout and your next pull. The README acknowledges its own rough edges: the Slurm executor ships with one cluster profile, the multinode path was lightly used for the paper, and the video pipeline depends on an external binary documented in a separate README. Those are the areas where a fork is likely to diverge from upstream. On licensing, Apache-2.0 permits commercial use and modification with the usual notice and patent terms, but the licence covers this repository's code only; the paper, the linked website and any model weights you obtain elsewhere carry their own terms, and the README's citation section exists for that reason. This is not legal advice, and if you plan to ship anything derived from the code, read the licence text and check the terms attached to the datasets and checkpoints you use.

Editorial conclusion

Adopt this repository if you already run multi-GPU PyTorch Lightning jobs on a Slurm cluster and want to reproduce or extend the EBT training recipe; the job scripts under job_scripts/nlp, job_scripts/img and job_scripts/vid are the entry points, and the minimal example at example_code/minimal_nlp_training_loop.py is the only piece that runs without that infrastructure. Do not adopt it if you need a packaged library, a pip-installable model class, or single-GPU fine-tuning of a released checkpoint, because the repository ships training code and configuration rather than a model distribution. Before committing, verify three things: that your cluster matches the reference_a100 header or that you can write a replacement, that the dataset paths and FFPROBE dependency in /data/vid/ are satisfied if you train on video, and that the checkpoint you intend to load is a real .ckpt file, since the inference scripts ship with a placeholder path.

Official sources

1. alexiglad/EBT on GitHub
2. Issues
3. License: Apache-2.0
4. Project website
5. README