Resumen Técnico: Los Registros Importan para los Transformadores de Difusión en el Espacio de Píxeles

Enunciado del Problema

Se sabe que los Transformadores de Visión (ViTs) exhiben valores atípicos de tokens de parche con normas altas que degradan la calidad de los mapas de características y crean artefactos en los mapas de atención. Este problema se mitiga eficazmente en los ViTs mediante la introducción de "tokens de registro", que actúan como sumideros para estos valores atípicos. A medida que los Transformadores de Difusión (DiTs) adoptan cada vez más arquitecturas de transformadores y se desplazan hacia el entrenamiento directamente en el espacio de píxeles (alejándose del espacio latente), se asemejan estructuralmente más a los ViTs. Esto plantea dos preguntas críticas:

1. ¿Heredan los DiTs los valores atípicos de tokens de parche con normas altas observados en los ViTs?
2. ¿Son beneficiosos los tokens de registro para los DiTs, particularmente en regímenes de entrenamiento en el espacio de píxeles?

Aunque los sumideros de atención y los comportamientos de tokens especiales han sido estudiados en Modelos de Lenguaje Grandes (LLMs) y modelos de difusión de lenguaje, su presencia y función en los transformadores de difusión de imágenes permanecen poco explorados.

Metodología

Los autores realizan un análisis sistemático de los tokens de registro en diferentes espacios de modelos de difusión (espacio de píxeles, espacio VAE y espacio RAE) y los comparan con ViTs estándar (específicamente DINOv2).

1. Análisis Empírico de Valores Atípicos: El estudio investiga las normas de características por token y los mapas de atención en DiTs (específicamente pDiTs entrenados con ajuste de flujo en ImageNet) frente a ViTs. Analizan si los tokens de parche exhiben valores atípicos de normas altas en ausencia de registros.
2. Integración de Tokens de Registro: Los tokens de registro se implementan como tokens adicionales aprendibles añadidos a la secuencia de tokens de parche, siguiendo el enfoque utilizado en ViTs. Estos tokens se excluyen del objetivo de entrenamiento (función de pérdida).
3. Análisis de Representaciones Internas: Los autores analizan los mapas de características intermedios utilizando la Variación Total (TV) para medir la suavidad espacial y el sondeo lineal para evaluar el contenido semántico de los tokens de registro. Examinan cómo los registros afectan las normas de características y la coherencia de las representaciones en diferentes pasos de difusión (niveles de ruido).
4. Exploración Arquitectónica: Motivados por los roles distintos de los registros y los tokens de parche, el artículo propone una arquitectura de doble flujo eficiente en parámetros. Este diseño procesa los tokens de registro y de parche a través de vías especializadas (específicamente desacoplando las capas RMSNorm, MLP y adaLN) mientras comparte los mecanismos de atención y la mayoría de los parámetros del transformador.
5. Estudios de Ablación: El estudio abla el número de registros, la profundidad a la que se introducen y los componentes específicos a duplicar para determinar configuraciones óptimas.

Hallazgos Clave y Contribuciones

1. Los DiTs Difieren Fundamentalmente de los ViTs en Relación con los Valores Atípicos

Contrario a los ViTs, los DiTs estándar (tanto en espacio latente como en espacio de píxeles) no exhiben valores atípicos de tokens de parche con normas altas. En los ViTs, los valores atípicos se localizan en tokens específicos (a menudo en regiones de baja información), mientras que en los DiTs, las normas de los tokens de parche permanecen relativamente uniformes y los mapas de atención carecen de los artefactos de baja información vistos en los ViTs.

2. Los Registros Mejoran los DiTs en Espacio de Píxeles a pesar de la Ausencia de Valores Atípicos Preexistentes

Sorprendentemente, añadir tokens de registro a los DiTs mejora significativamente la convergencia y la calidad de generación, particularmente en modelos en espacio de píxeles.

• Mecanismo: A diferencia de los ViTs, donde los registros absorben valores atípicos preexistentes, los registros en los DiTs crean tokens de normas altas dentro de sí mismos.
• Efecto: La presencia de registros reduce consistentemente las normas de características de todos los tokens de parche. Esto conduce a mapas de características intermedios más limpios y suaves, especialmente en niveles altos de ruido (pasos tempranos de difusión).
• Doble Función: Los tokens de registro se especializan en dos roles distintos:
  • Sumideros de Normas: Tokens con normas extremadamente altas que absorben magnitud de los tokens de parche.
  • Portadores de Información Global: Tokens con normas moderadas que codifican información semántica diversa sobre la imagen (por ejemplo, objetos en primer plano, contexto de fondo), verificados mediante sondeo lineal y visualización de mapas de atención.

3. Eficacia Específica del Espacio

Los beneficios de los tokens de registro dependen altamente del espacio de entrenamiento:

• Espacio de Píxeles: Los registros producen las mejoras más significativas en la calidad de generación (FID). Los modelos en espacio de píxeles producen inherentemente representaciones intermedias más ruidosas con normas de características más altas, lo que hace que el efecto de regularización de los registros sea crucial.
• Espacio Latente (VAE/RAE): Los registros proporcionan ganancias moderadas o, en el caso de modelos basados en RAE, degradan ligeramente el rendimiento. Los espacios latentes son más estructurados y de menor dimensión, reduciendo la necesidad de la fuerte regularización proporcionada por los registros.

4. Registros Implícitos en Arquitecturas Existentes

Los autores observan que las arquitecturas recientes de DiT en espacio de píxeles (por ejemplo, JiT) que utilizan condicionamiento de clase en contexto (añadiendo incrustaciones de clase a la secuencia de tokens) incorporan implícitamente mecanismos similares a los registros. Estos tokens en contexto exhiben un comportamiento similar a los registros explícitos: algunos actúan como sumideros de normas, mientras que otros codifican información semántica global. Esto sugiere que las ganancias de rendimiento en tales modelos pueden derivar en gran medida de este comportamiento similar a los registros en lugar de la información de clase en sí misma.

5. Arquitectura de Doble Flujo Eficiente en Parámetros

Para aprovechar estos conocimientos, los autores proponen una arquitectura compacta de doble flujo que especializa el procesamiento de los tokens de registro.

• Diseño: Los registros se introducen solo en capas más profundas (comenzando desde el bloque 4), donde se establece la estructura semántica. La arquitectura utiliza atención compartida pero desacopla RMSNorm, MLP y adaLN para registros y parches.
• Eficiencia: Este enfoque aumenta el tamaño del modelo solo en un ~14% (en comparación con ~77% para una duplicación completa ingenua) mientras mantiene costos de inferencia similares (GFLOPs).
• Rendimiento: Esta arquitectura mejora consistentemente la calidad de generación a través de escalas de modelos (B, L, H) y resoluciones (256x256, 512x512) en comparación con las líneas base de flujo único.

Resultados

• Calidad de Generación: En ImageNet 256x256, los pDiTs con registros superan consistentemente a aquellos sin ellos. Por ejemplo, un modelo pDiT-B/16 mejoró su FID de 7.39 a 5.30 a los 200 epochs.
• Suavidad de Características: Los modelos con registros muestran ratios de Variación Total (TV) más bajos en características intermedias en niveles altos de ruido, lo que indica representaciones más suaves y coherentes.
• Ablación: Los registros son más efectivos cuando se introducen después de las primeras 4 capas y cuando se utiliza un mayor número de tokens (por ejemplo, 32 registros para pDiT-B, en comparación con 4 para ViTs).

Significado y Afirmaciones

El artículo afirma que los tokens de registro son un componente crítico, aunque previamente pasado por alto, para la generación de difusión de alta calidad en espacio de píxeles. El significado radica en:

1. Redefinir el Rol de los Registros: Demostrar que los registros no son meramente una solución para valores atípicos específicos de ViT, sino que sirven una función de regularización más amplia en modelos de difusión al gestionar las normas de características y mejorar la coherencia espacial.
2. Explicar Éxitos Recientes: Proporcionar una justificación para el fuerte rendimiento empírico de los DiTs recientes en espacio de píxeles que utilizan condicionamiento en contexto, atribuyendo parte de su éxito al comportamiento implícito de los registros.
3. Diseño de Arquitectura Eficiente: Ofrecer un diseño de doble flujo eficiente en parámetros que capitaliza los roles distintos de los registros y los parches, logrando calidad de generación de vanguardia con una sobrecarga de tiempo de ejecución insignificante.

Los autores concluyen que, aunque sus diseños capturan la promesa de la especialización de tokens, el potencial completo de utilizar registros en transformadores de difusión en espacio de píxeles sigue siendo una pregunta abierta para futuras investigaciones.