Technische Samenvatting: SAGE – Surrogate-gradient Adaptatie via Attention-Guided Entropie voor Spiking Transformers

1. Probleemstelling

Spiking Neural Networks (SNNs) bieden een energiezuinig alternatief voor conventionele diepe neurale netwerken door gebruik te maken van ijle (sparse), event-gestuurde berekeningen. Het trainen van diepe Spiking Vision Transformers (ViTs) blijft echter uitdagend vanwege de niet-differentieerbare aard van de spike-generatiefunctie (Heaviside-stapfunctie). Dit vereist het gebruik van surrogaatgradiënten tijdens backpropagation om afgeleiden te benaderen.

Huidige state-of-the-art benaderingen maken doorgaans gebruik van vaste of globaal uniforme surrogaatgradiëntfuncties (bijv. statische arctangens of Fast Sigmoid) over alle ruimtelijke tokens en trainingsfasen heen. Het artikel betoogt dat deze ruimtelijk uniforme benadering een kritieke foutmodus creëert in diepe Spiking Transformers:

• Token-niveau Dynamiek Genegeerd: Visuele tokens vertonen variërende niveaus van semantische onzekerheid. Tokens met een hoge mate van zekerheid (bijv. duidelijke objectgrenzen) vereisen precieze gradiëntupdates, terwijl tokens met een hoge onzekerheid (bijv. ambigue achtergronden) baat hebben bij een bredere, exploratieve gradiëntstroom.
• Optimalisatiefout: Het afdwingen van een statisch surrogaatvenster voor alle tokens faalt in het bieden van contextuele credit assignment. In diepe lagen leidt dit tot ernstige gradiëntverdwijning (gradient vanishing) en veroorzaakt het dat attention heads instorten naar "dead heads" (waarbij membraanpotentialen onder de drempelwaarde blijven en gewichtupdates stagneren).

De kernuitdaging is het ontwikkelen van een adaptief surrogaatmechanisme dat de gradiëntstroom dynamisch afstemt op de lokale token-onzekerheid zonder de deterministische, laagvermogens binaire executie die vereist is voor inferentie in gevaar te brengen.

2. Methodologie: SAGE Framework

De auteurs stellen SAGE voor (Surrogate-gradient Adaptation via Attention-Guided Entropy), een onzekerheid-bewust framework dat de helling van de surrogaatgradiënt tijdens de training moduleert op basis van de dispersie van de attention-entropie.

2.1 Kernconcept

SAGE opereert uitsluitend tijdens de trainingsfase. Het schat de onzekerheid van elk transformer-blok door de dispersie van de attention-entropie over de self-attention heads te analyseren. Dit signaal wordt gebruikt om de helling ($\alpha$) van de surrogaatgradiëntfunctie adaptief aan te passen. Cruciaal is dat het inferentiemodel ongewijzigd blijft; de onzekerheidscontroller wordt na de training verwijderd.

2.2 Attention Onzekerheid Schatting

Het framework maakt gebruik van de Multi-Head Self-Attention (MHA) module van de Spikformer-architectuur:

1. Ruwe Attention Scores: De standaard Spikformer berekent ruwe scores $S = s QK^\top$ zonder softmax-normalisatie.
2. Auxiliary Distributie: Voor enkel de onzekerheidsschatting construeert SAGE een tijdelijke waarschijnlijkheidsdistributie $A$ door een detach-operatie toe te passen en een Softmax met een vaste temperatuur ($T_{ent} = 0.25$):
   $$A = \text{Softmax}\left(\frac{\text{detach}(S)}{T_{ent}}\right)$$
3. Entropie Berekening: De genormaliseerde Shannon-entropie $E_i$ wordt berekend voor elke attention head $i$.
4. Onzekerheidssignaal: In plaats van het gemiddelde van de entropie te gebruiken, berekent SAGE de standaarddeviatie (dispersie) van de entropieën over alle $H$ heads:
   $$D = \text{Std}(E_1, E_2, \dots, E_H)$$
   Een hoge dispersie duidt op onenigheid tussen de heads (hoge onzekerheid), terwijl een lage dispersie consistente attention-patronen impliceert.

2.3 Onzekerheid-gestuurde Surrogaat Modulatie

De dispersiestatistiek $D$ wordt verwerkt om een blok-specifieke surrogaathelling $\alpha$ te genereren:

1. Smoothing: De instantane dispersie wordt gladgestreken met behulp van een Exponential Moving Average (EMA) om stabiliteit te garanderen.
2. Normalisatie: De gladgestreken waarde wordt genormaliseerd met behulp van running mean en standaarddeviatie statistieken, en vervolgens gecentreerd over alle transformer-blokken om een relatieve onzekerheidsscore $\tilde{z}$ te verkrijgen.
3. Adaptieve Mapping: Een dead-zone controller onderdrukt kleine fluctuaties. De helling wordt gemoduleerd via een hyperbolische tangensfunctie:
   $$\hat{\alpha} = \begin{cases} 4.0 & \text{indien } |\tilde{z}| < 0.25 \\ 4.0 + 0.5 \tanh(\tilde{z}) & \text{anders} \end{cases}$$
   De uiteindelijke helling wordt afgekapt binnen het interval $[3, 5]$.
4. Integratie: Deze adaptieve $\alpha$ wordt alleen toegepast tijdens de backward pass. De forward pass en de inferentie-architectuur blijven identiek aan de baseline Spikformer.

2.4 Computationele Overhead

De extra berekening omvat het berekenen van entropie ($O(HN)$) en standaarddeviatie ($O(H)$), wat verwaarloosbaar is vergeleken met de $O(N^2)$ complexiteit van het attention-mechanisme. Profiling toont een overhead van ongeveer 0,03 ms per mini-batch tijdens de training en nul overhead tijdens de inferentie.

3. Belangrijkste Bijdragen

Het artikel schetst drie primaire bijdragen:

1. Eerste Onzekerheid-bewuste Framework: Introductie van het eerste adaptieve surrogaatgradiënt-framework specifiek ontworpen voor spiking transformers dat gebruikmaakt van attention-afgeleide onzekerheid.
2. Entropy Dispersie als Signaal: Demonstratie dat de dispersie (standaarddeviatie) van de attention-entropie over de heads een meer discriminerend en betrouwbaar online onzekerheidssignaal is voor surrogaat-adaptatie dan alternatieve metrieken (zoals Gini-onzuiverheid of gemiddelde entropie).
3. Prestaties en Efficiëntie: Validatie op concurrerende datasets die consistente verbeteringen laten zien ten opzien van vaste en leerbare baselines met verwaarloosbare trainingsoverhead en geen inferentiekosten.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op CIFAR-10, CIFAR-100 en ImageNet-200 met de Spikformer-4-384 backbone.

• Nauwkeurigheidswinst: SAGE behaalde consistente verbeteringen ten opzien van vaste surrogaat-baselines op de CIFAR-datasets.
  • CIFAR-10: 96,32% Top-1 nauwkeurigheid (vs. 95,51% voor baseline Spikformer).
  • CIFAR-100: 78,69% Top-1 nauwkeurigheid (vs. 78,21% voor de baseline).
  • ImageNet-200: Op deze benchmark behaalde SAGE een Top-1 nauwkeurigheid van 60,60%, wat een afname is vergeleken met de 70,24% van de baseline Spikformer (Tabel 2). Het artikel merkt op dat hoewel SAGE de prestaties op de CIFAR-datasets verbetert, de resultaten op ImageNet-200 lager waren dan die van de baseline, hoewel de methode wordt benadrukt vanwege de effectiviteit in low-latency regimes en het vermogen om de originele architectuur te behouden.
• Temporele Robuustheid: SAGE vertoonde de meest uitgesproken verbeteringen bij lagere simulatietijdstappen ($T \le 4$), met name bij de standaardinstelling van $T=4$. Naarmate de tijdstappen toenamen ($T=8, 12$), kromp het prestatieverschil, wat suggereert dat de methode het meest effectief is voor low-latency toepassingen.
• Vergelijking: SAGE presteerde beter dan zowel de vaste surrogaatgradiënten als een "Learnable" configuratie (waarbij één globale parameter wordt getraind) over meerdere datasets en tijdstappen.

5. Betekenis en Claims

Het artikel positioneert SAGE als een plug-and-play optimalisatiestrategie die de specifieke beperkingen van statische surrogaatgradiënten in diepe Spiking Transformers aanpakt.

• Architecturale Preservatie: De methode behoudt de originele netwerkarchitectuur, modelparameters en inferentiepipeline. Het introduceert geen extra leerbare parameters voor het model zelf en veroorzaakt geen extra inferentielatentie.
• Contextuele Credit Assignment: Door de gradiënthelling dynamisch af te stemmen op de lokale token-onzekerheid, mitigeert SAGE het "dead head"-fenomeen en de gradiëntverdwijning in diepe lagen, waardoor het netwerk tijdens de training beter kan omgaan met ambigue kenmerken.
• Praktische Bruikbaarheid: De auteurs beweren dat SAGE een lichtgewicht oplossing is die naadloos kan worden geïntegreerd in bestaande training-pipelines, waardoor het direct toepasbaar is op een breed scala aan SNN- en Transformer-gebaseerde SNN-modellen zonder dat er wijzigingen in de deployment nodig zijn.

De auteurs concluderen dat onzekerheid-gestuurde surrogaatoptimalisatie een effectieve richting is voor het verbeteren van de training van diepe SNN's, met name voor toepassingen die een lage latentie vereisen.