Riepilogo Tecnico: TAS-LoRA

Enunciato del Problema

La Ricerca di Architetture Transformer (TAS) mira a scoprire automaticamente architetture ottimali per Vision Transformer (ViT), riducendo lo sforzo manuale richiesto per la progettazione. I metodi TAS allo stato dell'arte si basano principalmente su approcci one-shot, che addestrano una "supernet" contenente tutti i possibili subnet. In queste supernets con pesi intrecciati, tutti i subnet condividono un insieme comune di pesi.

Il documento identifica una limitazione critica in questo paradigma: il problema del collasso delle caratteristiche. Poiché la supernet è ottimizzata per minimizzare una perdita globale su tutti i subnet, i pesi condivisi tendono ad apprendere caratteristiche generiche che funzionano adeguatamente per architetture diverse ma non riescono a catturare le caratteristiche specifiche del subnet. Di conseguenza, i singoli subnet non possono sfruttare appieno le loro proprietà architetturali uniche (ad esempio, numeri specifici di teste di attenzione o rapporti MLP), portando a prestazioni subottimali rispetto ai subnet addestrati indipendentemente da zero. L'analisi empirica nel documento conferma che i subnet all'interno di supernets con pesi intrecciati esibiscono rappresentazioni di caratteristiche quasi identiche (alta similarità del coseno), mentre i subnet addestrati indipendentemente sviluppano rappresentazioni distinte.

Metodologia: TAS-LoRA

Per affrontare il collasso delle caratteristiche mantenendo l'efficienza computazionale, gli autori propongono TAS-LoRA (Ricerca di Architetture Transformer con Mixture-of-LoRA Experts). Il metodo integra l'Adattamento a Basso Rango (LoRA) nel processo di addestramento della supernet per abilitare l'apprendimento di caratteristiche specializzate per ogni subnet senza incorrere nel costo proibitivo di assegnare parametri unici a ogni possibile configurazione di subnet.

1. Mixture-of-LoRA-Experts (MoLE)

Invece di assegnare un modulo LoRA unico a ogni subnet (il che è computazionalmente non fattibile a causa dell'enorme spazio di ricerca), TAS-LoRA impiega una strategia Mixture-of-LoRA-Experts (MoLE).

• Architettura: Gli esperti LoRA sono applicati ai livelli lineari della supernet (inclusi i livelli di proiezione QKV e i livelli MLP).
• Instradamento Dinamico: Un router leggero seleziona e combina dinamicamente un insieme limitato di esperti LoRA in base all'architettura specifica del subnet campionata.
• Efficienza nell'Inferenza: Al momento dell'inferenza, i parametri LoRA selezionati vengono fusi nei pesi originali della supernet. Poiché la configurazione del subnet è fissa durante l'inferenza, il router viene eseguito una sola volta e i pesi fusi introducono nessun sovraccarico computazionale aggiuntivo.

2. Progettazione del Router

Il router è progettato per catturare le dipendenze sequenziali dei blocchi transformer:

• Input: Accetta sia attributi a livello di blocco (ad esempio, numero di teste di attenzione, rapporto MLP) sia attributi a livello di subnet (ad esempio, dimensione dell'embedding, profondità) come input.
• Meccanismo: Questi attributi sono codificati da un livello di embedding apprendibile e elaborati da un LSTM per catturare le dipendenze sequenziali tra i blocchi. L'LSTM produce rappresentazioni nascoste, che vengono inviate attraverso un "testa del router" completamente connessa per generare una distribuzione di probabilità sugli esperti LoRA disponibili per ogni livello.

3. Inizializzazione del Router per Gruppi

Una sfida chiave identificata è che durante le prime fasi dell'addestramento, il router tende ad assegnare pesi simili a tutti gli esperti, portando a ridondanza e impedendo la specializzazione. Per mitigare ciò, gli autori introducono una tecnica di inizializzazione del router per gruppi:

• Raggruppamento: I blocchi transformer tra subnet diverse sono raggruppati in base a similarità architetturali (specificamente numero di teste di attenzione, rapporto MLP e dimensione dell'embedding).
• Inizializzazione: Il numero di esperti LoRA è impostato per corrispondere al numero di gruppi. Il router è inizializzato in modo che ogni esperto sia distorto verso un gruppo distinto. Ciò è ottenuto impostando termini di bias specifici nella testa del router su un valore positivo $\beta$ per la coppia gruppo-esperto corrispondente, mentre gli altri vengono azzerati.
• Effetto: Questo costringe il router ad assegnare esperti distinti a diversi gruppi architetturali nelle prime fasi dell'addestramento, incoraggiando un apprendimento di caratteristiche diversificato. Man mano che l'addestramento procede, il router aggiorna dinamicamente queste assegnazioni in base alle proprietà del subnet.

Contributi Chiave

1. Identificazione del Collasso delle Caratteristiche in TAS: Il documento identifica e analizza formalmente il problema del collasso delle caratteristiche nelle supernets con pesi intrecciati, dimostrando che i pesi condivisi impediscono ai subnet di apprendere rappresentazioni specializzate.
2. Framework TAS-LoRA: Un metodo innovativo che combina LoRA con un meccanismo di instradamento dinamico (MoLE) per apprendere in modo efficiente caratteristiche specifiche del subnet. Consente ai subnet di raffinare le caratteristiche condivise in base alle loro proprietà architetturali senza riaddestrare l'intera supernet da zero.
3. Inizializzazione del Router per Gruppi: Una tecnica per prevenire la ridondanza degli esperti durante le prime fasi dell'addestramento, garantendo che gli esperti LoRA apprendano rappresentazioni di caratteristiche distinte.
4. Prestazioni allo Stato dell'Arte: Esperimenti estensivi dimostrano che TAS-LoRA supera significativamente i metodi TAS esistenti (ad esempio AutoFormer, DYNAS, FocusFormer) su vari spazi di ricerca e dataset.

Risultati Sperimentali

Gli autori hanno valutato TAS-LoRA sul dataset ImageNet e su diversi benchmark di transfer learning (CIFAR-10/100, Flowers, CARS, INAT-19).

• Prestazioni su ImageNet: TAS-LoRA ha migliorato costantemente l'accuratezza top-1 delle supernets addestrate con AutoFormer e DYNAS negli spazi di ricerca Tiny, Small e Base.
  • Ad esempio, nello spazio AutoFormer-T, TAS-LoRA ha raggiunto un'accuratezza del 75,7%, superando AutoFormer originale (74,7%) e altri metodi SOTA come PreNAS-Ti (77,1% ma con meno flessibilità).
  • Nello spazio AutoFormer-B, TAS-LoRA-B ha raggiunto un'accuratezza del 82,6% con 54M di parametri, superando AutoFormer-B (82,4%) e eguagliando o superando modelli più grandi come GLiT-B (82,3% con 96M di parametri).
• Transfer Learning: TAS-LoRA ha dimostrato una generalizzazione superiore, ottenendo risultati competitivi o allo stato dell'arte su CIFAR-10/100, Flowers, CARS e INAT-19.
• Studi di Ablazione:
  • Inizializzazione: L'inizializzazione per gruppi ha superato costantemente l'inizializzazione casuale, confermando il suo ruolo nel ridurre la ridondanza degli esperti.
  • Similarità delle Caratteristiche: L'analisi della similarità del coseno ha mostrato che TAS-LoRA con inizializzazione per gruppi ha prodotto rappresentazioni di caratteristiche significativamente più diversificate tra gli esperti rispetto all'inizializzazione casuale o ai metodi standard con pesi intrecciati.
  • Confronto con il Fine-tuning: TAS-LoRA ha superato una supernet completamente fine-tunata (dove tutti i pesi vengono aggiornati) nonostante l'addestramento per soli 50 epoche sui livelli MoLE. Ciò indica che i guadagni derivano dalla capacità di apprendimento di caratteristiche specifiche del subnet di MoLE, non semplicemente da un addestramento prolungato.

Significato e Affermazioni

Il documento afferma che TAS-LoRA rappresenta un avanzamento significativo nella Ricerca di Architetture Transformer risolvendo efficacemente il problema del collasso delle caratteristiche intrinseco alle supernets con pesi intrecciati.

• Efficienza vs Prestazioni: A differenza dei metodi che richiedono l'addestramento di più supernets (ad esempio MoS) o il riaddestramento da zero quando cambiano i vincoli hardware (ad esempio PreNAS), TAS-LoRA addestra una singola supernet con livelli di adattamento specifici per subnet. Ciò consente una flessibile adattamento a nuovi vincoli hardware (ad esempio FLOPs, latenza) senza riaddestrare l'intero modello.
• Scalabilità: Utilizzando un insieme limitato di esperti e un router dinamico, il metodo scala a supernets con enormi numeri di possibili configurazioni di subnet senza un aumento lineare dei parametri o dei costi computazionali.
• Generalizzabilità: L'approccio è presentato come una soluzione generale applicabile a qualsiasi supernet con pesi intrecciati, migliorando le prestazioni indipendentemente dalla strategia di addestramento sottostante (ad esempio AutoFormer o DYNAS).

Gli autori concludono che TAS-LoRA mitiga con successo il collasso delle caratteristiche, permettendo ai subnet di apprendere rappresentazioni specializzate e raggiungendo prestazioni allo stato dell'arte sia nelle attività di ricerca architetturale che nel transfer learning.