Riepilogo Tecnico: Allineamento Multi-Granulare delle Traiettorie (MTA) per la Distillazione di LLM

Enunciato del Problema

I metodi esistenti di Distillazione della Conoscenza (KD) per i Modelli Linguistici di Grande Dimensione (LLM) si concentrano principalmente sull'allineamento delle distribuzioni di probabilità di output (logit) o degli stati nascosti intermedi a livelli fissi e statici. Sebbene efficaci nel trasferire segnali locali, questi approcci impiegano spesso una strategia di allineamento uniforme a livello di token attraverso tutte le profondità della rete.

Il documento sostiene che questo trattamento uniforme è in disallineamento con la natura gerarchica del linguaggio umano e l'organizzazione funzionale dei modelli Transformer. La ricerca linguistica e gli studi sull'interpretabilità dei Transformer indicano che:

1. I livelli inferiori enfatizzano caratteristiche di superficie, lessicali e di recupero fattuale.
2. I livelli superiori codificano progressivamente informazioni semantiche astratte e ragionamento composizionale.

Trattando tutti i token in isolamento e applicando lo stesso obiettivo di allineamento in modo uniforme, i metodi attuali non riescono a catturare la traiettoria evolutiva progressiva delle rappresentazioni mentre si trasformano da unità lessicali in strutture semantiche astratte. Questa limitazione restringe la capacità del modello studente di generalizzare verso compiti downstream.

Metodologia: Allineamento Multi-Granulare delle Traiettorie (MTA)

Per affrontare queste limitazioni, gli autori propongono MTA, un framework modulare che allinea la traiettoria rappresentazionale dei modelli insegnante e studente utilizzando una strategia adattiva ai livelli. MTA opera sull'ipotesi che la distillazione debba rispettare i ruoli semantici dipendenti dalla profondità dei livelli Transformer.

Componenti Principali

1. Span Multi-Granulari Adattivi ai Livelli:

   • Livelli Inferiori: Gli allineamenti sono eseguiti a livello di parola (Span di Parole) per preservare l'ancoraggio lessicale e le dipendenze fini.
   • Livelli Superiori: Gli allineamenti si spostano a livello di frase (Frasi Nominali e Frasi Verbal) per catturare la semantica composizionale e le relazioni strutturali di ordine superiore.
   • Costruzione degli Span: Le unità semantiche vengono estratte utilizzando un parser sintattico (ad es. spaCy). I pesi di importanza dei token sono calcolati tramite un meccanismo di auto-attenzione bidirezionale (escludendo i loop su se stessi) per aggregare le rappresentazioni dei token in rappresentazioni degli span.

2. Funzione di Perdita di Allineamento Strutturale Dinamico ($L_{DSA}$):

   • Questo obiettivo impone la coerenza geometrica tra insegnante e studente.
   • Invece di corrispondere vettori di token individuali, $L_{DSA}$ minimizza la discrepanza nelle distanze relazionali a coppie tra span semantici (parole o frasi) all'interno di un livello specifico.
   • Prioritizza gli span semanticamente salienti ponderando l'allineamento in base all'importanza aggregata dei token costituenti.
   • Matematicamente, allinea la struttura topologica dello spazio delle rappresentazioni:
     $$L_{DSA} = \frac{1}{\|L_{key}\|} \sum_{l \in L_{key}} \sum_{i,j} w_{ij} \left( d(U^S_{i,l}, U^S_{j,l}) - d(U^T_{i,\phi(l)}, U^T_{j,\phi(l)}) \right)^2$$
     dove $d(\cdot)$ è la distanza coseno, e $U$ rappresenta la rappresentazione dello span ponderata.

3. Funzione di Perdita di Allineamento delle Rappresentazioni Nascoste ($L_{Hid}$):

   • Per completare i vincoli strutturali, una perdita a livello di feature allinea stati nascosti specifici.
   • Poiché le dimensioni dello studente sono tipicamente inferiori a quelle dell'insegnante, un proiettore lineare apprendibile mappa gli stati dello studente nello spazio dell'insegnante.
   • La perdita minimizza la distanza coseno ponderata tra gli stati dello studente proiettati e gli stati dell'insegnante, concentrandosi sui token più informativi.

4. Ottimizzazione:

   • MTA è progettato come un modulo plug-in. La perdita totale combina la perdita di distillazione di base ($L_{Base}$) dei framework esistenti (ad es. DistiLLM, FDD) con le perdite strutturali e nascoste proposte:
     $$L_{Total} = L_{Base} + \lambda_{DSA}L_{DSA} + \lambda_{Hid}L_{Hid}$$

Contributi Chiave

1. Identificazione delle Limitazioni: Gli autori identificano che l'allineamento intermedio uniforme ignora l'evoluzione gerarchica delle rappresentazioni, portando a un trasferimento di conoscenza subottimale.
2. Framework Adattivo ai Livelli: Introducono MTA, che transita dinamicamente dall'allineamento a livello di parola nei livelli inferiori all'allineamento strutturale a livello di frase nei livelli superiori, rispecchiando la natura composizionale dal basso verso l'alto del linguaggio.
3. Modulo Generalizzato: MTA è dimostrato essere un layer di ottimizzazione generalizzato compatibile con metodi di distillazione all'avanguardia (DistiLLM, DistiLLM-2, FDD) che condividono la tokenizzazione.
4. Validazione Empirica: Esperimenti estensivi mostrano che MTA migliora costantemente le prestazioni attraverso diverse architetture di modelli (GPT-2, Qwen1.5, OPT) e benchmark.

Risultati Sperimentali

Gli autori hanno valutato MTA su compiti di follow-up alle istruzioni utilizzando tre coppie studente-insegnante:

• GPT-2 1.5B $\to$ 120M
• Qwen1.5 1.8B $\to$ 0.5B
• OPT 6.7B $\to$ 1.3B

Metriche di Prestazione:

• ROUGE-L: MTA ha costantemente superato le linee di base (FDD, DistiLLM, DistiLLM-2) su tutti i dataset (Dolly, SelfInst, Vicuna, S-NI). Ad esempio, nell'impostazione GPT-2 1.5B $\to$ 120M, l'aggiunta di MTA a DistiLLM ha migliorato il ROUGE-L medio da 20.21 a 21.45.
• LLM-as-a-Judge: Utilizzando GPT-4o-mini per valutare accuratezza fattuale, coerenza e utilità, i modelli potenziati da MTA hanno mostrato miglioramenti costanti dei punteggi rispetto alle linee di base.

Studi di Ablazione:

• Componenti della Perdita: Sia $L_{DSA}$ (strutturale) che $L_{Hid}$ (feature) hanno contribuito positivamente, con la loro combinazione che ha prodotto i migliori risultati.
• Granularità: La strategia "Full-level" (adattiva ai livelli) ha superato le strategie statiche (solo Parole o solo Frasi), confermando che allineare l'ancoraggio lessicale nei livelli inferiori e la semantica composizionale nei livelli superiori è cruciale.
• Selezione del Livello: Il metodo è robusto rispetto a diversi programmi di livelli, sebbene un'allocazione ibrida (1 livello a livello di parola : 2 livelli a livello di frase) sia risultata ottimale per le configurazioni testate.

Efficienza:

• MTA introduce un sovraccarico di addestramento principalmente dovuto all'estrazione sintattica degli span (ridotta a ~0.27s/passo tramite ottimizzazione) e al matching geometrico.
• I costi di inferenza rimangono invariati poiché l'estrazione degli span non viene utilizzata al momento del test.
• Confronti abbinati temporalmente hanno confermato che i guadagni di prestazioni sono dovuti all'obiettivo di allineamento e non a un budget di addestramento aumentato.

Significato e Affermazioni

Il documento afferma che MTA offre un approccio semanticamente fedele alla distillazione della conoscenza modellando esplicitamente la traiettoria rappresentazionale gerarchica degli LLM. Spostandosi oltre la corrispondenza statica a livello di token verso un allineamento strutturale dinamico e multi-granulare, MTA permette ai modelli studente di replicare meglio le dinamiche interne di elaborazione delle informazioni dell'insegnante.

Gli autori posizionano MTA non come un sostituto dei framework di distillazione esistenti, ma come un modulo complementare, plug-and-play che migliora la generalizzazione e la robustezza degli LLM compressi. Il lavoro sottolinea l'importanza di rispettare la gerarchia linguistica e l'evoluzione delle feature dipendente dalla profondità nella progettazione di future tecniche di compressione dei modelli.