Technische Zusammenfassung: FAR (Funktionserhaltender Attention-Ersatz für IMC-freundliche Inferenz)

1. Problemstellung

Transformer-Architekturen, insbesondere solche, die auf Self-Attention basieren, stoßen beim Einsatz auf neuartigen In-Memory-Computing (IMC)-Hardware, wie ReRAM-Crossbars, auf erhebliche Effizienzengpässe. Während Self-Attention auf parallelen digitalen Prozessoren (GPUs/TPUs) durch die Ermöglichung einer globalen Token-Integration glänzt, stehen ihre Rechencharakteristika im fundamentalen Konflikt mit dem IMC-Ausführungsmodell:

• Aktivierung-zu-Aktivierung-Multiplikationen: Self-Attention stützt sich stark auf Operationen wie $QK^\top$ und Softmax-Normalisierung, die eine dynamische, nicht-lokale Mischung von Aktivierungen erfordern, anstatt der gewichtsstationären Vektor-Matrix-Multiplikationen (VMM), die IMC-Hardware optimiert.
• Speicher-Overhead: Die quadratische Komplexität ($O(T^2)$) der Token-Interaktionen und die Notwendigkeit, räumlich verteilte Aktivierungen wiederholt auszulesen, führen zu fragmentierten analogen Operationen, umfangreichen Analog-Digital/Digital-Analog (ADC/DAC)-Konversionen und einer schlechten Crossbar-Auslastung.
• Aktuelle Einschränkungen: Bestehende IMC-Beschleuniger verlagern Attention oft auf digitale Einheiten, wodurch der quadratische Engpass ungelöst bleibt, oder versuchen, Attention direkt auf analoge Arrays abzubilden, ohne die inhärenten Datenbewegungskosten zu adressieren.

Gleichzeitig deuten neuere Analysen darauf hin, dass Attention-Schichten während der Inferenz eine erhebliche Redundanz aufweisen, wobei Attention-Karten über die Tiefe hinweg glatter und komprimierbarer werden. Dies motiviert die Hypothese, dass die funktionale Rolle von Attention durch sequenzielle Module angenähert werden kann, die besser mit IMC-Lokalität und Gewichtsstationarität übereinstimmen.

2. Methodik: Das FAR-Framework

Die Autoren schlagen FAR (Function-preserving Attention Replacement) vor, ein Framework, das alle Self-Attention-Blöcke in einem vortrainierten Vision Transformer (speziell der DeiT-Familie) durch sequenzielle Module ersetzt, konkret Multi-head Bidirectional LSTMs (BiLSTMs). Der Ansatz ist als „Drop-in"-Architekturkonversion konzipiert, die kein Training von Grund auf erfordert.

2.1 Ersetzungsstrategie

• Blockweise Substitution: Jeder Attention-Block $A_l$ in einem vortrainierten Transformer wird durch einen lernfähigen Ersatz $A'_l$ (ein Multi-head BiLSTM) ersetzt. Alle anderen Komponenten (MLP-Blöcke, Patch-Embeddings, Layernorms) bleiben eingefroren.
• Funktionale Äquivalenz durch Distillation: Die Ersetzung wird durch schichtenweise Distillation überwacht. Für jede Schicht $l$ minimiert ein Ähnlichkeitsverlust ($L_{sim}$) den Unterschied zwischen dem Ausgang des ursprünglichen Attention-Kopfes und dem neuen BiLSTM-Kopf:
  $$L_{sim}^{(l)} = ||A_l(x_l) - A'_l(x_l)||_2^2$$
  Dies stellt sicher, dass der Ersatz die repräsentationale Struktur des Lehrers repliziert.
• Trainingsprotokoll: Der Prozess umfasst zwei Phasen:
  1. Distillationsphase: Nur die ersetzenden BiLSTM-Module werden unter Verwendung einer Kombination aus Ähnlichkeitsverlust und Aufgabenverlust (Klassifizierung) trainiert, während der Rest des Modells eingefroren bleibt.
  2. Feinabstimmungsphase: Das gesamte Modell wird gemeinsam feinabgestimmt, um etwaige verbleibende Genauigkeitsverluste auszugleichen.

2.2 Multi-head BiLSTM-Architektur

Um den induktiven Bias von Multi-head-Attention zu bewahren, ist der Ersatz als Multi-head BiLSTM strukturiert:

• Projektion und Aufteilung: Das Eingabe-Embedding wird projiziert und in $N$ Teilräume aufgeteilt (entsprechend der Anzahl der Attention-Köpfe des Lehrers).
• Bidirektionale Verarbeitung: Jeder Teilraum wird von einem separaten BiLSTM-Kopf verarbeitet, der sowohl vorwärts- als auch rückwärtsgerichtete Token-Abhängigkeiten erfasst.
• Verkettung und Projektion: Die Ausgaben werden verkettet und zurück auf die ursprüngliche Dimension projiziert.
• IMC-Kompatibilität: Im Gegensatz zu Attention stützt sich diese Struktur auf wiederkehrende Gewichts-Wiederverwendung und sequenziellen Datenfluss, wodurch Interaktionen zwischen allen Aktivierungspaaren eliminiert werden. Dies bildet sich natürlich auf ReRAM-Crossbars ab, bei denen die Gewichte stationär sind und der Datenfluss sequenziell erfolgt.

2.3 Strukturiertes Pruning für IMC-Beschränkungen

Um sich weiter an ressourcenbeschränkte IMC-Arrays anzupassen, wenden die Autoren strukturiertes Pruning auf die verborgenen Dimensionen der BiLSTM-Köpfe an:

• Erweiterung von DeepHoyer: Sie erweitern das DeepHoyer-Framework, um auf verborgene LSTM-Einheiten abzuzielen. Da LSTM-Tore (Input, Forget, Cell, Output) verborgene Einheiten teilen, erfordert das Pruning einer einzelnen Einheit eine koordinierte Entfernung über alle zugehörigen Gewichtsmatrizen hinweg ($W_{ih}$, $W_{hh}$ und Projektionsgewichte).
• Gruppenspärlichkeit: Ein differenzierbares Hoyer-Spärlichkeitsmaß wird auf Gruppen von Gewichten angewendet, die jeder verborgenen Einheit entsprechen, und ermutigt das Netzwerk, zu lernen, welche Einheiten redundant sind.
• Schichtenadaptive Kompression: Das Pruning wird unabhängig pro Kopf und Richtung durchgeführt, wodurch sich das Modell an die spezifische Redundanz jeder Schicht anpassen kann, was zu kompakten Modulen führt, die in IMC-Array-Größen passen.

3. Wichtige Ergebnisse

Das Framework wurde an der DeiT-Familie (Tiny, Small, Base) auf ImageNet und verschiedenen Downstream-Aufgaben (CIFAR-10/100, Cars, Flowers, iNaturalist) evaluiert.

• Genauigkeitsbewahrung:
  • DeiT-Tiny: FAR-Tiny übertroffen tatsächlich das Basis-DeiT-Tiny um +1,2 % auf ImageNet (73,3 % vs. 72,2 %), was darauf hindeutet, dass der sequenzielle Ersatz als vorteilhafter induktiver Bias für kleinere Modelle wirkt.
  • Größere Modelle: Für DeiT-Small und Base behielt FAR eine wettbewerbsfähige Genauigkeit mit einer moderaten Lücke von etwa 1–2 % im Vergleich zum Lehrer bei, was zeigt, dass die funktionale Rolle von Attention weitgehend durch sequenzielle Rekurrenz reproduzierbar ist.
  • Downstream-Transfer: FAR-Modelle behielten eine starke Übertragbarkeit auf Downstream-Klassifizierungsaufgaben bei, was darauf hindeutet, dass die gelernten Sequenz-zu-Sequenz-Mappings effektiv generalisieren.
• Robustheit gegenüber Pruning:
  • Die Anwendung von strukturiertem Pruning reduzierte die durchschnittliche verborgene Dimension um 40–60 % bei nur einem zusätzlichen Genauigkeitsverlust von 1–2 %.
  • Die Pruning-Muster waren nicht einheitlich, mit einer höheren Beibehaltung in mittleren Schichten und einem aggressiveren Pruning an den Netzwerkrändern, was die interne Heterogenität der repräsentationalen Redundanz widerspiegelt.
• Kopfweise Spezialisierung: Die Visualisierung von Token-Interaktionskarten zeigte, dass FAR-Köpfe erfolgreich die distincten semantischen Rollen der ursprünglichen Attention-Köpfe übernahmen und sich auf spezifische Objektregionen konzentrierten, anstatt in redundante Muster zu kollabieren.

Hardware-Effizienz-Analyse

Unter Verwendung eines analytischen ReRAM-IMC-Modells schätzten die Autoren die Latenz und den Energieverbrauch:

• Latenz: FAR auf ReRAM IMC wird als ~18-mal schneller als DeiT auf ReRAM IMC geschätzt (aufgrund der Vermeidung des quadratischen Attention-Overheads) und ~400-mal schneller als DeiT auf einer GPU.
• Energie: FAR auf ReRAM IMC wird als ~3-mal energieeffizienter als DeiT auf ReRAM IMC und ~150-mal effizienter als DeiT auf einer GPU geschätzt.
• Mechanismus: Diese Gewinne resultieren aus der Eliminierung von Aktivierung-zu-Aktivierung-Multiplikationen, der Reduzierung von ADC/DAC-Konversionen und der Ermöglichung einer gewichtsstationären Ausführung. Die Autoren stellen fest, dass das Ausklammern der Softmax-Kosten (die signifikante Engpässe für Attention auf IMC darstellen) diese Schätzungen für FAR konservativ macht.

4. Bedeutung und Behauptungen

Die Arbeit behauptet, dass FAR eine praktische, hardware-effiziente Alternative für den Einsatz vortrainierter Transformer auf IMC-basierten Edge-Beschleunigern bietet, ohne dass Modelle von Grund auf neu trainiert werden müssen.

• Durchführbarkeit einer vollständigen Ersetzung: Die Arbeit zeigt, dass Attention während der Inferenz end-to-end ersetzt werden kann, während eine praktische Genauigkeit erhalten bleibt, und stellt damit die Vorstellung in Frage, dass Interaktionen zwischen allen Token-Paaren für hochleistungsfähige Vision-Aufgaben strikt notwendig sind.
• Hardware-Ausrichtung: Durch den Ersatz von Attention durch sequenzielle, gewichtsstationäre Module richtet FAR den rechnerischen Datenfluss an den physischen Stärken von ReRAM-Crossbars aus und reduziert den Speicherverkehr und die Latenz erheblich.
• Bewahrung vortrainierten Wissens: Das Framework beweist, dass die reichen semantischen Darstellungen, die während des großskaligen Pretrainings gelernt wurden, bewahrt und durch blockweise Distillation auf strukturell unterschiedliche, IMC-freundliche Architekturen übertragen werden können.

Zusammenfassend bietet FAR einen Weg, die Transformer-Inferenz um IMC-freundliche sequenzielle Module neu zu strukturieren und so die Kluft zwischen der Dominanz von Attention-basierten Modellen und den Beschränkungen neuartiger speicherzentrierter Hardware zu überbrücken.