💻 computer science

Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

Das Paper stellt Reg4Pru vor, eine Regularisierungstechnik, die den Leistungsverlust beim Token-Pruning für Vision Transformer mildert und dabei eine absolute Verbesserung der durchschnittlichen Präzision um 46 % sowie eine Beschleunigung um 29 % auf dem FIVES-Blutgefäß-Segmentierungsdatensatz erzielt.

Ursprüngliche Autoren: Julian Wyatt, Ronald Clark, Irina Voiculescu

Veröffentlicht 2026-02-04

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Julian Wyatt, Ronald Clark, Irina Voiculescu

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Das große Problem: Die „überfüllte Party“

Stellen Sie sich einen Vision Transformer (das KI-Gehirn hinter moderner Bilderkennung) wie eine riesige, hochkarätige Party vor. Die „Gäste“ sind Tokens – winzige Teile des Bildes, die die KI analysiert.

In einem Standard-Setup spricht jeder einzelne Gast mit jedem anderen Gast, um herauszufinden, was das Bild darstellt. Das funktioniert großartig für die Genauigkeit, ist aber unglaublich langsam. Wenn man die Anzahl der Gäste verdoppelt, vervierfacht sich der Gesprächsbedarf. Es ist, als würde man versuchen, ein bedeutungsvolles Gespräch in einem Stadion voller Menschen zu führen; es dauert eine Ewigkeit.

Um dies zu beheben, erfanden Forscher das Token Pruning. Dies ist wie ein Türsteher an der Tür, der die „langweiligen“ Gäste (redundante Teile des Bildes) hinauswirft, damit die Party schneller laufen kann.

Der Fehler: Der „Amnesie“-Effekt

Die Arbeit identifiziert ein großes Problem mit dieser Türsteher-Strategie.

Während des Trainings: Die KI lernt, indem sie die ganze Party betrachtet und dann gelegentlich Leute rauswirft, um Effizienz zu üben.
Während des Tests (Inferenz): Die KI wirft Leute dauerhaft raus, um Zeit zu sparen. Aber hier ist der Haken: Wenn sie ein detailliertes Ergebnis liefern muss (wie etwa die Zeichnung einer detaillierten Karte von Blutgefäßen), muss sie die rausgeworfenen Gäste zurückholen, um die fehlenden Details zu ergänzen.

Das Problem: Die KI wird verwirrt. Die „rausgeworfenen“ Gäste saßen in der Dunkelheit (sie haben Schichten übersprungen), während die „behaltenen“ Gäste feiern und lernen konnten. Wenn die KI versucht, die ausgeschlossenen Gäste wieder ins Licht zu holen, passen sie nicht mehr hinein. Sie haben eine „Amnesie“ darüber, was während ihrer Abwesenheit passiert ist. Dies führt dazu, dass die Leistung der KI einbricht, besonders in den tieferen, komplexeren Schichten des Netzwerks. Es ist, als würde man versuchen, ein komplexes Puzzle mit Teilen zu vollenden, die auf der Hälfte des Weges vergessen haben, wie das Bild aussah.

Die Lösung: Reg4Pru (Die „Zufällige Probe“)

Die Autoren schlagen eine neue Trainingsmethie namens Reg4Pru (Regularisation Through Random Token Routing) vor.

Betrachten Sie dies als ein Probespiel für die KI vor der eigentlichen Show.

Anstatt einfach nur Leute rauszuschmeißen und auf das Beste zu hoffen, spielt die KI während des Trainings ein Spiel, bei dem sie zufällig verschiedenen Gruppen von Gästen sagt, bestimmte Teile der Party für eine zufällige Zeit zu „überspringen“.

Zufälligkeit ist der Schlüssel: Manchmal überspringt Gast A zwei Räume; manchmal überspringt Gast B fünf Räume. Die „Übersprung-Zone“ ändert sich jedes Mal.
Das Ergebnis: Da die Gäste ständig angewiesen werden, zufällige Abschnitte zu überspringen und dann wieder einzusteigen, lernen sie, anpassungsfähig zu sein. Sie lernen, dass es keine Rolle spielt, wo sie in die Party eintreten oder wie lange sie weg waren – sie können immer noch hineinpassen und zum fertigen Bild beitragen.

Diese „zufällige Probe“ wirkt wie ein Stabilisator. Sie lehrt die KI, dass das Zurückholen von Gästen aus der „Dunkelheit“ sicher ist, was die Verwirrung und den Leistungsabfall verhindert, die bei früheren Methoden auftraten.

Die Ergebnisse: Schneller und Schlauer

Das Team testete dies auf einem medizinischen Datensatz namens FIVES, bei dem es darum geht, Blutgefäße in Augenbildern zu finden (eine Aufgabe, die hohe Detailgenauigkeit erfordert).

Der Vergleich: Sie verglichen ihre Methode mit einem Standard-„Türsteher“ (Pruning ohne diese neue Methode) und einer „No-Pruning“-Baseline.
Der Sieg:
- Genauigkeit: Die Standard-Pruning-Methode war in der Endphase schrecklich (der Präzisionsverlust war massiv). Reg4Pru behob dies und verbesserte die durchschnittliche Präzision im Vergleich zur Standard-Pruning-Methode um gewaltige 46 %.
- Geschwindigkeit: Sie erreichten diese hohe Genauigkeit und waren dennoch 29 % schneller als die langsame, nicht-geprunte Version.
- Visuelles: Die Arbeit zeigt Heatmaps, in denen die Standard-Pruning-Methode unscharfe, unordentliche Ergebnisse liefert, während Reg4Pru scharfe, klare Bilder von Blutgefäßen erzeugt, die fast so gut sind wie die der langsamen, nicht-geprunten Version.

Zusammenfassung

Reg4Pru ist ein Trainings-Trick, der verhindert, dass die KI verwirrt wird, wenn sie versucht, die Verarbeitung von Bildteilen zu beschleunigen, indem sie diese ignoriert. Durch das zufällige Üben von „Überspringen“ und „Wiedereinstiegen“ während des Trainings lernt die KI, stabil und präzise zu bleiben. Dies ermöglicht es ihr, hochauflösende medizinische Bilder viel schneller zu verarbeiten, ohne die feinen Details zu verlieren, die Ärzte benötigen.

Technisches Resümee: Reg4Pru – Regularisierung durch zufälliges Token-Routing für das Token-Pruning

Problemstellung
Vision Transformer (ViTs) sind aufgrund ihrer Skalierbarkeit und Generalisierungsfähigkeit zu einer Grundlage moderner Computer-Vision-Modelle geworden. Ihre Abhängigkeit von Self-Attention-Mechanismen führt jedoch zu einer Rechenkomplexität, die quadratisch ( $O(N^2)$ ) mit der Anzahl der Token $N$ skaliert. Dieser Engpass ist besonders gravierend in der medizinischen Bildsegmentierung, wo hochauflösende Eingaben erforderlich sind, um kritische diagnostische Details zu bewahren.

Obwohl Token-Pruning-Methoden (Entfernen redundanter Token) und Reaktivatierungsstrategien (Wiederaufnahme geprunter Token während der Inferenz für dichte pixelweise Vorhersagen) einen Weg zur Effizienz bieten, beobachten die Autoren eine signifikante Leistungsverschlechterung in tiefen Schichten. Speziell wenn geprunte Token während der Inferenz reaktiviert werden, weisen sie drastisch unterschiedliche Repräsentationen auf als in ihrem Zustand während des Trainings. Dieser Train-Test-Verteilungsunterschied führt zu Instabilität und schlechter dichter Vorhersageleistung, was kontraintuitiv dazu führt, dass die Leistung mit zunehmender Modelltiefe sinkt.

Methodik
Das Paper schlägt Reg4Pru vor, eine Regularisierungstechnik zur Trainingszeit, die darauf ausgelegt ist, die mit Token-Pruning und -Reaktivierung verbundene Instabilität zu mildern. Der Kernmechanismus besteht aus zufälligem Token-Routing während der Trainingsphase, um die Effekte des Prunings zu simulieren, ohne die Token dauerhaft zu verwerfen.

Zufälliges Token-Routing: Während des Forward-Passes des Trainings wählt das Modell einen zufälligen Routing-Bereich $[l, n]$ zwischen den intermediären Transformer-Blöcken. Eine zufällige Teilmenge von Token ( $\tau_r$ ) wird so gewählt, dass sie die Berechnung der Blöcke von $l$ bis $n$ überspringt, während die verbleibenden Token ( $\tau_k$ ) normal verarbeitet werden.
Gradientenfluss und Repräsentation: Die gerouteten Token umgehen spezifische Blöcke, was effektiv deren Gradienten in Bezug auf diese Schichten unterbricht. Da jedoch sowohl die Teilmenge der gerouteten Token als auch der Routing-Bereich $[l, n]$ in jedem Schritt randomisiert werden, tragen unterschiedliche Token in verschiedenen Iterationen Gradienten zu unterschiedlichen Blöcken bei. Dies zwingt das Modell dazu, tiefen-invariante Repräsentationen zu lernen, wodurch sichergestellt wird, dass Token in jeder Tiefe während der Inferenz reaktiviert werden können, ohne unter einem Repräsentations-Mismatch zu leiden.
Inferenz-Strategie: Bei der Inferenz sagt ein separates Policy-Netzwerk voraus, welche Token zu prunen sind, basierend auf einem hierarchischen Keep-Ratio-Schema. Geprunte Token werden vorübergehend entfernt, um Rechenleistung zu sparen, und werden am finalen Block reaktiviert, um dichte Segmentierungsmasken zu generieren.
Verlustfunktionen: Das Trainingsziel kombiniert einen Standard-Ratio-basierten Verlust (um Token-Budgets durchzusetzen) und optional einen informierten Policy-Verlust ( $L_{pol}$ ), der Ground-Truth-Segmentierungsmasken verwendet, um Entscheidungen zum Pruning zu leiten. Der Routing-Mechanismus fungt dabei als Regularisierer neben diesen Verlusten.

Wesentliche Beiträge

Tiefen-Instabilitätsanalyse: Die Autoren identifizieren und charakterisieren eine spezifische Instabilität beim Token-Pruning für die dichte Vorhersage, bei der die Leistung in tieferen Schichten aufgrund der Reaktivierung von Token sinkt, die während des Trainings nicht der vollen Tiefe des Netzwerks ausgesetzt waren.
Reg4Pru-Framework: Sie führen eine Token-Routing-Strategie zur Trainingszeit mit zufälligen Grenzen ein. Im Gegensatz zu bisherigen Routing-Methoden, die feste Grenzen oder gelernte Routing-Policies verwenden, randomisiert Reg4Pru den Routing-Bereich, um die Feature-Verteilung über alle Pruning-Tiefen hinweg zu approximieren und dadurch das Modell zu regularisieren.
Empirische Validierung: Die Methode wird auf dem FIVES-Datensatz zur Segmentierung von Netzhaut-Blutgefäßen validiert, wobei sie zeigt, dass zufälliges Routing die Leistung im Vergleich zu festem Routing oder keinem Routing signifikant stabilisiert.

Ergebnisse
Die Experimente wurden auf dem FIVES-Datensatz (2048 $\times$ 2048 Bilder, zur Durchsatzprüfung auf Quadranten skaliert) unter Verwendung einer mit DINOv2-small initialisierten EoMT-Architektur durchgeführt.

Leistungsgewinne: Im Vergleich zu einem geprunten Modell, das ohne Routing trainiert wurde, verbesserte Reg4Pru die Average Precision (AP) absolut um 46 % (von 25,60 % auf 71,67 % für die $L_{ratio}$ -Konfiguration).
Effizienz: Die Reg4Pru-Konfiguration erreichte eine relative Beschleunigung von 29 % in der Wandlaufzeit gegenüber der nicht-geprunten Baseline, während sie eine hohe Segmentierungsqualität beibehielt.
Stabilität: Die Tiefen-Analyse (Abbildung 5) zeigte, dass während nicht-geroutete geprunte Modelle einen starken Abfall der AP über tiefere Blöcke hinweg erlitten, Reg4Pru eine stabile Leistung beibehielt. Fixed-Bound-Routing milderte die Instabilität teilweise ab, konnte jedoch den Leistungsabfall nach Ende des festen Bereichs nicht verhindern, was die Notwendigkeit der Randomisierung bestätigt.
Vergleich: Reg4Pru übertraf Token-Merging (ToMe)-Methoden in Bezug auf die AP bei vergleichbaren Durchsatz-Konfigurationen, insbesondere wenn ToMe über sowohl MHSA- als auch MLP-Schichten angewendet wurde.

Bedeutung und Ansprüche
Das Paper behauptet, dass Reg4Pru ein wertvoller Regularisierer für Token-Reduktionsstrategien bei Aufgaben der dichten Vorhersage ist. Die Autoren argumentieren, dass die beobachtete Instabilität beim Pruning ein Resultat des Train-Test-Verteilungsunterschieds hinsichtlich der Token-Repräsentationen bei verschiedenen Tiefen ist. Durch die Randomisierung des Routing-Bereichs während des Trainings erzwingt Reg4Pra, dass das Modell Repräsentationen lernt, die robust gegenüber dem Überspringen an verschiedenen Tiefen sind, wodurch die Lücke zwischen Training und Inferenz effektiv geschlossen wird.

Die Autoren bleiben bescheiden hinsichtlich der Generalisierbarkeit ihrer Ergebnisse und merken an, dass die Experimente auf einem einzelnen hochauflösenden medizinischen Datensatz durchgeführt wurden. Sie legen nahe, dass die Methode zwar effektiv für die Segmentierung ist, weitere Evaluierungen jedoch erforderlich sind, um ihre Anwendbarkeit auf andere Aufgaben der dichten Vorhersage (z. B. Tiefenschätzung) und ihr Potenzial als allgemeiner Regularisierungsmechanismus in Kombination mit Techniken wie Dropout oder Stochastic Depth zu bewerten. Sie räumen zudem ein, dass sie zwar informierte Policies untersucht haben, Ratio-basierte Policies in ihrem spezifischen Setup derzeit eine überlegene Leistung erbringen.