MaskAlign: Выравнивание представлений подмножества токенов для эффективного обучения диффузионных моделей

Аннотация

Выравнивание представлений с предобученными моделями зрения недавно продемонстрировало высокий потенциал для ускорения обучения диффузионных трансформеров. Выравнивая промежуточные признаки диффузии с представлениями чистых изображений от самоконтролируемых кодировщиков зрения, существующие методы улучшают сходимость и качество генерации. Однако такое выравнивание также вводит нетривиальное ограничение: диффузионные модели работают с зашумленными входными данными, полезная информация в которых варьируется по временным шагам, в то время как эталонные признаки извлекаются из чистых изображений. В этой статье мы пересматриваем это несоответствие с точки зрения токенов. Мы обнаружили, что при полном выравнивании представлений токенов токены с большими нормами градиентов выравнивания демонстрируют стабильное пространственное предпочтение, что предполагает, что цель выравнивания не влияет на все токены равномерно и может побуждать модель полагаться на полный набор токенов чистого изображения. Чтобы решить эту проблему, мы предлагаем MaskAlign — метод выравнивания представлений на подмножестве токенов, который применяет выравнивание к случайно выбранным подмножествам токенов во время обучения. Подвергая модель различным подмножествам токенов на разных итерациях, MaskAlign уменьшает зависимость выравнивания представлений от полного набора токенов и способствует поведению выравнивания, которое является более устойчивым к возмущениям подмножества токенов. Чтобы смягчить потерю информации, вызванную прямым отбрасыванием токенов, мы дополнительно вводим легковесный блок смешивания токенов перед маскированием, который обменивается информацией между токенами до маскирования.

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.