Resumen Técnico: Más allá de la Última Capa: Fusión de Representaciones Multi-Capa para la Tokenización Visual

1. Declaración del Problema

Los tokenizadores visuales más avanzados actuales, particularmente aquellos que aprovechan Modelos Fundacionales de Visión (VFMs) preentrenados y congelados como DINOv2 como backbones, sufren de un cuello de botella de información de capa única crítico. Los métodos existentes extraen exclusivamente características de la última capa del codificador para definir el espacio latente. Si bien estas características de la última capa están optimizadas para semánticas de alto nivel, descartan sistemáticamente los detalles visuales de bajo nivel (texturas, bordes, gradientes de color) a través de la abstracción semántica progresiva.

Un análisis reciente indica que la información de bajo nivel sobrevive en la última capa solo como residuos atenuados a través de conexiones residuales pasivas, volviéndose cada vez más pérdida con cada capa sucesiva. En consecuencia, los tokenizadores de capa única descartan la rica información jerárquica distribuida a través de las capas intermedias, limitando tanto la fidelidad de reconstrucción como la calidad de la generación aguas abajo. Además, sigue siendo una pregunta abierta si aumentar la "riqueza de representación" de los tokenizadores visuales sigue una ley de escalado predecible análoga al escalado del tamaño del vocabulario en los tokenizadores de PLN.

2. Metodología: DRoRAE

Los autores proponen DRoRAE (Autoencoder de Representación con Enrutamiento por Profundidad), un módulo de fusión ligero diseñado para agregar características de todas las capas del codificador manteniendo la compatibilidad con decodificadores congelados y modelos generativos aguas abajo.

Arquitectura Central

DRoRAE inserta un módulo de fusión entrenable entre el backbone del VFM congelado y el decodificador. El proceso involucra tres mecanismos clave:

1. Expertos por Capa: Cada una de las $L$ capas del codificador está asociada con una red experta dedicada (un MLP de dos capas). Estos expertos proyectan características heterogéneas de las capas sobre una escala común, normalizada por la normalización de capas del backbone para manejar las disparidades de varianza.
2. Enrutamiento con Restricción de Energía: A diferencia de la Mezcla Estándar de Expertos (MoE) que utiliza normalización softmax, DRoRAE emplea un mecanismo de enrutamiento con restricción de energía. Una red enrutadora produce pesos por token a través de todas las capas. Crucialmente, esta formulación permite pesos negativos, permitiendo que el enrutador suprima activamente contribuciones de capas perjudiciales (eliminación de ruido) en lugar de solo seleccionar ganadores. La salida se normaliza mediante la norma $\ell_2$ del vector de pesos para acotar la energía de salida.
3. Corrección Incremental: En lugar de reemplazar la salida original de la última capa ($z_{base}$) con la representación fusionada ($z_{fuse}$), DRoRAE formula la salida como una corrección incremental:
   $$z_{final} = LN_{bb}(z_{base} + \beta \cdot (z_{fuse} - z_{base}))$$
   Aquí, $\beta$ controla la fuerza de fusión. Esta formulación residual asegura que el latente fusionado permanezca anclado a la distribución original, evitando una deriva arbitraria que podría romper la compatibilidad con la generación aguas abajo.

Estrategia de Entrenamiento Desacoplado en Tres Fases

Para abordar el desafío de desplazar la distribución latente lejos de lo que un decodificador congelado espera, los autores proponen una estrategia de entrenamiento progresiva:

• Fase 1 (Entrenamiento del Decodificador): Entrenar el decodificador sobre la línea base estándar RAE de capa única para establecer una fuerte "capacidad de decodificación" y definir la distribución objetivo.
• Fase 2 (Entrenamiento del Módulo de Fusión): Congelar el backbone y el decodificador de la Fase 1. Entrenar solo el módulo de fusión. El decodificador congelado actúa como una restricción distribucional implícita, forzando al módulo de fusión a aprender información complementaria que el decodificador ya puede invertir, evitando la deriva distribucional.
• Fase 3 (Ajuste Fino del Decodificador): Congelar el módulo de fusión optimizado y desbloquear el decodificador. Ajustar finamente el decodificador para aprovechar plenamente el espacio latente enriquecido, maximizando la calidad de reconstrucción sin dañar la generación.

3. Contribuciones Clave

1. Identificación del Cuello de Botella: El artículo identifica que los autoencoders de representación existentes descartan información visual jerárquica al depender únicamente de la última capa del codificador, donde los detalles de bajo nivel son meramente residuos atenuados.
2. Arquitectura DRoRAE: Un módulo de fusión con enrutamiento por profundidad ligero (~29M de parámetros) que agrega adaptativamente todas las capas del codificador mediante enrutamiento con restricción de energía (permitiendo pesos negativos) y corrección incremental.
3. Estrategia de Entrenamiento: Un protocolo de entrenamiento desacoplado en tres fases que asegura la compatibilidad de generación aprendiendo la fusión bajo la restricción de un decodificador congelado antes de la co-adaptación del decodificador.
4. Descubrimiento de Ley de Escalado: El artículo establece una ley de escalado log-lineal ($R^2=0.86$) entre la capacidad de fusión (definida por la capacidad de expertos y el número de capas fusionadas) y la calidad de reconstrucción. Esto identifica la "riqueza de representación" como una nueva dimensión predeciblemente escalable para tokenizadores visuales, análoga al tamaño del vocabulario en PLN.

4. Resultados Experimentales

Los experimentos se realizaron en ImageNet-256 (256x256) para reconstrucción y generación condicional a clases, y en CC12M-LLaVA-Next para síntesis de texto a imagen.

• Reconstrucción: DRoRAE mejora significativamente las métricas de reconstrucción sobre la línea base RAE. El rFID disminuye de 0.57 a 0.29, el PSNR aumenta de 18.8 a 24.32 dB, y el LPIPS mejora de 0.256 a 0.134.
• Generación Condicional a Clases: Utilizando un generador DiTDH-XL con AutoGuidance, DRoRAE mejora el FID de generación (gFID) de 1.74 a 1.65 mientras mantiene o mejora ligeramente la Puntuación de Incepción (IS) y la Precisión/Recall.
• Síntesis de Texto a Imagen: Las ganancias se transfieren a tareas de texto a imagen, con DRoRAE alcanzando puntuaciones GenEval comparables o ligeramente mejores (0.60 vs. 0.56 para RAE) en las dimensiones de objeto, conteo y espacial.
• Estudios de Ablación:
  • Enrutamiento con Restricción de Energía: Supera al enrutamiento softmax al permitir la supresión activa de capas ruidosas.
  • Corrección Incremental: Esencial para la compatibilidad de generación; sin ella, la reconstrucción mejora ligeramente, pero la pérdida de entrenamiento de difusión aguas abajo aumenta significativamente, indicando una incompatibilidad distribucional.
  • Estrategia de Entrenamiento: El entrenamiento conjunto sin la restricción de la Fase 2 falla en lograr la misma calidad de generación, confirmando la necesidad del enfoque desacoplado.

5. Significado y Afirmaciones

El artículo afirma que DRoRAE demuestra que la fusión multi-capa es una actualización práctica y efectiva para tokenizadores basados en representación, recuperando detalles de alta frecuencia perdidos en la extracción de capa única.

Un significado principal de este trabajo es el descubrimiento de una ley de escalado predecible para tokenizadores visuales. Al mostrar que aumentar el "presupuesto de representación" (mediante más capas fusionadas o mayores capacidades de expertos) produce mejoras log-lineales en la calidad de reconstrucción, los autores argumentan que la riqueza de representación es una dimensión escalable comparable al tamaño del vocabulario en PLN. Esto sugiere que las mejoras futuras en la tokenización visual pueden cosecharse sistemáticamente escalando la capacidad de fusión, en lugar de depender únicamente de cambios arquitectónicos o conjuntos de datos más grandes.

El trabajo concluye que, si bien los experimentos actuales utilizan DINOv2-B, la metodología es generalizable a codificadores más grandes y potencialmente a la tokenización de video, posicionando la fusión multi-capa como una dirección fundamental para mejorar los sistemas de generación visual.