Adaptive MLP Pruning for Large Vision Transformers

Este artículo presenta un método de poda adaptativa de MLP (AMP) que reduce aproximadamente un 40% de los parámetros y FLOPs en grandes transformadores de visión mediante un criterio de entropía de información libre de etiquetas y una búsqueda binaria, logrando una compresión casi sin pérdida de rendimiento sin necesidad de ajuste fino.

Lo esencial

¡Claro! Imagina que los Transformers de Visión (como CLIP o DINOv2) son como gigantes de la inteligencia artificial que pueden ver y entender el mundo con una precisión increíble. Pero hay un problema: son tan grandes y pesados que requieren una computadora enorme y costosa para funcionar, como intentar correr una maratón con una mochila llena de ladrillos.

Los autores de este paper (Chengchao Shen y su equipo) han encontrado una forma de quitarle esos "ladrillos" sin que el gigante pierda su agilidad ni su inteligencia. Aquí te explico cómo lo hacen, usando analogías sencillas:

1. El Problema: El "Gordito" del Equipo

Dentro de estos gigantes de IA, hay una parte llamada MLP (Perceptrón Multicapa). Piensa en el MLP como el estómago del gigante: es donde procesa la información y "digiere" lo que ve.

• El hallazgo: Los investigadores descubrieron que este "estómago" es el que más pesa. En algunos modelos, el 81% del peso total del cerebro de la IA está en este estómago.
• La solución: Si queremos aligerar al gigante, no debemos tocar sus ojos (la atención) ni sus músculos, sino adelgazar su estómago.

2. La Herramienta 1: La Brújula de la "Entropía" (No solo mirar la respuesta correcta)

Antes, para decidir qué neuronas (células del estómago) podían ser eliminadas, los científicos usaban una regla simple: "¿Esta neurona ayudó a acertar la respuesta correcta?".

• El error: Es como si un profesor solo mirara si un alumno acertó la respuesta de "A" en un examen de opción múltiple, ignorando si el alumno también entendía por qué las opciones "B", "C" y "D" eran incorrectas. Se perdía mucha información.
• La innovación (Entropía de Información): El nuevo método usa una brújula más inteligente. En lugar de mirar solo la respuesta correcta, mira todas las posibilidades que el gigante considera.
  • Analogía: Imagina que el gigante está adivinando qué animal es una foto. En lugar de solo mirar si dijo "Perro", miramos si también pensó en "Lobo" o "Zorro" y cómo de seguro estaba en cada opción. Esto les da una brújula mucho más precisa para saber qué neuronas son realmente importantes y cuáles son solo "ruido" o redundancia. Además, esta brújula funciona sin necesidad de las respuestas correctas (etiquetas), lo que es genial porque a veces no tenemos el manual de respuestas de modelos muy nuevos.

3. La Herramienta 2: El "Corte Inteligente" (Búsqueda Binaria)

Una vez que tienen la lista de neuronas ordenadas de "más importantes" a "menos importantes", no cortan de golpe.

• El método antiguo: Era como decir: "¡Cortemos el 40% del estómago!" sin saber si eso mataría al gigante.
• El método nuevo (Adaptativo): Es como un juego de adivinar el peso.
  1. Empiezan cortando un poco.
  2. Ponen al gigante a caminar (evalúan si sigue entendiendo bien).
  3. Si camina bien, cortan un poco más.
  4. Si tropieza, vuelven un paso atrás.
  • Usan un algoritmo llamado búsqueda binaria (como buscar una palabra en un diccionario abriéndolo por la mitad, luego por la mitad de esa mitad, etc.) para encontrar el punto exacto donde el gigante está lo más ligero posible sin perder su inteligencia. No hay una regla fija; se adapta a cada modelo.

4. El Remate: El "Entrenador" (Distilación de Conocimiento)

Después de cortar, el gigante está un poco aturdido. Para recuperarlo, usan al gigante original (antes de cortarle el estómago) como un entrenador personal.

• El gigante original le enseña al gigante "adelgazado" cómo pensar de nuevo.
• Analogía: Es como si un maestro de ajedrez (el modelo original) le enseñara a un estudiante (el modelo cortado) a jugar las mismas partidas. Gracias a que la estructura es muy similar, el estudiante aprende muy rápido y recupera su nivel casi al instante.

¿Qué lograron? (Los Resultados)

• Reducción masiva: Lograron quitarle un 40% de peso (parámetros y cálculos) a estos gigantes de IA.
• Velocidad: Ahora son un 50% más rápidos (casi el doble de velocidad).
• Sin dolor: Lo increíble es que, incluso sin volver a entrenarlos mucho, siguen funcionando casi igual de bien que antes. En muchos casos, después del entrenamiento con el "entrenador", ¡incluso funcionan mejor que el original!
• Versatilidad: Funciona con modelos famosos como CLIP (que entiende imágenes y texto) y DINOv2 (que solo entiende imágenes), incluso si no tenemos acceso a todas sus partes internas.

En resumen

Este paper nos dice que no necesitamos gigantes pesados para tener inteligencia. Con una brújula más inteligente para saber qué cortar y un corte a medida (no a lo loco), podemos tener modelos de visión por computadora que son ligeros, rápidos y baratos, pero que siguen siendo geniales entendiendo el mundo. ¡Es como convertir a un elefante en un guepardo sin perder su fuerza!

Resumen técnico

Resumen Técnico: Poda Adaptativa de MLP para Grandes Transformadores de Visión

1. El Problema

Los grandes transformadores de visión (Large Vision Transformers) han demostrado una excelente capacidad de escalado, mejorando su rendimiento a medida que aumenta la capacidad del modelo. Sin embargo, esto conlleva un costo prohibitivo en términos de parámetros, consumo computacional y requisitos de memoria, lo que dificulta su despliegue eficiente en aplicaciones prácticas.

El análisis de la arquitectura de estos modelos revela que los módulos Multicapa Perceptrón (MLP) representan la mayor parte de los parámetros del modelo (por ejemplo, el 81.1% en EVA-CLIP-E). Aunque existen métodos de poda basados en Taylor para comprimir modelos, presentan limitaciones críticas:

• Dependencia de etiquetas: Utilizan la pérdida de entropía cruzada one-hot, que ignora las predicciones potenciales sobre otras categorías, degradando la calidad de la evaluación de importancia.
• Falta de generalidad: Requieren el acceso a la función de pérdida o módulos adicionales (como el cabezal DINO en DINOv2 o el codificador de texto en CLIP) que a menudo no están disponibles públicamente.
• Rigidez: Suelen depender de una tasa de compresión predefinida, lo que impide una poda adaptativa basada en la redundancia real de cada módulo.

2. Metodología Propuesta (AMP)

Los autores proponen AMP (Adaptive MLP Pruning), un método que reduce significativamente los parámetros de los transformadores de visión sin degradación notable del rendimiento. El proceso consta de tres etapas principales:

A. Evaluación de Importancia de Neuronas (Criterio de Entropía de Información)
En lugar de usar la pérdida de entropía cruzada one-hot, AMP introduce un criterio libre de etiquetas basado en la entropía de información:

• Se calcula una matriz de similitud entre instancias en un mini-batch utilizando las representaciones de salida del modelo ($z_{cls}$).
• Se aplica una operación softmax con un coeficiente de temperatura ($\tau$) para obtener una matriz de probabilidad de predicción.
• Se calcula la entropía de esta distribución.
• Ventaja: Este criterio modela completamente la distribución de predicciones del modelo original, capturando información que la entropía cruzada ignora. Además, no requiere etiquetas ni módulos adicionales (como el cabezal DINO), permitiendo la compresión de modelos con pesos no totalmente publicados.

B. Poda Adaptativa mediante Búsqueda Binaria
Una vez evaluada la importancia de las neuronas ocultas del MLP:

• Las neuronas se clasifican según su puntuación de importancia.
• Se utiliza un algoritmo de búsqueda binaria para determinar el número óptimo de neuronas a podar en cada módulo MLP.
• El algoritmo ajusta dinámicamente el tamaño oculto: si la variación de la entropía de información tras la poda supera un umbral ($\Delta E$), se reduce la cantidad de neuronas podadas; de lo contrario, se poda más.
• Resultado: Se evita la necesidad de definir una tasa de compresión fija, adaptándose a la redundancia específica de cada módulo.

C. Recuperación de Rendimiento (Distillation)
Para recuperar el rendimiento tras la poda:

• Se utiliza Knowledge Distillation (Distilación de Conocimiento).
• El modelo original actúa como "maestro" y el modelo podado como "estudiante".
• Se minimiza el error cuadrático medio entre las salidas de los tokens de clase ($z_{cls}$) y los tokens de parche ($z_{patch}$) de ambos modelos.
• Gracias a la afinidad estructural (las dimensiones de salida se mantienen idénticas), no se requieren módulos de alineación adicionales.

3. Contribuciones Clave

1. Criterio de Entropía Libre de Etiquetas: Permite una evaluación de importancia más precisa y generalizable, funcionando incluso cuando no se dispone de los pesos completos del modelo de entrenamiento o de la función de pérdida original.
2. Poda Adaptativa: Elimina la necesidad de tasas de compresión predefinidas mediante el uso de búsqueda binaria, optimizando la reducción de parámetros según la redundancia real de cada capa.
3. Alta Eficiencia y Rendimiento: Logra una aceleración casi sin pérdida ("near lossless") con una reducción de aproximadamente el 40% en parámetros y FLOPs, superando significativamente a otros métodos de poda cuando no se realiza fine-tuning.

4. Resultados Experimentales

El método fue evaluado en transformadores de visión de última generación, incluyendo CLIP (OpenCLIP-g, OpenCLIP-G, EVA-CLIP-E, EVA-CLIP-8B) y DINOv2.

• Reducción de Recursos: Se logró una reducción de ~40% en parámetros y FLOPs, con un aumento en el throughput de inferencia de aproximadamente 1.5x.
• Rendimiento en Clasificación Zero-Shot:
  • Los modelos podados y distilados recuperaron el rendimiento de los modelos originales, superándolos ligeramente en algunos casos (ej. EVA-CLIP-8B distilado superó al original en 0.4% en Recall Medio).
  • Sin fine-tuning: Los modelos podados sin distilación superaron a otros métodos de poda (como Taylor pruning, SAViT, NViT) por un margen enorme (ej. +42.7% en OpenCLIP-g).
• Comparativa: En tareas de clasificación en ImageNet-1K y ObjectNet, así como en recuperación de imágenes/texto en Flickr30K y COCO, AMP demostró ser superior a la poda aleatoria, basada en normas L2, y otros métodos basados en Taylor.
• Evaluación kNN: En modelos puramente visuales como DINOv2-g, la versión podada y distilada alcanzó el mismo rendimiento que el original con solo el 54.4% de los parámetros.

5. Significado e Impacto

Este trabajo es significativo porque aborda la eficiencia de los grandes modelos de visión de una manera que es robusta, generalizable y adaptable.

• Democratización: Al no depender de etiquetas ni de módulos internos no públicos, permite comprimir modelos de vanguardia que de otro modo serían inaccesibles para la optimización.
• Eficiencia Operativa: Ofrece una ruta viable para desplegar modelos masivos en entornos con recursos limitados sin sacrificar significativamente la precisión.
• Futuro: Los autores planean extender esta metodología a la reducción adaptativa de los módulos de atención multi-cabeza y su aplicación en Grandes Modelos de Lenguaje (LLMs).

En conclusión, AMP representa un avance crucial en la compresión de modelos de visión, demostrando que es posible reducir drásticamente la carga computacional de los transformadores modernos manteniendo su capacidad de generalización y rendimiento.