MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training
Abstract Overview
This paper studies representation alignment for diffusion transformers from a token-level perspective and argues that aligning all diffusion tokens to clean-image encoder features creates a mismatch because diffusion inputs are noisy and informative content varies by timestep. The authors show that, under full-token alignment, alignment gradients are not uniformly distributed across spatial tokens and instead exhibit stable spatial preferences, suggesting reliance on specific token patterns. To address this, they propose MaskAlign, which performs representation alignment on randomly sampled token subsets during training while always keeping the class token. The method also adds a lightweight pre-mask token mixing block to reduce information loss before masking. Experiments on ImageNet 256x256 with SiT backbones evaluate convergence speed, alignment stability, generation quality, and per-step training cost.
Novelty
The distinctive contribution is reframing representation alignment in diffusion training as a token-subset problem rather than a full-token matching objective. MaskAlign combines random subset alignment with a lightweight pre-mask token mixer, motivated by a token-level analysis showing concentrated alignment gradients and instability under token-subset perturbations.
Results
On ImageNet 256x256, MaskAlign improves convergence and final sample quality relative to alignment baselines such as REG and also strengthens REPA when added to it. For SiT-XL/2, it reaches the 8.3 FID level about 77x faster than vanilla SiT-XL/2 and the 5.9 FID level about 30x faster than SiT-XL/2 + REPA in terms of training iterations, while improving FID from 3.4 to 2.8 at 400K iterations and from 2.7 to 2.4 at 1M iterations versus REG. It also reduces per-step training time by 11.6% relative to REG and shows a much smaller alignment-loss gap under masking, indicating more stable alignment behavior.
Key Points
- Token-level analysis shows that full-token representation alignment produces concentrated, spatially stable high-gradient tokens rather than uniform alignment across patch tokens.
- MaskAlign aligns randomly sampled token subsets during training and uses pre-mask token mixing to preserve information before dropping patch tokens.
- Experiments on ImageNet 256x256 report better FID than REG across multiple training budgets, faster convergence to matched FID levels, lower per-step cost, and improved stability under token-subset perturbations.