ADSEQ: A delay-aware autograd-compatible framework for spike-event delivery in SNNs

Resumo Técnico: ADSEQ – Um Framework de Filas de Eventos de Spike Compatível com Autograd e Consciente de Atraso para SNNs

Declaração do Problema
A simulação eficiente e o treinamento baseado em gradiente de Redes Neurais Pulsadas (SNNs) e modelos cerebrais exigem o tratamento da transmissão esparsa de spikes, particularmente ao incorporar atrasos sinápticos heterogêneos. As abordagens atuais enfrentam uma dicotomia:

1. Ineficiência de Memória: Simuladores baseados em bibliotecas de Aprendizado de Máquina (ML) frequentemente dependem de gradientes substitutos (ex: SuperSpike) para lidar com as descontinuidades do spike. Embora eficazes para modelos simplificados, os gradientes substitutos convertem eventos de spike esparsos em sinais contínuos e densos. Quando atrasos são introduzidos, esses sinais densos devem ser armazenados em buffers, anulando os benefícios da esparsidade e criando gargalos de memória.
2. Falta de Generalização: Métodos de gradiente exato existentes (ex: EventProp) são amplamente restritos a modelos de neurônios simplificados (como o Leaky Integrate-and-Fire) e conjuntos de parâmetros específicos. Eles frequentemente falham ao generalizar para modelos biofisicamente realistas ou estruturas de fila arbitrárias necessárias para simulações de cérebro completo.
3. Desajuste de Hardware: Há uma falta de compreensão sobre quais estruturas de dados de fila de eventos de spike (ex: ring buffers, heaps, FIFOs) apresentam desempenho otimável através de diversas arquiteturas de aceleradores de IA (CPUs von Neumann, GPUs, TPUs e LPUs de dataflow determinístico), especialmente quando combinadas com operações diferenciáveis.

Metodologia
Os autores introduzem o ADSEQ (AutoDifferentiable Spike-Event Queues), um framework projetado para permitir o cálculo de gradiente consciente de atraso de forma exata dentro de bibliotecas modernas de ML (especificamente o JAX).

• Formulação Matemática: A contribuição central é a derivação de derivadas implícitas para a entrega de eventos de spike. Ao aplicar o Teorema da Função Implícita à condição de detecção de spike ($v_{pre}(t_{pre}) = v_{th}$), os autores derivam o gradiente do tempo de spike em relação aos parâmetros do modelo ($\theta$) e aos atrasos sinápticos ($d$).
  • Eles generalizam as equações do EventProp para incluir gradientes de atraso: $dt_{post}/d\theta = dt_{pre}/d\theta + dd/d\theta$.
  • Eles derivam a atualização do gradiente para variáveis de estado pós-sinápticas ($x$) após a chegada do spike, contabilizando o atraso: $\partial x^+/\partial \theta = \partial x^-/\partial \theta + (1/\tau) \cdot \partial t_{post}/\partial \theta$.
• Implementação via Gradientes Customizados: Para integrar essas derivações ao sistema de diferenciação automática (autograd) do JAX, os autores implementam operadores de Produto Jacobiano-Vetor (JVP) customizados. Esses gradientes customizados são injetados nas operações de enqueue() e pop() de várias estruturas de fila. Isso permite que o mecanismo de autograd propague corretamente os gradientes através dos eventos de spike com atraso temporal sem a necessidade de derivações manuais de backpropagation para modelos de neurônios complexos (ex: Hodgkin-Huxley).
• Variantes de Fila: O framework implementa um conjunto de estruturas de fila diferenciáveis, incluindo:
  • Ring Buffers: Variantes padrão e com perda (somando spikes ou descartando-os).
  • Estruturas FIFO: Filas do tipo First-In-First-Out com limites de capacidade.
  • Estruturas de Árvore/Heap: Heaps binários e árvores de torneio para lidar com atrasos heterogêneos.
  • Arrays Ordenados: Arrays que mantêm a ordem via operações de ordenação.
  • Estruturas de Nível de Bit: Otimizadas para cenários de spike único.
• Mecanismos de Perda (Lossy): O framework inclui implementações de fila "lossy" que deliberadamente descartam spikes quando a capacidade é excedida, permitindo um compromisso ajustável entre desempenho e precisão.

Principais Contribuições

1. Formulação Autodiferenciável: Uma formulação matemática generalizada para gradientes exatos de eventos de spike que suporta tanto parâmetros neurais quanto atrasos sinápticos, compatível com modelos de neurônios arbitrários.
2. Biblioteca ADSEQ: Uma coleção de estruturas de fila diferenciáveis implementadas em JAX com gradientes customizados, permitindo o uso de armazenamento de spike esparso no treinamento baseado em gradiente.
3. Benchmarking Cross-Accelerator: Uma avaliação abrangente dessas estruturas de fila através de quatro arquiteturas de hardware distintas:
   • CPU (Intel i7): Arquitetura von Neumann.
   • GPU (NVIDIA H100): Processamento paralelo com acesso à memória coalescido.
   • TPU (Google TPU v4): Cores de Tensor e Sparse com unidades de ordenação dedicadas.
   • LPU (Groq LPU): Arquitetura de dataflow determinístico sem ramificação de fluxo de controle.
4. Análise de Queda de Spike (Spike Drop): Quantificação das taxas de queda de spike para filas lossy, fornecendo dados para o design de simuladores com restrições de precisão controláveis.

Resultos
Os resultados do benchmarking revelam que a estrutura de fila ideal é altamente dependente da arquitetura do hardware e das características da carga de trabalho:

• CPU: Filas baseadas em Árvore (Binary Heap) e em FIFO apresentam melhor desempenho, provavelmente devido ao menor overhead e ao tratamento eficiente de lógica condicional em arquiteturas von Neumann.
• GPU: Ring buffers se destacam para redes menores e cargas de trabalho que cabem no cache. No entanto, para redes maiores, as estruturas FIFO escalam de forma mais favorável devido à redução da pressão de memória.
• TPU: Implementações baseadas em ordenação (Sorted Arrays) superam significativamente outros métodos. Isso é atribuído aos Sparse Cores dedicados do TPU, que contêm unidades de ordenação aceleradas por hardware.
• LPU: O agendamento de dataflow determinístico favorece estruturas que evitam a divergência de fluxo de controle.
• Escalabilidade: À medida que o número de filas aumenta, o desempenho da CPU degrada rapidamente em comparação aos aceleradores. Contudo, para simulações de pequena escala onde os dados cabem nos caches locais, as CPUs permanecem competitivas.
• Treinamento vs. Inferência: O treinamento (ajuste de parâmetros) é significativamente mais intensivo em termos de computação e memória do que a inferência. O Forward AD é notavelmente mais lento que o Reverse AD em certas configurações, sugerindo a necessidade de métodos híbridos como checkpointing.

Significância e Alegações
O artigo afirma que o ADSEQ estabelece uma base para componentes reutilizáveis e escaláveis para simulações de SNN compatíveis com autograd que incluem atrasos, aplicáveis tanto a redes artificiais quanto biologicamente realistas.

• Generalização: Diferente de métodos anteriores de gradiente exato limitados a modelos LIF, o ADSEQ suporta modelos biofísicos complexos (ex: Hodgkin-Huxley) ao tratar a fila como uma primitiva diferenciável.
• Consciência de Hardware: O trabalho demonstra que estratégias de "tamanho único" para filas são subótimas. A escolha de uma estrutura de dados deve alinhar-se com a arquitetura do acelerador subjacente (ex: usar estruturas baseadas em ordenação para TPUs e heaps para CPUs).
• Integração Futura: Os autores sugerem que sua abordagem poderia permitir que simuladores cerebrais tradicionais (que atualmente carecem de suporte a gradiente) incorporem gradientes de nível de rede ao modificar modelos de neurônios para enviar mensagens de gradiente codificadas em float junto com os spikes, preenchendo a lacuna entre simuladores de neurociência tradicional e frameworks modernos de ML.

O artigo conclui que, embora as bibliotecas de ML frequentemente negligenciem estruturas de dados tradicionais como FIFOs e árvores em favor de arrays simplificados, essas estruturas possuem mérito significativo para cenários específicos de hardware e restrição de atraso, oferecendo um caminho para o treinamento mais eficiente e realista de SNNs.