# Just Align bmx: Aligning Predictions, Not Representations

arXiv: https://arxiv.org/abs/2610.00600 · HTML: https://mlsanity.com/paper/2610.00600/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 32 of its 35 references that are in the map.
- First posted on arXiv 30 Sep 2026.

## Builds on (most cited first, 30 of 32)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers](https://mlsanity.com/paper/2401.08740/) (2024)
- [Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think](https://mlsanity.com/paper/2410.06940/) (2024)
- [Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models](https://mlsanity.com/paper/2501.01423/) (2025)
- [Diffusion Transformers with Representation Autoencoders](https://mlsanity.com/paper/2510.11690/) (2025)
- [Applying Guidance in a Limited Interval Improves Sample and Distribution Quality in Diffusion Models](https://mlsanity.com/paper/2404.07724/) (2024)
- [REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion Transformers](https://mlsanity.com/paper/2504.10483/) (2025)
- [PixelFlow: Pixel-Space Generative Models with Flow](https://mlsanity.com/paper/2504.07963/) (2025)
- [Masked Autoencoders Are Effective Tokenizers for Diffusion Models](https://mlsanity.com/paper/2502.03444/) (2025)
- [What matters for Representation Alignment: Global Information or Spatial Structure?](https://mlsanity.com/paper/2512.10794/) (2025)
- [PixelDiT: Pixel Diffusion Transformers for Image Generation](https://mlsanity.com/paper/2511.20645/) (2025)
- [Diffuse and Disperse: Image Generation with Representation Regularization](https://mlsanity.com/paper/2506.09027/) (2025)
- [DeCo: Frequency-Decoupled Pixel Diffusion for End-to-End Image Generation](https://mlsanity.com/paper/2511.19365/) (2025)
- [Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be Consistent](https://mlsanity.com/paper/2302.09057/) (2023)
- [No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves](https://mlsanity.com/paper/2505.02831/) (2025)
- [Latent Forcing: Reordering the Diffusion Trajectory for Pixel-Space Image Generation](https://mlsanity.com/paper/2602.11401/) (2026)
- [Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis](https://mlsanity.com/paper/2603.06507/) (2026)
- [Improved Baselines with Representation Autoencoders](https://mlsanity.com/paper/2605.18324/) (2026)
- [Representation Alignment for Just Image Transformers is not Easier than You Think](https://mlsanity.com/paper/2603.14366/) (2026)
- [Stable Velocity: A Variance Perspective on Flow Matching](https://mlsanity.com/paper/2602.05435/) (2026)
- [SRA 2: Variational Autoencoder Self-Representation Alignment for Efficient Diffusion Training](https://mlsanity.com/paper/2601.17830/) (2026)

## Abstract (arXiv)

Representation alignment has become an effective way to accelerate diffusion training, but its benefits do not transfer reliably to pixel-space clean-image prediction. In JiT, we find that auxiliary feature alignment can improve access to semantic features while reducing access to image variation needed for clean-image prediction, creating a mismatch between the auxiliary objective and the denoising task. This suggests a different principle: auxiliary supervision should improve the prediction target itself rather than impose a separate representation target. We introduce JAx (Just Align x), a prediction-supervision method that aligns clean-image predictions across noise levels. JAx couples a noisier student observation with a cleaner observation through a Markov degradation that preserves the original JiT input distribution. Under this coupling, the oracle prediction from the cleaner state has the same conditional mean as the optimal JiT target, while its conditional target covariance is no greater. Thus, oracle prediction alignment preserves the population JiT objective up to a constant while providing a lower-variance training target. To make this construction practical with an imperfect EMA teacher, JAx combines ground-truth supervision with a reliability-gated coupling band that selects nearby teacher states based on prediction risk. On ImageNet 256x256, JAx consistently improves FID and accelerates convergence across JiT-B/16, L/16, and H/16, without an external encoder or changes to the architecture or sampling procedure. Gradient diagnostics further show reduced minibatch gradient variance, while ablations demonstrate that the gains cannot be explained by time reweighting alone. These results show that prediction-space supervision provides a simple and principled alternative to representation alignment for pixel-space generative models.