Energy-Based Models (EBMs) are a flexible class of generative models that decouple probabilistic modeling from the generation process. This decoupling directly enables compositional generation and inference-time scaling, where sample quality can be refined simply by increasing sampling steps. However, severe challenges in high-dimensional sampling and training instabilities with deep architectures have restricted EBMs to low-dimensional regimes. To resolve these bottlenecks, we introduce the Transformer Energy-Based Model (TEM), a scalable transformer energy function operating entirely within a latent space. On top of this backbone, we extract energies at multiple layers and sum them, composing complementary constraints from local texture to global semantics. Furthermore, two components make this high-quality at scale: Positive Representation Alignment (PRA), which supervises TEM's intermediate representations with a vision foundation model, and advanced sampling strategies that improve sample quality at inference. TEM-XL achieves an FID-50K of 3.82 on class-conditional ImageNet 256×256 without classifier-free guidance, showing that a contrastive EBM can generate high-fidelity images. On ImageNet 512×512, TEM achieves an FID-50K of 8.24, demonstrating that the approach scales to higher resolution. We further validate the core promises of EBMs: compositional generation, zero-shot inpainting, and out-of-distribution detection.
Figures & tables
Figure 1: High-fidelity ImageNet 256 × 256, 512 × 512 generation and compositional synthesis by TEM-XL/2. Left: uncurated class-conditional samples from 80 Langevin steps. The two large images are 512 × 512 samples and the remaining tiles are 256 × 256. Right: compositional generation by summing the energy functions of two independent classes, producing a coherent scene that satisfies both concepts without retraining.
Figure 2: Proposed Architecture and Training Strategy. Left : the Transformer Energy-Based Model processes VAE-encoded latent tokens through a stack of TEM blocks. Energy heads at multiple intermediate layers produce scalar energies whose sum defines the total energy. Right : contrastive training pushes down the energy of data samples x+ and pushes up the energy of MCMC negatives x− , while Positive Representation Alignment (PRA) aligns the model’s internal features with a frozen VFM.
ImageNet 256 × 256
ImageNet 512 × 512
Method
Params
NFE
FID ↓
FID ↓ (CFG)
NFE
FID ↓
FID ↓ (CFG)
GAN
BigGAN-deep
112M
1
6.95
–
1
8.43
–
GigaGAN
569M
1
3.45
–
–
–
–
Autoregressive / Masked
MaskGIT
227M
8
6.18
–
12
7.32
–
Table 1: Class-conditional generation on ImageNet-256 × 256 and 512 × 512 (FID-50K). Params refer to the 256 × 256 configuration unless a resolution-specific model is used. Baseline values are taken as reported in the original papers; “–” indicates that the corresponding setting is or was not reported.
Figure 3: Scaling behavior of TEM on class-conditional ImageNet 256 × 256. Left : FID-50K as a function of training steps across model sizes. Larger models consistently achieve lower FID at every checkpoint; 40 sampling steps are used for sampling during training. Middle : FID-50K as a function of cumulative training GFLOPs. All four model sizes follow a consistent compute–quality trend (dashed line). Right : FID-50K versus training GFLOPs for TEM-XL compared against DiT-XL, SiT-XL and SiT-XL+REPA showing that TEM tracks, and eventually achieves lower FID than unguided diffusion baselines with less training compute.
Figure 4: Effect of PRA on training. FID-50K across training steps for TEM-B with no representation alignment, Pos+Neg RA, and PRA. PRA consistently outperforms both, reaching 19.9 FID at 50K steps versus 42.8 without alignment and 36.6 with Pos+Neg RA. This supports aligning only positive samples in our method.
Figure 5: Ablations on Multi-energy Head and Loss Terms. Left : generated images from different energy heads, where (a), (b), (c) decompose the energy into early, middle, and final heads respectively, and (d) uses the full energy head. Middle : cosine similarity between energy head gradients across Langevin steps. Right : FID-50K at 50K steps for TEM-B. The upper block adds each loss term cumulatively with three heads; the lower block varies the number of energy heads under the full loss, against the three-head configuration above.
Figure 7: Inference-time scaling pseudocode and quality–diversity trade-off. Left: MCMC sampling procedure of TEM. Chains can run beyond 40 steps, but only the step-40 state is inserted into the replay buffer, so the buffer stays matched to the training dynamics. Right: FID-50K vs. Inception Score for TEM-XL, increasing the number of sampling steps (blue, η=500 ) or the step size (green, T=40 ) from the training configuration ( T=40 , η=500 ). Moderately increasing the number of steps improves both metrics (FID 3.82 at 60 steps). Beyond that, IS keeps rising while FID degrades as samples concentrate on high-density modes. This mirrors the fidelity–diversity trade-off controlled by classifier-free guidance, but here it comes from a single network with no retraining.
Figure 8: Application of TEM for Image Composition and Inpainting. Left : composed ImageNet images, car mirror + white shark and border collie + rapeseed (top), strawberry + petri dish and lemon + fountain (bottom). Right : inpainting of a meerkat (top) and a movie theater (bottom) across sampling steps. Both applications require no retraining or architectural modification.
Figure 9: CLEVR 256 × 256 compositional generation. Top: conditioning inputs, yellow boxes indicate specified object positions; second row shows TEM-B samples, green boxes indicate detected objects. Bottom: compositional generation accuracy ( ↑ ), object detection success rate by a trained classifier over 1000 scenes for 1-, 2-, and 3-object compositions.
OOD Dataset
DiT
SiT
SiT-REPA
Ours
iNaturalist
0.777
0.733
0.780
0.904
Textures
0.439
0.426
0.438
0.687
Places365
0.669
0.621
0.481
0.636
SUN
0.505
0.513
0.500
0.640
Average
0.598
0.573
0.549
0.717
Figure 10: Out-of-distribution detection on ImageNet benchmarks (AUROC ↑ ). Comparison of DiT-XL/2, SiT-XL/2, SiT-XL/2+REPA and our TEM-XL.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Model
TEM-S
TEM-B
TEM-L
TEM-XL
Params
33M
122M
415M
609M
Depth (layers)
12
12
24
28
Hidden size
384
768
1024
1152
Attention heads
6
12
16
16
Appendix
Table 2: TEM architecture configurations. Model variants from TEM-S to TEM-XL. All variants use a patch size of 2×2 and operate in the latent space of a pretrained Stable Diffusion VAE. Scaling from S to XL increases both depth and hidden dimension, yielding monotonic FID improvements (see Figure 3 ).
Model
# Steps
FID ↓
IS ↑
Precision ↑
Recall ↑
TEM-XL
40
7.63
106.0
0.65
0.65
TEM-XL
60
3.82
176.2
0.76
0.57
TEM-XL
80
4.72
221.1
0.81
0.51
DiT-XL ( g =1.0)
250
9.62
121.5
0.67
0.67
DiT-XL ( g =1.25)
250
3.22
201.8
0.76
0.62
DiT-XL ( g =1.5)
250
2.27
278.2
0.83
0.57
Appendix
Table 3: Fidelity-diversity trade-off. Top: TEM-XL with varying numbers of sampling steps, precision/recall with k=3 . Bottom: DiT-XL/2 with varying classifier-free guidance scale g results from Peebles and Xie [33]
Model
Training steps
Training FLOPs
FID ↓
DiT-XL/2
7M
638 E
9.62
SiT-XL/2
7M
638 E
8.30
SiT-XL/2+REPA
4M
393 E
5.90
TEM-XL
50K
146 E
3.82
Appendix
Table 4: Training cost on ImageNet 256 × 256. All values use the same accounting convention. VFM forwards are included for TEM and SiT+REPA.
Model
Steps / Sample
FLOPs / Sample
FID ↓
DiT-XL/2
250
29.66 T
9.62
SiT-XL/2
250
29.66 T
8.30
TEM-L
60
9.87 T
4.74
TEM-XL
60
14.48 T
3.82
Appendix
Table 5: Sampling cost on ImageNet 256 × 256. All results are without classifier-free guidance.