Riepilogo Tecnico: LEDE – Uscite Dinamiche Guidate dall'Esperienza per LLM con Reinforcement Learning

1. Definizione del Problema

I Large Language Models (LLM) soffrono di una latenza significativa durante l'inferenza autoregressiva perché ogni token generato deve attraversare l'intero stack di layer del transformer. Sebbene la Self-Speculative Decoding (SSD) sia stata proposta per accelerare questo processo utilizzando i propri layer iniziali come un "drafter" interno per generare token candidati in parallelo, gli attuali metodi SSD sono fondamentalmente limitati da configurazioni statiche.

Gli approcci esistenti si affidano a profondità di draft fisse e lunghezze di speculazione fisse applicate uniformemente durante l'intero processo di generazione. Questa rigidità non tiene conto della sparsità contestuale, dove la difficoltà di predire i token varia significativamente a seconda del contesto locale. Di conseguenza, i metodi statici non possono bilanciare in modo ottimale il trade-off tra qualità del draft (tasso di accettazione) e velocità di speculazione, portando a prestazioni di inferenza subottimali.

2. Metodologia: Il Framework LEDE

Gli autori propongono LEDE (Learning-based Dynamic Exit), un framework che inquadra il controllo dinamico della SSD come un Processo Decisionale di Markov (MDP). Invece di utilizzare regole rigide e predefinite, LEDE impiega il reinforcement learning (RL) offline per addestrare una policy che selezioni dinamicamente sia lo strato di uscita ottimale che la lunghezza della speculazione in base al contesto locale della sequenza generata.

2.1 Formulazione del Processo Decisionale di Markov (MDP)

Il framework modella il processo decisionale durante un singolo step di generazione come un MDP con i seguenti componenti:

• Spazio degli Stati ($S$): L'agente osserva un vettore di stato $s_l$ ad ogni layer candidato $l$. Questo vettore è derivato dalla distribuzione interna dei token e include:
  • Confidenza del Token ($C_k^l$): Il log-prob negativo medio dei top-$k$ token.
  • Entropia del Token ($H^l$): L'incertezza della distribuzione di predizione del token successivo.
  • Traccia Media della Confidenza del Token ($C_{trace}$): Il log-prob negativo medio del top-1 token su una finestra scorrevole di step precedenti.
• Spazio delle Azioni ($A$): Ad ogni layer candidato, la policy (una rete Q) decide tra due azioni:
  • Continua ($a=0$): Procedi al layer candidato successivo.
  • Esci ($a=1$): Designa l'attuale layer come lo strato di uscita ottimale ($l^*$) per la fase di draft.
• Funzione di Ricompensa ($R$): Progettata per ottimizzare l'efficienza dell'inferenza.
  • Ricompensa Positiva: Viene data se il token di draft viene accettato ($x_t^{l^*} = x_t^L$), scalata sulla profondità risparmiata ($1 + \frac{L-l^*}{L}$).
  • Ricompensa Negativa: Viene data se il token di draft viene rifiutato ($-2 + \frac{L-l^*}{L}$).
  • Penalità: Una piccola penalità ($-0.01$) viene applicata per l'azione di continuazione, incoraggiando le uscite precoci quando appropriato.

2.2 Ottimizzazione della Policy tramite RL Offline

LEDE utilizza una Deep Q-Network (DQN) addestrata offline utilizzando un replay buffer di passate esperienze di inferenza:

1. Raccolta Dati: Una fase iniziale di esplorazione raccoglie tuple di esperienza diversificate $(s_t, a_t, r_t, s_{t+1})$ in cui l'agente interagisce con l'ambiente.
2. Addestramento Offline: La rete Q viene addestrata minimizzando l'errore di Differenza Temporale (TD) n-step. Il valore target è calcolato utilizzando una rete target stabile, bilanciando la ricompensa immediata di un'uscita anticipata rispetto al valore a lungo termine dello stato.
3. Meccanismo di Doppio Controllo: Durante l'inferenza online, l'agente addestrato esegue due controlli adattivi:
   • Selezione Dinamica dello Strato di Uscita: Determina la profondità ottimale $l^*$ per iniziare il drafting.
   • Lunghezza del Draft Adattiva: Dopo aver selezionato $l^*$, l'agente rivaluta lo stato dopo ogni token di draft. Un'azione di "uscita" continua il drafting (indicando alta confidenza), mentre un'azione di "continua" (indicando bassa confidenza) termina immediatamente il drafting per la verifica.

3. Contributi Chiave

Il documento delinea tre contributi primari:

1. Formulazione MDP: È il primo lavoro a formulare il controllo dinamico della profondità di draft e della lunghezza della speculazione nella SSD come un problema MDP, applicando il reinforcement learning per derivare una policy di controllo.
2. Framework LEDE: La progettazione e l'implementazione di un framework che sfrutta il RL offline e l'esperienza di replay per apprendere una policy consapevole del contesto, superando i limiti delle configurazioni statiche e dei semplici metodi basati su euristiche.
3. Validazione Empirica: Esperimenti esaustivi che dimostrano incrementi di velocità significativi rispetto alla generazione autoregressiva e sostanziali guadagni di efficienza rispetto ai baseline statici.

4. Risultati Sperimentali

Il framework è stato valutato su modelli Llama-2 (7B, 13B) e Llama-3 (1B), nonché CodeLLaMA (7B, 34B), attraverso diversi compiti tra cui l'instruction following, la summarization, il language modeling e la generazione di codice.

• Prestazioni di Accelerazione (Speedup):
  • LEDE raggiunge un'accelerazione media di 2.32× rispetto alla decodifica autoregressiva standard.
  • I picchi di accelerazione raggiungono 2.72× (ad esempio, su Llama-2-7B per il language modeling).
  • LEDE fornisce un ulteriore incremento del 17% rispetto al miglior baseline speculativo statico (LayerSkip).
• Confronto con i Baseline:
  • LayerSkip (Statico): Le prestazioni sono altamente sensibili alle configurazioni fisse e non generalizzano bene tra i diversi task.
  • Metodi basati su Regole (LITE, Draft & Verify): Vincolati da soglie predefinite, portando a trade-off subottimali tra profondità e tassi di accettazione.
  • LEDE: Apprendendo una policy di controllo fine, supera costantemente tutti i baseline in termini di accelerazione pur mantenendo una qualità di generazione comparabile (misurata tramite Rouge-L).
• Studio di Ablazione: Rimuovere o la selezione dinamica dello strato di uscita o la lunghezza del draft adattiva degrada significativamente le prestazioni (riducendo l'accelerazione a circa ~2.0×), confermando che entrambi i meccanismi sono sinergici ed essenziali.

5. Significato e Rivendicazioni

Il paper sostiene che LEDE rappresenti uno slittamento concettuale dall'inferenza rigida guidata da euristiche verso un'inferenza adattiva guidata da policy. Imparando a co-ottimizzare la profondità computazionale e la lunghezza della speculazione in base al contesto in tempo reale, il framework consente una decodifica più robusta ed efficiente attraverso diversi modelli e task.

Gli autori posizionano questo lavoro come un passo verso la costruzione di sistemi LLM più intelligenti e computazionalmente efficienti. Notano che il lavoro futuro si concentrerà sulla scalabilità di questi benefici a modelli più grandi (scala 70B e oltre), ma il contributo attuale stabilisce la validità dell'uso del reinforcement learning offline per gestire dinamicamente la pipeline di inferenza dei grandi modelli linguistici.