Source-linked AI summary
MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training
Lianyu Pang, Tianlin Pan, Cheng Da, Changqian Yu, Huan Yang, Kun Gai, Song Guo, Wenhan Luo
TL;DR
Representation alignment must reconcile clean-image reference features with noisy diffusion features whose usable information varies across timesteps. MaskAlign applies alignment to random token subsets with pre-mask token mixing, achieving faster convergence and more stable alignment under token-subset perturbations. Its evaluation is concentrated on ImageNet 256 × 256 with SiT backbones and DINOv2 features.
Problem
Clean-image reference features can mismatch noisy diffusion features, while full-token alignment produces stable spatial concentration among high-gradient tokens.
Method
MaskAlign aligns randomly sampled token subsets and uses lightweight pre-mask token mixing to reduce complete-token dependence and information loss.
Results
MaskAlign reaches 8.3 FID about 77× faster than vanilla SiT-XL/2 and 5.9 FID about 30× faster than SiT-XL/2 + REPA, while improving alignment stability.
Takeaways & Limitations
Random token-subset alignment encourages behavior that is less sensitive to token-subset perturbations while improving convergence and training efficiency.
Takeaways & Limitations
Evaluation is mainly limited to ImageNet 256 × 256, SiT-based backbones, and pretrained DINOv2 features; broader generality remains unexplored.
Abstract
from arXiv · showhide
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.
1 Introduction
Diffusion representation alignment can mismatch noisy diffusion features with clean-image references and concentrate gradients on stable spatial token positions. MaskAlign addresses this with random token-subset alignment and pre-mask mixing, improving alignment stability and convergence efficiency.
- Training diffusion models at scale remains computationally expensive despite advances such as latent diffusion and transformer architectures.
- Clean-image reference features may mismatch noisy diffusion features because usable information changes across denoising timesteps.
- Full-token alignment produces non-uniform gradients, with high-gradient tokens exhibiting a stable spatial preference.
- At 8.3 FID, MaskAlign converges about 77× faster than vanilla SiT-XL/2; at 5.9 FID, it is about 30× faster than SiT-XL/2 + REPA.
- MaskAlign reduces concentrated token-gradient patterns and improves alignment stability under token-subset perturbations.
- MaskAlign applies alignment to randomly sampled token subsets and adds lightweight pre-mask token mixing to reduce shortcuts and information loss.
2 Related Work
Diffusion training has progressed from latent-space and transformer architectures to token masking and external representation alignment. MaskAlign differs by studying alignment at the token level and using random token subsets to improve stability.
- Latent diffusion uses VAEs to move generation into latent space, while DiT and SiT improve scalability through transformer and continuous-time formulations.
- Token-masking methods reduce training tokens, and lightweight mixers can aggregate token information before masking to mitigate information loss.
- Representation-alignment methods match intermediate diffusion features with pretrained vision-model features to improve training efficiency and generation performance.
- MaskAlign studies representation alignment from a token-level perspective and uses random token subsets to improve alignment stability during training.
3 Preliminaries
Diffusion models train by predicting injected noise or matching velocities along a continuous-time latent interpolation. The SiT framework uses VAE latents and a transformer to learn this velocity function.
- Denoising Diffusion Probabilistic Models: Under noise-prediction training, the network minimizes the distance between injected noise and its prediction from a corrupted latent input.
- Scalable Interpolant Transformers: SiT encodes clean images into latent representations with a pretrained VAE encoder before constructing a continuous-time interpolation process.
- Scalable Interpolant Transformers: As time increases, the interpolation coefficient for clean content decreases while the noise coefficient increases.
- Scalable Interpolant Transformers: SiT uses stacked Transformer blocks to learn a velocity function and trains it with a velocity-matching objective.
- Scalable Interpolant Transformers: The implementation uses α_t = 1 − t and σ_t = t, yielding constant derivatives α̇_t = −1 and σ̇_t = 1.
4 Token-level Analysis
Representation alignment bridges diffusion features and clean-image references, but noisy inputs and timestep-dependent information create a mismatch. Token-level analysis finds stable spatial concentration in alignment gradients, motivating random token-subset alignment to reduce dependence on complete token sets.
- Representation alignment matches intermediate diffusion features with clean-image representations, despite diffusion inputs varying in usable information across noise levels.The alignment projector injects supervision at the alignment layer, where hidden states are compared with reference features.
- Full-token analysis aligns every patch token with its corresponding clean-image reference feature while excluding the non-spatial class token from spatial heatmaps.The heatmap analysis therefore focuses on patch-token gradient behavior.
- For each image, selecting top-k patch tokens by alignment-gradient norm and aggregating their spatial frequencies reveals non-uniform token influence.The analysis computes the probability that each spatial position appears in the top-k set across multiple images.
- The largest spatial top-k probability is about 21× the smallest, indicating a stable preference rather than minor random fluctuation.Certain positions remain more likely to produce high-gradient tokens after averaging over many images.
- This concentration may reinforce feature-fitting shortcuts that reduce loss for complete token sets but are not consistently useful under noisy denoising conditions.The concern follows from repeatedly reinforcing high-gradient positions while reference features remain clean-image representations.
- MaskAlign randomly samples patch-token subsets during training so alignment signals remain stable across changing visible subsets.Varying subsets makes shortcuts dependent on the complete token set less consistently reinforced.
5 MaskAlign
MaskAlign preserves the class token, randomly masks patch tokens during training, and mixes token information before masking. It trains prediction and representation-alignment losses on the visible sequence while measuring sensitivity to full-token versus masked inputs.
- Framework: A lightweight pre-mask token mixer shares information across tokens before patch masking to mitigate disruption from dropping tokens.The mixed class and visible patch tokens are then fed into the diffusion transformer.
- Framework: MaskAlign always preserves the class token while randomly masking only patch tokens during training.The visible sequence contains the class token and a sampled subset of patch tokens.
- Training losses: The final objective combines prediction loss with representation-alignment loss, whose strength is controlled by λ.Alignment is computed over the class token and visible patch tokens at the alignment layer.
- Alignment stability: MaskAlign defines the alignment-loss gap as the masked-input alignment loss minus the full-token alignment loss.The full-token loss uses all patch tokens, whereas the masked loss uses a randomly sampled patch-token subset.
- Alignment stability: A smaller alignment-loss gap indicates alignment is less sensitive to token-subset perturbations and more stable across random subsets.At 200K steps with a 25% mask ratio, MaskAlign’s gap was 13.8% of REG’s, indicating substantially lower sensitivity.
6 Experiments
Experiments on ImageNet 256 × 256 show that MaskAlign accelerates convergence, improves generation quality, and reduces per-step computational cost. Ablations indicate that random masking, pre-mask token mixing, and moderate masking are important to its performance.
- MaskAlign consistently achieves the best FID among evaluated methods at matched training iterations across backbones and budgets.
- 77× faster convergence than vanilla SiT-XL/2 at 8.3 FID and 30× faster than SiT-XL/2 + REPA at 5.9 FID.Speedup is measured by the number of iterations required to reach the same FID level.
- MaskAlign improves REG from 15.2 to 14.8 FID on SiT-B/2 at 400K iterations and from 5.0 to 4.0 on SiT-XL/2 at 200K iterations.On SiT-XL/2, gains continue through longer training budgets, including a reduction from 2.2 to 2.1 FID at 2.4M iterations.
- At 80 epochs, MaskAlign improves REG from 1.86 to 1.82 FID and precision from 0.76 to 0.81 while maintaining recall.The model already achieves lower FID than vanilla SiT-XL/2 trained for 1,400 epochs.
- MaskAlign reduces input tokens by 24.9% and training time per step by 11.6% despite introducing about 8% more parameters.The comparison uses the same SiT-XL/2 backbone and GPU hardware at 400K iterations.
- Removing either token mixing or random masking degrades performance, showing that the two components are complementary.Token mixing reduces information loss from dropping tokens, while random masking supplies the token-subset training signal.
- A mask ratio of 0.25 reduces FID from 3.52 without masking to 2.84, while ratios of 0.5 and 0.75 weaken or reverse the improvement.The results support preserving enough input information while applying token-subset perturbations.
- Two pre-mask mixing layers achieve 2.84 FID, compared with 3.23 for one layer; three layers also degrade performance.The results indicate that a lightweight mixing block is sufficient.
7 Conclusion
MaskAlign addresses full-token alignment by training on randomly sampled token subsets and using pre-mask token mixing. ImageNet experiments show improved stability, faster convergence, better generation quality, and lower per-step computational cost, while broader generality remains open.
- MaskAlign applies representation alignment to random token subsets and adds lightweight pre-mask token mixing to reduce information loss.
- Experiments on ImageNet 256 × 256 show improved alignment stability, faster convergence, better generation quality, and lower per-step computational cost.
- The study is mainly evaluated on ImageNet 256 × 256 with SiT-based backbones and pretrained DINOv2 features.Generality to higher-resolution generation, text-to-image generation, and other teacher representations remains to be explored.
- Overly aggressive masking or excessive token mixing can degrade performance, making mask ratio and mixing depth important design choices.
A Experimental Setup
The experiments use DINOv2-B-based alignment with fixed MaskAlign hyperparameters across SiT-B/2 and SiT-XL/2 model scales. The setup uses two mixing layers and a 25% mask ratio.
- MaskAlign uses DINOv2-B, cosine similarity, two pre-mask token mixing layers, and a 25% mask ratio across both model scales.The alignment weight is λ = 0.5.
- The reported hyperparameter settings cover both SiT-B/2 and SiT-XL/2 model scales.
B Additional Token-Level Alignment Heatmaps
Additional heatmaps show that top-gradient alignment tokens retain non-uniform spatial patterns across timesteps and training checkpoints. This further supports the claim that full-token alignment affects patch tokens unevenly.
- Across timesteps and checkpoints, high-gradient tokens exhibit non-uniform spatial patterns.Each heatmap reports the probability that a spatial position appears among the top-10% tokens ranked by alignment-gradient norm.
- The heatmaps further support that full-token representation alignment does not affect all patch tokens uniformly.
C Additional Results on ImageNet
Additional ImageNet results show that MaskAlign improves generation quality with training and outperforms REG at the same 1M-iteration budget. The accompanying heatmaps and metrics further characterize convergence and alignment behavior.
- Quantitative Results: FID decreases from 22.36 at 50K iterations to 2.38 at 1M iterations.These results are reported without classifier-free guidance on ImageNet 256 × 256.
- Quantitative Results: At 1M iterations, MaskAlign achieves better FID, sFID, and IS than REG while maintaining comparable precision and recall.
- Alignment Analysis: Table 8 visualizes the spatial probability of top-10% alignment-gradient tokens across training iterations and timesteps.All heatmaps use the same color range, [0, 0.8].
- Quantitative Results: Table 9 reports FID, sFID, IS, precision, and recall across different training iterations.
F More Visualization Results
Figures 4–8 provide additional MaskAlign samples from SiT-XL/2 across animal and architectural ImageNet classes.
- Visualization Results: Figure 4 shows generated samples for the “great grey owl” class.
- Visualization Results: Figure 5 shows generated samples for the “golden retriever” class.
- Visualization Results: Figure 6 shows generated samples for the “arctic wolf” class.
- Visualization Results: Figure 7 shows generated samples for the “polecat” class.
- Visualization Results: Figure 8 shows generated samples for the “castle” class.