Variable-Length Tokenization via Learnable Global Merging for Diffusion Transformers

Resumen Técnico: Tokenización de Longitud Variable mediante Fusión Global Aprendible para Transformers de Difusión

1. Planteamiento del Problema

Los Modelos de Difusión Latente (LDM) dependen de tokenizadores para comprimir imágenes en espacios latentes, lo que crea un compromiso entre la calidad de generación y el coste computacional. Los tokenizadores convencionales utilizan una relación de compresión fija, lo que requiere múltiples variantes del modelo para adaptarse a diferentes presupuestos computacionales o requisitos de calidad.

Los Tokenizadores de Longitud Variable (VLT) se han propuesto para abordar esto ajustando adaptativamente el número de tokens basándose en condiciones específicas (por ejemplo, el presupuesto computacional). El método predominante, el dropout anidado (nested dropout), entrena modelos para descartar los tokens de la cola de una secuencia ordenada, situando la semántica de alto nivel en los primeros tokens y los detalles de bajo nivel en los últimos.

Sin embargo, los autores identifican un fallo crítico en el dropout anidado cuando se aplica a los Transformers de Difusión (DiT):

• Desplazamiento de Longitud Cruzada (Cross-Length Shift): Truncar una secuencia ordenada altera la estructura de similitud por pares (alineación representacional) del espacio latente. Las secuencias más cortas se centran en la semántica central, mientras que las más largas incluyen detalles residuales.
• Desajuste Distribucional: Este desplazamiento provoca que la distribución latente y la función de puntuación (score function) correspondiente varíen significativamente entre longitudes de token. En consecuencia, un único modelo de difusión entrenado conjuntamente a través de diversas longitudes no generaliza eficazmente, rindiendo a menudo peor que los modelos entrenados para longitudes específicas.
• Limitación de las Soluciones Existentes: Las pérdidas de alineación de representación estándar requieren que las dimensiones de los tokens coincidan (N = M), lo cual es imposible en entornos de longitud variable donde la cardinalidad difiere.

2. Metodología

El artículo propone un nuevo tokenizador de longitud variable basado en la fusión (merging-based) que modula la longitud combinando tokens en lugar de truncarlos. La innovación central es la Fusión Global Aprendible (LGM), diseñada para asegurar la alineación representacional a través de diferentes conteos de tokens y, al mismo tiempo, ser compatible con modelos generativos.

A. Modulación de Longitud Basada en la Fusión
En lugar de descartar tokens, el método fusiona tokens basándose en una matriz de asignación $\Gamma$.

• Atención Proporcional: Para gestionar los tokens fusionados, el decodificador y el transformer de difusión utilizan atención proporcional, donde la puntuación de atención se ajusta mediante el tamaño efectivo del clúster ($m_j$) de cada token fusionado.
• Representación Equivalente de Longitud Completa: Los autores demuestran que procesar latentes modulados con atención proporcional es matemáticamente equivalente a procesar latentes de longitud completa proyectados ($Wz$) con atención estándar. Esta equivalencia permite medir y minimizar directamente la diferencia entre los latentes originales ($z$) y los latentes fusionados ($Wz$), minimizando eficazmente el desplazamiento de representación al fomentar la fusión de tokens similares.

B. Fusión Global Aprendible (LGM)
Surge un desafío porque la fusión estándar (por ejemplo, el agrupamiento aglomerativo) depende de los datos, derivando los patrones de fusión a partir de la imagen de entrada. Esto es incompatible con la generación de difusión, donde la imagen objetivo no está disponible.

• Independencia de los Datos: Los autores introducen Embeddings Globales Aprendibles ($e$) que son independientes de los datos de entrada. Estos embeddings determinan el patrón de fusión para un determinado objetivo de reducción.
• Optimización: El patrón de fusión se deriva mediante agrupamiento aglomerativo sobre los embeddings globales aprendibles. Dado que el agrupamiento implica operaciones discretas, se aplica un truco de paso directo (straight-through trick) para permitir la optimización basada en gradientes de los embeddings.
• Pérdida de Alineación: Se introduce una pérdida de alineación ($L_{align}$) para asegurar que la estructura de similitud de los embeddings globales aprendibles ($e$) coincida con la estructura de similitud de los tokens latentes reales ($z$). Esto guía el patrón de fusión global para agrupar tokens similares, minimizando así el desplazamiento de representación definido en el objetivo de fusión.

C. Integración con el Transformer de Difusión

• Embeddings Posicionales: El método emplea embeddings posicionales fusionados, computados de la misma manera que los latentes, para gestionar la información posicional de los tokens agregados.
• Post-entrenamiento: Se introduce un paso de post-entrenamiento ligero mediante LoRA (Low-Rank Adaptation). Tras el entrenamiento conjunto en diversas longitudes, se entrenan parámetros LoRA específicos para cada longitud durante unos pocos ciclos (epochs) para adaptar aún más el modelo a conteos de tokens específicos con un coste adicional insignificante.

3. Contribuciones Clave

1. Tokenización Basada en la Fusión: Una nueva arquitectura VLT que modula la longitud mediante la fusión de tokens, permitiendo la alineación directa de la representación a través de diferentes conteos de tokens al minimizar el desplazamiento entre los latentes originales y los fusionados.
2. Fusión Global Aprendible: Un mecanismo de fusión independiente de los datos que resuelve la incompatibilidad entre la fusión dependiente de los datos y la generación de difusión, asegurando que el patrón de fusión sea accesible durante la inferencia.
3. Mejora en el Compromiso Calidad-Computación: La integración de este tokenizador con un Transformer de Difusión logra un rendimiento superior al equilibrar la calidad de generación (gFID) y el coste computacional en comparación con los métodos VLT existentes.

4. Resultados Experimentales

Los experimentos se realizaron en la generación condicional de clase ImageNet 256×256.

• Alineación de la Representación: El método propuesto logra puntuaciones de Alineación de Vecino Cercano de Kernel Centrado (CKNNA) significativamente más altas a través de diferentes longitudes de token en comparación con el dropout anidado. Notablemente, la alineación es a menudo superior a la de los tokenizadores de longitud fija entrenados de forma independiente.
• Rendimiento del Entrenamiento Conjunto:
  • Dropout Anidado: Muestra una gran brecha (0.8–1.7 gFID) entre los modelos entrenados conjuntamente en diversas longitudes y aquellos entrenados por separado.
  • Método Propuesto: Reduce esta brecha a menos de 0.3 gFID. Con el post-entrenamiento LoRA, el modelo conjunto alcanza un rendimiento dentro de 0.05 gFID de los modelos entrenados por separado.
• Comparación a Nivel de Sistema:
  • El método supera a los VLT existentes (Semanticist, FlexTok) y a las líneas base de longitud fija (SoftVQ, MAETok) en términos del compromiso gFID–computación.
  • Logra una calidad de generación competitiva en altas tasas de compresión (por ejemplo, 32 tokens) sin requerir un decodificador de tokenizador basado en difusión, a diferencia de algunas líneas base.
• Calidad de Reconstrucción: El tokenizador mantiene métricas de reconstrucción sólidas (rFID, PSNR, SSIM) a través de diversas longitudes de token, particularmente en altas tasas de compresión.

5. Significado y Reivindicaciones

El artículo sostiene que la barrera principal para un modelo de difusión de longitud variable eficaz es el desplazamiento de la distribución latente por longitud cruzada causado por los métodos de truncamiento convencionales. Al cambiar hacia un enfoque basado en la fusión con patrones aprendibles e independientes de los datos, los autores logran mantener la alineación representacional con éxito.

La importancia radica en permitir que un único modelo de difusión opere eficazmente a través de diversos conteos de tokens, lo que permite un control flexible de la calidad y la computación sin necesidad de múltiples variantes del modelo. Los autores posicionan esto como un paso hacia el despliegue más eficiente de modelos generativos en escenarios con diversas restricciones computacionales, señalando que, aunque persiste una brecha entre los tokenizadores 1D y 2D, su enfoque ofrece una dirección prometedora para arquitecturas híbridas. El trabajo se enmarca como una mejora de eficiencia en lugar de una nueva capacidad generativa, sin reivindicar preocupaciones éticas.