Technische Samenvatting: FAR (Functiebehoudende Aandachtvervanging voor IMC-vriendelijke Inferentie)

1. Probleemstelling

Transformer-architecturen, met name die welke vertrouwen op zelfaandacht, ondervinden aanzienlijke efficiëntieflessenhalsen bij implementatie op opkomende In-Memory Computing (IMC)-hardware, zoals ReRAM-kruisbalken. Hoewel zelfaandacht uitstekend presteert op parallelle digitale processors (GPU's/TPU's) door wereldwijde token-integratie mogelijk te maken, staan de rekenkundige kenmerken ervan fundamenteel haaks op het IMC-uitvoeringsmodel:

• Activering-naar-activering vermenigvuldigingen: Zelfaandacht is sterk afhankelijk van bewerkingen zoals $QK^\top$ en softmax-normalisatie, die dynamische, niet-lokale menging van activeringen vereisen in plaats van de gewichtsstationaire vector-matrixvermenigvuldigingen (VMM) waar IMC-hardware op is geoptimaliseerd.
• Geheugenoverhead: De kwadratische complexiteit ($O(T^2)$) van tokeninteracties en de noodzaak om ruimtelijk verspreide activeringen herhaaldelijk te lezen, leiden tot gefragmenteerde analoge bewerkingen, uitgebreide Analog-naar-Digitaal/Digitaal-naar-Analoge (ADC/DAC) conversies en een slechte benutting van kruisbalken.
• Huidige beperkingen: Bestaande IMC-accelerators verplaatsen aandacht vaak naar digitale eenheden, waardoor de kwadratische bottleneck onopgelost blijft, of proberen aandacht direct op analoge arrays toe te passen zonder rekening te houden met de inherente kosten van dataverplaatsing.

Tegelijkertijd suggereren recente analyses dat aandachtlagen tijdens inferentie aanzienlijke redundantie vertonen, waarbij aandachtskaarten over de diepte heen gladder en comprimeerbaar worden. Dit motiveert de hypothese dat de functionele rol van aandacht kan worden benaderd door sequentiële modules die beter aansluiten bij IMC-localiteit en gewichtsstationariteit.

2. Methodologie: Het FAR-kader

De auteurs stellen FAR (Functiebehoudende Aandachtvervanging) voor, een kader dat alle zelfaandachtblokken in een vooraf getrainde Vision Transformer (specifiek de DeiT-familie) vervangt door sequentiële modules, namelijk Multi-head Bidirectionele LSTMs (BiLSTMs). De aanpak is ontworpen als een "drop-in" architecturale conversie die geen training vanaf nul vereist.

2.1 Vervangingsstrategie

• Bloksgewijze substitutie: Elk aandachtsblok $A_l$ in een vooraf getrainde transformer wordt vervangen door een leerbare vervanger $A'_l$ (een Multi-head BiLSTM). Alle andere componenten (MLP-blokken, patch-embeddings, lagen-normalisatie) blijven bevroren.
• Functionele equivalentie via distillatie: De vervanging wordt begeleid door lagen-gewijze distillatie. Voor elke laag $l$ minimaliseert een gelijkaardigheidsverlies ($L_{sim}$) het verschil tussen de output van de originele aandachtskop en de nieuwe BiLSTM-kop:
  $$L_{sim}^{(l)} = ||A_l(x_l) - A'_l(x_l)||_2^2$$
  Dit zorgt ervoor dat de vervanger de representatieve structuur van de leraar nabootst.
• Trainingsprotocol: Het proces omvat twee fasen:
  1. Distillatiefase: Alleen de vervangende BiLSTM-modules worden getraind met een combinatie van gelijkaardigheidsverlies en taakverlies (classificatie), terwijl de rest van het model bevroren blijft.
  2. Finetuning-fase: Het volledige model wordt gezamenlijk gefinetuned om eventuele resterende nauwkeurigheidsdalingen te herstellen.

2.2 Multi-head BiLSTM-architectuur

Om de inductieve bias van multi-head-aandacht te behouden, is de vervanger gestructureerd als een Multi-head BiLSTM:

• Projectie en splitsing: De input-embedding wordt geprojecteerd en opgesplitst in $N$ subruimten (overeenkomend met het aantal aandachtskoppen van de leraar).
• Bidirectionele verwerking: Elke subruimte wordt verwerkt door een afzonderlijke BiLSTM-kop, wat zowel voorwaartse als achterwaartse token-afhankelijkheden vastlegt.
• Concatenatie en projectie: Outputs worden samengevoegd en geprojecteerd terug naar de oorspronkelijke dimensie.
• IMC-compatibiliteit: In tegenstelling tot aandacht, vertrouwt deze structuur op recurrente gewichtsrehergebruik en sequentiële dataflow, waardoor alle-paren activeringsinteracties worden geëlimineerd. Dit past natuurlijk op ReRAM-kruisbalken waar gewichten stationair zijn en data sequentieel stroomt.

2.3 Gestructureerde pruning voor IMC-beperkingen

Om verder aan te sluiten bij resource-beperkte IMC-arrays, passen de auteurs gestructureerde pruning toe op de verborgen dimensies van de BiLSTM-koppen:

• Uitbreiding van DeepHoyer: Ze breiden het DeepHoyer-kader uit om LSTM-verborgen eenheden te targeten. Omdat LSTM-poorten (input, forget, cell, output) verborgen eenheden delen, vereist het verwijderen van een enkele eenheid gecoördineerde verwijdering over alle bijbehorende gewichtsmatrices ($W_{ih}$, $W_{hh}$ en projectiegewichten).
• Groeps-sparseheid: Een differentieerbare Hoyer-sparseheidsmaat wordt toegepast op groepen gewichten die overeenkomen met elke verborgen eenheid, waardoor het netwerk wordt aangemoedigd te leren welke eenheden redundant zijn.
• Lagen-adaptieve compressie: Pruning wordt onafhankelijk per kop en richting uitgevoerd, waardoor het model zijn capaciteit kan aanpassen aan de specifieke redundantie van elke laag, wat resulteert in compacte modules die passen in IMC-array-groottes.

3. Belangrijkste Resultaten

Het kader werd geëvalueerd op de DeiT-familie (Tiny, Small, Base) op ImageNet en diverse downstream-taken (CIFAR-10/100, Cars, Flowers, iNaturalist).

• Behoud van nauwkeurigheid:
  • DeiT-Tiny: FAR-Tiny presteerde op ImageNet zelfs beter dan de baseline DeiT-Tiny met +1,2% (73,3% versus 72,2%), wat suggereert dat de sequentiële vervanger fungeert als een gunstige inductieve bias voor kleinere modellen.
  • Grotere modellen: Voor DeiT-Small en Base behield FAR concurrerende nauwkeurigheid met een bescheiden gat van ongeveer 1–2% ten opzichte van de leraar, wat aantoont dat de functionele rol van aandacht grotendeels reproduceerbaar is via sequentiële recurrentie.
  • Downstream-overdracht: FAR-modellen behielden sterke overdraagbaarheid naar downstream-classificatietaken, wat aangeeft dat de geleerde sequentie-naar-sequentie-mapping effectief generaliseert.
• Robuustheid tegen pruning:
  • Het toepassen van gestructureerde pruning verlaagde de gemiddelde verborgen dimensie met 40–60% met slechts een extra nauwkeurigheidsdaling van 1–2%.
  • Pruningpatronen waren niet-uniform, met hogere retentie in middenlagen en agressievere pruning aan de uiteinden van het netwerk, wat interne heterogeniteit in representatieve redundantie weerspiegelt.
• Kopgewijze specialisatie: Visualisatie van token-interactiekarten toonde aan dat FAR-koppen met succes de distincte semantische rollen van de originele aandachtskoppen erfden, met concentratie op specifieke objectregio's in plaats van te instorten in redundante patronen.

Hardware-efficiëntieanalyse

Met behulp van een analytisch ReRAM IMC-model schatten de auteurs latentie en energie:

• Latentie: FAR op ReRAM IMC wordt geschat op ~18x sneller dan DeiT op ReRAM IMC (door het vermijden van kwadratische aandachtoverhead) en ~400x sneller dan DeiT op een GPU.
• Energie: FAR op ReRAM IMC wordt geschat op ~3x energie-efficiënter dan DeiT op ReRAM IMC en ~150x efficiënter dan DeiT op een GPU.
• Mechanisme: Deze winsten vloeien voort uit het elimineren van activering-naar-activering vermenigvuldigingen, het verminderen van ADC/DAC-conversies en het mogelijk maken van gewichtsstationaire uitvoering. De auteurs merken op dat het uitsluiten van softmax-kosten (die aanzienlijke bottlenecks vormen voor aandacht op IMC) deze schattingen conservatief maakt voor FAR.

4. Betekenis en Claims

Het artikel claimt dat FAR een praktisch, hardware-efficiënt alternatief biedt voor het implementeren van vooraf getrainde transformers op IMC-gebaseerde edge-accelerators zonder de noodzaak om modellen vanaf nul te trainen.

• Haalbaarheid van volledige vervanging: Het werk toont aan dat aandacht end-to-end kan worden vervangen tijdens inferentie terwijl praktische nauwkeurigheid wordt behouden, wat de notie uitdaagt dat alle-paren tokeninteracties strikt noodzakelijk zijn voor hoogwaardige visietaken.
• Hardware-uitlijning: Door aandacht te vervangen door sequentiële, gewichtsstationaire modules, brengt FAR de rekenkundige dataflow in lijn met de fysieke sterktes van ReRAM-kruisbalken, wat het dataverkeer en de latentie aanzienlijk verlaagt.
• Behoud van vooraf getrainde kennis: Het kader bewijst dat de rijke semantische representaties die tijdens grootschalige pretraining zijn geleerd, kunnen worden behouden en overgedragen naar structureel verschillende, IMC-vriendelijke architecturen via bloksgewijze distillatie.

Kortom, FAR biedt een weg om transformer-inferentie te herstructureren rond IMC-vriendelijke sequentiële modules, waardoor de kloof wordt overbrugd tussen de dominantie van op aandacht gebaseerde modellen en de beperkingen van opkomende geheugencentrische hardware.