Sharpness Aware Surrogate Training for Spiking Neural Networks

1. Problema: La Brecha de Transferencia en SNNs

Las Redes Neuronales de Spikes (SNN) son prometedoras para hardware neuromórfico debido a su computación esparsa y basada en eventos. Sin embargo, su entrenamiento supervisado es difícil porque la función de disparo (spike) es no diferenciable.

• Solución actual: Se utilizan gradientes de sustituto (surrogate gradients), donde se reemplaza la derivada no suave de la función de disparo por una aproximación suave durante la retropropagación.
• El desafío principal: Existe una brecha de transferencia (transfer gap) significativa entre el rendimiento durante el entrenamiento (usando el modelo de sustituto suave) y el despliegue (usando el modelo de umbral duro o "hard spike").
  • Durante el entrenamiento, las activaciones son graduales.
  • En la inferencia, al cambiar a umbral duro, las decisiones de disparo pueden cambiar drásticamente si los potenciales de membrana están cerca del umbral.
  • Esto provoca que un modelo que parece preciso en la evaluación de sustituto tenga un rendimiento muy inferior en la inferencia real de spikes.
• Fragilidad: El entrenamiento de SNNs es sensible al ruido y a las perturbaciones debido a la recurrencia temporal, lo que a menudo lleva a soluciones "puntas" (sharp) en el espacio de parámetros que no generalizan bien.

2. Metodología: SAST (Sharpness Aware Surrogate Training)

El autor propone SAST, una integración de la Minimización Consciente de la Agudeza (SAM) en el entrenamiento de SNNs, pero con un enfoque teórico específico:

• Enfoque del Modelo: A diferencia de los métodos tradicionales que usan un paso forward de "spike duro" con un backward de sustituto (lo cual genera un estimador de gradiente sesgado), SAST entrena un modelo forward de sustituto suave.
  • Se utiliza una función de activación suave (ej. arctan o sigmoide rápida) tanto en el forward como en el backward.
  • Esto permite calcular el gradiente exacto de un objetivo empírico suave, eliminando la discrepancia entre el forward y el backward.
• Aplicación de SAM: Se aplica SAM al objetivo de pérdida suave del modelo de sustituto.
  • El objetivo es minimizar la pérdida en el peor caso dentro de una vecindad de radio $\rho$ alrededor de los parámetros actuales: $\min_w \max_{\|\epsilon\| \le \rho} \tilde{L}(w + \epsilon)$.
  • Esto busca soluciones más planas (flat) en el espacio de parámetros, lo que teóricamente mejora la robustez y la generalización.
• Implementación:
  • Utiliza un algoritmo de dos pasos: primero calcula el gradiente, luego perturba los parámetros en la dirección del gradiente, y finalmente calcula un nuevo gradiente en un segundo minibatch independiente para actualizar los parámetros.
  • Se reinician los estados de la SNN en cada paso de SAM para evitar sesgos de estado.

3. Contribuciones Clave

A. Contribuciones Teóricas

El artículo establece garantías rigurosas bajo supuestos explícitos de acotación y contracción:

1. Estabilidad de Estado y Lipschitz: Se derivan límites compactos para la estabilidad del estado interno y la continuidad Lipschitz de la entrada, demostrando que el modelo de sustituto es estable bajo ciertas condiciones de contracción ($\gamma < 1$).
2. Suavidad del Objetivo: Se prueba que el objetivo empírico del modelo de sustituto es suave ($\beta$-smooth), lo cual es un requisito previo para la convergencia.
3. Convergencia No Convexa: Se proporciona una garantía de convergencia para el entrenamiento estocástico de SAST con un segundo minibatch independiente.
4. Proposición de Mecanismo Local: Se propone una relación local (no universal) que vincula el control de los gradientes de parámetros por muestra con la reducción de las normas de los gradientes de entrada, bajo condiciones de buen acondicionamiento del Jacobiano.

B. Contribuciones Empíricas y Protocolo

• Evaluación Rigurosa: Se define un protocolo que incluye métricas de precisión en sustituto, precisión en spikes duros, la brecha de transferencia, robustez a corrupción y controles de costo computacional.
• Análisis de Brecha de Transferencia: El objetivo principal no es solo mejorar la precisión, sino minimizar la diferencia entre el rendimiento de entrenamiento y el de despliegue.

4. Resultados Experimentales

Los experimentos se realizaron en N-MNIST y DVS Gesture.

• Reducción de la Brecha de Transferencia:
  • N-MNIST: La precisión en "spikes duros" aumentó drásticamente de 65.7% a 94.7% (con $\rho=0.30$), mientras que la precisión del modelo de sustituto se mantuvo alta (~97%). La brecha de transferencia se redujo de 0.3034 a 0.0248 (una reducción relativa del ~91.8%).
  • DVS Gesture: La precisión en "spikes duros" mejoró de 31.8% a 63.3% (con $\rho=0.40$), reduciendo la brecha de 0.4320 a 0.1358.
• Robustez: Los modelos entrenados con SAST mostraron una degradación más suave ante corrupciones (como la caída de eventos o ruido temporal) en comparación con el entrenamiento base.
• Costo Computacional: SAST requiere aproximadamente el doble de tiempo por época (debido a las dos evaluaciones de gradiente), pero la mejora en la transferencia justifica este costo en escenarios de despliegue.
• Observaciones: El beneficio principal no es una mayor precisión en el modelo de sustituto, sino una mayor fidelidad al cambiar al modelo de spikes duros.

5. Significado y Conclusión

• Validación Teórica-Práctica: El trabajo es notable por vincurar directamente la teoría de optimización suave (SAM) con la dinámica específica de las SNNs (recurrencia temporal, reset, umbrales), proporcionando garantías teóricas que a menudo faltan en la literatura de SNN.
• Solución al Problema de Despliegue: SAST aborda directamente el cuello de botella más grande en la adopción de SNNs: la incapacidad de los modelos entrenados con sustitutos para funcionar bien con hardware de spikes reales.
• Limitaciones y Futuro: El autor reconoce que la teoría asume ciertas condiciones (como pooling promedio y normalización afín) que pueden no cumplirse en todas las implementaciones prácticas (ej. MaxPool). Además, se señala que se necesitan más comparaciones con baselines de "costo computacional emparejado" y calibración de umbrales para determinar si SAST es la solución más eficiente en términos de costo-beneficio.

En resumen, SAST demuestra que optimizar un objetivo suave consciente de la agudeza en el espacio de parámetros permite entrenar SNNs que son robustas y transfieren su rendimiento casi perfectamente al hardware de spikes duros, cerrando la brecha que históricamente ha limitado la utilidad práctica de estas redes.