MaskAlign : Alignement de représentations de sous-ensembles de tokens pour un entraînement efficace de diffusion

Résumé

L'alignement des représentations avec des modèles de vision pré-entraînés a récemment montré un fort potentiel pour accélérer l'entraînement des transformers de diffusion. En alignant les caractéristiques intermédiaires de diffusion avec les représentations d'images propres issues d'encodeurs visuels auto-supervisés, les méthodes existantes améliorent la convergence et la qualité de génération. Cependant, un tel alignement introduit également une contrainte non triviale : les modèles de diffusion opèrent sur des entrées bruitées dont les informations utilisables varient selon les pas de temps, tandis que les caractéristiques de référence sont extraites d'images propres. Dans cet article, nous réexaminons ce décalage d'un point de vue au niveau des tokens. Nous constatons que, sous un alignement des représentations sur tous les tokens, les tokens présentant de grandes normes de gradient d'alignement affichent une préférence spatiale stable, ce qui suggère que l'objectif d'alignement n'affecte pas tous les tokens de manière uniforme et pourrait encourager le modèle à se fier à l'ensemble complet des tokens d'images propres. Pour résoudre ce problème, nous proposons MaskAlign, une méthode d'alignement des représentations sur un sous-ensemble de tokens, qui applique l'alignement à des sous-ensembles de tokens échantillonnés aléatoirement lors de l'entraînement. En exposant le modèle à différents sous-ensembles de tokens au fil des itérations, MaskAlign réduit la dépendance de l'alignement des représentations vis-à-vis de l'ensemble complet de tokens et encourage un comportement d'alignement plus stable face aux perturbations des sous-ensembles de tokens. Pour atténuer la perte d'information causée par la suppression directe de tokens, nous introduisons en outre un bloc léger de mélange de tokens avant masquage, qui partage l'information entre les tokens avant l'application du masque.

English

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.