Sintesi Tecnica: LARK (Learnability-Grounded Trajectory Selection)

Definizione del Problema

Il documento affronta la sfida della distillazione del ragionamento (reasoning distillation), un paradigma in cui un modello linguistico (LLM) studente più piccolo viene sottoposto a fine-tuning per imitare le traiettorie di ragionamento generate da modelli insegnanti più grandi. Mentre i metodi esistenti di selezione dei dati operano spesso a livello di domanda o si affidano a euristiche come la qualità della traiettoria, la correttezza o la confidenza del modello, essi trascurano frequentemente un fattore critico: la learnability (capacità di apprendimento).

Nello specifico, per una singola domanda, possono esistere molteplici traiettorie candidate (generate da diversi insegnanti, campionamenti o stili). Anche se queste traiettorie conducono allo stesso risultato finale corretto, esse forniscono segnali di addestramento profondamente diversi. Una traiettoria che appare "di alta qualità" o "naturale" non è necessariamente quella dalla quale il modello studente può apprendere in modo più efficiente. La domanda centrale posta è: È possibile progettare un criterio di selezione delle traiettorie basato su principi che garantisca che le traiettorie di ragionamento distillate siano apprendibili dal modello studente?

Metodologia: LARK

Gli autori propongono LARK (Learnability-grounded Anchor-time Ranking), un metodo che seleziona le traiettorie basandosi su una definizione di learnability derivata dalla teoria dell'ottimizzazione, piuttosto che su metriche di qualità statiche.

1. La Learnability come Obiettivo di Ottimizzazione

LARK inquadra la selezione delle traiettorie come un problema di ottimizzazione di policy. Definisce un fattore di learnability $\rho(\theta; q)$, che caratterizza il tasso con cui diminuisce la perdita di addestramento dello studente. Questo fattore è definito come il rapporto tra la norma del gradiente al quadrato e la perdita:
$$\rho(\theta; q) \triangleq \frac{\|\nabla_\theta \mathcal{L}(\theta; q)\|^2}{\mathcal{L}(\theta; q)}$$
dove $q$ è il vettore dei pesi su $K$ traiettorie candidate. Sotto la condizione di Polyak–Łojasiewicz (PL), un $\rho$ più elevato implica una convergenza più rapida.

2. Approssimazione Anchor-Time

L'ottimizzazione diretta dell'obiettivo di learnability integrato nel tempo sull'intero processo di fine-tuning è intrattabile. LARK introduce la condizione relativa all'ancora (anchor-relative condition), la quale postula che, per modelli sovra-parametrizzati nel regime di "lazy training", il tasso di apprendimento rimanga proporzionale al suo valore ai parametri iniziali di riferimento $\theta_{ref}$. Ciò consente di approssimare l'obiettivo massimizzando $\rho(\theta_{ref}; q)$.

3. Linearizzazione Locale e Proxy del Forward-Pass

Massimizzare direttamente $\rho(\theta_{ref}; q)$ è problematico perché è una funzione quasi-convessa che collassa in una soluzione "one-hot" degenere (selezionando una singola traiettoria). Per evitare ciò, LARK:

• Linearizzazione Locale: Espande $\rho$ attorno alla distribuzione uniforme $p$ (SFT standard senza selezione) utilizzando un'espansione di Taylor del primo ordine.
• Proxy del Forward-Pass ($\hat{g}_k$): Invece di calcolare i gradienti esatti (che richiederebbero passaggi all'indietro per ogni candidato), LARK deriva un proxy computabile $\hat{g}_k$ utilizzando solo i passaggi in avanti (forward passes). Questo proxy stima il gradiente locale del tasso di apprendimento basandosi sui residui dei token dello studente e sulla perdita di cross-entropy.
  • Nello specifico, $\hat{g}_k$ favorisce le traiettorie in cui lo studente è "sicuramente errato" in modo strutturato: la perdita $\ell_k$ è alta (indicando una scarsa predizione), ma la massa residua è concentrata su pochi token "quasi corretti" (fornendo un segnale di correzione chiaro).

4. Politica di Selezione con Regolarizzazione $\chi^2$

Per evitare che il selettore "inganni" il criterio di learnability allontanandosi troppo dalla distribuzione originale dei dati (il che danneggerebbe la generalizzazione), LARK impone un termine di regolarizzazione $\chi^2$. Ciò porta a una regola di selezione a forma chiusa e parametrizzata in base al budget:

• Le traiettorie sono classificate in base al loro punteggio $\hat{g}_k$.
• Vengono selezionate le prime $B$ traiettorie.
• A differenza della semplice troncamento top-$B$, LARK assegna pesi soft proporzionali al margine del punteggio rispetto alla $(B+1)$-esima traiettoria. Questa regola di pesatura è derivata per massimizzare il limite inferiore del miglioramento della learnability controllando al contempo la divergenza $\chi^2$ dalla distribuzione uniforme di partenza.

Contributi Chiave

1. Formulazione: Gli autori formulano la selezione delle traiettorie di ragionamento come un problema di ottimizzazione di policy basato sulla learnability, introducendo il tasso di apprendimento anchor-time $\rho$ come obiettivo rigoroso.
2. Algoritmo: Derivano una regola di selezione a forma chiusa e parametrizzata in base al budget che utilizza un proxy del forward-pass $\hat{g}_k$ e la regolarizzazione $\chi^2$. Questo evita costosi passaggi all'indietro (backward passes) per ogni traiettoria, garantendo teoricamente limiti sull'errore di stima e sulla generalizzazione.
3. Validazione Empirica: Il documento dimostra che LARK supera costantemente i baseline esistenti (inclusi metodi basati su euristiche, sulla qualità o sul punteggio lato studente) attraverso molteplici modelli base (Qwen-2.5-7B/1.5B, Llama-3.2-3B) e benchmark di ragionamento (AIME, AMC, GPQA, MATH-500).

Risultati Sperimentali

• Performance: LARK ha ottenuto la media Acc@5 più alta in tutti i setting di budget del modello testati. Ad esempio, su Qwen-2.5-7B con un budget di $B=3$, LARK ha migliorato le prestazioni di 7,07 punti percentuali rispetto al più forte baseline (RSR).
• Efficienza: LARK è un metodo basato esclusivamente sul forward-pass, eguagliando il costo computazionale dei baseline lato studente a minor costo, mentre i metodi basati su gradienti esatti richiederebbero molta più computazione.
• Analisi Diagnostica:
  • Decadimento della Perdita: Le traiettorie selezionate da LARK hanno indotto la riduzione più rapida della perdita SFT durante l'addestramento.
  • Validità Predittiva: Il punteggio LARK $\hat{g}_k$ ha mostrato una forte correlazione positiva (coefficiente di Pearson $r=0,69$) con l'accuratezza a valle quando aggregato per modello insegnante.
  • Comportamento a livello di Token: Le traiettorie selezionate esibivano probabilità di token corretti inferiori ma una maggiore "concentrazione di token errati", confermando che LARK seleziona esempi in cui lo studente sbaglia, ma l'errore è strutturato e correggibile.

Significato e Affermazioni

Il documento afferma che la learnability è un principio distinto ed efficace per l'efficiente distillazione del ragionamento, separato dalla qualità o dalla difficoltà della traiettoria. Fondando la selezione sulla dinamica di ottimizzazione del modello studente, LARK identifica le traiettorie che forniscono il segnale di apprendimento più efficiente.

Gli autori sottolineano che il loro metodo:

• Preserva la Generalizzazione: Attraverso la regolarizzazione $\chi^2$, impedisce alla politica di selezione di andare in overfitting su un set ristretto di traiettorie "facili" o "manipolate".
• È Teoricamente Fondato: La regola di selezione è derivata dai primi principi della teoria dell'ottimizzazione (condizioni PL, flusso del gradiente) con espliciti limiti di errore.
• È Pratico: Non richiede il tuning degli iperparametri per la regola di selezione stessa (la forza della regolarizzazione è determinata dal budget $B$) ed evita l'overhead computazionale dei passaggi all'indietro.

Il lavoro suggerisce che futuri sviluppi potrebbero applicare questa prospettiva basata sulla learnability a pool di candidati più rumorosi e a domini di ragionamento più ampi, ma l'attuale ambito è strettamente limitato alla distillazione del ragionamento nella risoluzione di problemi matematici.