Energy-Based Transformers

What the architecture inherits from classical energy models, and where that leaves it

Introduction

In July 2025, Alexi Gladstone and his coauthors at the University of Virginia, the University of Illinois, Amazon GenAI, Stanford University, and Harvard University published a paper introducing Energy-Based Transformers. The reported results are strong. Faster scaling than the standard transformer on six axes. Better behavior on data that differs from the training distribution. An inference procedure that allocates more compute to harder predictions.

The architecture revives an old framework. Reading it against the classical energy-based literature pays off in both directions. The same literature predicts where the architecture should struggle and where it should do well. This post explains what an energy-based model is, then what the paper does. It then walks through four limitations that matter in practice and the two places where the framework gives the architecture an advantage. The aim is to set the limitations against the strengths and let both sides count.

What an energy-based model is

A regular neural network maps an input to an output directly. Pass an image through a classifier and the network returns a label. Pass a sequence of tokens through a transformer and the network returns a probability distribution over the next token. The output is a function of the input through a single forward pass.

An energy-based model works differently. It does not produce an output directly. Given an input and a candidate output, it returns a single number called an energy. The energy says how compatible the input and the candidate output are. Low energy means the pair fits well. High energy means the pair fits poorly.

To make a prediction, the model has to find the candidate output with the lowest energy for the given input. It does this by gradient descent on the output. Start from a random candidate. Compute the gradient of the energy with respect to the candidate. Take a step downhill. Repeat until the energy stops decreasing. The output is the candidate at convergence.

The tradeoff differs from a standard model. A standard forward pass trades flexibility for speed. One pass, one output. An energy-based model trades speed for two things. It gets variable compute per prediction, since harder predictions take more optimization steps. It also gets a verification signal, since the energy itself is a score that says how good a prediction is.

The framework comes from physics and dates back to Ludwig Boltzmann in the 1870s. It entered neural networks in 1982 through John Hopfield’s associative memory network and was extended to learning in 1985 with the Boltzmann machine. The central obstacle is the partition function. To turn an energy into a probability that sums to one, the model divides by a normalizing constant that sums the exponentiated negative energy across every possible output. For any model with meaningful capacity, that sum has too many terms to compute. Most of the energy-based literature is a search for ways around this sum.

The framework never went away. Yann LeCun’s 2006 tutorial recast a wide range of models as energy-based. His later writing has argued for energy as the foundation of future architectures. Score matching, noise contrastive estimation, and the score-based view of diffusion models are all modern descendants that handle the partition function in different ways. The Gladstone paper is the latest entry in this continuing line of work. What it adds is scale. The framework had not been pushed to the size of a modern transformer before.

What the Gladstone paper proposes

The paper takes one of the standard routes around the partition function. It avoids the problem by never computing a probability. The energy function returns an unnormalized score. Lower is better. Higher is worse.

Training works as follows. Take an input from the dataset. Initialize a candidate output as Gaussian noise. Run 2 or 3 steps of gradient descent on the energy function, updating the candidate to lower its energy. Compare the final candidate to the ground truth through a standard loss function. Backpropagate through the entire optimization sequence to update the model parameters. The procedure needs gradients of gradients, since the loss depends on a gradient step. These can be computed efficiently through Hessian-vector products at cost linear in model size.

The forward pass of an Energy-Based Transformer is therefore not a single computation. It is a small loop of gradient descent on the energy. The number of steps in the loop is a hyperparameter. At inference time the loop can be run for more steps than during training, which lets the model spend extra compute on harder predictions.

Three regularization techniques shape the energy landscape during training. Langevin noise, a small random perturbation added at each gradient step, encourages the optimizer to explore the landscape away from any single descent path. A replay buffer stores intermediate states from earlier training and reuses them, which extends the effective length of the optimization trajectories the model sees. The step size and number of optimization steps are randomized across training examples. Together these techniques produce an energy surface that is approximately convex around the training examples.

The reported results cover both discrete and continuous modalities. On language modeling, the architecture scales faster than the standard transformer recipe across data, batch size, depth, parameters, total compute, and embedding dimension. On image denoising, the architecture matches or exceeds Diffusion Transformers while using 99 percent fewer forward passes. On data far from the training distribution, the gain from extra optimization steps grows roughly linearly with the distributional shift. The largest model the paper trains has 800 million parameters. The four limitations below sit inside this setup.

Figure 1 of the paper places the architecture next to feedforward transformers, RNNs, and Diffusion Transformers. The last comparison is the most informative. Both the Diffusion Transformer and the Energy-Based Transformer take a candidate next token as part of their input. They differ in what they return. The Diffusion Transformer predicts the noise to remove from the candidate. The Energy-Based Transformer predicts a single energy that scores the candidate. That one difference is what lets the energy model act as a verifier. The same difference shapes both the limitations and the advantages discussed below.

Limitation 1: The convex landscape and multimodal data

The paper acknowledges in Section B.2 that text-to-image generation fails on the COCO dataset, a standard benchmark of natural images paired with captions. When a single caption maps to many plausible images, such as “a giraffe standing in a field,” the model produces blurred averages instead of any single distinct image. The authors attribute the failure to the convex energy landscape assumption made during training. The acknowledgment stops there.

The classical energy literature has been clear on this point since the Boltzmann machine. The energy landscape that matches a one-to-many conditional distribution has multiple basins, one per mode. A landscape that is convex around the training data, with a single broad minimum, flattens that structure into one. The three regularization techniques the paper introduces, Langevin noise, replay buffer, and randomized step size, all reward smooth energy surfaces that gradient descent can navigate reliably. Smoothness and multimodality pull against each other. The smoother the learned surface, the harder it is for that surface to keep distinct modes apart.

A careful reader can push back on how far this goes. Local convexity along the training trajectories does not prove global convexity everywhere. Stronger conditioning might collapse some of the multimodality before it becomes a problem. A different initialization scheme might let the optimizer land in different basins on different runs and recover distinct modes that way. Each of these counterarguments stays open. What remains is the narrower claim. The current training procedure is biased against highly multimodal conditional tasks, and the COCO failure is what that bias looks like in practice.

The practical reading follows from the bias. Tasks with strong one-to-many structure are where this architecture is most likely to disappoint as currently trained. Conditional image generation and open-ended text generation are the clearest cases of one-to-many structure. Multi-target retrieval and multimodal policy learning are the same problem in different settings. The paper suggests stronger conditioning as a future fix. The current architecture leaves the problem open.

Limitation 2: Energies are not comparable across contexts

Figure 8 in the paper plots the energy of different tokens on the same vertical axis. Function words such as “the” and “but” sit at low energy across the inference steps. Content words such as “quick” and “problem” sit higher. The caption reads the plot as evidence that the model has learned to express uncertainty. The intuition is that low energy means the model is confident.

The plot misleads. Comparing energies across contexts is like comparing two temperatures that were each measured on their own hidden scale. Each reading is internally consistent. Putting both numbers on one axis says nothing, because the zero point and the spacing differ from one scale to the next. The same thing happens with energies. The model produces unnormalized energies, and an unnormalized energy means something only next to other energies that share the same partition function. Within a single context the same normalizing constant applies, so comparing the energy of two candidate predictions is well posed. Across contexts, comparing the energy of “problem in sentence A” to “but in sentence B” mixes scores from different normalizing constants.

The within-context limit has practical consequences. Best-of-N selection works as expected, because the model generates several candidates for the same input and picks the one with the lowest energy. That is a within-context comparison. Selective prediction is a different case. A model doing selective prediction refuses to answer on inputs it is unsure about, which requires comparing confidence across different inputs. The energy cannot do this, because energies from different contexts sit on different scales. The paper does not flag this gap.

Limitation 3: Training procedure bias

Training an Energy-Based Transformer involves unrolling 2 or 3 gradient descent steps and backpropagating through them. The procedure resembles contrastive divergence, the approximate training method Geoffrey Hinton introduced in 2002 for restricted Boltzmann machines. Contrastive divergence replaces an expensive expectation over the model distribution with a few steps of Markov chain sampling started from a training example. Carreira-Perpiñán and Hinton showed in 2005 that the truncation introduces a bias that does not vanish with more data. Persistent contrastive divergence, introduced by Tieleman in 2008, keeps a long-running chain across training steps, much as the Gladstone paper keeps a replay buffer. Each variant has carried some form of bias.

The resemblance to contrastive divergence does not automatically mean the Gladstone procedure inherits its bias. A skeptical reader is right to say that similar training structure does not guarantee the same convergence problem. The contrastive divergence bias comes from a specific source. The chain samples from the model distribution and stops before it mixes, so the negative-phase gradient is computed from the wrong distribution. The unrolled gradient procedure in the new paper never samples from the model distribution. It runs deterministic gradient descent on the energy and compares the result to the ground truth. The contrastive divergence bias, in its precise form, does not obviously carry over.

Still, three things could go wrong, and the paper and the analogy tend to blur them. The first is truncated optimization. The forward pass stops after a few steps. The prediction sits short of an energy minimum, and the gradient that trains the energy function is taken at a point that has not converged. The second is short-trajectory approximation in general. Any method that replaces a long process with a few steps risks shaping the model to behave well only on those few steps. The third is replay-buffer staleness. Reusing old states ties the gradient to the history of training and brings the same staleness that persistent contrastive divergence is known for. These are distinct concerns. Only the third has a direct match in the classical literature. The first two are plausible on general grounds, and no one has characterized them for this exact procedure.

So the honest claim is modest. The classical literature does not predict a specific bias for unrolled energy descent. It does say that every approximate energy training method invented so far has carried some bias, and it points at three places to look. The procedure works empirically up to 800 million parameters. Whether it stays well-behaved at larger scale depends on whether any of these three effects compounds as the model grows. The paper does not measure them. For a practitioner weighing a larger investment, the open question is that missing measurement, not a predicted failure.

Limitation 4: The diffusion comparison is asymmetric

The paper reports that Energy-Based Transformers match or exceed Diffusion Transformers, often abbreviated DiT, on image denoising while using 99 percent fewer forward passes. Denoising is the standard objective for diffusion training. The model receives an image corrupted with Gaussian noise and is asked to predict the clean image. Ho, Jain, and Abbeel introduced the procedure in 2020. The comparison number is striking. It is also not symmetric in training compute.

Energy-Based Transformers train with 2 or 3 unrolled optimization steps per sample. Diffusion Transformers train with the full denoising schedule, often dozens or hundreds of steps. The training compute per sample differs by an order of magnitude or more. The reported 99 percent inference-time gap is real, and it sits inside an uneven setup. A practitioner would want to know whether the gap survives when both models get the same total training budget. A practitioner would also want to know whether the diffusion model can match the EBT inference cost once it trains with a shorter schedule. The deeper comparison between these two is the geometry, not the benchmark, and the upside section takes that up.

The paper does not run this experiment, and that is the gap. The energy framework may be more inference-efficient than diffusion. It may not. The answer turns on engineering choices, and the classical theory is silent on engineering questions. Read the 99 percent as a statement about inference at a fixed setting, and wait for a matched-compute comparison before reading it as anything more.

Where the architecture has a real edge

The same classical framework that predicts the failures also predicts two strengths, and the paper shows evidence for both. Each deserves stating as plainly as the limitations. The strengths matter as much as the weaknesses.

The first is adaptive computation that tracks difficulty. Because prediction is an optimization loop, the model can run the loop longer when a prediction is hard and stop early when it is easy. The paper measures this. The gain from extra optimization steps grows roughly linearly with how far a test input sits from the training distribution. That is the behavior a model meant to handle unfamiliar inputs should show. A standard transformer spends the same compute on every token, easy or hard, novel or familiar. An energy model can spend more where spending more helps. The property is structural, not a tuning trick. It pays off on tasks where the hard cases are rare but costly, and where inference-time compute can be traded for accuracy.

The second is verification built into the model. The energy is a score that says how good a candidate prediction is, produced by the same network that makes the prediction. The word “verification” is narrower than it sounds. The model verifies consistency with its own learned energy surface, not external correctness. The property is what Figure 1 shows. A diffusion model predicts a denoising direction and gives no scalar that says how good its current output is. An energy model gives that scalar for free. Within a single context the score is well posed, which makes it usable for selection. Generate several candidates, keep the lowest-energy one. The paper uses exactly this and reports gains from it. Low-entropy structured prediction is the natural fit for this, where one input has essentially one correct output and the task is to find it. It also fits any task where checking a candidate matters more than producing many different outputs.

These two strengths come with a boundary worth naming. Both depend on the prediction being close to single-valued. Adaptive refinement converges cleanly when there is one basin to converge to. The verification score is trustworthy when low energy points to one right answer and not to an average of several. The same single-mode assumption that creates the multimodal failure in the first limitation is what makes these two strengths work. Single-valued means low output entropy. The input maps to one answer or to a narrow set of close ones, and the conditional distribution has few modes. A task can still carry real uncertainty as long as the mode count is small. The architecture is matched to problems with low output entropy and mismatched to problems with high output entropy. The limitations and the advantages are two readings of one design choice.

A sharp reader will object here, and the objection is the most useful question in the whole discussion. Diffusion models also refine step by step, and they handle multimodal generation well. If step-by-step refinement and multimodality coexist in diffusion, why not here? The answer is that the two methods refine over different objects. A diffusion model learns a score field, the gradient of log-density across the whole space, conditioned on a noise level. Run from different starting noise, that field can carry samples toward different modes, because it encodes the full shape of the distribution and not a single low point. An Energy-Based Transformer descends one learned energy surface, and the regularization has smoothed that surface toward a single basin. Same word, refinement. Different geometry underneath. Diffusion keeps every mode because its target is a density gradient. The energy model drops modes because its target has been shaped into a bowl with one bottom. The contrast is one of bias, not of guarantees. Diffusion can collapse too. Strong guidance, aggressive step reduction, and distillation all flatten its outputs in practice. The point is the direction each method leans by default. Diffusion leans toward keeping modes and has to be pushed off that. The energy model, as regularized here, leans toward dropping them. This is why diffusion, and not the standard transformer, is the right conceptual comparison for this architecture. The two are the step-by-step generative methods. The place they split is the place the energy model loses multimodality in order to gain its verifier.

Beyond these two, the paper invites a longer list of hopes. Planning, theorem proving, molecular design, and iterative scientific optimization all come up. Each is a low-entropy search problem in spirit, so each is plausible on the same grounds. The paper demonstrates none of them. They are reasonable extensions, and the paper has not yet shown them. The same caution the limitations call for applies to them.

Conclusion

The paper contributes something genuine. It trains an energy-based transformer at 800 million parameters, which had not been done before. It introduces three landscape regularization techniques that hold training stable. It reports scaling rates above the standard transformer recipe on six axes. It demonstrates two strengths the classical framework predicts. Adaptive compute grows with difficulty, and the verification signal is usable in practice.

The training procedure appears poorly matched to highly multimodal conditional tasks, and the COCO failure is what that looks like. The energy is a within-context score, so it cannot serve as cross-context confidence. The training procedure resembles methods whose bias has been studied for two decades, and three distinct bias sources remain unmeasured for this exact procedure. The diffusion comparison is striking but asymmetric in training compute.

The pattern underneath is the useful part. Three of the four limitations were predictable from the classical energy literature, and so were both strengths. The convex landscape failure, the cross-context score, and the single-mode assumption all follow from how energy models work. The three together explain the architecture’s good and bad behavior at the same time. The diffusion comparison is the only open item that is experimental. The other three are theoretical.

For a practitioner the reading is not a yes or a no. It is a map. The most useful next experiments are the ones the paper did not run. A matched-compute diffusion comparison. A direct measurement of training bias at larger scale. A multimodal task with conditioning strong enough to test whether the single-mode limit can be pushed back. The architecture is matched to low-entropy structured prediction, adaptive-depth reasoning, and tasks where verifying a candidate matters more than producing many. It is mismatched to high-entropy conditional generation.

References:

Carreira-Perpiñán, M. Á., & Hinton, G. E. (2005). On contrastive divergence learning. Proceedings of the Tenth International Workshop on Artificial Intelligence and Statistics, 33–40. https://proceedings.mlr.press/r5/carreira-perpinan05a.html

Gladstone, A., Nanduru, G., Islam, M. M., Han, P., Ha, H., Chadha, A., Du, Y., Ji, H., Li, J., & Iqbal, T. (2025). Energy-Based Transformers are Scalable Learners and Thinkers. arXiv preprint, arXiv:2507.02092. https://arxiv.org/abs/2507.02092

Hinton, G. E. (2002). Training products of experts by minimizing contrastive divergence. Neural Computation, 14(8), 1771–1800. https://doi.org/10.1162/089976602760128018

Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33, 6840–6851. https://proceedings.neurips.cc/paper_files/paper/2020/hash/4c5bcfec8584af0d967f1ab10179ca4b-Abstract.html

LeCun, Y., Chopra, S., Hadsell, R., Ranzato, M., & Huang, F. (2006). A tutorial on energy-based learning. In Predicting Structured Data. MIT Press. http://yann.lecun.com/exdb/publis/pdf/lecun-06.pdf

Tieleman, T. (2008). Training restricted Boltzmann machines using approximations to the likelihood gradient. Proceedings of the 25th International Conference on Machine Learning, 1064–1071. https://doi.org/10.1145/1390156.1390290