SpikeDecoder: Realizing the GPT Architecture with Spiking Neural Networks

Resumen Técnico: SpikeDecoder

Planteamiento del Problema
La arquitectura Transformer, particularmente los modelos GPT de solo decodificador, se ha establecido como el paradigma dominante para el procesamiento del lenguaje natural (NLP). Sin embargo, estos modelos dependen de Redes Neuronales Artificiales (ANN) que realizan operaciones de matrices complejas y de alto consumo energético (específicamente, operaciones de Multiplicación-Acumulación o MAC) durante la inferencia. Esto resulta en un alto consumo de energía, lo que limita su despliegue en dispositivos de borde (edge devices) o entornos con restricciones energéticas. Aunque las Redes Neuronales de Impulsos (SNN) ofrecen una alternativa prometedora y eficiente desde el punto de vista energético debido a su naturaleza dirigida por eventos y binaria, son notoriamente difíciles de entrenar directamente. Las adaptaciones de Transformers basadas en SNN existentes se han limitado en gran medida a tareas de visión computacional (por ejemplo, Spikformer) o arquitecturas de solo codificador. Además, los pocos intentos en NLP (por ejemplo, SpikeGPT) han dependido de estructuras de Redes Neuronales Recurrentes (RNN) como RWKV en lugar de la autoatención, o han utilizado la conversión de ANN a SNN en lugar del entrenamiento directo. Actualmente, no existe una implementación de bloque de decodificador Transformer de solo impulsos (spiking-only) y totalmente entrenable para NLP.

Metodología
Los autores proponen SpikeDecoder, una implementación completamente de impulsos del bloque decodificador de Transformer diseñada para la generación de lenguaje natural. El proceso de desarrollo consistió en una transición sistemática de un decodificador estándar basado en ANN a una arquitectura puramente de impulsos, analizando el impacto de cada cambio de componente.

1. Modelo de Neurona: La arquitectura utiliza el modelo de neurona de Leaky Integrate-and-Fire (LIF) de múltiples pasos, implementado mediante la librería SpikingJelly. Este modelo integra entradas a través de pasos de tiempo, aplica un umbral y realiza un reinicio, imitando la dinámica de las neuronas biológicas.
2. Estrategias de Embedding: El artículo evalúa varios métodos para proyectar caracteres de texto en formatos compatibles con impulsos. Las opciones incluyeron one-hot, binarios, estáticos (seno/coseno) y embeddings aprendidos. Los autores descubrieron que concatenar los embeddings de tokens aprendidos con la codificación posicional estática ofrecía el mejor equilibrio entre rendimiento y eficiencia de parámetros, aunque señalan que la codificación posicional estática requiere un manejo específico (por ejemplo, tablas de búsqueda o mapeo de Heaviside) para ajustarse a las restricciones binarias.
3. Conexiones Residuales: Un desafío crítico identificado fue la introducción de valores de punto flotante al sumar conexiones residuales entre capas de impulsos binarios. Los autores adoptaron una estrategia de reestructuración (inspirada en Zhou et al., 2023) donde las neuronas LIF se colocan inmediatamente después de los operadores de suma (tanto en el bloque de Multi-Head Self-Attention como en el MLP). Esto asegura que las entradas a las capas lineales subsiguientes sigan siendo binarias, eliminando las multiplicaciones de punto flotante causadas por los residuales.
4. Normalización: Se encontró que la Normalización de Capa (LN) estándar no es adecuada para modelos totalmente de impulsos debido a su cómputo dinámico de media y varianza durante la inferencia. Los autores evaluaron la Normalización de Lote (BN) y la Normalización de Potencia (PN). Mientras que la BN degradó el rendimiento, la Normalización de Potencia (PN) fue identificada como una alternativa viable que permite la fusión de capas lineales y de normalización durante el despliegue, aunque incurre en mayores tiempos de entrenamiento y uso de memoria.
5. Unificación de Pasos de Tiempo: Para manejar la dimensión temporal inherente a las SNN, los autores compararon cuatro estrategias para agregar las salidas de los pasos de tiempo antes del cabezal de clasificación: Promedio (Average), Acumulación (Accumulate), Concatenación (Concatenate) y Final (Final). Determinaron que la Concatenación de impulsos era la más adecuada para una arquitectura puramente de impulsos, ya que preserva la información y produce salidas compatibles con impulsos para la capa lineal final.
6. Transición Híbrida: Los autores introdujeron un marco de "Grado de Impulso" (Spike Degree) (0 a 4) para analizar la transición de ANN a SNN. Esto les permitió aislar las caídas de rendimiento en etapas específicas: conversión de embedding, conversión del mecanismo de atención, conversión de MLP y conversión del cabezal de clasificación.

Contribuciones Clave

• Arquitectura SpikeDecoder: La primera implementación de solo impulsos y totalmente entrenable de un bloque decodificador de Transformer para NLP, yendo más allá de los Transformers de solo codificador o basados en visión.
• Análisis de Transición Sistemático: Un desglose detallado de las compensaciones de rendimiento al convertir componentes específicos del Transformer (embedding, atención, MLP, residuales) a equivalentes de impulsos.
• Optimización de Componentes:
  • Validación de la Normalización de Potencia como una alternativa viable a la Normalización de Capa para Transformers totalmente de impulsos.
  • Demostración de que los umbrales de LIF aprendibles pueden mejorar el rendimiento (aumentando la precisión de ~78% a ~89% en configuraciones específicas), aunque señalaron inconsistencias biológicas (umbrales negativos) en algunas capas.
  • Identificación de la Concatenación como la estrategia óptima de unificación de pasos de tiempo para mantener la compatibilidad de impulsos.
• Análisis de Eficiencia Energética: Una estimación teórica del consumo de energía comparando el SpikeDecoder con su línea base ANN.

Resultados

• Rendimiento: El SpikeDecoder totalmente de impulsos (Grado de Impulso 4) logró una precisión del 87.0% en una tarea de predicción a nivel de carácter utilizando un subconjunto de "Guerra y Paz", en comparación con el 98.5% de la línea base ANN no impulsiva. El proceso de transición reveló que la conversión de la Autoatención de Múltiples Cabezales (SMHA) causó la mayor caída de rendimiento (aprox. 25%), la cual fue mitigada parcialmente aumentando los bloques del decodificador y reduciendo los cabezales de atención.
• Consumo de Energía: El análisis teórico estimó que el SpikeDecoder reduce el consumo de energía entre un 86.7% y un 92.3% en comparación con la línea base ANN. Esta reducción se atribuye al reemplazo de las operaciones MAC por Operaciones Sinápticas (SOPS) y la eliminación de la aritmética de punto flotante en los bloques de procesamiento centrales.
• Estudios de Ablación:
  • Eliminar las sumas de punto flotante en las conexiones residuales (haciéndolas puramente binarias) causó una caída catastrófica en el rendimiento (al 18.1%), lo que indica que los gradientes de impulsos actuales dependen fuertemente de los residuales de punto flotante para una retropropagación estable.
  • El uso de la Normalización de Potencia en lugar de la Normalización de Capa resultó en una precisión estable del 81.7% para la variante totalmente de impulsos, confirmando su viabilidad a pesar de los mayores tiempos de entrenamiento.

Significado y Reivindicaciones
El artículo posiciona a SpikeDecoder no como un modelo de lenguaje de vanguardia que supere a GPT en capacidad, sino como una prueba de concepto funcional y un punto de partida estable para el desarrollo de modelos Transformer de impulsos entrenables directamente. Los autores afirman modestamente que su trabajo cierra la breancia entre la poderosa arquitectura Transformer de solo decodificador y la eficiencia energética de la computación neuromórfica. Enfatizan que, si bien el rendimiento es actualmente inferior al de sus contrapartes ANN, la arquitectura demuestra con éxito que es factible un decodificador totalmente de impulsos y entrenable directamente. El trabajo proporciona un marco fundacional para la investigación futura sobre la optimización de modelos de lenguaje de impulsos, la exploración de conjuntos de datos más grandes y el refinamiento de las estrategias de codificación de impulsos.