Transformer energy models improve image synthesis quality and scale
Scaling Transformer Energy-based Models For High Quality Image Synthesis
Machine LearningComputer Vision and Pattern Recognition
Summary
Creating detailed images with computers can be tricky because it often requires balancing accuracy and speed. The authors developed a method using transformer-based energy models that work in a simplified space of image features. This approach allows the model to generate higher-quality images by refining the process with more steps. They also showed the model can mix parts from different images, fill in missing areas, and detect unusual inputs.
What this means in practice
- •For computer vision engineers: Create high-quality class-conditional images at 256x256 and 512x512 resolutions using energy-based transformer models.
- •For graphics software developers: Implement zero-shot inpainting features that fill missing parts of images without retraining the model on specific masks.
Authors
Jungsoo Lee, Alexi Gladstone, Yilun Du
Abstract
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.