# MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 24 of its 29 references that are in the map.
- First posted on arXiv 7 Jun 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 24 of 24)

- [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)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding](https://mlsanity.com/paper/2205.11487/) (2022)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models](https://mlsanity.com/paper/2112.10741/) (2021)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Improved Precision and Recall Metric for Assessing Generative Models](https://mlsanity.com/paper/1904.06991/) (2019)
- [eDiff-I: Text-to-Image Diffusion Models with an Ensemble of Expert Denoisers](https://mlsanity.com/paper/2211.01324/) (2022)
- [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)
- [Generating Images with Sparse Representations](https://mlsanity.com/paper/2103.03841/) (2021)
- [MDTv2: Masked Diffusion Transformer is a Strong Image Synthesizer](https://mlsanity.com/paper/2303.14389/) (2023)
- [REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion Transformers](https://mlsanity.com/paper/2504.10483/) (2025)
- [Fast Training of Diffusion Models with Masked Transformers](https://mlsanity.com/paper/2306.09305/) (2023)
- [Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think](https://mlsanity.com/paper/2507.01467/) (2025)
- [Dropout Training as Adaptive Regularization](https://mlsanity.com/paper/1307.1493/) (2013)
- [What matters for Representation Alignment: Global Information or Spatial Structure?](https://mlsanity.com/paper/2512.10794/) (2025)
- [Boosting Generative Image Modeling via Joint Image-Feature Synthesis](https://mlsanity.com/paper/2504.16064/) (2025)
- [REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training](https://mlsanity.com/paper/2505.16792/) (2025)
- [Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget](https://mlsanity.com/paper/2407.15811/) (2024)
- [TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training](https://mlsanity.com/paper/2501.04765/) (2025)
- [REGLUE Your Latents with Global and Local Semantics for Entangled Diffusion](https://mlsanity.com/paper/2512.16636/) (2025)

## Built on by (most cited first, 2 of 2)

- [Scaffold Then Internalize: Representation Injection for Diffusion Transformers](https://mlsanity.com/paper/2609.35292/) (2026)
- [GAD: Geometry Attention Distillation for Consistent Video Generation](https://mlsanity.com/paper/2610.18779/) (2026)

## Abstract (arXiv)

Representation alignment with pretrained vision models has recently shown strong potential for accelerating diffusion transformer training. By aligning intermediate diffusion features with clean-image representations from self-supervised vision encoders, existing methods improve convergence and generation quality. However, such alignment also introduces a non-trivial constraint: diffusion models operate on noisy inputs whose usable information varies across timesteps, while the reference features are extracted from clean images. In this paper, we revisit this mismatch from a token-level perspective. We find that, under full-token representation alignment, tokens with large alignment-gradient norms exhibit a stable spatial preference, suggesting that the alignment objective does not affect all tokens uniformly and may encourage the model to rely on the complete set of clean-image tokens. To address this issue, we propose MaskAlign, a token-subset representation alignment method that applies alignment to randomly sampled token subsets during training. By exposing the model to different token subsets across iterations, MaskAlign reduces the dependence of representation alignment on the complete token set and encourages alignment behavior that is more stable under token-subset perturbations. To mitigate the information loss caused by directly dropping tokens, we further introduce a lightweight pre-mask token mixing block that shares information across tokens before masking.