MaskAlign: Ausrichtung von Token-Teilmenge-Repräsentationen für effizientes Diffusionstraining

Zusammenfassung

Die Repräsentationsausrichtung mit vortrainierten visuellen Modellen hat sich in letzter Zeit als vielversprechend für die Beschleunigung des Trainings von Diffusionstransformatoren erwiesen. Durch die Angleichung intermediärer Diffusionsmerkmale an merkmalsfreie Repräsentationen aus selbstüberwachten visuellen Encodern verbessern bestehende Methoden die Konvergenz und die Generierungsqualität. Allerdings führt eine solche Ausrichtung auch eine nichttriviale Einschränkung ein: Diffusionsmodelle arbeiten auf verrauschten Eingaben, deren nutzbare Information über Zeitschritte variiert, während die Referenzmerkmale aus merkmalsfreien Bildern extrahiert werden. In diesem Beitrag beleuchten wir diese Diskrepanz aus der Perspektive der Token-Ebene. Wir stellen fest, dass bei einer vollständigen Token-Repräsentationsausrichtung Tokens mit großen Ausrichtungsgradientennormen eine stabile räumliche Präferenz aufweisen. Dies deutet darauf hin, dass das Ausrichtungsziel nicht alle Tokens gleichmäßig beeinflusst und das Modell dazu verleiten kann, sich auf den vollständigen Satz merkmalsfreier Bild-Tokens zu stützen. Um dieses Problem zu adressieren, schlagen wir MaskAlign vor, eine Methode zur Token-Teilmenge-Repräsentationsausrichtung, die während des Trainings eine Ausrichtung auf zufällig abgetastete Token-Teilmengen anwendet. Indem das Modell über Iterationen hinweg verschiedenen Token-Teilmengen ausgesetzt wird, reduziert MaskAlign die Abhängigkeit der Repräsentationsausrichtung vom vollständigen Tokensatz und begünstigt ein Ausrichtungsverhalten, das unter Token-Teilmenge-Perturbationen stabiler ist. Um den Informationsverlust durch direktes Verwerfen von Tokens zu mildern, führen wir zusätzlich einen leichten Vor-Maskierungs-Token-Mischblock ein, der vor der Maskierung Informationen über Tokens hinweg teilt.

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.