Riepilogo Tecnico: ARKD – Adaptive Reinforcement Learning-Guided Bidirectional KL Divergence Distillation

1. Definizione del Problema

La Distillazione della Conoscenza (Knowledge Distillation, KD) è una tecnica primaria per comprimere i Large Language Models (LLM) al fine di ridurre i costi di inferenza mantenendo le prestazioni. Tuttavia, i metodi tradizionali di KD che si affidano a un singolo obiettivo di divergenza di Kullback-Leibler (KL) affrontano un compromesso fondamentale nella generazione di testo open-ended:

• Forward KL (FKL): Incoraggia un comportamento di tipo "mode-covering" (copertura dei modi), in cui lo studente tenta di coprire tutti i modi della distribuzione del docente. Ciò spesso porta a una sovrastima delle probabilità nelle regioni a bassa confidenza e a coda lunga, degradando la qualità della generazione.
• Reverse KL (RKL): Incoraggia un comportamento di tipo "mode-seeking" (ricerca del modo), concentrandosi sui modi dominanti e ignorando le regioni a bassa probabilità. Sebbene ciò prevenga la sovrastima della coda lunga, compromette la diversità e la copertura, portando potenzialmente al collasso dei modi (mode collapse).

Gli esistenti metodi di interpolazione statica (ad esempio, pesatura fissa 0.5:0.5) non riescono a catturare la natura dinamica del divario tra studente e docente durante l'evoluzione dell'addestramento. Un coefficiente statico non può adattarsi ai requisiti mutevoli dell'allineamento della distribuzione, che richiedono equilibri diversi tra FKL e RKL in diverse fasi dell'addestramento.

2. Metodologia: ARKD

Gli autori propongono ARKD (Adaptive Reinforcement Learning-Guided Bidirectional KL Divergence Distillation), un framework che assegna dinamicamente i pesi a FKL e RKL utilizzando un agente di apprendimento per rinforzo (Reinforcement Learning, RL).

Fondamento Teorico

Il paper stabilisce innanzitutto la coerenza teorica di FKL e RKL. L'analisi mostra che entrambi gli obiettivi convergono alla stessa soluzione ($q_\theta = p$) quando i gradienti sono nulli. Ciò conferma che entrambe le divergenze sono valide per il matching della distribuzione, ma esibiscono comportamenti complementari durante la traiettoria di ottimizzazione, giustificando il loro uso congiunto.

Framework basato su RL

ARKD riformula il problema dell'assegnazione dei pesi come un processo decisionale sequenziale (Processo Decisionale di Markov):

• Stato ($s$): Un vettore a 6 dimensioni che cattura la dinamica attuale dell'addestramento, inclusi le entropie di docente/studente ($H_T, H_S$), le varianze ($\sigma^2_T, \sigma^2_S$) e i valori attuali delle perdite FKL/RKL ($L_{FKL}, L_{RKL}$).
• Azione ($\alpha$): Una rete di policy ($\pi_\phi$) produce un peso scalare $\alpha \in (0, 1)$, dove $\alpha$ pesa il termine FKL e $(1-\alpha)$ pesa il termine RKL.
• Ricompensa ($r$): Definita come la perdita di distillazione negativa ($r = -L_{distil}$) per il batch corrente. L'obiettivo è massimizzare la ricompensa cumulativa (minimizzare la perdita totale) lungo la traiettoria di addestramento.
• Ottimizzazione:
  • Aggiornamento dello Studente: Minimizza la perdita di distillazione pesata più una perdita di modellazione del linguaggio.
  • Aggiornamento della Policy: Utilizza l'algoritmo REINFORCE con una baseline di Media Mobile Esponenziale (EMA) e regolarizzazione dell'entropia per stabilizzare l'addestramento e prevenire la saturazione dei bordi.

Dinamiche di Addestramento

La rete di policy impara ad adattare il compromesso in modo dinamico:

1. Fase di Esplorazione: All'inizio dell'addestramento, la policy esplora un intervallo più ampio di $\alpha$ (ad esempio, 0.3–0.6), enfatizzando la FKL per garantire una copertura ampia e prevenire il collasso dei modi.
2. Fase di Convergenza: Man mano che la distribuzione dello studente si avvicina a quella del docente, la policy si sposta verso valori di $\alpha$ più bassi (dominanza RKL) per concentrarsi su un allineamento preciso ed evitare la sovrastima della coda lunga.
3. Fase di Stabilità: La policy si stabilizza a un $\alpha$ specifico (ad esempio, $\approx 0.18$) per il raffinamento finale.

3. Contributi Chiave

1. Analisi Teorica: Il paper fornisce un'analisi teorica che stabilisce la complementarità e la coerenza di convergenza di FKL e RKL, formando la base per la distillazione adattiva.
2. Framework ARKD: Introduce un nuovo framework di ottimizzazione congiunta basato su RL, in cui una rete di policy pesa dinamicamente FKL e RKL in base alle caratteristiche dipendenti dallo stato, andando oltre le euristiche statiche.
3. Validazione Empirica: Esperimenti estesi dimostrano che ARKD migliora costantemente la qualità della generazione, il fitting della distribuzione e la generalizzazione attraverso molteplici benchmark e scale di modelli.

4. Risultati Sperimentali

Gli autori hanno valutato ARKD su compiti di instruction tuning utilizzando configurazioni GPT-2 (1.5B $\to$ 120M/340M) e LLaMA (13B $\to$ 7B).

• Prestazioni In-Distribution (ID): Sul benchmark DollyEval, ARKD ha ottenuto il BertScore più alto e punteggi Rouge-L competitivi o superiori rispetto ai baseline (SFT, SeqKD, FKL, RKL, FKL+RKL e MiniLLM). Ad esempio, su GPT-2 1.5B $\to$ 120M, ARK ha ottenuto un Rouge-L di 24.5 e un BertScore di 0.815, superando il forte baseline MiniLLM.
• Generalizzazione Out-of-Distribution (OOD): ARKD ha dimostrato una generalizzazione superiore su dataset non visti (SelfInst, VicunaEval, S-NI, UnNI), superando costantemente tutti i baseline in entrambe le famiglie di modelli.
• Studi di Ablazione:
  • vs. Scheduling Statico/Lineare: ARKD ha superato significativamente la pesatura statica ( $\alpha$ fisso) e le strategie di scheduling lineare (decadimento/incremento basato sul tempo), che non riuscivano a catturare le complesse dinamiche di allineamento.
  • vs. Euristiche Greedy: ARKD ha superato le euristiche greedy (selezione della divergenza con perdita minore ad ogni step) di 0.4–0.6 punti. L'approccio greedy è risultato miope, degenerando spesso in pura RKL, mentre ARKD ha mantenuto una componente FKL benefica attraverso un mixing continuo.
  • Analisi della Traiettoria: La traiettoria di $\alpha$ appresa ha confermato una strategia intelligente: enfasi iniziale sulla FKL per la copertura, transizione verso la RKL per la precisione, un pattern non ottenibile tramite design manuale.

5. Significato e Rivendicazioni

Il paper afferma che ARKD fa progredire lo stato dell'arte della distillazione della conoscenza adattiva affrontando i limiti dei metodi basati su KL statici. Sfruttando l'apprendimento per rinforzo per bilanciare dinamicamente i comportamenti mode-covering e mode-seeking, ARKD ottiene un migliore allineamento della distribuzione che migliora sia la qualità della generazione che la generalizzazione.

Gli autori posizionano ARKD come un approccio rigoroso che valida la necessità di una pesatura dinamica e consapevole dello stato rispetto alle euristiche fisse. Notano che, sebbene l'attuale policy si basi su un vettore di stato costruito a mano e reward immediati per batch, il framework è compatibile con famiglie più ampie di $f$-divergence e paradigmi on-policy. Il lavoro fornisce una solida base per la compressione efficiente degli LLM, dimostrando che l'ottimizzazione adattiva guidata da RL può superare costantemente i baseline statici e greedy in diversi compiti e scale di modelli.