Energy-Based Models can now scale better than Diffusion models. Introducing TEM: Transformer Energy-based Models TLDR: • TEM outscales Diffusion without guidance • First scalable contrastive Energy-Based Models • Better generalization, test-time scaling, and OOD detection than Diffusion 🧵

[1/N] First, what are Energy-Based Models (EBMs)? EBMs learn a scalar energy E(x) for each data point, with low energy for real data, and high energy for fake data. This energy landscape defines generation, which starts from random noise and minimizes the energy until reaching a local minimum.

[2/N] Why formulate generative modeling with EBMs? It's easier to judge whether an image looks real, than to produce one in a single shot. EBMs learn this verifier (energy scalar) directly, then generate by optimizing against it, which enables: • Test-time scaling: more optimization steps → better samples • Compositional generalization: add energies to compose • Likelihoods and OOD detection: the energy acts as a verifier

[3/N] Despite these benefits, contrastive EBMs have struggled with scalability… so how can we scale them? To solve this we created Transformer Energy Models (TEMs). TEMs involve two key ingredients: • Multi-layer energies: sum energy predictions from early, middle, and final Transformer layers • Positive representation alignment: align EBM image features with a pretrained foundation model

[4/N] TEMs scale better than diffusion without guidance. TEM-XL's curve drops steeply with training compute and passes DiT, SiT and SiT+REPA, reaching a 3.82 FID with less compute than other diffusion baselines.

[5/N] Because TEM learns an energy landscape, they unlock capabilities diffusion models have traditionally struggled with. For example, TEMs enable compositional generation, OOD detection, and zero-shot inpainting. Here, we simply add the energies for compositional generation and TEM succeeds zero-shot over Compositional Diffusion.

[6/N] The energy also works as a verifier, enabling OOD detection that works better than Diffusion. Across four ImageNet OOD benchmarks, TEM averages 0.717 AUROC, vs. 0.55–0.60 for DiT, SiT, and SiT+REPA.

[7/N] Diffusion models struggle with generation tasks they weren’t trained on like inpainting. But TEMs handle these tasks zero-shot due to leveraging an energy landscape.

[8/N] TEMs can also do test-time scaling past the number of sampling steps they were trained with. Despite being trained with 40 sampling steps, we can run 120 inference steps and performance keeps improving.

[9/N] TEMs demonstrate the promise of Energy-Based Models as a generative modeling paradigm. More EBM work is coming soon so stay tuned! Thanks to the amazing @intlee0609 for leading this work and @du_yilun 🌐 Project page: 📄 Paper: 💻 Code: Jungsoo's thread:

how does the method without representation alignment perform relative to diffusion-based methods (without representation alignment)?

yah this is a great question and we will need to test it!

LeCun has been hammering on energy landscapes forever, cool to see contrastive EBMs actually scale with transformers

definitely! i am now extremely confident the future of EBMs will be a mix of reconstructive + contrastive :)

EBM FTW

cool work Alexi, looking forward to more EBM works :)

thanks jayden :) when can we get next EBM prediction?

really cool

thank you!

I have my doubts about its generality until I see a pixel-space version

im confident we can make it happen!

Doing god's work!

thank u for sharing, will definitely read it!!!

@choang333 Nice work Alexi. Looking forward to reading!

TEM-XL hitting 3.82 FID with less compute than the diffusion baselines is striking. A single energy score that also works as a verifier could make OOD detection feel built in, not bolted on. After years shipping large systems, how does that score hold up when real data drifts?

very cool!

So cool to finally see this out!! Congrats!

thanks aayush :)

Outscaling without guidance is the claim I want to see hold. Where does the compute go at sample time compared with a diffusion model of the same size?

mostly towards similar things as diffusion, denoising from the prior to the data distribution worth mentioning that more of the compute is gradients at inference-time than with diffusion (2/3rds here, vs none with diffusion)

That lines up. One third of 14.48T over 60 steps is about 0.08T of forward per step. A DiT-XL step is 0.12T. So the forward net is the cheaper part and the gradients are the surcharge.

@beffjezos

You became my bucket list author

thank you! although @intlee0609 deserves most of the credit here

@intlee0609 I guess my next plan to be pursue you and him for a paper in audio modality with energy models (because I never explored energy models and have no co-authors in this field)

super cool work!!!

thanks jonathan!!

@du_yilun 🚀🔥

@du_yilun EBM memory🤫

@du_yilun 💙

Interesting work!

thank you means a lot! big fan of your work on discrete diffusion :)

Awesome!

thanks lev :)

thank you.

EBM finally getting its scalability moment — the no-guidance test-time scaling is the killer feature. Curious whether the contrastive training cost stays competitive at frontier scale, or only at the current sizes.

Multi-layer energy still scales the EBM

appreciate your work

did you release product based on this? i wanna see

really cool work!

thank you!