Dual-Space Knowledge Distillation with Key-Query Matching for Large Language Models with Vocabulary Mismatch

1. Il Problema: Sfide nella Distillazione della Conoscenza Cross-Token

I Large Language Models (LLM) offrono prestazioni all'avanguardia ma sono costosi da distribuire a causa delle loro dimensioni e delle richieste di risorse. La Distillazione della Conoscenza (KD) risolve questo problema addestrando modelli studenti più piccoli a imitare modelli insegnanti più grandi.

Tuttavia, un ostacolo significativo sorge quando l'insegnante e lo studente utilizzano tokenizzatori e vocabolari distinti (ad esempio, modelli basati su BPE vs. modelli basati su WordPiece o vocabolari di dimensioni molto diverse). Questo crea due problemi principali:

1. Disallineamento delle sequenze: La stessa frase viene tokenizzata in sequenze di lunghezza diversa.
2. Disallineamento degli spazi di predizione: Le distribuzioni di probabilità sui token non sono direttamente confrontabili.

Il metodo stato dell'arte (SOTA) attuale per questo scenario è il DSKD-CMA (Dual-Space Knowledge Distillation with Cross-Model Attention), che utilizza meccanismi di attenzione per imparare implicitamente le allineamenti. Tuttavia, il funzionamento interno di DSKD-CMA rimane "opaco" (una scatola nera), rendendo difficile comprenderne i limiti e migliorarlo. Inoltre, le distribuzioni di Key (chiavi) e Query calcolate da modelli diversi possono non corrispondere, indebolendo l'efficacia dell'attenzione.

2. Metodologia Proposta

Gli autori hanno sviluppato un approccio in due fasi: un'analisi approfondita del meccanismo esistente e una nuova metodologia per correggerne le debolezze.

A. Analisi del DSKD-CMA

Per comprendere come DSKD-CMA allinea i token, gli autori hanno introdotto:

• Allineamento manuale dei chunk: Un algoritmo che segmenta le sequenze di token dell'insegnante e dello studente in "chunk" minimi che coprono lo stesso testo (basato sugli offset dei caratteri).
• Visualizzazione delle mappe di calore: Confronto tra le allineamenti appresi dall'attenzione del modello e l'allineamento manuale.
• Ablazioni: Sostituzione del meccanismo CMA con due varianti interpretabili:
  • Chunk-Level Projection (CLP): Proiezione uniforme tra chunk allineati senza attenzione adattiva.
  • Chunk-Level Attention (CLA): Attenzione limitata solo all'interno dei chunk allineati.

Risultati dell'analisi: Il DSKD-CMA apprende effettivamente strutture a livello di chunk, ma le sue mappature possono essere "rumorose" (specialmente con token ripetuti) e mancano di una localizzazione precisa. L'uso di CLA (attenzione solo locale) ha peggiorato le prestazioni, suggerendo che il modello beneficia di allineamenti "goffi" ma stabili a livello macro, piuttosto che di pesi granulari eccessivamente sensibili.

B. Nuova Metodologia: DSKD-CMA-GA e DSKD-CMA-CT

Per affrontare il disallineamento delle distribuzioni tra le Key (dall'insegnante) e le Query (dallo studente), gli autori propongono di integrare tecniche di matching Key-Query basate su apprendimento avversariale:

1. DSKD-CMA-GA (Generative Adversarial):

   • Utilizza un discriminatore che cerca di distinguere le Key dell'insegnante dalle Query dello studente.
   • Il proiettore delle Query dello studente agisce come generatore, cercando di ingannare il discriminatore per allineare le distribuzioni.
   • Obiettivo: Minimizzare la distanza tra le distribuzioni di input dell'attenzione.

2. DSKD-CMA-CT (Conditional Transport):

   • Utilizza una versione vincolata del Trasporto Ottimale per ridurre il costo computazionale.
   • Impiega un "critico" trainabile per massimizzare il costo di trasporto, mentre il proiettore cerca di minimizzarlo, creando un gioco avversariale simile al GA.

La funzione di perdita totale combina la distillazione originale con la perdita di allineamento Key-Query ($L_{train} = L_{DSKD-CMA} + L_{KQ}$).

3. Contributi Chiave

1. Framework di Analisi: Introduzione di sonde di allineamento a livello di chunk e visualizzazioni che rivelano i punti di forza (allineamento orientato ai chunk) e le debolezze (localizzazione limitata) del DSKD-CMA.
2. Nuovi Obiettivi di Distillazione: Proposta di DSKD-CMA-GA e DSKD-CMA-CT, che migliorano la qualità della generazione testuale allineando le distribuzioni di Key e Query tra modelli con vocabolari diversi.
3. Superamento del Gap Tokenizer: Dimostrazione che la KD cross-tokenizer può raggiungere, e talvolta superare, le prestazioni della KD con lo stesso tokenizzatore.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti utilizzando:

• Studente: GPT-2 (124M parametri).
• Insegnanti: GPT-2 XL (stesso tokenizzatore) e Qwen1.5 (1.8B parametri, vocabolario molto diverso da GPT-2).
• Dataset: Dolly 15K per l'addestramento; valutazione su dataset in-distribution (Dolly) e out-of-distribution (Self-Instruct, Vicuna, S-NI, UnNI).

Punti salienti dei risultati:

• Prestazioni Generali: Tutte le varianti DSKD superano significativamente i metodi baselines precedenti (MinED, ULD) e il semplice fine-tuning supervisionato (SFT).
• Miglioramento GA: DSKD-CMA-GA ha ottenuto i risultati migliori, mostrando guadagni consistenti di ROUGE-L (da +0.15 a +1.04 punti) rispetto al DSKD-CMA base.
• Dati Out-of-Distribution (OOD): Il miglioramento è particolarmente evidente sui dati OOD, con un guadagno medio di +0.37 in ROUGE-L.
• Chiusura del Gap: DSKD-CMA-GA riduce drasticamente il divario tra la KD cross-tokenizer e quella same-tokenizer, superando talvolta quest'ultima.
• Interazione Divergenza: È stato scoperto che la scelta della funzione di divergenza (es. Skewed KL) è cruciale e deve essere co-progettata con il metodo di allineamento (GA funziona meglio con SKL).

5. Significato e Implicazioni

Questo lavoro è significativo perché:

1. Rende trasparente la "scatola nera": Fornisce una comprensione empirica di come funziona l'attenzione incrociata tra modelli diversi, guidando future ricerche.
2. Migliora l'efficienza della KD: Offre un metodo pratico per distillare modelli con vocabolari incompatibili senza costi computazionali proibitivi (a differenza di metodi basati su Chain-of-Thought o OT completo).
3. Valida l'approccio Avversariale: Dimostra che l'apprendimento avversariale (GA) è uno strumento potente per allineare spazi latenti eterogenei in contesti di NLP.
4. Impatto Pratico: Consente di sfruttare modelli insegnanti potenti e proprietari (spesso con tokenizzatori chiusi o diversi) per addestrare studenti efficienti e open-source, democratizzando l'accesso alle capacità dei LLM avanzati.

In conclusione, la proposta DSKD-CMA-GA rappresenta un passo avanti verso la standardizzazione della distillazione della conoscenza in un panorama di LLM frammentato da diversi standard di tokenizzazione.