Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

Sintesi Tecnica: Reg4Pru – Regolarizzazione tramite Routing Casuale dei Token per il Pruning dei Token

Problema
I Vision Transformer (ViT) sono diventati fondamentali per i moderni modelli di visione grazie alle loro capacità di scalabilità e generalizzazione. Tuttavia, la loro dipendenza dai meccanismi di auto-attenzione comporta una complessità computazionale che scala quadraticamente ($O(N^2)$) con il numero di token $N$. Questo collo di bottiglia è particolarmente acuto nella segmentazione di immagini mediche, dove sono richiesti input ad alta risoluzione per preservare dettagli diagnostici critici.

Sebbene i metodi di pruning dei token (rimozione dei token ridondanti) e le strategie di riattivazione (reintroduzione dei token rimossi durante l'inferenza per predizioni dense pixel-wise) offrano una via verso l'efficienza, gli autori osservano una significativa degradazione delle prestazioni negli strati profondi. Nello specifico, quando i token rimossi vengono riattivati durante l'inferenza, essi esibiscono rappresentazioni drasticamente diverse rispetto al loro stato durante l'addestramento. Questo spostamento della distribuzione tra addestramento e test (train-test distribution shift) causa instabilità e scarse prestazioni nelle predizioni dense, provocando paradossalmente un calo delle prestazioni all'aumentare della profondità del modello.

Metodologia
Il paper propone Reg4Pru, una tecnica di regolarizzazione durante l'addestramento progettata per mitigare l'instabilità associata al pruning e alla riattivazione dei token. Il meccanismo centrale prevede il routing casuale dei token durante la fase di addestramento per simulare gli effetti del pruning senza scartare permanentemente i token.

1. Routing Casuale dei Token: Durante il passaggio in avanti (forward pass) dell'addestramento, il modello seleziona un intervallo di routing casuale $[l, n]$ tra i blocchi intermedi del transformer. Un sottoinsieme casuale di token ($\tau_r$) viene scelto per saltare il calcolo dei blocchi da $l$ a $n$, mentre i restanti token ($\tau_k$) vengono elaborati normalmente.
2. Flusso del Gradiente e Rappresentazione: I token instradati saltano blocchi specifici, tagliando efficacemente i loro gradienti rispetto a quegli strati. Tuttavia, randomizzando sia il sottoinsieme di token instradati che l'intervallo di routing $[l, n]$ ad ogni step, diversi token contribuiscono con gradienti a diversi blocchi attraverso le iterazioni. Ciò costringe il modello a imparare rappresentazioni invarianti rispetto alla profondità, garantendo che i token possano essere riattivati a qualsiasi profondità durante l'inferenza senza subire un disallineamento della rappresentazione.
3. Strategia di Inferenza: All'inferenza, una rete di policy separata predice quali token eliminare basandosi su uno schema gerarchico di rapporto di mantenimento (keep-ratio schedule). I token rimossi vengono temporaneamente eliminati per risparmiare computazione e vengono riattivati all'ultimo blocco per generare maschere di segmentazione dense.
4. Funzioni di Perdita: L'obiettivo di addestramento combina una funzione di perdita standard basata sul rapporto (per imporre i budget dei token) e, opzionalmente, una loss di policy informata ($L_{pol}$) che utilizza le maschere di segmentazione della ground truth per guidare le decisioni di pruning. Il meccanismo di routing agisce come un regolarizzatore insieme a queste perdite.

Contributi Chiave

• Analisi dell'Instabilità per Profondità: Gli autori identificano e caratterizzano una specifica instabilità nel pruning dei token per la predizione densa, in cui le prestazioni degradano negli strati più profondi a causa della riattivazione di token che non sono stati esposti alla piena profondità della rete durante l'addestramento.
• Framework Reg4Pru: Introducono una strategia di routing dei token durante l'addestramento con limiti casuali. A differenza dei precedenti metodi di routing che utilizzano limiti fissi o policy di routing apprese, Reg4Pru randomizza l'intervallo di routing per approssimare la distribuzione delle feature attraverso tutte le profondità di pruning, regolarizzando così il modello.
• Validazione Empirica: Il metodo è validato sul dataset di segmentazione dei vasi sanguigni retinici FIVES, dimostrando che il routing casuale stabilizza significativamente le prestazioni rispetto al routing a limiti fissi o all'assenza di routing.

Risultati
Gli esperimenti sono stati condotti sul dataset FIVES (immagini 2048$\times$2048, ridimensionate in quadranti per i test di throughput) utilizzando un'architettura EoMT inizializzata con DINOv2-small.

• Guadagni di Prestazioni: Rispetto a un modello con pruning addestrato senza routing, Reg4Pru ha migliorato l'Average Precision (AP) di un valore assoluto del 46% (dal 25,60% al 71,67% per la configurazione $L_{ratio}$).
• Efficienza: La configurazione Reg4Pru ha ottenuto un'accelerazione relativa del 29% nel tempo di esecuzione (wall-clock time) rispetto al baseline non-pruned, mantenendo un'alta qualità di segmentazione.
• Stabilità: L'analisi per profondità (Figura 5) ha mostrato che, mentre i modelli con pruning non-routed subivano un brusco calo dell'AP nei blocchi più profondi, Reg4Pru manteneva prestazioni stabili. Il routing a limiti fissi mitigava parte dell'instabilità ma non riusciva a prevenire i cali di prestazione dopo la fine dell'intervallo fisso, confermando la necessità della randomizzazione.
• Confronto: Reg4Pru ha superato i metodi di Token Merging (ToMe) in termini di AP a parità di configurazioni di throughput, specialmente quando ToMe era applicato sia ai livelli MHSA che MLP.

Significatività e Rivendicazioni
Il paper sostiene che Reg4Pru sia un prezioso regolarizzatore per le strategie di riduzione dei token in compiti di predizione densa. Gli autori argomentano che l'instabilità osservata nel pruning è il risultato di uno spostamento della distribuzione train-test riguardo alle rappresentazioni dei token a diverse profondità. Randomizzando l'intervallo di routing durante l'addestramento, Reg4Pru costringe il modello a imparare rappresentazioni che siano robuste all'essere saltate a varie profondità, colmando efficacementamente il divario tra addestramento e inferenza.

Gli autori rimangono modesti riguardo alla generalizzabilità dei loro risultati, notando che gli esperimenti sono stati condotti su un singolo dataset medico ad alta risoluzione. Suggeriscono che, sebbene il metodo sia efficace per la segmentazione, sia necessaria un'ulteriore valutazione per valutarne l'applicabilità ad altri compiti di predizione densa (es. stima della profondità) e il suo potenziale come meccanismo di regolarizzazione generale quando combinato con tecniche come il dropout o lo stochastic depth. Riconoscono inoltre che, sebbene abbiano esplorato policy informate, le policy basate sul rapporto attualmente producono prestazioni superiori nel loro specifico setup.