CompoDistill: Attention Distillation for Compositional Reasoning in Multimodal LLMs

Resumen Técnico: CompoDistill

1. El Problema: La Brecha en la Distilación de la Percepción Visual

Los Modelos de Lenguaje Grandes Multimodales (MLLMs) han demostrado un rendimiento superior en tareas de visión y lenguaje, pero su alto costo computacional limita su despliegue en aplicaciones del mundo real. La Distilación de Conocimiento (KD) se ha utilizado como solución para transferir el conocimiento de un modelo "maestro" grande a un modelo "estudiante" más pequeño.

Sin embargo, los autores identifican una limitación crítica en los métodos de KD existentes:

• Éxito en Reconocimiento Visual: Los métodos actuales mejoran significativamente las tareas de Preguntas de Respuesta Visual (VQA), que requieren reconocimiento de objetos y atributos básicos.
• Fracaso en Percepción Visual: Fallan en transferir la percepción visual, una capacidad más compleja que implica entender relaciones entre objetos, atributos precisos y la composición de la escena (evaluada mediante Razonamiento Composicional - CR).
• Causa Raíz: A través de un análisis sistemático, el paper demuestra que el fallo se debe a la desalineación de la atención visual. Mientras que el modelo maestro se enfoca en las regiones relevantes de la imagen para responder a una consulta, el modelo estudiante (distilado) atiende a regiones irrelevantes, fallando en heredar el mecanismo de atención visual del maestro.

2. Metodología: CompoDistill

Para abordar la desalineación de la atención, los autores proponen CompoDistill, un nuevo marco de distilación que consta de dos módulos centrales y una estrategia de entrenamiento en tres etapas.

A. Módulos Clave:

1. Módulo de Alineación de ATención Visual (VAT - Visual ATtention alignment):

   • Objetivo: Alinear explícitamente la distribución de atención del estudiante con la del maestro en las capas de "comprensión visual" (capas intermedias del LLM).
   • Estrategia de Emparejamiento de Capas (Group Layer Matching): Dado que el maestro suele tener más capas que el estudiante, en lugar de un emparejamiento uno a uno (que es subóptimo), el método empareja cada capa del estudiante con un grupo de capas consecutivas del maestro (relación uno a muchos). Esto permite al estudiante capturar un conocimiento de atención más rico y distribuido.
   • Función de Pérdida: Minimiza la distancia coseno entre la matriz de atención del estudiante y el promedio de las matrices de atención del grupo de capas del maestro correspondiente.

2. Módulo de Recuperación del Adaptador del Maestro (TAF - Teacher Adapter Fetch):

   • Problema: La atención del maestro está optimizada para su propio espacio de características visuales (salida de su adaptador). Transferir esta atención directamente a un estudiante con un espacio de características incompatible genera conflictos.
   • Solución: El módulo TAF utiliza el adaptador preentrenado y congelado del maestro para proyectar las características visuales al espacio del lenguaje, añadiendo solo una pequeña capa MLP entrenable para la alineación de dimensiones. Esto asegura que el estudiante procese la entrada visual a través de la misma "lente" que el maestro, haciendo que la transferencia de atención sea efectiva.

B. Estrategia de Entrenamiento en Tres Etapas:

1. Pre-entrenamiento Distilado (DPT): Se alinea el espacio de características visuales con el espacio lingüístico utilizando el módulo TAF y pérdidas de modelado de lenguaje y divergencia KL.
2. Ajuste Fino Distilado (DFT): Se activa el módulo VAT para alinear la atención visual entre maestro y estudiante, refinando la percepción visual.
3. Ajuste Fino Supervisado (SFT): Se entrena el modelo final solo con pérdida de modelado de lenguaje para consolidar el conocimiento y mejorar la capacidad de seguir instrucciones.

3. Contribuciones Clave

1. Identificación del Problema: Es el primer trabajo que identifica y analiza sistemáticamente que los métodos de KD actuales fallan en distilar la percepción visual debido a la desalineación de la atención visual, no solo por falta de datos o capacidad.
2. Nueva Arquitectura de KD: Propone CompoDistill, un marco que integra la alineación de atención (VAT) y la compatibilidad de espacios de características (TAF) para transferir tanto reconocimiento como percepción.
3. Estrategia de Emparejamiento: Introduce una estrategia de "Emparejamiento de Capas por Grupos" que supera a los métodos de emparejamiento simple o adaptativo en la transferencia de conocimiento de atención.

4. Resultados Experimentales

Los experimentos se realizaron comparando CompoDistill (estudiante de 2B parámetros) con otros métodos de KD y modelos de base, utilizando un maestro de 4B parámetros.

• Rendimiento en Razonamiento Composicional (CR): CompoDistill supera drásticamente a los métodos de KD existentes (como LLaVA-KD, LLaVA-MoD) y al modelo SFT estándar. Logra un rendimiento en tareas CR comparable al del modelo maestro de 4B, algo que ningún otro método de distilación de 2B había logrado antes.
• Rendimiento en VQA: Mantiene un rendimiento competitivo y de vanguardia en tareas de VQA (reconocimiento visual), demostrando que la mejora en la percepción no sacrifica el reconocimiento.
• Eficiencia de Datos: Logra estos resultados con solo 1.2M de muestras de entrenamiento, siendo significativamente más eficiente que otros modelos que requieren datasets masivos (ej. 5M o más).
• Generalización: El método funciona bien con diferentes arquitecturas de LLM (Qwen y MobileLLaMA) y escala positivamente cuando se usa un maestro más grande (7B) o se aumenta el tamaño de los datos.
• Reducción de Alucinaciones: El modelo muestra una mejora significativa en la reducción de alucinaciones relacionales en benchmarks como R-Bench y Reefknot, gracias a una mejor comprensión de las relaciones entre objetos.

5. Significado e Impacto

Este trabajo es fundamental para el desarrollo de MLLMs eficientes. Demuestra que la simple transferencia de logits o características no es suficiente para modelos multimodales; la atención visual es el componente crítico que debe alinearse para que los modelos pequeños puedan realizar razonamiento complejo.

Al resolver la desalineación de la atención, CompoDistill permite crear modelos pequeños que no solo "ven" objetos, sino que "comprenden" la escena, cerrando la brecha de rendimiento entre modelos pequeños y grandes en tareas de razonamiento de alto nivel. Esto abre la puerta a la implementación práctica de MLLMs avanzados en dispositivos con recursos limitados sin sacrificar la capacidad de razonamiento composicional.