Resumen Técnico: Replanteando la Poda de Profundidad para Vision Transformers: Una Perspectiva Consciente de la Heterogeneidad

1. Planteamiento del Problema

Los Vision Transformers (ViTs) han alcanzado un rendimiento de vanguardia en diversos dominios, pero sufren de altos costos computacionales y grandes conteos de parámetros, lo que conduce a una latencia de inferencia extendida. Aunque la poda estructurada es una estrategia de compresión efectiva, los métodos existentes enfrentan limitaciones significativas:

• Poda de Ancho (Width Pruning): Reduce canales o cabezales de atención dentro de las capas, pero a menudo produce aceleraciones limitadas en el hardware convencional.
• Poda de Profundidad (Depth Pruning): Elimina capas completas, ofreciendo un potencial de aceleración significativamente mayor (como se muestra en la Figura 2a). Sin embargo, es notoriamente difícil recuperar la precisión, especialmente con una eliminación agresiva de capas.
• El Cuello de Botella Central: El artículo argumenta que el fracaso de los métodos de poda de profundidad existentes se debe a dos problemas pasados por alto:
  1. Negligencia de la Heterogeneidad: Los ViTs contienen tipos de capas distintos (por ejemplo, capas de Atención frente a capas de Función de Activación dentro de las Redes de Feed-Forward). Los métodos existentes suelen tratar estos tipos de capas de manera uniforme o dependen de métricas basadas en gradientes que fallan debido a la disparidad significativa en las escalas de gradiente entre estos tipos.
  2. Desajuste de Dimensiones: Podar simultáneamente las capas de Atención y las capas Lineales en la Red de Feed-Forward (FFN) crea incompatibilidades estructurales. Dado que las FFN utilizan un mecanismo de expansión-contracción (por ejemplo, $D \to 4D \to D$), eliminar una capa Lineal rompe la continuidad de la forma del tensor requerida por las capas subsiguientes (ilustrado en la Figura 3).

2. Metodología: HetDPT

Los autores proponen HetDPT (Heterogeneity-aware DePTh pruning), un marco de dos etapas diseñado para abordar la heterogeneidad y evitar el desajuste de dimensiones.

Ideas Clave

El método se fundamenta en dos observaciones críticas respecto a la heterogeneidad de las capas:

1. Disparidad de Gradientes: Las capas de Atención y las capas de Activación exhiben magnitudes de gradiente muy diferentes durante la retropropagación. El uso de métricas de importancia estándar basadas en gradientes conduce a decisiones de poda sesgadas, favoreciendo consistentemente la poda de un tipo sobre el otro.
2. Asimetría de Recuperación:
   • La poda de capas de Atención causa caídas de precisión inicial moderadas pero requiere un reentrenamiento extensivo para su recuperación.
   • La poda de capas de Activación causa una degradación inicial severa (a menudo $\ge 90\%$ de pérdida de precisión) pero permite una recuperación rápida durante el ajuste fino (fine-tuning).
   • Implicación: Las métricas cortoplacistas que evalúan la precisión inmediata post-poda son ineficaces.

El Pipeline de Dos Etapas

Etapa 1: Identificación de Capas Redundantes
Esta etapa determina cuántas y qué capas específicas podar sin comparar directamente los tipos heterogéneos de capas.

• Paso 1: Asignación del Presupuesto de Poda (Model Accuracy Predictor - MAP): En lugar de comparar capas directamente, el método utiliza un proceso de recolección de datos ligero (poda iterativa $\to$ ajuste fino rápido $\to$ evaluación) para construir un conjunto de datos. Una función polinómica aproxima el Model Accuracy Predictor (MAP), el cual predice la precisión final de un modelo podado dado los ratios de poda de las capas de Atención ($\tilde{m}_a$) y de Activación ($\tilde{m}_g$). Esto permite al sistema encontrar la asignación óptima del presupuesto de poda entre los dos tipos de capas.
• Paso 2: Eliminación de Capas Redundantes Específicas: Una vez asignado el presupuesto, el método identifica las capas específicas a eliminar dentro de cada grupo homogéneo (por ejemplo, solo entre las capas de Atención). Utiliza parámetros de importancia aprendibles y propagación de gradientes para clasificar las capas, asegurando que no ocurra una comparación entre tipos.

Etapa 2: Optimización del Modelo Podado

• Ajuste Fino (Fine-tuning): El modelo podado se ajusta finamente, utilizando opcionalmente auto-destilación guiada por el ViT original no podado para restaurar la precisión.
• Fusión (Merging): Para resolver el desajuste de dimensiones y acelerar la inferencia, las capas Lineales adyacentes que fueron separadas por una capa de Activación podada se fusionan. Por ejemplo, si se elimina una capa de Activación, $L_{l,1}$ y $L_{l,2}$ se fusionan en una única capa lineal $L_l$. Esto reduce la profundidad del modelo mientras mantiene un alineamiento dimensional riguroso.

HetDPT+: Los autores integran además HetDPT con métodos establecidos de poda de ancho para crear HetDPT+, un pipeline para la compresión extrema de ViT que aborda tanto las dimensiones de profundidad como de ancho.

3. Contribuciones Clave

1. Poda de Funciones de Activación: El artículo propone la primera estrategia para eliminar las capas de Función de Activación (GELU) en los ViTs. Al eliminar estas no linealidades, las capas Lineales adyacentes pueden fusionarse naturalmente, resolviendo el problema de desajuste de dimensiones inherente a la poda de profundidad heterogénea.
2. Análisis y Solución de la Heterogeneidad: Los autores revelan y abordan la disparidad de gradientes y la asimetría de recuperación. Introducen un marco de dos etapas (HetDPT) equipado con un Predictor de Precisión del Modelo (MAP) para gestionar estas heterogeneidades, respaldado por un análisis teórico que garantiza la existencia de una configuración óptima y la consistencia del predictor bajo un ajuste fino rápido y ruidoso.
3. Récords de Vanguardia (SOTA):
   • Poda de Profundidad: HetDPT logra una aceleración de 1.58$\times$ para DeiT-B y 1.39$\times$ para DeiT-S con una precisión casi sin pérdidas, superando los trabajos de poda de profundidad existentes.
   • Compresión Extrema: HetDPT+ establece un nuevo referente para la compresión extrema de ViT. Para la configuración Isomorphic-Pruning-2.6G, mejora la ratio de aceleración de 4.24$\times$ a 5.19$\times$ manteniendo una precisión casi sin pérdidas.

4. Resultados Experimentales

El método fue evaluado en ImageNet-1K, CIFAR-100, COCO y ADE20K.

• Compromiso entre Precisión y Aceleración: En ImageNet-1K, HetDPT supera consistentemente la poda aislada de capas de atención o de activación. Para DeiT-B, alcanza un 81.8% de precisión Top-1 (igualando la línea base) con una aceleración de 1.58$\times$.
• Comparación con SOTA: HetDPT+ supera a los métodos de compresión extrema existentes (por ejemplo, Isomorphic Pruning, NViT). Con ~81% de precisión, HetDPT+ logra una aceleración de 5.19$\times$ comparado con 4.24$\times$ de Isomorphic Pruning. En configuraciones ultra-extremas (1.1G MACs), logra una aceleración de 9.19$\times$.
• Transferibilidad: El método demuestra un aprendizaje de transferencia robusto en CIFAR-100 y mantiene el rendimiento en tareas de predicción densa (detección de objetos COCO y segmentación semántica ADE20K) con una degradación de precisión mínima.
• Escalabilidad: El enfoque escala a modelos de miles de millones de parámetros, demostrado en DINO-V2-Giant, donde poda 15 capas con solo una caída de precisión del 0.06%.
• Ortogonalidad: HetDPT es ortogonal a los métodos de poda de tokens (por ejemplo, GTP-ViT, ToMe), proporcionando aceleraciones adicionales cuando se combinan.

5. Significado y Reivindicaciones

El artículo afirma que la contribución principal es permitir un despliegue más eficiente de los Vision Transformers al abordar los desafíos estructurales y dinámicos específicos de la poda de profundidad.

• Democratización: Al reducir sustancialmente los requisitos computacionales (hasta 9.19$\times$ de aceleración), el método hace que los modelos de visión de vanguardia sean accesibles para investigadores y profesionales con recursos limitados.
• Sostenibilidad: La reducción del consumo energético en los modelos comprimidos contribuye a un desarrollo de IA más sostenible desde el punto de vista ambiental.
• Despliegue en el Borde (Edge Deployment): Las ganancias de eficiencia facilitan el despliegue en tiempo real en dispositivos de borde para aplicaciones que requieren baja latencia, como diagnósticos médicos y sistemas autónomos.

Los autores señalan una limitación: el Model Accuracy Predictor (MAP) se reajusta actualmente para cada familia de arquitectura, y extender el método a modelos de lenguaje u otras modalidades sigue siendo un trabajo futuro. Sin embargo, este trabajo establece un nuevo fundamento teórico y práctico para la poda de profundidad heterogénea, yendo más allá de los enfoques de "talla única" de la literatura previa.