MaskAlign: Alineación de Representaciones de Subconjuntos de Tokens para el Entrenamiento Eficiente de Difusión

Resumen

La alineación de representaciones con modelos visuales preentrenados ha demostrado recientemente un gran potencial para acelerar el entrenamiento de transformers de difusión. Al alinear las características intermedias de difusión con representaciones de imágenes limpias obtenidas de codificadores visuales autosupervisados, los métodos existentes mejoran la convergencia y la calidad de generación. Sin embargo, dicha alineación también introduce una restricción no trivial: los modelos de difusión operan sobre entradas ruidosas cuya información utilizable varía a lo largo de los pasos temporales, mientras que las características de referencia se extraen de imágenes limpias. En este artículo, revisitamos este desajuste desde una perspectiva a nivel de tokens. Descubrimos que, bajo una alineación de representaciones con todos los tokens, aquellos con grandes normas de gradiente de alineación muestran una preferencia espacial estable, lo que sugiere que el objetivo de alineación no afecta a todos los tokens de manera uniforme y puede incentivar al modelo a depender del conjunto completo de tokens de imágenes limpias. Para abordar este problema, proponemos MaskAlign, un método de alineación de representaciones por subconjuntos de tokens que aplica alineación a subconjuntos de tokens muestreados aleatoriamente durante el entrenamiento. Al exponer el modelo a diferentes subconjuntos de tokens a lo largo de las iteraciones, MaskAlign reduce la dependencia de la alineación de representaciones del conjunto completo de tokens y fomenta un comportamiento de alineación más estable bajo perturbaciones de subconjuntos de tokens. Para mitigar la pérdida de información causada por eliminar tokens directamente, introducimos además un bloque ligero de mezcla de tokens previa a la máscara que comparte información entre los tokens antes del enmascaramiento.

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.