MaskAlign: Alinhamento de Representação de Subconjunto de Tokens para Treinamento Eficiente de Difusão

Resumo

O alinhamento de representações com modelos de visão pré-treinados tem demonstrado recentemente um forte potencial para acelerar o treinamento de transformers de difusão. Ao alinhar características intermediárias de difusão com representações de imagens limpas provenientes de codificadores visuais auto-supervisionados, métodos existentes melhoram a convergência e a qualidade da geração. No entanto, tal alinhamento também introduz uma restrição significativa: modelos de difusão operam sobre entradas ruidosas, cuja informação utilizável varia ao longo dos passos de tempo, enquanto as características de referência são extraídas de imagens limpas. Neste artigo, revisitamos essa incompatibilidade sob uma perspectiva de nível de token. Descobrimos que, sob o alinhamento de representações de tokens completos, tokens com grandes normas de gradiente de alinhamento exibem uma preferência espacial estável, sugerindo que o objetivo de alinhamento não afeta todos os tokens de maneira uniforme e pode incentivar o modelo a depender do conjunto completo de tokens de imagens limpas. Para abordar esse problema, propomos o MaskAlign, um método de alinhamento de representações por subconjuntos de tokens que aplica o alinhamento a subconjuntos de tokens amostrados aleatoriamente durante o treinamento. Ao expor o modelo a diferentes subconjuntos de tokens ao longo das iterações, o MaskAlign reduz a dependência do alinhamento de representações em relação ao conjunto completo de tokens e incentiva um comportamento de alinhamento mais estável sob perturbações de subconjuntos de tokens. Para mitigar a perda de informação causada pelo descarte direto de tokens, introduzimos ainda um bloco leve de mistura de tokens pré-máscara que compartilha informações entre tokens antes da mascaramento.

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.