Resumen Técnico: Regularización de Concentración de Pesos para Mejorar la Robustez de la Poda bajo Alta Esparsidad

Declaración del Problema

Las redes neuronales profundas (DNN) han logrado un rendimiento de vanguardia en tareas de visión y lenguaje, pero a menudo requieren recursos computacionales sustanciales, lo que limita su implementación en entornos con recursos limitados. La poda de modelos tiene como objetivo reducir el tamaño del modelo eliminando parámetros menos importantes. Si bien la poda iterativa con reentrenamiento es efectiva, resulta computacionalmente costosa. La poda de un solo paso (one-shot pruning), que elimina parámetros redundantes en una sola etapa sin reentrenamiento, ofrece una alternativa más eficiente. Sin embargo, los modelos entrenados con objetivos estándar suelen sufrir caídas significativas de precisión bajo esparsidad agresiva.

Trabajos anteriores intentan mitigar esto a través de dos direcciones principales:

1. Enfoques basados en regularización (por ejemplo, $\ell_1$, DeepHoyer) que moldean las distribuciones de pesos durante el entrenamiento. Sin embargo, $\ell_1$ reduce todos los pesos uniformemente, mientras que DeepHoyer induce una esparsidad invariante a la escala. Ninguno de los dos concentra explícitamente la energía de los pesos en un pequeño subconjunto de parámetros informativos.
2. Optimizadores robustos a la poda (por ejemplo, SAM, CrAM, S2SAM) que aplanan el paisaje de pérdida para mejorar la generalización.

La limitación central identificada es que los métodos estándar que inducen esparsidad no abordan explícitamente la compensación entre esparsidad y preservación de información. Cuando la importancia de los pesos está distribuida de manera difusa, eliminar pesos de pequeña magnitud aún puede degradar el rendimiento, ya que la energía "importante" no está suficientemente concentrada en unos pocos parámetros dominantes.

Metodología: Regularizador de Concentración de Pesos (WCR)

Los autores proponen el Regularizador de Concentración de Pesos (WCR), un regularizador diseñado para el tiempo de entrenamiento que remodela la distribución de la magnitud de los pesos. A diferencia de los métodos convencionales que reducen uniformemente los pesos, WCR busca concentrar la energía de los pesos en un pequeño subconjunto de parámetros, mientras impulsa el resto hacia cero. Esto crea una distribución de "cola pesada" donde unos pocos parámetros albergan la mayor parte de la energía de los pesos, permitiendo que la poda por magnitud elimine los parámetros insignificantes con una pérdida mínima de rendimiento.

Formulación Matemática

WCR opera penalizando el recíproco de la varianza de los pesos absolutos. Para una capa seleccionada con pesos $w^{(\ell)}$, los autores primero aplican una transformación suave de valor absoluto para garantizar la diferenciabilidad en el origen:
$$\tilde{w}^{(\ell)}_i = \sqrt{(w^{(\ell)}_i)^2 + r} \approx |w^{(\ell)}_i|, \quad r = 10^{-8}$$
La varianza se calcula por capa para manejar las diferencias de escala:
$$\text{Var}(\tilde{w}^{(\ell)}) = \frac{1}{n_\ell} \sum_{i=1}^{n_\ell} \left( \tilde{w}^{(\ell)}_i - \frac{1}{n_\ell} \sum_{j=1}^{n_\ell} \tilde{w}^{(\ell)}_j \right)^2$$
La penalización de regularización $\psi(w)$ es la suma de las varianzas recíprocas a través de todas las capas seleccionadas:
$$\psi(w) = \sum_{w^{(\ell)} \in \mathcal{W}} \frac{1}{\text{Var}(\tilde{w}^{(\ell)}) + \epsilon}$$
donde $\epsilon = 10^{-8}$ asegura la estabilidad numérica. El objetivo total de entrenamiento se convierte en:
$$\mathcal{L}_{total}(w) = \mathcal{L}(w) + \lambda \psi(w)$$
donde $\lambda > 0$ es un coeficiente de equilibrio.

Análisis Teórico

El artículo proporciona un análisis de convergencia que demuestra que WCR preserva la tasa de convergencia estándar del Descenso de Gradiente Estocástico (SGD). Bajo suavidad $\beta_1$ de la función de pérdida y suavidad $\beta_2$ del regularizador, el objetivo combinado satisface las garantías estándar de convergencia a la estacionariedad. Dado que el regularizador es una función determinista de los pesos, no introduce varianza adicional en el gradiente estocástico. En consecuencia, la tasa de convergencia asintótica permanece en $O(1/\sqrt{T})$, coincidiendo con el SGD estándar.

Contribuciones Clave

1. Introducción de WCR: Un nuevo regularizador para el tiempo de entrenamiento que concentra la energía de los pesos en un pequeño subconjunto de parámetros, fomentando explícitamente una distribución de pesos de cola pesada en lugar de una reducción uniforme.
2. Garantías Teóricas: Demostración de que WCR mantiene la tasa de convergencia estándar del SGD bajo supuestos de suavidad moderados.
3. Validación Empírica: Evaluación exhaustiva en diversos escenarios, incluyendo:
   • Clasificación de Imágenes: ResNet-18/50, WideResNet-28-10 en CIFAR-10/100, SVHN, y ViT-B/32 en Tiny-ImageNet.
   • Segmentación Médica: ResNet-50–UNet en el conjunto de datos de RM cerebral LGG.
   • Ajuste Fino de LLM: Qwen-2.5-1.5B y LLaMA-3.2-1B utilizando LoRA en Commonsense-170K.
4. Compatibilidad: Demostración de que WCR es compatible con y mejora los optimizadores existentes robustos a la poda (SAM, CrAM, S2SAM) y puede combinarse con ellos para obtener mayores beneficios.

Resultados Experimentales

Los autores evaluaron WCR bajo poda por magnitud de un solo paso sin reentrenamiento.

• Tareas de Visión (CNNs y ViTs): WCR mejoró consistentemente la precisión posterior a la poda, particularmente en niveles altos de esparsidad (por ejemplo, 96% de poda).
  • En CIFAR-10 con ResNet-18, añadir WCR a SGD mejoró la precisión al 96% de poda del 19.69% al 25.70%.
  • Cuando se combinó con SAM, WCR produjo mejoras dramáticas: al 96% de poda, la precisión aumentó del 48.36% (solo SAM) al 92.08% (SAM + WCR).
  • Se observaron tendencias similares con S2SAM (72.51% $\to$ 93.44%) y en SVHN (9.74% $\to$ 95.99%).
  • Para ViT-B/32, WCR mejoró la precisión al 80% de poda en CIFAR-100 del 64.11% (SAM) al 84.95%.
• Segmentación Médica: En el conjunto de datos de RM cerebral LGG, WCR mejoró significativamente las puntuaciones F1 y redujo la distancia de Hausdorff bajo podas del 85% y 88%. Cabe destacar que, para SGD al 88% de poda, la línea base falló en producir segmentaciones válidas, mientras que SGD + WCR logró una puntuación F1 de 0.5280. Los resultados cualitativos mostraron que WCR preservó mejor los límites de las lesiones en comparación con las líneas base.
• Ajuste Fino de LLM: En el ajuste fino basado en LoRA, WCR superó a los regularizadores $\ell_1$ y DeepHoyer. En Qwen-2.5-1.5B, WCR logró la mayor precisión promedio al 30% de poda (69.54%) entre los regularizadores comparados, demostrando robustez en escenarios de ajuste fino eficiente en parámetros.
• Generalización: El análisis del valor propio principal de Hessian indicó que WCR reduce la curvatura de la superficie de pérdida, sugiriendo que guía la optimización hacia mínimos más planos, lo cual se correlaciona con una mejor generalización.

Significado y Afirmaciones

El artículo afirma que moldear la distribución de la magnitud de los pesos durante el entrenamiento es un enfoque efectivo para mejorar la robustez bajo alta esparsidad. El significado de WCR radica en su capacidad para:

• Crear una separación más clara entre parámetros de gran magnitud (informativos) y parámetros de pequeña magnitud (insignificantes).
• Permitir que la poda basada en magnitud elimine una gran fracción de parámetros con un impacto funcional mínimo.
• Operar como un regularizador independiente o como un componente complementario a métodos de suavizado del paisaje de pérdida como SAM.

Los autores concluyen que WCR ofrece una solución práctica para el tiempo de entrenamiento a la caída de precisión asociada con la poda agresiva de un solo paso, aplicable a CNNs, Transformers y Modelos de Lenguaje Grandes, sin requerir cambios arquitectónicos ni reentrenamiento después de la poda.