Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

Technische Samenvatting: Reg4Pru – Regularisatie door Willekeurige Token-routing voor Token Pruning

Probleemstelling
Vision Transformers (ViTs) zijn fundamenteel geworden voor moderne visiemodellen vanwege hun schaalbaarheid en generalisatievermogen. Hun afhankelijkheid van self-attention mechanismen resulteert echter in een computationele complexiteit die kwadratisch ($O(N^2)$) schaalt met het aantal tokens $N$. Deze bottleneck is bijzonder acuut in medische beeldsegmentatie, waar hoge-resolutie inputs vereist zijn om kritieke diagnostische details te behouden.

Hoewel methoden voor token pruning (het verwijderen van redundante tokens) en reactivatiestrategieën (het herintroduceren van geprunede tokens tijdens inferentie voor dichte pixel-gewijze voorspellingen) een weg naar efficiëntie bieden, observeren de auteurs een significante prestatieafname in diepe lagen. Specifiek wanneer geprunede tokens tijdens inferentie worden gereactiveerd, vertonen zij drastisch andere representaties vergeleken met hun staat tijdens de training. Deze verschuiving in de distributie tussen training en test (train-test distribution shift) leidt tot instabiliteit en slechte prestaties bij dichte voorspellingen, wat contra-intuïtief ertoe leidt dat de prestaties afnemen naarmate de diepte van het model toeneemt.

Methodologie
Het paper stelt Reg4Pru voor, een regularisatietechniek tijdens de training die ontworpen is om de instabiliteit die geassocieerd wordt met token pruning en reactivatie te mitigeren. Het kernmechanisme bestaat uit willekeurige token-routing tijdens de trainingsfase om het effect van pruning te simuleren zonder tokens permanent te verwijderen.

1. Willekeurige Token-routing: Tijdens de forward pass van de training selecteert het model een willekeurig routingbereik $[l, n]$ tussen intermediaire transformer-blokken. Een willekeurige subset van tokens ($\tau_r$) wordt gekozen om de berekening van de blokken van $l$ tot $n$ over te slaan, terwijl de overige tokens ($\tau_k$) normaal worden verwerkt.
2. Gradiëntstroom en Representatie: De gerouteerde tokens omzeilen specifieke blokken, wat effectief hun gradiënten met betrekking tot die lagen doorsnijdt. Echter, door zowel de subset van gerouteerde tokens als het routingbereik $[l, n]$ bij elke stap te randomiseren, dragen verschillende tokens in verschillende iteraties bij aan verschillende gradiënten voor verschillende blokken. Dit dwingt het model om diepte-invariante representaties te leren, waardoor tokens op elke diepte tijdens inferentie kunnen worden gereactiveerd zonder last te hebben van een mismatch in representatie.
3. Inferentiestrategie: Tijdens de inferentie voorspelt een apart beleidsnetwerk (policy network) welke tokens gepruned moeten worden op basis van een hiërarchisch keep-ratio schema. Geprunede tokens worden tijdelijk verwijderd om berekening te besparen en worden gereactiveerd bij het laatste blok om dichte segmentatiemaskers te genereren.
4. Loss-functies: De trainingsdoelstelling combineert een standaard ratio-gebaseerde loss (om token-budgetten af te dwingen) en, optioneel, een geïnformeerde beleids-loss ($L_{pol}$) die de grondwaarheid (ground truth) segmentatiemaskers gebruikt om pruning-beslissingen te sturen. Het routingmechanisme fungeert als een regularisator naast deze losses.

Belangrijkste Bijdragen

• Diepte-wijze Instabiliteitsanalyse: De auteurs identificeren en karakteriseren een specifieke instabiliteit in token pruning voor dichte voorspellingen, waarbij de prestaties afnemen in diepere lagen door de reactivatie van tokens die tijdens de training niet aan de volledige diepte van het netwerk zijn blootgesteld.
• Reg4Pru Framework: Zij introduceren een training-time token routing strategie met willekeurige grenzen. In tegenstelling tot eerdere routingmethoden die vaste grenzen of geleerde routing-policies gebruiken, randomiseert Reg4Pru het routingbereik om de feature-distributie over alle pruning-dieptes te benaderen, waardoor het model wordt geregulariseerd.
• Empirische Validatie: De methode wordt gevalideerd op de FIVES retina bloedvat segmentatie dataset, waarbij wordt aangetoond dat willekeurige routing de prestaties aanzienlijk stabiliseert vergeleken met vaste routing of geen routing.

Resultaten
Experimenten werden uitgevoerd op de FIVES dataset (2048$\times$2048 afbeeldingen, geresized naar kwadranten voor throughput-testen) met behulp van een DINOv2-small geïnitialiseerde EoMT architectuur.

• Prestatiewinst: Vergeleken met een gepruned model getraind zonder routing, verbeterde Reg4Pru de Average Precision (AP) met een absolute 46% (van 25,60% naar 71,67% voor de $L_{ratio}$ configuratie).
• Efficiëntie: De Reg4Pru-configuratie behaalde een 29% relatieve versnelling in wall-clock tijd vergeleken met de niet-gepruned baseline, terwijl een hoge segmentatiekwaliteit behouden bleef.
• Stabiliteit: Diepte-wijze analyse (Figuur 5) toonde aan dat terwijl niet-gerouteerde geprunede modellen leden onder een scherpe daling in AP bij diepere blokken, Reg4Pru een stabiele prestatie behield. Fixed-bound routing mitigeerde enige instabiliteit, maar slaagde er niet in de prestatiedalingen na het einde van het vaste bereik te voorkomen, wat de noodzaak van randomisering bevestigt.
• Vergelijking: Reg4Pru presteerde beter dan Token Merging (ToMe) methoden in termen van AP bij vergelijkbare throughput-configuraties, met name wanneer ToMe werd toegepast op zowel MHSA als MLP lagen.

Significantie en Claims
Het paper claimt dat Reg4Pru een waardevolle regularisator is voor token reductie strategieën in taken voor dichte voorspellingen. De auteurs stellen dat de geobserveerde instabiliteit bij pruning een gevolg is van de verschuiving in de distributie tussen training en test met betrekking tot token-representaties op verschillende dieptes. Door het routingbereik tijdens de training te randomiseren, dwingt Reg4Pru het model om representaties te leren die robuust zijn tegen het overslaan op diverse dieptes, waardoor de kloof tussen training en inferentie effectief wordt overbrugd.

De auteurs blijven bescheiden over de generaliseerbaarheid van hun bevindingen, waarbij zij opmerken dat de experimenten werden uitgevoerd op een enkele hoog-resolutie medische dataset. Zij suggereren dat hoewel de methode effectief is voor segmentatie, verdere evaluatie vereist is om de toepasbaarheid op andere taken voor dichte voorspellingen (bijv. diepte-inschatting) en het potentieel als een algemene regularisatiemechanisme in combinatie met technieken zoals dropout of stochastic depth te beoordelen. Zij erkennen ook dat hoewel zij geïnformeerde policies onderzochten, ratio-gebaseerde policies momenteel superieure prestaties leveren in hun specifieke opzet.