STARFISH: faST Accuracy Recovery in pruned networks From Internal State Healing

Resumen Técnico: STARFISH

Planteamiento del Problema
La poda de redes neuronales es una técnica estándar para reducir el tamaño del modelo, la huella de memoria y la latencia de inferencia mediante la eliminación de pesos, neuronas o canales. Sin embargo, la poda agresiva a menudo conduce a una degradación sustancial en la precisión predictiva. Aunque se emplean procesos de recuperación (curación o healing) para recuperar el rendimiento perdido, los métodos existentes enfrentan limitaciones significativas. Los enfoques tradicionales suelen requerir un reentrenamiento completo o un ajuste fino (fine-tuning) utilizando el conjunto de datos de entrenamiento original, lo cual es computacionalmente costoso y frecuentemente imposible en entornos modernos donde los datos de entrenamiento son propietarios o inaccesibles. Además, los métodos de recuperación de vanguardia suelen tener dificultades para restaurar la precisión bajo condiciones de alta dispersión (por ejemplo, eliminar >75% de los pesos), particularmente para las arquitecturas modernas de Vision Transformer (ViT).

Metodología: STARFISH
Los autores proponen STARFISH (faST Accuracy Recovery in pruned networks From Internal State Healing), un método de recuperación post-poda que opera sin acceso a los datos de entrenamiento originales ni a las etiquetas. La hipótesis central es que optimizar la alineación de las representaciones intermedias (estados internos) entre la red podada y la red densa original es un proxy más efectivo para la restauración de la precisión que optimizar únicamente sobre las predicciones de salida finales.

El algoritmo procede en dos etapas utilizando un pequeño conjunto de calibración no etiquetado ($S_{cal}$) muestreado de la distribución de los datos pero distinto al conjunto de entrenamiento:

1. Almacenamiento de Representaciones Objetivo: Se ejecuta el modelo denso original sobre $S_{cal}$ para almacenar en caché las salidas vectorizadas de cada bloque computacional (por ejemplo, después de la normalización, las no linealidades y las adiciones residuales).
2. Optimización de la Alineación: El modelo podado ($f_{\theta \odot m}$) se entrena para minimizar la pérdida de alineación entre sus representaciones intermedias y los objetivos almacenados en caché. La función de pérdida se define como la distancia de coseno promedio entre los estados ocultos del modelo podado ($f^\ell_{\theta \odot m}(x)$) y los estados ocultos del modelo denso ($r^\ell_x$) a través de todos los bloques $\ell$:
   $$\mathcal{L}_{align}(\theta \odot m, x) = \frac{1}{L} \sum_{\ell \in [L]} \left( 1 - \frac{\langle f^\ell_{\theta \odot m}(x), r^\ell_x \rangle}{\|f^\ell_{\theta \odot m}(x)\| \|r^\ell_x\|} \right)$$
   Los pesos restantes del modelo podado se actualizan para minimizar la pérdida empírica sobre el conjunto de calibración.

Se proporciona una motivación teórica demostrando que la distancia entre las representaciones ocultas finales (entradas al cabezal de clasificación) sirve como un límite superior para la divergencia de Kullback–Leibler (KL) de las distribuciones de salida. Esto sugiere que alinear los estados internos restringe inherentemente el error de salida.

Contribuciones Clave

1. Recuperación Basada en Representaciones: El artículo introduce un método de recuperación independiente de los tipos de poda, máscaras o arquitecturas específicos. Demuestra empíricamente que la alineación de la representación intermedia supera a los objetivos directos basados en la salida (como el Cross-Entropy o la coincidencia de la divergencia KL) a medida que aumentan los niveles de dispersión.
2. Eficiencia en Alta Dispersión: El método es el primero en restaurar eficientemente altos porcentajes de la precisión original en niveles de dispersión extrema (por ejemplo, 85% de dispersión). Logra esto utilizando un conjunto de calibración diminuto (tan solo 1,000 ejemplos, o el 0.4% del tamaño del conjunto de entrenamiento) y no requiere acceso a los datos de entrenamiento originales.
3. Modularidad y Velocidad: El enfoque es computacionalmente ligero, requiriendo solo unas pocas pasadas hacia adelante (forward) y hacia atrás (backward) en un conjunto de datos pequeño. Es aplicable a patrones de poda estructurados, semiestructurados y no estructurados.

Resultados Experimentales
Los autores evaluaron STARFISH en ImageNet-1K utilizando diversas arquitecturas de Vision Transformer (DeiT, ViT) y Redes Neuronales Convolucionales (MobileNet).

• Alta Dispersión (0.75–0.85): En un modelo DeiT-B podado a un 75% de dispersión, STARFISH recuperó el 82% de la precisión original del modelo denso, mientras que los métodos competidores (como CORP) recuperaron solo el 40%. En un modelo DeiT3-H al 85% de dispersión, STARFISH retuvo el 92.08% de la precisión del modelo denso, mejorando el rendimiento del modelo podado en 72.36 puntos porcentuales.
• Dispersión Moderada (0.5): A través de varias arquitecturas DeiT y alcances de poda (MLP, QK, y conjunta MLP+QK), STARFISH superó consistentemente a los métodos de vanguardia. Por ejemplo, en un modelo DeiT-T con poda estructurada del 50% en MLP+QK, STARFISH logró una retención del 81.95% de la precisión densa, una mejora de 31.54 puntos porcentuales sobre el método CORP.
• Comparación: STARFISH demostró un rendimiento superior frente a métodos como SNOWS (que utiliza la aproximación de la Hessiana) y CORP (que utiliza la recuperación afín de forma cerrada), particularmente en regímenes donde la poda degrada severamente el modelo.

Significancia y Reivindicaciones
El artículo afirma que STARFISH representa un avance significativo en el campo de la compresión de modelos al desplazar el paradigma de recuperación del emparejamiento de salidas hacia la curación de estados internos. Los autores enfatizan que este enfoque es particularmente vital para los modelos modernos a gran escala donde los datos de entrenamiento no están disponibles y el reentrenamiento completo es inviable. Al aprovechar el principio de que las representaciones internas capturan el conocimiento del modelo de manera más robusta que las salidas finales bajo alta dispersión, STARFISH permite el despliegue de redes podadas altamente eficientes sin sacrificar el rendimiento predictivo. El método se presenta como una solución general de tipo plug-and-play que puede aplicarse a cualquier red podada independientemente de la estrategia de poda inicial.