HiAP: A Multi-Granular Stochastic Auto-Pruning Framework for Vision Transformers

El artículo presenta HiAP, un marco de auto-poda estocástica multi-granular que optimiza de forma continua y en una sola fase de entrenamiento Vision Transformers mediante el uso de puertas Gumbel-Sigmoid para eliminar selectivamente tanto componentes macro (cabezas de atención y bloques FFN) como micro (dimensiones intra-cabeza y neuronas), logrando arquitecturas eficientes con alto rendimiento en dispositivos de borde sin necesidad de heurísticas manuales o procesos multi-etapa.

Lo esencial

¡Hola! Imagina que tienes un gigante intelectual (un modelo de Inteligencia Artificial llamado "Vision Transformer") que es increíblemente inteligente, pero tan pesado y lento que no cabe en tu teléfono móvil ni en un reloj inteligente. Necesita demasiada energía y memoria para pensar.

El problema es que, hasta ahora, para hacerlo más ligero, los ingenieros tenían que actuar como arquitectos muy estrictos y manuales: tenían que decidir a mano qué partes del cerebro del gigante cortar, probarlo, ver si seguía funcionando, y repetir el proceso muchas veces. Era un trabajo tedioso, costoso y a veces resultaba en un gigante que ya no sabía ver bien.

Aquí es donde entra HiAP (el marco de "Poda Automática Jerárquica").

La Analogía: El Gran Hotel y el Gerente Inteligente

Imagina que el modelo de IA es un gigantesco hotel con muchas plantas (capas), cada una con muchas habitaciones (cabezas de atención) y dentro de cada habitación, muchos muebles (neuronas y dimensiones).

1. El problema de los métodos antiguos:
   Los métodos anteriores eran como un gerente que entra y dice: "¡Quiero quitar el 30% de los muebles de todas las habitaciones!" o "¡Quiero cerrar el 20% de las plantas!".

   • Si quitas demasiados muebles, la habitación se siente vacía y el cliente (la imagen) no se siente cómodo (baja la precisión).
   • Si solo quitas muebles pero dejas todas las plantas abiertas, el hotel sigue gastando mucha energía en mantener las puertas y los pasillos abiertos (problema de memoria), aunque las habitaciones estén vacías.
   • Además, este gerente necesitaba un manual de instrucciones muy complejo y tenía que hacer varias rondas de inspección.

2. La solución de HiAP (El Gerente que Aprende):
   HiAP es como un gerente inteligente y autónomo que vive dentro del hotel y aprende a gestionar el edificio mientras el hotel está abierto y funcionando. No necesita un manual ni un jefe que le diga qué hacer.

   HiAP tiene dos niveles de control (como dos tipos de llaves maestras):

   • Las Llaves Macro (Macro-gates): Son las llaves que cierran plantas enteras o alas completas del hotel. Si una planta entera no es necesaria, HiAP la apaga por completo. Esto ahorra una cantidad enorme de energía porque no tienes que caminar por esos pasillos (ahorro de memoria).
   • Las Llaves Micro (Micro-gates): Son las llaves que quitan muebles específicos de las habitaciones que sí se quedan abiertas. Si una habitación es necesaria, HiAP decide quitar solo la cama de invitados o la mesa de noche, pero deja la cama principal. Esto hace que la habitación sea más ligera pero sigue siendo funcional (ahorro de cálculos).

¿Cómo funciona la magia? (El proceso de "Aprendizaje")

En lugar de cortar y pegar manualmente, HiAP usa un truco matemático llamado "Gumbel-Sigmoid".

• Imagina que es como un termostato: Al principio del entrenamiento, el "termostato" está muy suave. Las puertas y los muebles están "semi-abiertos" o "semi-cerrados". El hotel está experimentando: "¿Qué pasa si cierro esta planta? ¿Qué pasa si quito este mueble?".
• El aprendizaje conjunto: Mientras el hotel experimenta con estas puertas semi-abiertas, los empleados (los pesos de la red neuronal) aprenden a trabajar mejor con lo que tienen. Se adaptan.
• El enfriamiento (Annealing): A medida que pasa el tiempo, el termostato se va enfriando. Las puertas "semi-abiertas" se deciden: o se abren del todo (se quedan) o se cierran del todo (se eliminan).
• El resultado final: Al final del proceso, el hotel se ha reorganizado solo. Ha encontrado su propia forma ideal: tiene las plantas exactas que necesita y los muebles justos. Lo mejor de todo es que no hubo que cerrar el hotel para remodelarlo; todo ocurrió mientras funcionaba.

¿Por qué es tan genial?

1. Todo en uno: No necesitas dos pasos (primero buscar qué cortar, luego volver a entrenar). HiAP hace todo en un solo viaje.
2. Equilibrio perfecto: Entiende que a veces es mejor cerrar una planta entera (para ahorrar memoria) que dejarla abierta con solo un mueble (que sigue costando memoria). Y a veces es mejor dejar la planta abierta pero quitar muebles. Lo decide solo.
3. Listo para usar: Al final, obtienes un edificio físico, real y compacto. No es un plano teórico; es un hotel listo para recibir huéspedes en un teléfono móvil, sin necesidad de herramientas especiales.

En resumen

Piensa en HiAP como un arquitecto autodidacta que entra en un edificio gigante, camina por él mientras vive en él, y decide: "Esta planta no hace nada, la cierro. En esta otra, solo necesito la cocina, así que quito los dormitorios".

Al final, tienes un edificio pequeño, eficiente y rápido, que sigue funcionando perfectamente, y todo esto ocurrió sin que nadie tuviera que dibujar planos ni hacer cálculos manuales complicados. ¡Es la IA aprendiendo a ser más eficiente por sí misma!

Resumen técnico

Aquí tienes un resumen técnico detallado del artículo "HiAP: A Multi-Granular Stochastic Auto-Pruning Framework for Vision Transformers" en español:

1. El Problema

Los Vision Transformers (ViT) han demostrado un rendimiento superior en visión por computadora, pero su alto costo computacional y de ancho de banda de memoria dificulta su despliegue en dispositivos con recursos limitados (edge devices).

• Limitaciones de los métodos actuales:
  • Granularidad única: La mayoría de los métodos de poda estructurada operan en un solo nivel. La poda "micro" (dimensiones internas) reduce FLOPs pero no mitiga la sobrecarga de memoria (carga de matrices completas), mientras que la poda "macro" (bloques completos) reduce memoria pero puede degradar drásticamente la capacidad de representación.
  • Procesos complejos: Los métodos existentes suelen requerir pipelines multi-etapa, heurísticas manuales de importancia, umbralización post-hoc y fases de ajuste fino (fine-tuning) separadas, lo que complica la implementación.
  • Dependencia de hardware: La poda no estructurada requiere hardware especializado, mientras que la estructurada busca sub-redes densas compatibles con hardware estándar.

2. Metodología: HiAP (Hierarchical Auto-Pruning)

HiAP propone un marco de relajación continua que descubre la sub-red óptima en una sola fase de entrenamiento end-to-end, sin necesidad de objetivos de dispersión predefinidos ni intervención manual.

A. Mecanismo de Puertas Jerárquico (Gating)

HiAP introduce puertas estocásticas basadas en la distribución Gumbel-Sigmoid en dos niveles de granularidad simultáneos:

1. Nivel Macro: Controla la retención de unidades completas.
   • Puertas para cabezas de atención completas ($g_{l,h}$).
   • Puertas para bloques enteros de la red neuronal feed-forward (FFN) ($b_l$).
   • Si una puerta es 0, la unidad se omite completamente, reduciendo la sobrecarga de memoria.
2. Nivel Micro: Opera dentro de las estructuras macro activas.
   • Puertas para dimensiones intra-cabeza ($d_{l,h,j}$).
   • Puertas para neuronas individuales del FFN ($c_{l,k}$).
   • Permite ajustar el ancho de la red de forma heterogénea.

B. Modelado de Costos Diferenciables

El método formula un costo exacto y diferenciable de las operaciones Multiplicación-Acumulación (MACs) que se incluye en la función de pérdida:

• Descomposición lineal: Separa el costo en sobrecarga macro (cálculo de proyecciones Q, K, V y mapas de atención) y costos micro (salidas de atención y neuronas activas).
• Penalización de viabilidad: Se introducen penalizaciones explícitas para evitar el colapso estructural (ej. eliminar capas enteras prematuramente), garantizando un flujo de gradiente estable y un paso hacia adelante válido.

C. Entrenamiento y Extracción

• Una sola fase: A diferencia de métodos que separan la búsqueda de arquitectura del ajuste fino, HiAP entrena la red densa y las puertas simultáneamente.
• Recocido de Temperatura (Annealing): La temperatura $\tau$ de la distribución Gumbel-Sigmoid se reduce gradualmente durante el entrenamiento. Esto permite que las puertas pasen de ser estocásticas (exploración) a determinísticas (0 o 1), convergiendo suavemente a una sub-red discreta sin "choque de gradiente".
• Extracción Física: Al final del entrenamiento, las puertas se endurecen (umbral > 0.5) y la arquitectura se extrae físicamente (truncando matrices), resultando en un modelo ViT comprimido listo para inferencia sin necesidad de fine-tuning adicional.

3. Contribuciones Clave

1. Unificación de Granularidades: HiAP es el primer marco que unifica la poda macro (bloques/cabezas) y micro (dimensiones/neuronas) en un único sistema de puertas diferenciable.
2. Descubrimiento Autónomo: Elimina la necesidad de heurísticas manuales, métricas de ranking proxy o objetivos de dispersión predefinidos. La red descubre su propia arquitectura óptima basada en la eficiencia.
3. Eficiencia en el Pipeline: Logra un estado del arte en la frontera de eficiencia-precisión (Pareto) con un proceso de entrenamiento simplificado de una sola etapa, eliminando la necesidad de fases de ajuste fino costosas.

4. Resultados Experimentales

Los experimentos se realizaron en CIFAR-10 y ImageNet-1K (usando DeiT-Small).

• Rendimiento en ImageNet (DeiT-Small):
  • HiAP logró reducir los MACs de 4.6G a 3.1G (una reducción del ~33%) manteniendo una precisión Top-1 de 79.1% (solo -0.75% respecto al modelo denso).
  • Superó o igualó a métodos complejos como ViT-Slim, GOHSP y S2ViTE, que requieren pipelines multi-etapa, pero con una implementación mucho más simple.
• Rendimiento en CIFAR-10 (ViT-Tiny):
  • HiAP superó a las heurísticas manuales (como el ranking por norma L1) y a la poda uniforme, mejorando la precisión en un +0.93% bajo el mismo presupuesto de MACs.
• Eficiencia de Hardware:
  • Al extraer físicamente la sub-red, se obtuvo una aceleración de inferencia de 1.44x (de 5.57 ms a 3.86 ms) en una GPU, confirmando que las reducciones teóricas de MACs se traducen directamente en velocidad real sin depender de motores de convolución dispersa.
• Análisis de la Estructura Descubierta:
  • El modelo aprendió a eliminar bloques FFN completos en capas finales (poda de profundidad) y a reducir agresivamente las dimensiones internas en capas profundas, demostrando una distribución heterogénea óptima de la dispersión.

5. Significado e Impacto

HiAP representa un cambio de paradigma en la optimización de ViT:

• Simplificación: Transforma la poda de un problema de ingeniería manual y multi-etapa a un problema de aprendizaje automático unificado y consciente del presupuesto.
• Adaptabilidad: Al abordar simultáneamente los cuellos de botella de memoria (poda macro) y de cómputo (poda micro), genera arquitecturas que son realmente eficientes en hardware moderno.
• Listo para Producción: Al generar sub-redes densas y físicamente extraídas sin necesidad de ajuste fino posterior, HiAP facilita enormemente el despliegue de modelos de visión en dispositivos con recursos limitados.

En resumen, HiAP demuestra que es posible descubrir arquitecturas de alto rendimiento y alta eficiencia de manera automática y end-to-end, superando las limitaciones de los métodos de poda estructurada tradicionales.