MaskAlign: Token-subset representatie-afstemming voor efficiënte diffusietraining

Samenvatting

Het afstemmen van representaties op voorgetrainde visiemodellen heeft recentelijk veel potentieel getoond voor het versnellen van de training van diffusietransformers. Door tussenliggende diffusiekenmerken af te stemmen op representaties van schone beelden uit zelfgecontroleerde visie-encoders, verbeteren bestaande methoden de convergentie en de generatiekwaliteit. Echter, een dergelijke afstemming introduceert ook een niet-triviale beperking: diffusiemodellen werken op ruizige invoer waarvan de bruikbare informatie varieert over tijdsstappen, terwijl de referentiekenmerken worden geëxtraheerd uit schone beelden. In dit artikel bekijken we deze mismatch opnieuw vanuit een token-niveau perspectief. We vinden dat, onder volledige token-representatie-afstemming, tokens met grote afstemmingsgradiëntnormen een stabiele ruimtelijke voorkeur vertonen, wat suggereert dat het afstemmingsdoel niet alle tokens gelijkmatig beïnvloedt en het model kan aanmoedigen om te vertrouwen op de volledige set van schone beeldtokens. Om dit probleem aan te pakken, stellen we MaskAlign voor, een token-subset representatie-afstemmingsmethode die afstemming toepast op willekeurig geselecteerde tokensubsets tijdens de training. Door het model bloot te stellen aan verschillende tokensubsets over iteraties heen, vermindert MaskAlign de afhankelijkheid van representatie-afstemming van de volledige tokenset en stimuleert het afstemmingsgedrag dat stabieler is onder tokensubset-verstoringen. Om het informatieverlies door het direct weglaten van tokens te beperken, introduceren we verder een lichtgewicht pre-mask token-mixblok dat informatie deelt over tokens heen voordat masking wordt toegepast.

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.