Riepilogo Tecnico: Draft-OPD

Enunciato del Problema

La decodifica speculativa (SD) accelera l'inferenza dei Modelli Linguistici di Grande Dimensione (LLM) utilizzando un modello bozza leggero per proporre blocchi di token, che vengono poi verificati in parallelo da un modello target più grande. Sebbene i metodi mainstream attuali (ad esempio EAGLE-3, DFlash) addestrino questi modelli bozza tramite Affinamento Supervisionato (SFT) su traiettorie fisse generate dal modello target, questo approccio ha raggiunto un limite di prestazioni.

Gli autori identificano un fondamentale mismatch offline-inferenza:

1. Fase di Addestramento: Il modello bozza apprende da traiettorie statiche in cui ogni prefisso è generato dal modello target.
2. Fase di Inferenza: Il modello bozza viene valutato su blocchi che propone esso stesso. La lunghezza di accettazione (il numero di token accettati per passo di verifica) è determinata dagli stati indotti dalla propria politica del modello bozza, non dagli stati statici osservati durante l'SFT.

Di conseguenza, l'SFT continuato raggiunge rapidamente un plateau, e il modello bozza non riesce a migliorare la propria lunghezza di accettazione sui dati di test. Inoltre, la Distillazione On-Policy (OPD) standard è difficile da applicare direttamente ai modelli bozza perché:

• I rollout solo-bozza (dove il modello bozza genera sequenze in modo indipendente) sono instabili e spesso degenerano in output ripetitivi.
• I rollout assistiti dal target (dove il modello target corregge la bozza) producono sequenze di alta qualità ma eliminano il segnale "on-policy". Il modello target scarta gli errori della bozza, il che significa che il modello bozza impara solo dalla distribuzione del target, fallendo nell'affrontare gli errori specifici che limitano l'accettazione speculativa.

Metodologia: Draft-OPD

Per colmare questo divario, gli autori propongono Draft-OPD, un framework di distillazione on-policy progettato specificamente per l'addestramento di modelli bozza. Il metodo permette al modello bozza di apprendere dai propri stati al momento dell'inferenza, inclusi sia le proposte riuscite che gli errori specifici che causano fallimenti nella verifica.

Il framework è composto da tre componenti coordinate:

1. Rollout con Raccolta delle Posizioni di Errore

Invece di costringere il modello bozza a generare intere sequenze da solo, Draft-OPD utilizza rollout assistiti dal target per garantire continuazioni stabili e di alta qualità.

• Durante la decodifica speculativa, il modello bozza propone un blocco di token, e il modello target li verifica.
• Il sistema registra la posizione di inizio di ogni blocco bozzato come un ancoraggio.
• Questo processo genera un rollout valido (seguendo la distribuzione del target) preservando allo stesso tempo gli stati specifici in cui il modello bozza ha agito.

2. Replay per il Calcolo delle Log-Probabilità

Per recuperare il segnale "on-policy" perso durante il rollout assistito dal target, il sistema riproduce il processo di bozzatura a partire dagli ancoraggi registrati.

• Partendo da uno stato di ancoraggio (il prefisso verificato dal target), il sistema riesegue la simulazione della proposta del blocco successivo da parte del modello bozza.
• Calcola le log-probabilità sia per il modello bozza (studente) che per il modello target (insegnante) su questi prefissi generati dalla bozza.
• Crucialmente, questo permette al sistema di valutare la distribuzione del modello target sui token che sono stati rifiutati durante la verifica iniziale, catturando gli errori specifici commessi dal modello bozza.

3. Obiettivo di Distillazione Consapevole dell'Accettazione

L'obiettivo di addestramento tratta i token accettati e quelli rifiutati in modo diverso per ottimizzare l'accettazione speculativa:

• Token Accettati: Per i token che superano la verifica, il metodo utilizza la Divergenza KL Forward ($D_{KL}(P_{target} || P_{draft})$). Questo incoraggia la distribuzione della bozza a coprire la distribuzione del target dove sono già d'accordo, rafforzando comportamenti affidabili.
• Token Rifiutati: Per i token che falliscono la verifica, il metodo utilizza la Divergenza KL Reverse ($D_{KL}(P_{draft} || P_{target})$). Questo penalizza il modello bozza per assegnare alta probabilità a token che il modello target rifiuta, insegnando efficacemente al modello bozza a evitare i propri errori ad alta confidenza.
• Ponderazione: I token rifiutati sono ponderati esponenzialmente in base alla loro posizione all'interno del blocco ($w_k = \gamma^{k-1}$). Gli errori che si verificano all'inizio di un blocco sono ponderati più pesantemente perché invalidano l'intero suffisso successivo, avendo un impatto maggiore sulla lunghezza complessiva di accettazione.

Contributi Chiave

1. Identificazione del Limite: Il documento dimostra che l'SFT offline per i modelli bozza raggiunge rapidamente un plateau a causa del mismatch tra gli stati di addestramento (generati dal target) e gli stati di inferenza (generati dalla bozza).
2. Analisi delle Sfide OPD: Spiega perché l'OPD standard fallisce per i modelli bozza: i rollout indipendenti sono instabili, mentre i rollout assistiti dal target in modo ingenuo cancellano il segnale di errore on-policy.
3. Framework Draft-OPD: Introduce un nuovo framework che combina rollout assistiti dal target con il replay delle posizioni di errore, abilitando un efficiente post-addestramento sugli errori al momento della verifica.
4. Obiettivo Consapevole dell'Accettazione: Propone un obiettivo ibrido di divergenza KL che differenzia tra token accettati e rifiutati per colpire specificamente gli errori che limitano l'accelerazione speculativa.

Risultati Sperimentali

Gli autori hanno valutato Draft-OPD sulla famiglia Qwen3 (4B, 8B e 30B-A3B-Thinking) su ragionamento matematico (GSM8K, MATH-500, AIME), generazione di codice (MBPP, HumanEval, SWE-bench Lite) e chat generale (MT-Bench).

• Miglioramenti delle Prestazioni: Con un budget di FLOPs di addestramento equivalente, Draft-OPD ha costantemente superato EAGLE-3 e DFlash.
  • Ha raggiunto un'accelerazione senza perdita superiore a 5× per i modelli di ragionamento ("thinking").
  • Ha migliorato la lunghezza media di accettazione e l'accelerazione del 23% rispetto a EAGLE-3 e del 13% rispetto a DFlash.
• Throughput: Nei benchmark di distribuzione SGLang, Draft-OPD ha mostrato guadagni costanti nel throughput a tutti i livelli di concorrenza (da 1 a 32), con miglioramenti della lunghezza media di accettazione di circa l'11,2% rispetto a DFlash.
• Studi di Ablazione:
  • L'applicazione semplice di SFT su dati generati da OPD (senza il meccanismo di replay) non ha migliorato le prestazioni, confermando la necessità del segnale di errore on-policy.
  • L'uso di una singola direzione KL (tutta forward o tutta reverse) è stato inferiore all'obiettivo misto e consapevole dell'accettazione.
  • La selezione casuale degli ancoraggi di replay è stata meno efficace rispetto alla selezione basata sulle posizioni di errore, evidenziando l'importanza di concentrare l'addestramento sugli stati di fallimento.

Significato e Affermazioni

Il documento afferma che Draft-OPD rappresenta un significativo passo avanti nell'ottimizzazione dei modelli bozza basati sull'addestramento. Spostandosi da una supervisione statica e offline a una distillazione dinamica e on-policy che prende di mira esplicitamente gli errori al momento della verifica, il metodo rompe il plateau di prestazioni dell'SFT tradizionale.

Gli autori sottolineano che il loro approccio preserva la natura senza perdita della decodifica speculativa (mantenendo la distribuzione di output del modello target) migliorando sostanzialmente l'efficienza. Posizionano questo lavoro come una soluzione pratica per accelerare i modelli di ragionamento ("thinking") e i compiti complessi, dove i modelli bozza standard spesso faticano a causa dell'alta incertezza delle lunghe tracce di ragionamento. Il documento conclude che il post-addestramento dei modelli bozza sui propri errori al momento dell'inferenza è una direzione critica per migliorare la decodifica speculativa, offrendo una via verso lunghezze di accettazione più elevate e inferenze più rapide senza compromettere la qualità della generazione.