TT-SNN: Tensor Train Decomposition for Efficient Spiking Neural Network Training

Dit paper introduceert TT-SNN, een methode die Tensor Train-decompositie en een parallelle rekenpijplijn combineert om de opslag, berekeningskosten en trainingstijd van Spiking Neural Networks aanzienlijk te verminderen met verwaarloosbaar verlies aan nauwkeurigheid.

De kern

Stel je voor dat je een heel slimme robot wilt bouwen die kan zien en leren, net als een mens. Maar er is een probleem: deze robot is momenteel nogal een "energieverslinder". Hij heeft een enorme geheugenbank nodig en doet er lang over om te leren. Dit is wat er gebeurt met de huidige Spiking Neural Networks (SNN's). Ze zijn wel heel energiezuinig als ze werken (zoals een menselijk brein dat alleen "aan" gaat als er iets te doen is), maar het leren van deze systemen is nog steeds erg zwaar en traag.

De onderzoekers van Yale University hebben een nieuwe oplossing bedacht, genaamd TT-SNN. Laten we dit uitleggen met een paar simpele vergelijkingen.

1. Het Probleem: De Zware Koffer

Stel je voor dat je een SNN wilt trainen. Dit is alsof je een student moet leren een complexe puzzel op te lossen.

• Huidige methode: De student moet elke stap van de puzzel één voor één doen, en elke stap moet hij onthouden om later terug te kunnen gaan en fouten te corrigeren. Hij moet een enorme koffer meenemen met al zijn aantekeningen (geheugen) en doet er eeuwen over (rekenkracht).
• Het gevolg: De computer wordt heet, de batterij is snel leeg en het duurt lang voordat de robot slim is.

2. De Oplossing: De "TT-SNN" (De Slimme Opbergtruc)

De onderzoekers gebruiken een techniek uit de wiskunde genaamd Tensor Train Decomposition. Klinkt ingewikkeld? Laten we het zo zien:

Stel je voor dat je een gigantische, zware muur van bakstenen (de data van de robot) moet verplaatsen.

• De oude manier: Je probeert de hele muur in één keer te tillen. Zwaar!
• De TT-methode: Je breekt die muur op in vier kleinere, lichtere stapels bakstenen. Je kunt deze stapels nu veel makkelijker en sneller verplaatsen.
• Het resultaat: De robot heeft veel minder "kofferinhoud" nodig om te leren. De onderzoekers zeggen dat ze de grootte van het model met bijna 8 keer hebben verkleind en de rekenwerk met 9 keer hebben versneld.

3. Twee Nieuwe Trucs: Parallel en Halfweg

De onderzoekers hebben niet alleen de bakstenen opgesplitst, ze hebben ook de manier waarop de robot werkt veranderd.

A. De "Parallelle Truc" (PTT)

• Oude manier: De robot doet de vier stapels bakstenen één voor één. Stap 1, dan stap 2, dan stap 3... Dit duurt lang.
• Nieuwe manier (PTT): De robot doet stap 2 en stap 3 tegelijkertijd. Het is alsof je twee handen hebt in plaats van één. Je pakt twee bakstenen tegelijk.
• Vergelijking: In plaats van een lange rij mensen die een emmer water doorgeven, heb je nu twee rijen mensen die tegelijkertijd water dragen. Het resultaat is hetzelfde, maar het gaat veel sneller.

B. De "Halfweg-Truc" (HTT)

• Het idee: Soms hoeft de robot niet alles te doen. Net als wanneer je een film kijkt: in het begin moet je goed opletten om te begrijpen wat er gebeurt, maar later, als je het verhaal al kent, kun je soms even "wegkijken" of minder aandacht hebben.
• Hoe het werkt: De robot doet in het begin van het leerproces alles volledig (alle bakstenen). Maar in de latere stappen gebruikt hij maar de helft van de bakstenen.
• Voordeel: Dit bespaart nog meer energie, vooral bij statische beelden (zoals een foto). Bij bewegende beelden (zoals een video) werkt dit iets minder goed, omdat je daar elke frame nodig hebt.

4. De Speciale Werkbank (De Hardware)

Je kunt een snelle auto niet rijden op een slechte weg. Omdat deze nieuwe methode (PTT en HTT) dingen tegelijkertijd doet, past het niet goed op de oude "werkbanken" (chips) die er nu zijn. Die zijn gemaakt voor één ding tegelijk.

De onderzoekers hebben daarom een nieuwe werkbank ontworpen.

• Stel je voor dat je een fabriek hebt met vier aparte teams die tegelijkertijd werken, in plaats van één team dat alles moet doen.
• Deze nieuwe chip is speciaal gebouwd om de "Parallelle Truc" van de robot te ondersteunen. Hierdoor wordt het trainen niet alleen sneller, maar ook 28% tot 43% energiezuiniger.

Wat betekent dit voor de toekomst?

Kortom, deze paper zegt: "We hebben een manier gevonden om slimme robots (SNN's) veel sneller en goedkoper te leren."

• Minder geheugen: De robot past op kleinere apparaten.
• Sneller leren: Wat nu dagen duurt, kan nu in uren.
• Minder stroom: Ideaal voor apparaten die op batterijen werken, zoals draagbare medische apparaten of zelfrijdende auto's.

Het is alsof je een zware, traag lopende olifant hebt omgezet in een snelle, wendbare gazelle die net zo slim is, maar veel minder te eten nodig heeft. En dat is een enorme stap vooruit voor de robotica van de toekomst!

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "TT-SNN: Tensor Train Decomposition for Efficient Spiking Neural Network Training" in het Nederlands.

Probleemstelling

Spiking Neural Networks (SNNs) worden gezien als een energie-efficiënt alternatief voor traditionele Artificial Neural Networks (ANNs) vanwege hun event-gestuurde, binaire activatie. Echter, het trainen van SNNs via backpropagation (BP) met gesurrogateerde gradiënten brengt aanzienlijke uitdagingen met zich mee:

1. Hoge geheugen- en rekenkosten: Door de ruimtetijddynamiek en de noodzaak om gradiënten over meerdere tijdstappen (timesteps) en lagen te berekenen, moeten veel tussenliggende activaties worden opgeslagen.
2. Inefficiëntie van bestaande methoden: Bestaande optimalisatietechnieken zoals kwantisatie, kennisdistillatie en pruning richten zich voornamelijk op het versnellen van de inference, niet op het trainingsproces zelf.
3. Hardware beperkingen: Bestaande SNN-training accelerators zijn vaak ontworpen voor sequentiële verwerking van lagen en kunnen de parallelle potentie van geavanceerde decompositiemethoden niet optimaal benutten.

Methodologie: TT-SNN

De auteurs introduceren TT-SNN, een methode die Tensor Train (TT) decompositie toepast op SNN's om de modelgrootte te verkleinen en de training te versnellen. De kern van de aanpak omvat drie belangrijke innovaties:

1. Tensor Train Decompositie voor Convoluties:
   In plaats van een enkele grote convolutielaag te gebruiken, worden de gewichten ontbonden in kleinere tensoren (TT-kernen). Een standaard $3\times3$ convolutie wordt opgesplitst in vier sub-convoluties met asymmetrische kernen.

2. Parallel TT (PTT) Module:
   Traditionele TT-methoden (Sequential TT of STT) berekenen deze sub-convoluties achter elkaar, wat leidt tot verlies van informatie door de asymmetrie van de kernen (bijv. alleen verticale of horizontale extractie).

   • De auteurs stellen een Parallel TT (PTT) pipeline voor waarbij de tweede en derde sub-convolutie parallel worden berekend op basis van de output van de eerste laag.
   • Dit benadert een $3\times3$ kernel (zonder de vier hoekwaarden) en zorgt voor gelijktijdige extractie van verticale en horizontale kenmerken, wat de nauwkeurigheid verbetert ten opzichte van STT.

3. Half TT (HTT) Module:
   Gebaseerd op het inzicht dat SNN's in de vroege tijdstappen meer informatie vastleggen dan in latere stappen, introduceert de HTT-module een strategie voor tijdsredundantie.

   • In de vroege tijdstappen worden volledige sub-convoluties gebruikt.
   • In latere tijdstappen worden slechts halve sub-convoluties gebruikt.
   • Dit creëert een "half-diagonale" berekening in het ruimtetijddiagram, wat de berekeningslast verder verlaagt.

4. Hardware Accelerator Ontwerp:
   Omdat PTT en HTT parallelle verwerking vereisen die niet goed past bij bestaande sequentiële accelerators, stellen de auteurs een multi-cluster systolic-array accelerator voor.

   • Het ontwerp gebruikt vier rekenclusters: Cluster 1 voor de eerste sub-convolutie, Clusters 2 en 3 werken parallel voor de middelste stappen, en Cluster 4 voor de laatste stap.
   • Dit ontwerp maximaliseert de parallelle verwerking en minimaliseert de data-overdracht naar het DRAM, wat essentieel is voor energie-efficiëntie.

Belangrijkste Bijdragen

1. TT-SNN Module: De eerste toepassing van tensor decompositie specifiek voor SNN-training, met een focus op het verminderen van trainingskosten.
2. Parallelle Verwerking: Introductie van PTT om de informatieverliesproblemen van sequentiële TT-methoden op te lossen en de snelheid te verhogen.
3. HTT Strategie: Een nieuwe methode om tijdsredundantie in SNN's te exploiteren voor verdere compressie.
4. Specifieke Hardware Accelerator: Een nieuw accelerator-ontwerp dat volledig is afgestemd op de parallelle aard van PTT en HTT, in tegenstelling tot bestaande oplossingen.
5. Flexibiliteit: De modules kunnen eenvoudig worden geïntegreerd in bestaande SNN-architecturen (zoals ResNet) zonder significante nauwkeurigheidsverlies.

Resultaten

De methode is getest op statische datasets (CIFAR-10/100) en dynamische event-datasets (N-Caltech101) met ResNet18 en ResNet34 architecturen.

• Compressie en Snelheid:

  • Op de N-Caltech101 dataset werd een reductie van 7,98x in het aantal parameters en 9,25x in FLOPs bereikt.
  • De trainingstijd werd met 17,7% verkort (PTT) en tot 19,3% (HTT).
  • Er was sprake van een verwaarloosbaar verlies in nauwkeurigheid (soms zelfs een lichte verbetering ten opzichte van de baseline).

• Energie-efficiëntie:

  • Op bestaande accelerators (ontworpen voor sequentiële verwerking) presteerde PTT slechter qua energie dan STT door extra DRAM-toegang.
  • Op de nieuwe voorgestelde accelerator werd echter een energiebesparing van 28,3% (voor PTT) en 43,5% (voor HTT) bereikt ten opzichte van de sequentiële STT-methode.

• Compatibiliteit:

  • De PTT-module werd succesvol geïntegreerd in andere bestaande SNN-algoritmen (zoals tdBN, TEBN, TET, NDA), wat resulteerde in trainingstijdsreducties van 9% tot 25% zonder significante nauwkeurigheidsdaling.

Betekenis en Conclusie

Dit paper biedt een doorbraak in de efficiëntie van SNN-training. Door Tensor Train decompositie te combineren met parallelle verwerking en een daarop afgestemde hardware-architectuur, lossen de auteurs het probleem van hoge trainingskosten op.

• Technisch: Het bewijst dat SNN's niet alleen energie-efficiënt zijn tijdens inferentie, maar ook efficiënter getraind kunnen worden, wat een cruciale stap is voor de adoptie van neuromorfe computing.
• Praktisch: De methode is compatibel met zowel statische als dynamische (event-based) datasets, wat het zeer waardevol maakt voor toepassingen in robotica en sensoren waar real-time verwerking en lage energieverbruik essentieel zijn.

De code is beschikbaar op GitHub, wat de reproduceerbaarheid en verdere ontwikkeling van deze techniek stimuleert.