Layerwise Progressive Freezing: A Training Scaffold for Depth-Scalable Binary Networks

Resumen Técnico: Congelación Progresiva por Capas para Redes Binarias Escalables en Profundidad

Declaración del Problema

Entrenar Redes Neuronales Binarias (BNN), donde tanto los pesos como las activaciones están restringidos a $\{-1, +1\}$, desde cero sigue siendo un desafío, particularmente a medida que aumenta la profundidad de la red. El enfoque de entrenamiento dominante se basa en el Estimador de Paso Directo (STE, por sus siglas en inglés), que sustituye una función de signo no diferenciable en el pase hacia adelante con un gradiente sustituto (típicamente la identidad) en el pase hacia atrás. Aunque es efectivo para redes poco profundas, el STE sufre un desajuste fundamental entre el pase hacia adelante y el de retroceso. A medida que las redes se profundizan, este desajuste conduce a una degradación severa de la precisión.

La literatura existente intenta mitigar esto principalmente refinando el propio STE (por ejemplo, mejores gradientes sustitutos, modelado de gradientes o escalado aprendido). Sin embargo, estos métodos mantienen la estructura central de un operador de pase hacia adelante rígido emparejado con un pase hacia atrás sustituto. Este artículo identifica un eje de falla diferente: bloqueos de gradiente inducidos por la activación. En las BNN completas, una vez que una activación se "congela" (se compromete con un valor binario rígido), su derivada se vuelve cero en casi todas partes. Si las activaciones se congelan en profundidades arbitrarias durante el entrenamiento, el flujo de gradiente hacia las capas anteriores se corta, impidiendo un aprendizaje efectivo en arquitecturas profundas. Este problema es específico de las BNN completas; las Redes de Pesos Binarios (BWN), donde las activaciones permanecen continuas, no sufren este bloqueo.

Metodología: StoMPP

Los autores introducen StoMPP (Stochastic Masked Partial Progressive Binarization - Binarización Progresiva Parcial con Máscara Estocástica), un andamiaje de entrenamiento que controla cuándo y dónde se impone la binarización, en lugar de mejorar el estimador de gradiente en sí mismo.

Mecanismo Central

StoMPP mantiene variables de valores reales subyacentes para los pesos y las pre-activaciones. Una máscara binaria determina si cada entrada es tratada como discreta (congelada) o continua (no congelada). El método emplea dos estrategias clave:

1. Programación Progresiva por Capas: A diferencia de la cuantización progresiva global (donde los parámetros se congelan aleatoriamente en toda la red), StoMPP congela las capas secuencialmente desde la entrada hacia la salida.

   • Prefijo Congelado: Las capas $1$a$\ell-1$ están totalmente binarizadas.
   • Capa de Transición: La capa $\ell$ está parcialmente congelada, evolucionando de continua a binaria.
   • Sufijo No Congelado: Las capas $\ell+1$ a $N$ permanecen totalmente continuas.
   • Racional: Esto asegura que la capa que actualmente está experimentando la binarización siempre tenga un sufijo descendente no congelado, preservando un camino de gradiente válido hacia la función de pérdida. Esto evita los bloqueos de gradiente causados por el congelamiento temprano de las activaciones descendentes.

2. Congelación Parcial con Máscara Estocástica y Refresco Suave: Dentro de la capa de transición, la máscara no se establece de forma determinista. En su lugar, una fracción $p(\tau)$ de las entradas se congela según un esquema (por ejemplo, cúbico). Crucialmente, el conjunto congelado se "refresca suavemente" en cada paso: un pequeño subconjunto de índices ($1/r$) se vuelve a muestrear según una distribución de Bernoulli. Esto evita el compromiso prematuro con una configuración congelada específica y permite que el modelo se adapte al patrón binario emergente.

Manejo de Gradientes

El método es agnóstico al tratamiento del gradiente de las entradas congeladas, admitiendo dos regímenes:

• StoMPP libre de STE: Las entradas congeladas reciben un gradiente cero ($Surr(u) = 0$). Esto proporciona un procedimiento de entrenamiento completamente libre de STE.
• StoMPP+STE: Las entradas congeladas reciben un gradiente de identidad estándar de STE ($Surr(u) = 1$). Esto combina la estructura de progresión con gradientes sustitutos.

Los autores realizan estudios de ablación en el régimen libre de STE para aislar la contribución del orden de progresión de los beneficios de los gradientes sustitutos.

Contribuciones Clave

1. Marco de Trabajo StoMPP: Un procedimiento de congelación progresiva por capas con máscara estocástica y refresco suave. Soporta tanto el entrenamiento totalmente libre de STE como el entrenamiento híbrido con STE en las entradas congeladas.
2. Papel Decisivo del Orden de Progresión: El artículo establece que, para las BNN completas, el orden de la binarización es crítico.
   • Progresión hacia Adelante (Entrada $\to$ Salida): Tiene éxito en redes profundas.
   • Progresión hacia Atrás (Salida $\to$ Entrada): Causa un colapso catastrófico hacia un rendimiento cercano al azar.
   • Orden Global: Se degrada significamente con la profundidad.
   • Insensibilidad de BWN: Las Redes de Pesos Binarios son mayormente insensibles al orden, confirmando que el modo de falla es específico a las activaciones binarias que bloquean los gradientes.
3. Análisis de la Dinámica de Entrenamiento: Los autores caracterizan la dinámica de entrenamiento, observando un patrón de convergencia de "diente de sierra" donde la precisión cae cuando una nueva capa entra en la fase de congelación y se recupera a medida que el sufijo no congelado se adapta. Demuestran que la programación por capas es robusta al tipo de esquema y a las tasas de refresco, a diferencia de la programación global.

Resultados Experimentales

La evaluación utiliza un protocolo controlado con una receta de entrenamiento mínima (tasa de aprendizaje constante, sin decaimiento de pesos) para aislar el efecto de la regla de binarización.

• Escalabilidad de Profundidad: En BNN de ResNet-50, el rendimiento de vanilla STE cae significativamente a medida que aumenta la profundidad (por ejemplo, de 41.9% en ResNet-18 a 30.8% en ResNet-50 en ImageNet).
  • StoMPP (Libre de STE): Mejora sobre vanilla STE en +18.0% (CIFAR-10), +13.5% (CIFAR-100) y +3.8% (ImageNet) en ResNet-50.
  • StoMPP+STE: Logra ganancias aún mayores, mejorando sobre vanilla STE en +27.1% (CIFAR-10), +19.8% (CIFAR-100) y +17.7% (ImageNet) en ResNet-50.
• Generalización: El patrón se mantiene en ResNet-18/34, MobileNetV2 y el ajuste fino de BERT en SST-2. Notablemente, en SST-2, vanilla STE colapsa a casi el azar (50.9%), mientras que StoMPP+STE alcanza el 82.8%.
• Ablaciones:
  • Orden: El orden de progresión hacia atrás causa el colapso de las BNN (por ejemplo, 8.6% en CIFAR-100 ResNet-50), mientras que el orden hacia adelante mantiene el rendimiento.
  • Hiperparámetros: El método es robusto a la elección del esquema (cúbico, lineal, etc.) y la tasa de refresco, siempre que se utilice una progresión por capas.
  • Arquitecturas: StoMPP se transfiere a arquitecturas específicas de binarización como Bi-Real Net sin modificaciones, aunque las interacciones de ajuste fino varían.

Significancia y Reivindicaciones

El artículo sostiene que el cuello de botella principal en el entrenamiento de BNN profundas no es la calidad del gradiente sustituto, sino el momento y la ubicación de la binarización. Al imponer una progresión hacia adelante por capas, StoMPP evita la formación de bloqueos de gradiente inducidos por la activación que cortan las señales de aprendizaje en redes profundas.

Los autores enfatizan que este es un enfoque ortogonal a la mejora del STE; cambia el andamiaje de entrenamiento en lugar del estimador de gradiente. Los resultados sugieren que el "colapso de profundidad" observado en las BNN es una consecuencia estructural de los bloqueos de gradiente, que pueden mitigarse asegurando un camino de gradiente continuo a través de las capas descendentes no congeladas durante el proceso de binarización. El método se presenta como un procedimiento de entrenamiento que puede combinarse con mejoras arquitectónicas existentes o técnicas de gradiente sustituto, ofreciendo un camino hacia redes profundas totalmente binarias y escalables.