HiAP: A Multi-Granular Stochastic Auto-Pruning Framework for Vision Transformers

Die Arbeit stellt HiAP vor, ein neuartiges Framework zur stochastischen, mehrstufigen automatischen Pruning von Vision Transformern, das durch die gleichzeitige Optimierung von Makro- und Mikro-Granularitäten in einem einzigen End-to-End-Trainingsschritt effiziente Subnetzwerke ohne manuelle Heuristiken oder mehrstufige Pipelines erzeugt.

Das Wesentliche

Hallo! Hier ist eine einfache Erklärung des Papers „HiAP" auf Deutsch, voller anschaulicher Vergleiche.

Stell dir vor, du hast einen riesigen, hochmodernen Schwarm von 100 Robotern (das ist das „Vision Transformer"-Modell), die zusammen Bilder erkennen sollen. Jeder Roboter hat einen eigenen Kopf (um Details zu sehen) und einen eigenen Arbeitsbereich (um zu rechnen).

Das Problem: Dieser Schwarm ist zu groß, zu langsam und verbraucht zu viel Strom. Er passt nicht in das kleine Handy oder den kleinen Computer am Rand des Netzwerks (dem „Edge Device").

Früher haben Ingenieure versucht, diesen Schwarm zu verkleinern, indem sie manuell Roboter rausgeworfen haben. Das war wie ein Gärtner, der mit einer Schere blindlings Äste abschneidet. Oft schnitt er wichtige Äste ab oder ließ unnötige Zweige stehen. Außerdem musste er danach den ganzen Baum neu pflanzen und warten, bis er wieder wuchs (das war die langwierige „Nachjustierung").

HiAP ist wie ein intelligenter, selbstorganisierender Garten.

Hier ist, wie HiAP funktioniert, Schritt für Schritt:

1. Der „Geister-Schalter" (Stochastische Gatter)

Stell dir vor, jeder Roboter und jeder Teil eines Roboters hat einen unsichtbaren Schalter.

• Makro-Schalter (Groß): Diese Schalter entscheiden, ob ein ganzer Roboter (ein ganzer „Aufmerksamkeitskopf") oder sogar eine ganze Abteilung (ein ganzer „FFN-Block") im Raum bleibt oder komplett verschwindet.
• Mikro-Schalter (Klein): Wenn ein Roboter bleibt, entscheiden diese Schalter, ob er nur seine linke Hand benutzt oder ob er seine ganze Kraft auf die rechte Hand konzentriert. Sie verkleinern die inneren Arbeitsbereiche.

Das Geniale: Diese Schalter sind nicht fest verdrahtet. Sie sind wie Geister, die während des Trainings hin und her flackern. Manchmal sind sie zu 50 % an, manchmal zu 10 %. Das erlaubt dem System, „probeweise" zu lernen, was passiert, wenn man etwas weglässt, ohne es sofort endgültig zu löschen.

2. Der eine große Tanz (End-to-End Training)

Früher musste man erst die Roboter aussortieren (Phase 1) und dann den Rest neu trainieren (Phase 2). Das war wie ein zweistufiger Umzug.
HiAP macht alles in einem Schritt. Während das System lernt, Bilder zu erkennen, lernen die Schalter gleichzeitig, welche Roboter nicht gebraucht werden. Es ist, als würde der Tanzmeister (der Algorithmus) während des Tanzes entscheiden, wer aus dem Kreis ausscheidet, während die Musik (das Training) weiterläuft. Niemand stolpert, weil sich alle an die neuen Schritte anpassen.

3. Der „Gummibärchen-Temperatur"-Trick (Gumbel-Sigmoid)

Am Anfang des Trainings sind die Schalter wie schmelzende Gummibärchen – sie sind weich und unsicher. Das System probiert viel aus.
Je mehr das System lernt, desto mehr wird die „Temperatur" heruntergefahren. Die Gummibärchen werden hart. Am Ende sind die Schalter fest: Entweder ist ein Roboter da (1) oder weg (0).
Das Ergebnis ist ein fester, kleinerer Schwarm, der sofort einsatzbereit ist. Kein Nachjustieren nötig!

4. Warum ist das so clever? (Makro vs. Mikro)

Stell dir vor, du musst einen LKW entladen.

• Nur Mikro-Pruning: Du nimmst nur ein paar Kartons aus dem LKW, aber der LKW fährt trotzdem mit voller Größe zur Baustelle. Der Motor (der Prozessor) läuft, aber der LKW ist fast leer. Das spart wenig Zeit, weil der LKW immer noch schwer ist (Speicherbandbreite).
• Nur Makro-Pruning: Du nimmst den ganzen LKW weg, aber die verbleibenden Kartons sind riesig und schwer zu tragen.
• HiAP: HiAP entscheidet: „Wir brauchen nur 3 kleine Lieferwagen (Makro), und in jedem davon tragen wir nur die wichtigsten 50 % der Ware (Mikro)."
  Das Ergebnis: Der LKW ist klein, die Kartons sind leicht, und die Fahrt ist superschnell.

5. Das Ergebnis

HiAP hat auf großen Datensätzen (wie ImageNet) gezeigt, dass es Modelle findet, die:

1. Viel schneller sind (weniger Rechenaufwand).
2. Genau so gut funktionieren wie die riesigen Originalmodelle.
3. Keine manuelle Hilfe brauchen. Das System findet die perfekte Größe von selbst.

Zusammenfassend:
HiAP ist wie ein selbstoptimierender Architekt, der während des Baus eines Hauses entscheidet, welche Zimmer man streichen kann und welche Wände man dünner machen darf, ohne dass das Haus einstürzt. Am Ende hast du ein kleines, effizientes Haus, das sofort bezugsfertig ist – ohne dass du als Bauleiter mühsam Pläne zeichnen musst.

Technische Zusammenfassung

1. Problemstellung

Vision Transformers (ViT) sind zwar leistungsstark, aber aufgrund ihres hohen Rechen- und Speicherbedarfs schwer auf ressourcenbeschränkten Edge-Geräten einzusetzen.

• Herausforderung: Herkömmliche Methoden zur strukturierten Pruning (Entfernung ganzer Komponenten wie Attention-Heads oder Neuronen) arbeiten oft nur auf einer einzigen Granularitätsebene (entweder nur Mikro-Strukturen wie Dimensionen oder nur Makro-Strukturen wie ganze Blöcke).
  • Mikro-Pruning reduziert zwar FLOPs, ignoriert aber oft den Speicherbandbreiten-Overhead (DRAM/SRAM-Zugriffe), der durch das Laden aller Schichten entsteht.
  • Makro-Pruning kann Speicherzugriffe reduzieren, führt aber oft zu einem signifikanten Genauigkeitsverlust durch den Verlust von Repräsentationskapazität.
• Limitierung bestehender Ansätze: Viele aktuelle Differentiable Architecture Search (DAS)-Methoden erfordern mehrstufige Pipelines, manuelle Heuristiken für die Wichtigkeit, vordefinierte Sparsity-Ziele oder teure Nachjustierungen (Fine-Tuning) nach der Suche.

2. Methodik: HiAP (Hierarchical Auto-Pruning)

HiAP ist ein Framework, das das Pruning als ein einziges, end-to-end lernbares Problem formuliert, ohne manuelle Eingriffe.

• Hierarchische Gating-Mechanismen:
  Das System führt stochastische Gumbel-Sigmoid-Gates auf zwei Ebenen ein, die gleichzeitig optimiert werden:

  1. Makro-Ebene: Gates ($g_{l,h}, b_l$) steuern das Entfernen ganzer Attention-Heads oder ganzer Feed-Forward-Network (FFN)-Blöcke. Dies adressiert den speichergebundenen Overhead (Memory-Bound).
  2. Mikro-Ebene: Gates ($d_{l,h,j}, c_{l,k}$) selektiv entfernen interne Dimensionen innerhalb der Attention-Heads oder einzelne Neuronen in den FFN-Schichten. Dies adressiert die rechengebundenen Operationen (Compute-Bound).

• Differentiable Cost Modeling:
  Anstatt auf heuristischen Rankings zu basieren, integriert HiAP eine exakte, differenzierbare Berechnung der Multiply-Accumulate-Operationen (MACs) direkt in die Verlustfunktion.

  • Die Kosten werden in statische Overheads und dynamische Kosten (gesteuert durch die Gates) zerlegt.
  • Ein linearer Zerlegungsterm ($\mathbb{E}[C]$) ermöglicht es, Hardware-Strafen präzise einzelnen Strukturen zuzuordnen.

• Training und Konvergenz (Single-Phase):

  • Gumbel-Sigmoid Relaxation: Während des Trainings werden die binären Gates als kontinuierliche Werte $\hat{z} \in (0,1)$ modelliert, um Gradientenfluss zu ermöglichen.
  • Temperatur-Annealing: Die Temperatur $\tau$ wird während des Trainings von einem hohen Wert (stochastisches Verhalten, ähnlich Dropout) auf einen niedrigen Wert (deterministische binäre Entscheidung) abgekühlt.
  • Feasibility Penalties: Um das „Layer Collapse" (das vollständige Löschen ganzer Schichten vor der Anpassung der Gewichte) zu verhindern, werden explizite Strafen eingeführt, die eine Mindestanzahl an aktiven Heads, Dimensionen und Neuronen garantieren.
  • Ergebnis: Am Ende des Trainings werden die Gates bei einem Schwellenwert (z. B. 0.5) „gehärtet", wodurch ein physisch komprimiertes, dichtes Subnetzwerk entsteht, das sofort für die Inferenz bereit ist, ohne weiteres Fine-Tuning.

3. Schlüsselbeiträge

1. Einheitliches Framework: HiAP vereinigt Makro-Pruning (Blöcke/Heads) und Mikro-Pruning (Dimensionen/Neuronen) in einem einzigen differentierbaren Ansatz.
2. Budget-Bewusstes Lernen: Das Netzwerk entdeckt autonom seine optimale Sub-Architektur innerhalb eines einzigen Trainingsdurchlaufs. Es werden keine manuellen Wichtigkeitsmetriken, Proxy-Rankings oder teure Nachoptimierungsphasen benötigt.
3. Physische Extraktion: Im Gegensatz zu Methoden, die nur „weiche" Masken erzeugen, liefert HiAP ein physisch zugeschnittenes, hardware-freundliches Subnetzwerk, das native Beschleunigung auf Standard-Hardware ermöglicht.
4. Theoretische Fundierung: Das Paper liefert Beweise dafür, dass der Suchraum von HiAP eine strikte Obermenge von reinen Makro-Ansätzen ist und dass die erwarteten Kosten linear in die Gates zerlegbar sind.

4. Ergebnisse

Die Methode wurde auf CIFAR-10 (mit ViT-Tiny) und ImageNet-1K (mit DeiT-Small) evaluiert.

• ImageNet-1K (DeiT-Small):
  • HiAP reduzierte die MACs von 4,6G auf 3,1G (ca. 33% Reduktion) bei einer Top-1-Genauigkeit von 79,1% (nur -0,75% im Vergleich zum dichten Basismodell).
  • Im Vergleich zu State-of-the-Art-Methoden wie GOHSP oder ViT-Slim erreicht HiAP eine wettbewerbsfähige Genauigkeit-Effizienz-Pareto-Front, jedoch mit einem deutlich vereinfachten, einstufigen Prozess.
• CIFAR-10 (ViT-Tiny):
  • HiAP übertraf manuelle Heuristiken (wie $\ell_1$-Norm-Ranking) und Uniform-Ratio-Baselines bei gleicher MAC-Reduktion.
  • Bei einer 33%igen Reduktion wurde eine Genauigkeitssteigerung von +0,93% gegenüber der Uniform-Baseline erzielt.
• Hardware-Effizienz:
  • Auf einer GPU wurde eine Latenzreduktion von 5,57 ms auf 3,86 ms gemessen (ca. 1,44-fache Beschleunigung), was bestätigt, dass die gefundenen Subnetzwerke keine speziellen Sparse-Hardware-Engines benötigen.

5. Bedeutung und Fazit

HiAP stellt einen Paradigmenwechsel dar, indem es das Pruning von ViTs von manuell gestalteten Heuristiken und mehrstufigen Pipelines hin zu einem automatisierten, budgetbewussten Lernproblem führt.

• Praktische Relevanz: Durch die gleichzeitige Optimierung von Speicherbandbreite (durch Makro-Pruning) und Rechenleistung (durch Mikro-Pruning) entstehen Modelle, die effizient auf Edge-Geräten laufen.
• Vereinfachung: Die Eliminierung von Nachjustierungsphasen und manueller Eingriffe macht die Methode besonders attraktiv für den industriellen Einsatz.
• Zukunftsausblick: Während HiAP derzeit MACs optimiert, könnten zukünftige Arbeiten die Lücke zwischen theoretischen MACs und tatsächlicher Latenz/Energie durch plattformspezifische Kalibrierung schließen und die Methode mit Token-Pruning oder Quantisierung kombinieren.

Zusammenfassend demonstriert HiAP, dass komplexe, mehrstufige Suchverfahren durch einen eleganten, einstufigen stochastischen Ansatz ersetzt werden können, der sowohl die Architektur als auch die Gewichte gleichzeitig anpasst, um hochleistungsfähige und effiziente Vision Transformer zu erzeugen.