PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Transformer Energy-based Models For High Quality Image Synthesis

Jungsoo Lee, Alexi Gladstone, Yilun Du

arXiv:2610.31889v1cs.LGcs.CV

TL;DR

High-dimensional sampling and training instability have limited EBMs to smaller-scale benchmarks. TEM combines a scalable transformer EBM with representation alignment and sampling strategies, and reports competitive high-fidelity ImageNet generation at two resolutions. The paper also demonstrates compositional generation, inpainting, and out-of-distribution detection.

  • Problem

    EBMs face poor mixing and mode collapse in high-dimensional MCMC, and deeper architectures can be severely unstable, limiting evaluation on high-resolution datasets.

    Source: §1 · pp. 1–2

  • Method

    TEM uses a scalable Vision Transformer EBM with multiple energy heads, representation alignment, and advanced sampling strategies.

    Source: §1 · p. 2, §3.1 · p. 3, §3.2 · p. 4

  • Results

    TEM-XL reports competitive high-fidelity ImageNet synthesis at 256×256 and 512×512, with guidance-free FID-50K 3.82 and 8.24, respectively.

    Source: §4.1 · p. 6, §6 · p. 9

  • Takeaways & Limitations

    At high resolution, TEM retains compositional generation, inpainting, and out-of-distribution detection capabilities.

    Source: §6 · pp. 9–10

    Iterative sampling during training and inference is computationally costly, and longer trajectories can sharpen images while reducing diversity and worsening FID.

    Source: §6 · p. 9

Abstract

from arXiv · show

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$\times$256 without classifier-free guidance, showing that a contrastive EBM can generate high-fidelity images. On ImageNet 512$\times$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.

1 INTRODUCTION

EBMs separate probabilistic modeling from the generation process, supporting flexible sampling and composition. This work addresses their high-dimensional sampling and scaling challenges with a transformer-based approach and demonstrates high-quality generation alongside compositional capabilities.

1 INTRODUCTION

  • EBMs decouple learning from inference, enabling test-time compute scaling through more MCMC steps and composition of different distributions.These properties can support zero-shot generation and generalization.Source: §1 · p. 1
  • EBMs face poor high-dimensional MCMC mixing and mode collapse, while deeper conventional architectures suffer severe training instability.These challenges have limited EBM research to smaller-scale benchmarks and constrained evaluation on high-resolution datasets.Source: §1 · pp. 1–2
  • Figure 1 contrasts uncurated class-conditional samples at two ImageNet resolutions with a coherent two-class composition produced without retraining.The composition sums the energy functions of two independent classes.Source: Figure 1 · p. 2
  • Contrastive EBMs with their own MCMC negatives achieve high-quality 256×256 ImageNet generation through persistent sampling initialized from Gaussian noise.The work also reports compositional generalization and flexible inference-time control.Source: §1 · p. 2
  • The proposed contributions include a scalable Vision Transformer EBM with multiple energy heads and a training framework combining latent-space design, representation alignment, and advanced sampling.The work also evaluates compositional generation, zero-shot inpainting, and out-of-distribution detection.Source: §1 · p. 2

2 BACKGROUND

EBMs represent distributions with an energy function, but their normalizing constant makes exact likelihood training infeasible. The background then motivates short-run Langevin sampling as a finite-step approximation.

2.1 ENERGY-BASED MODELS

  • An EBM uses a neural-network energy function to define an unnormalized density, with low-energy regions corresponding to high data density.The normalizing constant is intractable, making exact maximum-likelihood training computationally infeasible.Source: §2.1 · p. 2
  • The likelihood gradient is tractable, but its model expectation requires sampling from the EBM distribution, which cannot be done exactly in practice.MCMC convergence would require infinitely long chains with infinitesimal steps.Source: §2.1 · p. 2

2.2 SHORT-RUN LANGEVIN MCMC

  • Short-run Langevin MCMC starts from a fixed distribution and runs a finite number of updates, producing an efficient but biased sampler.Figure 2 depicts contrastive training with data samples and MCMC negatives alongside the TEM architecture.Source: §2.2 · pp. 2–3, Figure 2 · p. 3

3 SCALABLE TRANSFORMER ENERGY-BASED MODELS

TEM uses a Vision Transformer energy function and sums energy heads from different depths, while training in latent space with contrastive learning and representation alignment. At inference, adjusting sampling can refine samples without retraining, with a trade-off in mode coverage.

3.1 TRANSFORMER ENERGY-BASED MODELS ARCHITECTURE

  • A ViT maps spatial patch embeddings through Transformer blocks, pools final tokens, and uses an MLP head to produce scalar energy.The design reports stable, effective scaling without EBM-specific components such as spectral normalization.Source: §3.1 · p. 3
  • Multiple heads at different Transformer depths sum energy terms, combining constraints from progressively abstract features within one model.Early layers capture low-level structure and deeper layers high-level semantics.Source: §3.1 · p. 3

3.2 TRAINING TRANSFORMER ENERGY-BASED MODELS

  • Training operates on autoencoder latents and contrasts data samples against sampler negatives, with an L2 penalty bounding energy magnitude.Sampler matching encourages low-energy, diverse outputs; PRA aligns positive-sample representations with vision foundation model features, and alignment is not applied to MCMC negatives.Source: §3.2 · p. 4

3.3 INFERENCE TIME SCALING

  • Increasing the sampling step size or extending sampling beyond the training horizon can sharpen samples and make them more class-typical, at the cost of mode coverage.The step-size adjustment keeps injected noise fixed and does not require retraining or additional unconditional network evaluations.Source: §3.3 · p. 5

4 EXPERIMENTS

Experiments evaluate TEM’s ImageNet quality and scaling, its training and sampling choices, and applications including composition, inpainting, and out-of-distribution detection. Results include guidance-free high-resolution generation and inference-time control, alongside ablations of multi-head energy and representation supervision.

4.1 MAIN RESULTS

  • FID-50K is 3.82 on ImageNet 256×256 and 8.24 on 512×512 for TEM-XL without classifier-free guidance; Figure 3 shows scaling against unguided diffusion baselines.The paper reports that TEM eventually achieves lower FID than those baselines with less training compute.Source: §4.1 · p. 6, Figure 3 · p. 6

4.2 ABLATIONS

  • At 50K steps, the three-head model reaches lower FID than single- or double-head baselines, while its heads contribute complementary texture, layout, and semantic constraints.The summed heads satisfy all three aspects; gradient similarity is high early in sampling and decreases later.Source: §4.2 · p. 7
  • At 50K steps, PRA reaches 19.9 FID-50K for TEM-B, versus 42.8 without alignment and 36.6 with positive-and-negative representation alignment.The authors attribute the weaker result from aligning negatives to their providing a wrong alignment target.Source: §4.2 · p. 7

4.3 TEST-TIME SCALING ANALYSIS

  • Beyond the 40-step training horizon, sampling produces sharper, more class-typical images, while longer trajectories can raise Inception Score at the cost of FID and diversity.The paper describes this fidelity–diversity trade-off as resembling classifier-free guidance, but achieved through inference-time scaling.Source: §4.3 · p. 8
  • Users can tune sampling length and step size for fidelity or diversity without retraining or architectural changes.Only the step-40 state is returned to the replay buffer, keeping it matched to training dynamics.Source: Figure 7 · p. 7, §4.3 · p. 8

4.4 APPLICATIONS OF TEM

  • Summing energy functions composes their distributions, and TEM-B outperforms Composable Diffusion at every tested CLEVR composition count.The reported counts are compositions of one, two, and three objects at specified positions.Source: §4.4 · p. 8
  • For inpainting, TEM refines missing latent regions while restoring observed-region latents after each sampling step; its scalar energy also provides an OOD score without task-specific training.The OOD evaluation uses ImageNet-1k validation data as the in-distribution reference and reports AUROC on four benchmarks.Source: §4.4 · pp. 8–9

5 RELATED WORK

Prior EBM work demonstrated practical sampling and latent-space control, but left scalability at high resolution unresolved. TEM uses a single Transformer energy function in a pretrained autoencoder’s frozen latent space, contrasting with earlier approaches that added auxiliary mechanisms or non-noise initialization.

Foundations of Energy-Based Models

  • Earlier EBMs used short-run Langevin MCMC and replay buffers for coherent CIFAR-10 samples, while follow-up methods remained at 32×32 or lower resolutions.The paper characterizes scalability of the EBM paradigm as unresolved.Source: §5 · p. 9

Latent-Space Energy-Based Models

  • Prior latent-space EBMs explored VAE priors, hierarchical models, and latent-manifold control; TEM instead trains one scalar-output Transformer energy in a frozen pretrained-autoencoder latent space.VAEBM uses EBM refinement in data space and reparameterizes its MCMC chain through VAE latents.Source: §5 · p. 9

Scaling Energy-Based Models

  • Prior high-resolution EBM efforts rely on diffusion denoising or auxiliary-image initialization; TEM is contrastively trained and generates high-fidelity 256×256 images, enabling inference-time scaling and native EBM applications.The adversarial sampler is initialized from an auxiliary image dataset at training and inference and is never seeded from noise at this scale.Source: §5 · p. 9

6 LIMITATIONS AND CONCLUSION

TEM’s primary limitation is the computational cost of iterative sampling, with longer trajectories trading diversity for sharper images. The conclusion reports high-fidelity ImageNet synthesis at 256×256 and 512×512 while retaining compositional generation, inpainting, and out-of-distribution detection.

6 LIMITATIONS AND CONCLUSION

  • Iterative sampling during both training and inference is TEM’s primary limitation.The paper suggests feed-forward proposals as a direction for reducing the sampling needed to refine candidates.Source: §6 · p. 9
  • Longer sampling trajectories sharpen images but can reduce diversity and worsen FID.Adaptive noise schedules and proposals that move between modes are suggested as possible ways to preserve diversity while improving sample quality.Source: §6 · p. 9
  • TEM synthesizes high-fidelity ImageNet images at 256×256 and 512×512 while retaining compositional generation, inpainting, and out-of-distribution detection at these resolutions.The paper presents these capabilities as properties that motivate energy-based modeling.Source: §6 · pp. 9–10

A THE SAMPLER-MATCHING LOSS

The sampler-matching loss addresses the gap between the short-chain sampler and the model distribution by encouraging sampler outputs to be low-energy and diverse. Its energy term is tractable through the sampling chain, while its entropy term uses a nearest-neighbor diversity regularizer.

A THE SAMPLER-MATCHING LOSS

  • The maximum-likelihood gradient requires model-distribution negatives, but short Langevin chains approximate them and can be biased when they miss low-energy regions.The sampler-matching loss is motivated by this mismatch and encourages low-energy, diverse sampler outputs.Source: A THE SAMPLER-MATCHING LOSS · p. 14

A.1 ESTIMATING THE SAMPLER-ENERGY TERM

  • The sampler-energy term is differentiable through sampling steps, but backpropagation is limited to the final step because differentiating through all steps is memory-prohibitive.The energy and entropy terms together form the sampler-matching objective.Source: §A.1 · p. 14

A.2 ESTIMATING THE ENTROPY TERM

  • Exact sampler entropy is intractable because its density marginalizes over all trajectories that could produce a sample.The paper instead uses a nearest-neighbor entropy estimator as a diversity regularizer, not as an accurate entropy estimate; its buffer samples are neither independent nor drawn from the current sampler distribution.Source: §A.2 · p. 14

B.1 TRAINING DETAILS.

The training and evaluation setup uses a latent-space Transformer EBM with persistent sampling. The reported sampling-step comparison quantifies a fidelity–diversity trade-off, and the appendix gives compute details for training and inference.

B.1 TRAINING DETAILS.

  • TEM variants use a patch size of 2 × 2 in a pretrained Stable Diffusion VAE latent space, with depth and hidden dimension increasing from S to XL alongside monotonic FID improvements.Latents are norm-clipped, and the reported training setup uses 40 sampling steps, a replay buffer, and DINOv2-large for representation alignment.Source: §B.1 · pp. 14–15, Table 3 · p. 15

B.2 INFERENCE DETAILS.

  • Increasing TEM-XL sampling from 40 to 80 steps raises precision from 0.65 to 0.81 while reducing recall from 0.65 to 0.51.Table 4 compares TEM-XL sampling steps with DiT-XL/2 classifier-free guidance settings; evaluation uses persistent chains, with later chains continuing from buffered states.Source: §B.2 · p. 15, Table 4 · p. 16

C COMPUTATIONAL COST

  • Sampling 50K images takes about an hour on one H200 GPU at 60 steps, while 512×512 training and inference take roughly 3–4 times longer than at 256×256.The paper also reports approximately six days to train TEM-XL for 50K steps on 256×256 ImageNet using four H200 GPUs.Source: C COMPUTATIONAL COST · pp. 15–16

C.1 TRAINING COST

TEM has per-step training and sampling overhead, but the paper reports lower guidance-free FID with less total training compute and comparable sampling cost. Its sampling-cost comparison excludes classifier-free guidance.

C.1 TRAINING COST

  • TEM reaches a lower guidance-free FID with less total training compute despite per-step overhead.Table 5 compares total training cost on ImageNet 256×256 using a common accounting convention, including VFM forwards for TEM and SiT+REPA.Source: §C.1 · p. 16, Table 5 · p. 16

C.2 INFERENCE COST

  • TEM achieves lower guidance-free FID at comparable sampling cost, with each sampling step costing about 2F for a K-step chain costing 2FK.Table 6 compares ImageNet 256×256 sampling costs without classifier-free guidance; guidance would improve baseline quality but double baseline sampling cost.Source: §C.2 · pp. 16–17, Table 6 · p. 17

D TRAINING ALGORITHM

The full TEM training procedure is provided in Algorithm 2.

D TRAINING ALGORITHM

E COMPOSITIONAL GENERATION IN 256X256 CLEVR DATASET

The section describes compositional-generation and OOD-evaluation setups, and reports that foundation-model alignment alone does not account for TEM’s OOD-performance gap. It also specifies the baselines used for comparison.

E COMPOSITIONAL GENERATION IN 256X256 CLEVR DATASET

  • For 256×256 CLEVR compositional generation, TEM-B is conditioned on normalized object locations through global AdaLN coordinates and spatial patch-grid bias.Coordinates use an 8-band Fourier encoding, producing a 34-dimensional code; training runs for 150k steps with batch size 64.Source: E COMPOSITIONAL GENERATION IN 256X256… · pp. 17–18
  • Compositional detection success is measured over 1,000 scenes for one-, two-, and three-object compositions using a classifier at 128×128.TEM’s 256×256 samples are downsampled for scoring; the classifier exceeds 99.9% detection accuracy on real CLEVR images.Source: E COMPOSITIONAL GENERATION IN 256X256… · p. 18

G.1 BASELINES IN TABLE 1

  • OOD baselines include DiT-XL/2, SiT-XL/2, and SiT-XL/2+REPA; their scores use native residuals and a single network evaluation.The reported comparison also draws on GAN, autoregressive, masked, diffusion, and flow baselines, with parameter counts for 256×256 configurations unless resolution-specific models were released.Source: F OUT-OF-DISTRIBUTION DETECTION · p. 18, §G.1 · pp. 18–19

F OUT-OF-DISTRIBUTION DETECTION

  • SiT-XL/2+REPA’s AUROC is close to SiT-XL/2’s, so DINOv2 alignment alone does not account for TEM’s OOD-performance gap.Source: F OUT-OF-DISTRIBUTION DETECTION · p. 18
  • The OOD protocol uses ImageNet-1k validation as in-distribution data and evaluates AUROC on iNaturalist, SUN, Places365, and Textures.TEM uses its unconditional branch to score inputs with S(x) = −Eθ(x) in one forward pass, without labels, post-hoc calibration, or OOD exposure.Source: F OUT-OF-DISTRIBUTION DETECTION · p. 18

G.2 SCALING ENERGY-BASED MODELS

Prior work has scaled energy landscapes using predictive, transport, or reconstruction-style objectives, while some explicit scalar-energy variants have shown quality or resolution limits. TEM instead trains a directly parameterized scalar energy contrastively with negatives from its own MCMC sampler, and the authors report scaling to ImageNet 256×256.

G.2 SCALING ENERGY-BASED MODELS

  • Prior approaches scale energy landscapes through denoising targets, equilibrium-gradient regression, or reconstruction loss through unrolled optimization.Denoising-score approaches can recover MCMC-based composition without a single global energy, while equilibrium matching defines energy implicitly through its gradient field.Source: §G.2 · p. 19
  • Explicit scalar-energy variants of Equilibrium Matching report FID 73.40 and 75.53 in the B/4 setting, while Energy Matching’s image-generation experiments are limited to 32×32.Energy Matching combines optimal-transport and contrastive energy-learning objectives.Source: §G.2 · p. 19
  • TEM directly parameterizes a scalar energy and trains contrastively with negatives from its own MCMC sampler, without a transport or flow-matching objective.The authors contrast this with prior scaling approaches and report that this explicit contrastive EBM scales to ImageNet 256×256.Source: §G.2 · p. 19

Related papers

Loading 2610.31889v1…