$\mathcal{X}$-KD: General Experiential Knowledge Distillation for Large Language Models

1. Il Problema

L'articolo affronta le limitazioni attuali della Knowledge Distillation (KD) per i Large Language Models (LLM). Le metodologie esistenti si concentrano principalmente sull'imitazione del comportamento del modello "insegnante" (ad esempio, attraverso la clonazione comportamentale o la minimizzazione della divergenza tra distribuzioni di token). Tuttavia, questi approcci trascurano l'ambiente di apprendimento originale che ha plasmato le conoscenze e le strategie dell'insegnante.

In particolare:

• Gli studenti imparano a imitare l'output, ma non necessariamente la funzione di ricompensa sottostante che ha guidato l'insegnante.
• Approcci recenti come MiniLLM hanno introdotto concetti di Inverse Reinforcement Learning (IRL), ma si discostano dal paradigma dell'apprendimento supervisionato, richiedendo tecniche di stabilizzazione complesse (come il policy gradient) e limitando la flessibilità nell'uso di diverse funzioni di divergenza o nella distillazione black-box.
• Manca un framework unificato che permetta allo studente di "vivere" e apprendere nell'ambiente originale dell'insegnante, mantenendo al contempo la semplicità e la flessibilità dei metodi di distillazione supervisionati.

2. Metodologia: X-KD (Experiential Knowledge Distillation)

Gli autori propongono X-KD, un framework generale basato sulla teoria dell'apprendimento esperienziale e sull'Approssimata Variational Reward Imitation Learning (AVRIL), un framework di IRL bayesiano.

Concetto Fondamentale

L'obiettivo è permettere al modello studente ($\pi_\theta$) di apprendere nell'ambiente di ricompensa originale in cui il modello insegnante ($\pi$) ha operato, piuttosto che limitarsi a copiare i suoi comportamenti. Poiché la funzione di ricompensa originale $R_\pi$ è sconosciuta, X-KD la stima e la approssima.

Componenti Chiave

1. Modellazione della Ricompensa Originale (ORM):

   • Viene formulato un problema di IRL bayesiano per stimolare la distribuzione a posteriori della funzione di ricompensa dell'insegnante.
   • Utilizzando il framework AVRIL, il sistema apprende congiuntamente una funzione di ricompensa (codificata da $q_\phi$) e una politica (decodificata da $Q_\theta$).
   • L'obiettivo include un termine di regolarizzazione esperienziale che assicura che la politica dello studente sia coerente con la funzione di ricompensa stimata, minimizzando l'errore Temporal Difference (TD).

2. Integrazione con Metodi KD Esistenti:
   X-KD è progettato per essere un "layer" esperienziale aggiunto a varie tecniche di distillazione:

   • Sequence-Level X-KD: Estende la distillazione a livello di sequenza (Kim & Rush, 2016) aggiungendo un termine di regolarizzazione esperienziale alla funzione di perdita.
   • Supervised X-KD: Si applica alla distillazione supervisionata (Forward-KL), combinando la divergenza KL standard con il termine esperienziale.
   • Generalized X-KD (G-XKD): Si applica alla distillazione basata su divergenze (come GKD e MiniLLM), permettendo l'uso di funzioni di divergenza flessibili (es. Jensen-Shannon) e dati on-policy.

3. Semplicità e Paradigma Supervisionato:
   A differenza di MiniLLM, X-KD mantiene il paradigma dell'apprendimento supervisionato. L'obiettivo finale è la somma della perdita KD originale e di un termine di regolarizzazione esperienziale. Questo rende l'integrazione nei pipeline di training esistenti molto più semplice.

   • Per i modelli Black-box (dove non si ha accesso alla distribuzione di probabilità dell'insegnante), X-KD utilizza direttamente l'obiettivo a livello di sequenza, rendendolo applicabile anche quando l'insegnante è un'API chiusa.

3. Contributi Principali

• Nuovo Framework: Introduzione di X-KD, il primo framework di distillazione che integra esplicitamente l'apprendimento esperienziale (modellazione della ricompensa dell'ambiente) nei LLM.
• Flessibilità e Semplicità: Il metodo è compatibile con diverse tecniche di distillazione (sequenziale, supervisionata, generalizzata) e tipi di modelli (white-box e black-box), mantenendo la struttura di training supervisionato.
• Validazione Empirica: Dimostrazione che X-KD supera gli stati dell'arte (baselines come GKD e MiniLLM) in termini di prestazioni, compromesso prestazioni-diversità ed efficienza dei dati.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su tre compiti di generazione del linguaggio naturale (NLG): Riassunto astrattivo (XSum), Traduzione automatica (WMT14 en-de) e Ragionamento aritmetico (GSM8k).

• Prestazioni Generali: X-KD (nelle sue varianti G-XKD, S-XKD, SeqXKD) ha superato sistematicamente le baselines (SFT, SeqKD, SKD, GKD, MiniLLM) in quasi tutti i setting di dimensione del modello (da Large a Small/Base).
• Compromesso Prestazioni-Diversità:
  • X-KD mostra un trade-off superiore tra qualità della generazione e diversità.
  • Mentre le baselines tendono a degradare rapidamente la qualità all'aumentare della diversità (temperatura di campionamento), X-KD mantiene prestazioni più stabili, suggerendo che lo studente ha appreso strategie di generalizzazione più robuste invece di imitare superficialmente l'insegnante.
• Efficienza dei Dati:
  • X-KD raggiunge prestazioni comparabili o superiori alle baselines utilizzando significativamente meno dati di training (es. 50% o 75% del dataset).
  • Questo indica che la modellazione della funzione di ricompensa permette allo studente di catturare i principi fondamentali del compito più rapidamente.
• Distillazione Black-Box:
  • Utilizzando Llama-2-7b come insegnante black-box, la variante SeqXKD ha ottenuto un tasso di vittoria superiore rispetto a SeqKD nelle valutazioni condotte da GPT-4, confermando l'efficacia del metodo anche senza accesso ai gradienti o alle distribuzioni interne dell'insegnante.

5. Significato e Impatto

Il lavoro di X-KD rappresenta un passo avanti significativo nella teoria della distillazione dei LLM:

1. Cambio di Paradigma: Sposta il focus dalla semplice "imitazione del comportamento" all'"apprendimento dell'ambiente", allineandosi meglio con le teorie cognitive umane (Teoria dell'Apprendimento Esperienziale di Kolb).
2. Praticità: Risolve la complessità di metodi precedenti basati su IRL (come MiniLLM) fornendo una soluzione che si integra nativamente nei flussi di lavoro di training supervisionato standard.
3. Robustezza: La capacità di funzionare sia in setting white-box che black-box rende X-KD una soluzione versatile per scenari reali dove i modelli insegnanti sono spesso proprietari o inaccessibili.

In sintesi, X-KD dimostra che incorporare la conoscenza dell'ambiente di apprendimento originale dell'insegnante permette di creare modelli studenti più efficienti, diversificati e capaci di generalizzare meglio, superando i limiti delle tecniche di distillazione puramente comportamentali.