Resumen Técnico: HiReLC – Aprendizaje por Refuerzo Jerárquico para la Compresión de Redes Neuronales

1. Declaración del Probleo

Las redes neuronales profundas, particularmente los Vision Transformers (ViTs), enfrentan desafíos significativos de despliegue en plataformas con recursos limitados debido a sus altos recuentos de parámetros, huellas de memoria y latencias de inferencia. Aunque existen técnicas de compresión actuales, como heurísticas estáticas, asignación de precisión mixta de un solo paso (one-shot) y optimización de un solo objetivo, estas suelen fallar al no capturar las dependencias entre capas y carecen de un refinamiento adaptativo durante el ajuste fino (fine-tuning). Además, los enfoques de Aprendizaje por Refuerzo (RL) para la compresión han luchado históricamente contra la alta varianza en espacios de acción extensos, la granularidad gruesa y la falta de una guía sensible a la importancia.

El artículo aborda la necesidad de un marco de trabajo de propósito general y agnóstico a la arquitectura capaz de optimizar conjuntamente la poda estructurada (structured pruning) y la cuantización de precisión mixta, navegando los complejos compromisos entre la reducción del tamaño del modelo y la retención de la precisión.

2. Metodología: El Marco de Trabajo HiReLC

HiReLC (Hierarchical Reinforcement Learning for Neural Network Compression) es un marco de aprendizaje por refuerzo de conjunto jerárquico de dos niveles diseñado para descomponer el espacio de búsqueda de compresión. El sistema es agnóstico a la arquitectura a nivel del controlador, basándose en una abstracción de capa modular para desacoplar el entorno de RL de las topologías de red específicas.

2.1 Estructura de Agentes Jerárquicos

El marco emplea dos niveles distintos de agentes:

• Agentes de Alto Nivel (HLAs): Estos agentes operan al nivel de la red para coordinar la asignación global del presupuesto. Asignan un "Presupuesto de Capa" (LayerBudget) a cada bloque compresible, determinando las fracciones de tamaño retenido objetivo ($R_i$), los límites de poda y los límites de ancho de bits. Los HLAs utilizan estimaciones de sensibilidad basadas en la Información de Fisher para guiar la asignación del presupuesto, asegurando que los bloques con alta sensibilidad reciban una compresión más ligera.
• Agentes de Bajo Nivel (LLAs): Operando de forma independiente por bloque, los LLAs refinan los presupuestos globales en configuraciones específicas por kernel. Seleccionan de un espacio de acción multidiscreto de cuatro dimensiones para cada kernel:
  1. Ancho de bits ($b_{i,k}$): Selección de niveles de precisión.
  2. Ratio de retención ($\rho_{i,k}$): Determinación de la fracción de canales de salida a retener (poda estructurada).
  3. Tipo de cuantización ($\tau_{i,k}$): Elección entre Entero (INT) y Punto Flotante (FLOAT).
  4. Granularidad ($\mu_{i,k}$): Selección entre esquemas uniformes, de dominio logarítmico, por canal o de magnitud-signo aprendida.

2.2 Aprendizaje de Conjunto y Consenso

Para mitigar la varianza inherente al RL en espacios discretos grandes, HiReLC utiliza un enfoque de conjunto (ensemble).

• Conjunto de LLAs: Múltiples políticas paralelas (combinando PPO y A2C) se entrenan con pesos de recompensa heterogéneos (variando el énfasis entre precisión vs. compresión). Se deriva una acción de consenso mediante votación ponderada, mezclando estrategias conservadoras, moderadas y agresivas para producir una configuración robusta.
• Conjunto de HLAs: De manera similar, un conjunto de HLAs propone asignaciones de presupuesto, determinando el presupuesto final mediante votación por mayoría o agregación de media aritmética.

2.3 Aprendizaje Activo y Modelado de Sustitutos (Surrogate Modeling)

Para abordar el costo computacional de evaluar cada configuración candidata mediante un ajuste fino completo post-compresión, HiReLC integra un bucle de aprendizaje activo iterativo:

• Proxy de Logit-MSE: Durante la fase de "arranque en frío" (antes de recolectar suficientes datos), se utiliza un proxy ligero basado en el Error Cuadrático Medio (MSE) entre los logits del modelo comprimido y el modelo base para estimar la retención de la precisión.
• Sustituto MLP: Una vez que se evalúa un número mínimo de configuraciones, se entrena un sustituto de Perceptrón Multicapa (MLP) ligero basado en los datos recolectados para predecir la precisión. Este sustituto se utiliza para el modelado de la recompensa (reward shaping) durante la optimización de RL, mientras que el ajuste fino completo se reserva para las evaluaciones finales de ciclo.
• Bucle Iterativo: El proceso alterna entre la optimización de la política de RL, la aplicación de la configuración, el ajuste fino y las actualizaciones del sustituto.

2.4 Formulación de la Recompensa

Las funciones de recompensa están diseñadas para equilibrar la precisión, la compresión y la estabilidad:

• Recompensas de LLA: Incluyen términos para la retención de precisión ($R_{acc}$), ganancia de compresión ($R_{comp}$), cumplimiento del presupuesto ($R_{budget}$) y estabilidad ($R_{stab}$). Se aplica una penalización de sensibilidad ($P_{sens}$) para desalentar la compresión agresiva en bloques altamente sensibles.
• Recompensas de HLA: Se centran en equilibrar la heterogeneidad entre bloques ($R_{bal}$), la asignación consciente de la sensibilidad ($R_{sens}$), la adherencia al objetivo global ($R_{global}$) y la retroalimentación dinámica basada en el rendimiento del ciclo anterior ($R_{dyn}$).

3. Contribuciones Clave

El artículo describe tres contribuciones primarias:

1. Marco de Trabajo HiReLC: Un marco de RL de conjunto jerárquico agnóstico a la arquitectura para la optimización conjunta de la poda estructurada y la cuantización de precisión mixta. Logra ratios de compresión de almacenamiento de parámetros efectivos de 5.99–6.72× en cuatro arquitecturas (DeiT-Small, CLIP ViT-B/32, ResNet18, MobileNetV2) con caídas de precisión que oscilan entre 0.55% y 5.62%, e incluso una ganancia de precisión del 3.83% en un escenario determinado.
2. Guía Consciente de la Sensibilidad: La integración de puntuaciones de sensibilidad de Información de Fisher en ambos niveles jerárquicos. Esto asegura que los bloques críticos reciban una compresión más ligera sin requerir evaluaciones adicionales durante la búsqueda, propagando la importancia de los bloques basada en gradientes a través de la formación de recompensas y la corrección de acciones.
3. Aprendizaje Activo Aumentado por Sustitutos: Un bucle iterativo que intercala la optimización de RL con el ajuste fino post-compresión. Al utilizar un sustituto MLP ligero y un proxy de logit-MSE, el marco amortiza las evaluaciones costosas, permitiendo estudios prácticos a nivel de marco de trabajo dentro de un número limitado de ciclos.

4. Resultados Experimentales

Los experimentos se realizaron en CIFAR-10/100 y Tiny ImageNet utilizando Vision Transformers y CNNs.

• Rendimiento de Compresión: El marco logró consistentemente ratios de compresión entre 5.99× y 6.72× (equivalente a una reducción del tamaño del modelo del 83.3%–85.1%).
• Impacto en la Precisión:
  • CLIP ViT-B/32 en CIFAR-100: Logró la caída más pequeña de 0.55% con 6.63× de compresión.
  • CLIP ViT-B/32 en CIFAR-10: Demostró una ganancia de precisión del 3.83% (de 89.22% a 93.05%), atribuido a que la compresión actúa como un regularizador para modelos sobre-parametrizados.
  • ResNet18 y MobileNetV2 en Tiny ImageNet: Mostraron caídas de 4.22% y 5.62% respectivamente, reflejando los desafíos de comprimir arquitecturas ya compactas y más pequeñas en conjuntos de datos más difíciles.
• Estudios de Ablación:
  • Eliminar el HLA (usando presupuestos uniformes) provocó una sobre-compresión y una caída de precisión significativamente mayor (6.80% frente a 1.72% para el sistema completo).
  • El uso de un conjunto homogéneo de A2C resultó en compromisos menos estables en comparación con el conjunto mixto de PPO/A2C.
  • La guía de sensibilidad de Fisher mejoró la retención de precisión en 1.30 puntos porcentuales y produjo asignaciones de presupuesto por bloque más diferenciadas.
  • Precisión Mixta: Los experimentos que permitieron la selección por kernel de tipos INT/FLOAT y la granularidad de escala logarítmica resultaron en menores caídas de precisión en comparación con los ajustes uniformes de INT, sugiriendo que la flexibilidad ayuda a la recuperación.

5. Significado y Reivindicaciones

El artículo posiciona a HiReLC como una elección de diseño práctica para la compresión conjunta de redes neuronales, enfatizando los siguientes puntos:

• Espacio de Búsqueda Unificado: A diferencia de métodos previos que apuntan a ejes de compresión únicos (poda o cuantización) o dependen de heurísticas estáticas, HiReLC optimiza conjuntamente ambos, aprovechando las sinergias entre la poda estructurada y la cuantización de precisión mixta.
• Agnóstico a la Arquitectura: El marco comprime con éxito diversas familias de modelos (Transformers, CLIP, ResNets, MobileNets) sin requerir el rediseño del controlador específico de la arquitectura.
• Conciencia de la Sensibilidad: El uso de la Información de Fisher permite al sistema identificar y proteger automáticamente los componentes sensibles de la red, mejorando el compromiso entre precisión y compresión.
• Practicidad: Al utilizar modelos sustitutos y aprendizaje activo, el marco reduce la carga computacional del RL, haciendo que sea factible explorar grandes espacios de configuración sin necesidad de un ajuste fino exhaustivo de cada candidato.

Los autores señalan que, si bien los ratios de compresión alcanzados son significativos, los resultados dependen del contexto (por ejemplo, evaluados en CIFAR/Tiny ImageNet en lugar de ImageNet-1K en las tablas principales) y asumen un almacenamiento estructurado empaquetado. El marco se presenta como un punto de partida para escenarios de despliegue donde el tamaño del modelo es la restricción principal, con trabajo futuro necesario para optimizar las métricas de latencia y energía específicas del hardware.