Riepilogo Tecnico: X-Token: Distillazione della Conoscenza Cross-Tokenizer Guidata da Proiezione

Enunciato del Problema

La Distillazione della Conoscenza (KD) richiede tipicamente un tokenizzatore condiviso tra i modelli insegnante e studente, bloccando di fatto i praticanti in specifiche famiglie di modelli. Questa restrizione impedisce di sfruttare insegnanti più potenti o specializzati con vocabolari incompatibili (ad esempio, distillare da un modello Qwen o Phi in uno studente Llama) e ostacola la distillazione multi-insegnante attraverso diverse famiglie di tokenizzatori.

I metodi KD cross-tokenizer esistenti rientrano in due categorie: basati sulla rappresentazione (allineamento degli stati nascosti) e basati sulla distribuzione dei logit (operanti sulle probabilità di output). Quest'ultima è preferita per il pre-addestramento continuo poiché non richiede modifiche architetturali. Tuttavia, il metodo basato sui logit più avanzato allo stato dell'arte, GOLD, soffre di due limitazioni strutturali critiche quando gestisce tokenizzatori incompatibili:

1. Fallimento dei Token Non Comuni: Quando i tokenizzatori frammentano il testo in modo diverso (ad esempio, un insegnante che divide i numeri multi-cifra mentre lo studente li impacchetta), i token critici finiscono in un sottoinsieme "non abbinato". Questi token vengono degradati dal rumore agnostico all'identità derivante dall'abbinamento basato sul rango e soppressi dai gradienti del termine KL comune, portando a drastici cali delle prestazioni (ad esempio, l'accuratezza su GSM8k scende da 12,89 a 2,56).
2. Abbinamento Eccessivamente Conservativo: La rigorosa uguaglianza stringa 1-a-1 esclude allineamenti semantici validi dove le forme superficiali differiscono (ad esempio, un token studente "Centinaia" corrispondente ai token insegnante "Hund" + "reds"), lasciando inutilizzati segnali di supervisione utili.

Metodologia: X-Token

X-Token affronta questi problemi attraverso un framework a tre componenti: Allineamento degli Span, una Matrice di Proiezione ($W$) e Formulazioni di Loss Complementari.

1. Allineamento delle Sequenze

Per gestire lunghezze di sequenza diverse e disallineamenti posizionali, X-Token impiega l'allineamento degli span. I token sono raggruppati in chunk che decodificano allo stesso testo sottostante. Viene applicata una fusione a catena di regole all'interno di ogni chunk per combinare le probabilità per-token in una singola distribuzione a livello di chunk, creando unità allineate per la distillazione.

2. Matrice di Proiezione ($W$)

Una matrice di proiezione sparsa $W \in \mathbb{R}^{|\mathcal{V}_S| \times |\mathcal{V}_T|}$ mappa le probabilità dei token studente nello spazio del vocabolario dell'insegnante. È costruita deterministicamente in due passaggi:

• Passaggio di Corrispondenza Esatta: Imposta $W[s, t] = 1$ per i token con corrispondenze di stringa canonizzate.
• Passaggio di Regola Multi-Token: Per i token studente non abbinati, il testo decodificato viene re-tokenizzato dall'insegnante. Il token studente è mappato alla risultante sequenza di token insegnante con pesi decrescenti (ad esempio, $\beta=0,9, \gamma=0,1$).
  La matrice è inizializzata dalle regole del tokenizzatore e può essere opzionalmente affinata durante l'addestramento.

3. Formulazioni di Loss Complementari

X-Token introduce due modalità di loss, selezionate in base a un audit di copertura dei token critici (ad esempio, i numeri) nel vocabolario dello studente:

• P-KL (Projection-KL): Progettata per scenari in cui i token critici rientrano nel set "non comune" (ad esempio, insegnante Qwen). Rimuove l'intera partizione. La distribuzione studente è proiettata tramite $W$ nello spazio del vocabolario insegnante e viene calcolata una divergenza KL diretta tra la distribuzione insegnante e la distribuzione studente proiettata. Questo elimina i gradienti basati sul rango errati e gli effetti soppressivi.
• H-KL (Hybrid-KL): Progettata per scenari in cui la partizione è strutturalmente solida (ad esempio, insegnante Phi-4-mini). Mantiene la struttura ibrida ma rilassa la definizione del set "comune". Invece della rigorosa uguaglianza di stringa, include coppie $(s, t^*)$ dove $t^*$ è il token insegnante di primo rango per il token studente $s$ sotto $W$. Questo cattura allineamenti quasi equivalenti preservando al contempo la supervisione più netta della KL diretta sulle coppie allineate per identità.

4. Distillazione Multi-Insegnante

X-Token si estende naturalmente a più insegnanti aggregando le loss per insegnante. Impiega una pesatura statica (che ha superato gli schemi adattivi negli esperimenti) per combinare i segnali da insegnanti eterogenei, sfruttando i loro punti di forza complementari (ad esempio, ragionamento vs. senso comune).

5. Scalatura Dinamica

Per stabilizzare l'addestramento, la loss di distillazione viene ridimensionata dinamicamente a ogni passo per corrispondere alla magnitudine della loss cross-entropy, prevenendo l'instabilità dell'ottimizzazione causata da scale di loss variabili.

Contributi Chiave

1. Allineamento Cross-Token Deterministico: Introduzione di una matrice di proiezione sparsa $W$ costruita tramite regole di stringa e decodifica multi-token, che consente un allineamento diretto del vocabolario senza componenti ausiliari addestrabili.
2. Dualità delle Formulazioni di Loss: La proposta di P-KL e H-KL, che si adattano a diversi regimi di disallineamento dei tokenizzatori. P-KL recupera il segnale per i token non abbinati critici, mentre H-KL affina l'abbinamento per i token ben allineati.
3. Capacità Multi-Insegnante: Dimostrazione che X-Token abilita una distillazione efficace da insegnanti eterogenei, con risultati che mostrano come le combinazioni complementari di insegnanti producano prestazioni superiori rispetto a quelle a insegnante singolo o accoppiamenti omogenei.
4. Allineamento di Sequenza Robusto: Un metodo di allineamento degli span scalabile basato su programmazione dinamica che garantisce unità coerenti con il testo per la distillazione.

Risultati Sperimentali

Gli esperimenti sono stati condotti utilizzando Llama-3.2-1B come studente e Qwen3-4B, Phi-4-mini-Instruct e Llama-3.2-3B come insegnanti.

• Prestazioni vs SOTA: X-Token supera lo stato dell'arte attuale (GOLD) di +3,82 punti medi quando si utilizza un insegnante Qwen3-4B (con un miglioramento di 6 volte su GSM8k, da 2,56 a 15,54) e di +0,5 punti con un insegnante Phi-4-mini.
• Superiorità Cross-Famiglia: Nel regime Qwen, X-Token (P-KL) ha superato la KD con lo stesso tokenizzatore da un insegnante più debole Llama-3.2-3B, dimostrando che la distillazione cross-tokenizer può superare i limiti della stessa famiglia quando la formulazione della loss è corretta.
• Guadagni Multi-Insegnante: Una configurazione a due insegnanti (Phi-4-mini + Llama-3.2-3B) ha migliorato la distillazione a insegnante singolo di +1,3 punti, evidenziando il valore della combinazione di punti di forza diversi degli insegnanti.
• Risultati di Ablazione:
  • La selezione della loss è critica: P-KL fallisce su Phi-4-mini e H-KL fallisce su Qwen3-4B.
  • L'apprendimento congiunto della matrice di proiezione $W$ produce guadagni modesti ma costanti rispetto a un'inizializzazione basata su regole e congelata.
  • La pesatura statica degli insegnanti ha superato gli schemi adattivi (basati su entropia o confidenza).

Significato

Il paper afferma che X-Token dimostra che un attento allineamento sia a livello di sequenza che di vocabolario, combinato con un design adattivo della loss, è la chiave per sbloccare il pieno potenziale della distillazione della conoscenza cross-tokenizer. Affrontando le modalità di fallimento specifiche dei metodi basati su partizione, X-Token permette ai praticanti di sfruttare insegnanti più potenti, specializzati o diversificati senza essere vincolati dalla compatibilità del tokenizzatore, rompendo efficacemente il "blocco del tokenizzatore" che attualmente limita il scaling e la specializzazione dei modelli.