Scaling Transformer Energy-based Models For High Quality Image Synthesis

Pixel-space generative models bypass lossy latent compression, yet necessitate joint learning of global structure and fine-grained details in a high-dimensional space. Standard flow matching interpolates noise toward a fixed clean-image endpoint, leaving the spectral evolution to be learned implicitly. In this paper, we introduce Energy-Guided Flow Matching(EG-FM) that explicitly models a coarse-to-fine generative trajectory by moving endpoint. Specifically, EG-FM replaces the fixed endpoint with a heat-kernel-filtered endpoint that evolves smoothly from low-frequency image to clean image. The fraction of high-frequency signal in moving endpoint is released by an image-specific energy-guided scheduling, leading to the re-targeting of velocity in flow matching. Our framework requires no adaptation of the backbone and training data, bringing negligible cost on the training and inference stages. In our experiment, EG-FM consistently achieves lower FID on the ImageNet class-conditional image generation task at $256 \times 256$ with fewer epochs, reaching an FID of 1.55 at 200 epochs and 1.45 at 600 epochs. We continue training the generation task on the setting of $512 \times 512$ resolution, yielding a FID of 1.58 after only 40 high-resolution adaptation epochs. Furthermore, we transfer EG-FM on text-to-image generation and achieve 0.85 on GenEval score and 83.9 on DPG-Bench. Code is available at https://github.com/ysng123/EG-FM.

Energy-Based Models (EBMs) provide a flexible framework for generative modeling by learning an energy landscape that assigns low energy values to realistic samples and higher energies to unlikely observations. Despite their theoretical appeal, training EBMs remains challenging due to the computational cost of Langevin sampling and the difficulty of efficiently exploring the learned data manifold. In this work, we propose a cooperative Energy-Based Model and Autoencoder (EBM-AE) framework that combines energy-based refinement with manifold projection. The proposed approach jointly trains an EBM with a denoising autoencoder and introduces an iterative EBM$\rightarrow$AE$\rightarrow$EBM sampling procedure in which Langevin dynamics and autoencoder projection alternately refine generated samples. Within this framework, the autoencoder acts as a manifold projection operator that regularizes sampling trajectories, while the EBM performs energy-based refinement toward low-energy regions of the learned distribution. Extensive experiments conducted on the MNIST dataset demonstrate that joint EBM-AE training substantially improves generation quality compared with a conventional autoencoder. Beyond unconditional generation, we evaluate the proposed framework on image inpainting tasks involving structured and random masks. The results show that manifold projection provides the majority of the reconstruction capability, whereas the final energy-based refinement becomes increasingly beneficial as the reconstruction problem becomes more challenging. Taken together, the results indicate that combining manifold projection and energy minimization provides an effective and interpretable framework for generation, reconstruction, and out-of-distribution detection, while offering new insights into the complementary roles of energy-based modeling and representation learning.

In diffusion transformers, a class label or a text prompt is embedded once, and the same condition is reused at every denoising step. We ask whether predicted embeddings can serve as this condition instead. Next-Embedding Predictive Autoregression (NEPA) trains a Transformer to predict the next continuous embedding in a sequence. In generation, the clean image follows the noisy image, so its embeddings are the next embeddings after the condition and the noisy image. We train a NEPA model to predict them all at once with Multi-Embedding Prediction, and in Embedding Conditioned Generation, a DiT generator is conditioned on these predictions, recomputed at every denoising step, so the conditioning signal adapts to the current noisy state. Experiments on class-conditional ImageNet $256\times256$ study the condition of the generator, the design of Multi-Embedding Prediction, and the scaling of both models. The NEPA model adds a second network to every sampling step; with it, and combined with REPA, our final model, NEPA-DiT-XL, reaches an FID of 1.32 using about a third of the training compute of REPA.