Accelerating Vision Transformers with Adaptive Patch Sizes

1. Il Problema

I Vision Transformer (ViT) hanno stabilito un nuovo standard nel riconoscimento visivo, ma la loro scalabilità è limitata dal costo quadratico dell'attenzione self-attention rispetto alla lunghezza della sequenza di input.

• Limitazione attuale: I ViT tradizionali dividono le immagini in patch di dimensioni fisse e uniformi, indipendentemente dal contenuto dell'immagine.
• Conseguenza: Per immagini ad alta risoluzione, questo genera sequenze di token estremamente lunghe e ridondanti. Ad esempio, un'area uniforme (come un cielo blu) richiede lo stesso numero di calcoli e token di una regione complessa (come il volto di un uccello), sprecando risorse computazionali.
• Limiti delle soluzioni esistenti: Metodi precedenti che riducono i token (pruning o merging) spesso utilizzano rapporti di riduzione fissi (non adattivi alla complessità dell'immagine) o introducono forme irregolari e padding che annullano i guadagni di velocità nella pratica.

2. Metodologia: Adaptive Patch Transformer (APT)

Gli autori propongono APT, un approccio che adatta la dimensione delle patch in base al contenuto dell'immagine, ispirandosi ai tokenizer adattivi usati nei modelli linguistici (come Byte-Pair Encoding).

A. Assegnazione Adattiva delle Patch (Deciding Patch Sizes)

Invece di usare una griglia fissa, APT divide l'immagine in patch di dimensioni variabili (es. $16\times16$, $32\times32$, $64\times64$) seguendo una struttura simile a un quadtree:

1. Calcolo dell'Entropia: Viene calcolata l'entropia di Shannon su patch a diverse scale per misurare la complessità e la ridondanza del contenuto.
2. Soglia Dinamica: Le regioni con bassa entropia (omogenee, poco informative) vengono rappresentate con patch grandi. Le regioni ad alta entropia (dettagliate, complesse) ricevono patch piccole.
3. Processo Gerarchico: L'algoritmo inizia con patch grandi e le suddivide ricorsivamente solo se l'entropia supera una soglia $\tau$ specifica per quel livello.

B. Aggregazione delle Patch (Patch Aggregation)

Per gestire patch di dimensioni diverse all'interno dello stesso modello, APT utilizza una strategia di embedding ibrida:

• Le patch grandi vengono ridimensionate alla dimensione base ($p \times p$) e passate attraverso il layer di embedding standard.
• Parallelamente, le patch grandi vengono divise nelle loro sottopatch originali, che vengono anch'esse embedded.
• Zero-Initialized MLP: Le informazioni dalle sottopatch (dettagli ad alta risoluzione) vengono aggregate tramite un layer convoluzionale e poi fuse con l'embedding della patch ridimensionata utilizzando un MLP inizializzato a zero.
  • Vantaggio: Questo permette al modello di apprendere gradualmente a incorporare i dettagli ad alta risoluzione senza degradare le prestazioni iniziali, facilitando una convergenza rapida (anche in 1 epoca) su checkpoint pre-addestrati.

C. Gestione delle Sequenze Variabili

Poiché il numero di token varia per ogni immagine:

• Sequence Packing: I token di un batch vengono concatenati in una singola sequenza lunga, utilizzando una maschera a blocchi diagonali per garantire che l'attenzione avvenga solo tra token della stessa immagine.
• Compatibilità: Questo approccio è nativamente supportato da backend efficienti come FlashAttention, evitando l'overhead del padding.
• Posizional Encoding: Vengono utilizzati encoding posizionali interpolati (simili a NaViT) per gestire le diverse scale spaziali mantenendo la coerenza spaziale.

3. Contributi Chiave

1. Introduzione di APT: Un metodo che accelera l'inferenza e l'addestramento dei ViT fino al 40-50% riducendo drasticamente il numero di token di input senza perdere accuratezza.
2. Adattabilità e Velocità di Convergenza: APT può essere applicato a ViT pre-addestrati (es. su ImageNet) e raggiungere le prestazioni originali in sole 1 epoca di fine-tuning, grazie all'uso dell'MLP inizializzato a zero.
3. Generalizzazione ai Task Densi: A differenza di molti metodi di pruning, APT funziona efficacemente anche su task di predizione densa (rilevamento oggetti, segmentazione semantica) e task vision-language (VQA), mantenendo la fedeltà spaziale necessaria.
4. Efficienza su Alta Risoluzione: I guadagni di velocità sono proporzionali alla risoluzione e alla dimensione del modello, offrendo i benefici maggiori su immagini ad alta risoluzione e modelli grandi (ViT-H).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet, COCO (rilevamento), ADE20K (segmentazione) e vari benchmark VQA (LLaVA).

• Classificazione (ImageNet):
  • Su ViT-L e ViT-H, APT ottiene un aumento del throughput del 40% e 50% rispettivamente.
  • Mantiene l'accuratezza del modello baseline (es. 86.1% su ViT-L/336) con un fine-tuning di 1 epoca.
  • Riduce il tempo di addestramento wall-clock fino al 30-60% a seconda della risoluzione.
• Task Densi (Rilevamento e Segmentazione):
  • Su COCO (1536x1536), APT riduce i token del 30% ottenendo un speedup del 30% nel rilevamento oggetti, con mAP quasi identico al baseline.
  • Su ADE20K, riduce i token del 28-32% mantenendo l'accuratezza IoU, dimostrando che la riduzione dei token non compromette la precisione a livello di pixel.
• Vision-Language (VQA):
  • Integrato in LLaVA-1.5, APT aumenta il throughput del 22-23% mantenendo o superando le prestazioni su benchmark come GQA e VQAv2.
• Confronto con Baseline:
  • APT supera metodi di merging a livello di layer (come ToMe, EViT) sia in termini di accuratezza che di throughput, specialmente quando i baselines vengono ri-implementati con FlashAttention per un confronto equo.

5. Significato e Impatto

Il lavoro di APT rappresenta un passo significativo verso l'efficienza dei modelli visivi:

• Riduzione dei Costi Computazionali: Dimostra che è possibile addestrare e inferire modelli SOTA con budget computazionali inferiori, rendendo accessibili modelli grandi su hardware più limitato.
• Paradigma Adattivo: Sposta il paradigma dai patch fissi a una tokenizzazione visiva adattiva, colmando il divario di efficienza tra i modelli linguistici e quelli visivi.
• Praticità: La capacità di adattarsi a modelli pre-addestrati esistenti con un minimo di ri-addestramento lo rende una soluzione immediatamente applicabile nell'industria e nella ricerca, senza la necessità di addestrare da zero.

In sintesi, APT risolve il problema della ridondanza nelle immagini ad alta risoluzione sfruttando la struttura intrinseca dei dati visivi, offrendo un compromesso ottimale tra velocità e accuratezza.