TiTok: Transfer Token-level Knowledge via Contrastive Excess to Transplant LoRA

1. Il Problema

I Large Language Models (LLM) sono ampiamente utilizzati, ma il loro fine-tuning completo comporta costi computazionali e di storage proibitivi. Metodi di Parameter-Efficient Fine-Tuning (PEFT) come LoRA (Low-Rank Adaptation) mitigano questi costi aggiornando solo un piccolo sottoinsieme di parametri. Tuttavia, sorge un problema fondamentale: i parametri adattati da LoRA sono strettamente dipendenti dal modello base su cui sono stati addestrati e non possono essere trasferiti direttamente su modelli diversi (backbone diversi o versioni diverse).

Le soluzioni esistenti affrontano questo limite in modi imperfetti:

• Knowledge Distillation (KD): Richiede solitamente l'accesso ai dati di addestramento originali del task target, che spesso non sono disponibili o sono costosi da ottenere.
• TransLoRA: Utilizza dati sintetici generati da LLM per trasferire la conoscenza, ma richiede l'addestramento di un modello discriminatore aggiuntivo per filtrare i dati sintetici di bassa qualità, introducendo complessità e overhead computazionale.

2. Metodologia: TITOK

Il paper propone TITOK (Transfer Token-Level Knowledge), un nuovo framework che abilita il trapianto efficace di LoRA attraverso il trasferimento di conoscenza a livello di token, senza richiedere modelli aggiuntivi o dati originali.

Il processo si articola in tre fasi principali:

A. Generazione di Dati Sintetici

Utilizzando un modello esperto sorgente (Modello Base + LoRA addestrato), il framework genera un dataset sintetico di coppie query-label partendo da un piccolo set di prompt few-shot. A differenza di TransLoRA, che usa il modello target non addestrato per generare le query, TITOK utilizza il modello esperto sorgente per generare sia le query che le etichette, garantendo una maggiore coerenza distributiva.

B. Calcolo del "Contrastive Excess Score" (Punteggio di Eccesso Contrastivo)

Questa è l'innovazione centrale. Per filtrare i dati sintetici, TITOK non usa un discriminatore esterno, ma sfrutta la differenza interna tra il modello esperto e il suo backbone.
Per ogni token $y_i$ in una risposta sintetizzata, viene calcolato un punteggio di eccesso $S(y_i)$:
$$S(y_i) = L_e(y_i) - L_a(y_i)$$
Dove:

• $L_e$ è la perdita (log-likelihood negativa) calcolata dal modello esperto (Base + LoRA).
• $L_a$ è la perdita calcolata dal modello "amatore" (solo Base).

Interpretazione: Un punteggio alto indica che il token è incerto per il modello base ma altamente probabile per il modello esperto. Questi token contengono la conoscenza specifica del task "iniettata" dal LoRA. Teoricamente, questo è collegato al rapporto di verosimiglianza (LLR) statistico, identificando i token dove la distribuzione predittiva cambia significativamente a causa dell'adattatore.

C. Filtraggio a Due Livelli e Addestramento

Una volta calcolati i punteggi, il dataset sintetico viene filtrato in due stadi:

1. Filtraggio dei Campioni (Sample Filtering): Vengono selezionati solo i campioni (query-label) con il punteggio medio di eccesso più alto, scartando quelli meno informativi.
2. Selezione dei Token (Token Selection): All'interno dei campioni selezionati, non tutti i token vengono usati per l'addestramento. Viene mantenuta solo una percentuale $k\%$ dei token con i punteggi di eccesso più alti.

Il modello target viene quindi addestrato con un nuovo LoRA inizializzato, utilizzando una funzione di perdita mascherata che considera solo questi token selezionati.

D. Allineamento dei Tokenizer

Per gestire il trasferimento tra modelli con tokenizer diversi (es. Mistral e Llama), TITOK introduce un algoritmo di allineamento a due puntatori. Questo mappatura propagando le maschere dei token dalla sequenza sorgente a quella target, gestendo casi di mappatura uno-a-uno, uno-a-molti e molti-a-uno tramite regole di copia, replicazione o media.

3. Contributi Chiave

• Trasferimento senza Discriminatori: Elimina la necessità di addestrare un modello discriminatore aggiuntivo (come in TransLoRA), riducendo significativamente l'overhead computazionale e la complessità del pipeline.
• Selezione Fine-Grained: Sposta il focus dal trasferimento di intere sequenze sintetiche alla selezione di token specifici che contengono la conoscenza trasferibile, rendendo il processo più efficiente e mirato.
• Robustezza ai Dati Esterni: Dimostra che il metodo funziona efficacemente anche con dati esterni reali (non sintetici) o provenienti da domini diversi, superando i limiti delle soluzioni basate esclusivamente sulla sintesi dei dati.
• Indipendenza dal Modello: Funziona efficacemente in scenari di trasferimento tra modelli della stessa famiglia, tra famiglie diverse, tra diverse dimensioni (es. 3B a 8B) e diverse versioni.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su tre benchmark principali: Big-Bench Hard (BBH), MMLU (compiti di ragionamento) e LaMP (compiti di personalizzazione/generazione).

• Performance Generali: TITOK ha superato costantemente le baseline (Vanilla, KD, TransLoRA). In media, ha ottenuto un miglioramento del +4-10% rispetto alle baseline su tutti i task e scenari di trasferimento.
  • Rispetto al modello target "Vanilla": +9.94%.
  • Rispetto alla Knowledge Distillation (KD): +8.5%.
  • Rispetto a TransLoRA: +4.4%.
• Efficienza Computazionale: Grazie all'eliminazione del training del discriminatore, TITOK riduce il tempo di calcolo totale di circa 1.5x - 2.5x rispetto a TransLoRA.
• Robustezza: Il metodo mantiene prestazioni elevate anche quando il modello sorgente è più debole del modello target o quando si utilizzano dati sintetici di qualità inferiore.
• Analisi Ablativa: Ha dimostrato che sia il filtraggio dei campioni che la selezione dei token sono componenti essenziali; l'uso combinato di entrambi produce i risultati migliori.

5. Significato e Impatto

TITOK rappresenta un paradigma metodologicamente semplice ma potente per il trasferimento di conoscenza tra modelli LLM. Risolve il collo di bottiglia della dipendenza dai dati di addestramento originali e riduce l'overhead computazionale delle soluzioni attuali.

La sua capacità di identificare e trasferire solo la conoscenza "essenziale" a livello di token lo rende particolarmente adatto per scenari reali dove:

1. I dati originali non sono accessibili (privacy o costi).
2. È necessario aggiornare o migrare modelli su architetture diverse.
3. Le risorse computazionali sono limitate.

In sintesi, TITOK offre una soluzione scalabile ed efficiente per democratizzare l'uso di adattatori LoRA su una vasta gamma di modelli, facilitando l'evoluzione e il riutilizzo delle competenze acquisite dai modelli LLM.