TrainDeeploy: Hardware-Accelerated Parameter-Efficient Fine-Tuning of Small Transformer Models at the Extreme Edge

TrainDeeploy is een framework dat hardware-versnelde, parameter-efficiënte fine-tuning van zowel CNN- als Transformer-modellen mogelijk maakt op extreem beperkte randapparatuur, waarbij het de eerste end-to-end on-device training van een Compact Convolutional Transformer op een RISC-V-SoC realiseert met aanzienlijke verbeteringen in geheugengebruik en prestaties.

De kern

Hier is een uitleg van het paper "TrainDeeploy" in gewoon Nederlands, met behulp van creatieve vergelijkingen om het begrijpelijk te maken.

De Grote Droom: Slimme Toestellen die Zelf Leren

Stel je voor dat je slimme horloge of je draagbare sensor niet alleen dingen weet (zoals "ik zie een hond"), maar ook zelf kan leren zonder dat je ze naar de cloud (het internet) hoeft te sturen. Dat is het idee van Edge AI.

Het probleem is echter: leren is zwaar werk.

• Inferentie (het toepassen van wat je al weet) is als het lezen van een boek. Dat is makkelijk.
• Training (leren van nieuwe dingen) is als het herschrijven van een heel boek terwijl je het leest, met alle notities en krabbels erbij. Dat kost enorm veel energie en geheugen.

Voor kleine, batterij-aangedreven apparaten (zoals een sensor in een bos of een medisch implantaat) is dit meestal onmogelijk. Ze hebben te weinig batterij en te weinig geheugen.

De Oplossing: TrainDeeploy

De auteurs van dit paper hebben TrainDeeploy bedacht. Dit is een soort "bouwplan" of "besturingssysteem" dat het mogelijk maakt om deze zware leeropdrachten uit te voeren op de kleinste, zuinigste computers die er zijn.

Hier is hoe het werkt, vertaald naar alledaagse beelden:

1. De "Zware Last" van Transformers

Moderne AI-modellen (zoals die in ChatGPT of geavanceerde beeldherkenning) worden vaak gebouwd met een architectuur die Transformers heet.

• Vergelijking: Stel je een Transformer voor als een gigantisch kantoor met duizenden werknemers die allemaal met elkaar praten. Om iets nieuws te leren, moet iedereen zijn notities bijwerken. Op een klein apparaat is er geen ruimte voor al die notities en niet genoeg kracht om iedereen tegelijk aan het werk te zetten.

2. De Slimme Truc: LoRA (Low-Rank Adaptation)

Om dit op te lossen, gebruiken ze een techniek genaamd LoRA.

• De Vergelijking: In plaats van het hele kantoor (het hele AI-model) te herschrijven, doen ze alsof het kantoor al perfect is opgeleid. Ze plakken er alleen een paar kleine post-its op de muren.
• Hoe het werkt: Ze laten de grote, dure werknemers (de zware onderdelen van het model) rustig slapen (vriezen in). Ze laten alleen een paar kleine, flexibele assistenten (de LoRA-matrices) werken om de nieuwe informatie te leren.
• Het Resultaat: In plaats van 15.000 pagina's notities te moeten opslaan, hebben ze nu maar 10 pagina's nodig. Dit bespaart enorm veel geheugen en energie.

3. De Krachtige Motor: Hardware-versnelling

Zelfs met die kleine post-its is het rekenwerk nog steeds zwaar. Daarom gebruiken ze speciale hardware-chips die in deze apparaten zitten.

• De Vergelijking: Stel je voor dat je een berg aardappelen moet schillen.
  • Normaal: Je doet het met je handen (de gewone processor). Dat duurt lang en je wordt moe.
  • TrainDeeploy: Je gebruikt een speciale aardappelroterende machine (de GEMM-accelerator op de chip) die de aardappelen in een seconde schilt.
• De auteurs hebben hun software zo ontworpen dat het de zware rekenwerk (het schillen) automatisch naar deze snelle machines stuurt, terwijl de gewone processor zich bezighoudt met het organiseren.

Wat hebben ze bewezen?

Ze hebben dit getest op een heel klein apparaatje (een RISC-V chip, vergelijkbaar met die in moderne sensoren) en een model dat CCT heet (een slimme mix van een camera en een taalmodel).

• Succes: Het was de eerste keer dat iemand een volledig Transformer-model op het apparaat zelf kon laten leren.
• Snelheid: Ze konden ongeveer 11 keer per seconde een nieuwe les leren (een "gradient update").
• Efficiëntie: Door de LoRA-truc en de snelle machine, gebruikten ze 23% minder geheugen en was het 2,3 tot 3,5 keer sneller dan zonder deze trucjes.
• Resultaat: Het apparaat leerde nieuwe dingen (zoals het herkennen van verschillende soorten bloemen of auto's) met bijna dezelfde nauwkeurigheid als een groot model, maar dan met een fractie van de energie.

Waarom is dit belangrijk?

Vroeger moest je je data naar een grote server sturen om te leren. Dat kostte internet, tijd en was minder privé.
Met TrainDeeploy kan je apparaatje:

1. Privé blijven: Je data blijft bij jou.
2. Zich aanpassen: Het kan leren van jouw specifieke situatie (bijvoorbeeld: "deze sensor werkt het beste als ik hem zo vastzet").
3. Lang meegaan: Het verbruikt zo weinig batterij dat het jarenlang kan werken.

Kortom: TrainDeeploy is de sleutel die de deur opent naar een wereld waar al onze kleine, slimme gadgets niet alleen slim zijn, maar ook kunnen groeien en leren waar ze staan, zonder dat ze een zware batterij of een internetverbinding nodig hebben.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "TrainDeeploy: Hardware-Accelerated Parameter-Efficient Fine-Tuning of Small Transformer Models at the Extreme Edge" in het Nederlands.

Probleemstelling

Het trainen van diepe neurale netwerken (DNN's) direct op apparaten ("on-device training") biedt grote voordelen voor privacy en langdurige adaptatie aan de rand van het netwerk (edge). Echter, het backpropagatie-proces voor het fine-tunen van modellen is extreem rekenintensief en geheugenafhankelijk. Dit vormt een onoverkomelijke barrière voor ultra-low-power System-on-Chips (SoCs) in de "extreme edge" (zoals sensoren en wearables), die vaak beperkt zijn tot enkele honderden kilobytes SRAM en een lage energiebudget.

Specifiek voor Transformer-modellen (die gebaseerd zijn op attention-mechanismen) is de uitdaging groter dan bij traditionele CNN's vanwege de complexe architectuur en de schaal. Bestaande oplossingen hebben beperkingen:

• Ze zijn vaak gericht op CNN's en niet op Transformers.
• Ze missen een end-to-end compilatieflow die rekening houdt met de geheugenhiërarchie (tiling en allocatie).
• Ze bieden geen ondersteuning voor heterogene hardware (CPU + versnellers) in combinatie met parameter-efficiënte technieken.

Methodologie: TrainDeeploy

De auteurs introduceren TrainDeeploy, een nieuw compilatie- en uitvoeringsframework dat efficiënte inferentie en on-device training unifyt op heterogene ultra-low-power SoCs.

1. Architectuur en Compiler-Flow:

• Basis: TrainDeeploy bouwt voort op Deeploy, een domeinspecifieke compiler voor energie-efficiënte inferentie.
• Extensie voor Training: Het framework voegt automatische differentiatie (autograd) toe. Het converteert een PyTorch-model (via ONNX) naar een statische graaf die zowel de forward- als backward-pass (gradiënten) omvat.
• Geheugenbeheer: Een centrale component is de "Midend", die een geïntegreerde optimalisatie uitvoert voor operator-tilling (het opsplitsen van grote matrices in kleinere blokken) en statische geheugenallocatie. Dit lost het probleem op van het opslaan van activaties voor gradiëntberekening binnen de beperkte SRAM.
• Heterogene Hardware: Het framework is ontworpen voor platforms met een host-processor (RISC-V) en versnellers (zoals GEMM-engine). Het mapt zware rekenwerklasten (GEMM en convoluties) naar deze versnellers, terwijl de host de controle en onregelmatige taken behoudt.

2. Parameter-Efficiënte Fine-Tuning (PEFT) met LoRA:
Om het geheugenvoetafdruk te verkleinen, implementeert TrainDeeploy Low-Rank Adaptation (LoRA).

• In plaats van alle gewichten bij te werken, worden de oorspronkelijke gewichten ($W_0$) bevroren.
• Er worden twee kleine, trainbare matrices ($A$ en $B$) toegevoegd met een lage rang $r$.
• Dit reduceert het aantal trainbare parameters en de benodigde opslag voor gradiënten drastisch, waardoor training mogelijk wordt binnen de strenge geheugengrenzen van microcontrollers.

3. Doelplatform:
De implementatie is getest op een RISC-V gebaseerde heterogene SoC (gebaseerd op de PULP-architectuur):

• Een host-core en een cluster van 8 RISC-V cores.
• Een RedMulE-versneller: een hardware-accelerator voor Floating Point GEMM-bewerkingen (geoptimaliseerd voor FP32 training).
• Hiërarchisch geheugen: 128 KB L1 (TCDM), 2 MB L2 SRAM, en 32 MB extern L3 (HyperRAM).

Belangrijkste Bijdragen

1. Eerste End-to-End Pipeline: TrainDeeploy is het eerste framework dat volledige end-to-end training van Transformers (specifiek Compact Convolutional Transformers, CCT) mogelijk maakt op ultra-low-power SoCs.
2. Hardware-Acceleratie: Het is de eerste implementatie van LoRA-training die gebruikmaakt van on-chip GEMM-versnellers op extreme edge-apparaten.
3. Unificatie: Het ondersteunt zowel CNN's als Transformers met parameter-efficiënte tuning, wat een breder toepassingsgebied biedt dan bestaande CNN-centric frameworks.
4. Geheugen- en Rekenoptimalisatie: Door het combineren van LoRA (voor reductie van trainbare parameters) en geavanceerde compiler-technieken (tiling en allocatie), wordt training haalbaar gemaakt binnen de beperkte resources.

Resultaten

De prestaties zijn gemeten op een Compact Convolutional Transformer (CCT) met 0,28 miljoen parameters.

• Snelheid:

  • TrainDeeploy bereikt een doorvoer van 11 getrainde afbeeldingen per seconde (gradient updates) bij single-sample training.
  • Het gebruik van de RedMulE-versneller levert een 2,3x tot 3,5x snelheidswinst op ten opzichte van een puur CPU-gebaseerde uitvoering.
  • FLOP/cycle: Het systeem bereikt tot 4,6 FLOP/cycle voor CCT en 13,4 FLOP/cycle voor kleinere netwerken (Deep-AE), wat aanzienlijk hoger is dan state-of-the-art frameworks voor vergelijkbare hardware.

• Geheugen en Efficiëntie (LoRA vs. Full Backpropagation):

  • Dynamisch geheugen: 23% reductie in piekgebruik van dynamisch geheugen.
  • Parameters: 15x minder trainbare parameters en gradiënten.
  • Gegevensoverdracht: 1,6x minder data-overdracht tussen on-chip en off-chip geheugen.
  • Accuraatheid: LoRA-2 (fine-tuning van twee attention-blokken) bereikte 96,0% nauwkeurigheid op MNIST (50-shot transfer), wat slechts 1% lager is dan full fine-tuning, maar met 15x minder parameters.

• Vergelijking met State-of-the-Art:

  • TrainDeeploy presteert aanzienlijk beter in reken-efficiëntie (FLOP/cycle) dan frameworks zoals PULP-TrainLib, POET, MiniLearn en TTE, terwijl het grotere modellen (Transformers) aankan zonder in te leveren op nauwkeurigheid.

Betekenis en Impact

Dit paper markeert een doorbraak in Edge AI. Het bewijst dat het mogelijk is om geavanceerde Transformer-modellen, die normaal gesproken alleen in de cloud of op krachtige servers kunnen worden getraind, lokaal te fine-tunen op energiezuinige, geheugenbeperkte apparaten.

De combinatie van hardware-versnelling (GEMM-cores) en software-efficiëntie (LoRA en geoptimaliseerde compilatie) opent de deur voor:

• Persoonlijke AI: Apparaten die zich continu aanpassen aan de gebruiker zonder data naar de cloud te sturen.
• Privacy: Gevoelige data blijft volledig lokaal.
• Duurzaamheid: Vermindering van de noodzaak voor cloud-infrastructuur en data-overdracht.

TrainDeeploy stelt een nieuwe standaard voor het trainen van AI-modellen op de "extreme edge", waarbij de beperkingen van energie en geheugen worden omzeild door slimme co-design van hardware en software.