Riepilogo Tecnico: Depth Exploration Decoding (DEX)

1. Definizione del Problema

La decodifica autoregressiva (AR) nei Large Language Models (LLM) è intrinsecamente costosa perché ogni token generato deve passare attraverso l'intero stack di layer prima che il token successivo possa essere prodotto. Studi recenti hanno identificato una significativa ridondanza lungo la dimensione della profondità: molti token diventano prevedibili a livelli intermedi, il che significa che non richiedono strettamente l'intera profondità del modello per essere accurati.

I metodi esistenti di adattamento della profondità senza perdita (lossless) tentano di sfruttare questo aspetto impiegando una strategia di selezione della profondità (depth-selection). Questi metodi propongono un token da un singolo livello di profondità intermedia pre-selezionato e lo verificano rispetto al modello a profondità finale. Tuttavia, gli autori identificano un collo di bottiglia della selezione (selection bottleneck):

• Uscita Tardiva (Late Exit): Se la profondità di uscita scelta è troppo profonda, la computazione viene sprecata su layer che erano già ridondanti.
• Uscita Anticipata (Early Exit): Se l'uscita scelta è troppo superficiale e la previsione non corrisponde all'output della profondità finale, il sistema attiva un fallback. Questo non solo scarta il token corrente, ma invalida anche qualsiasi "discendente" (draft) dipendente, portando a significative perdite di efficienza.

Il paper sostiene che il divario tra gli attuali metodi basati sulla selezione e un "oracolo" ideale (che conosce l'esatta profondità minima in cui un token diventa stabile) sia sostanziale. La domanda fondamentale è se questa ridondanza possa essere sfruttata più efficacementamente riducendo il rischio di una decisione errata sulla profondità di uscita.

2. Metodologia: Depth Exploration Decoding (DEX)

Gli autori propongono il Depth Exploration Decoding (DEX), un algoritmo lossless che sostituisce la selezione a profondità singola con l'esplorazione parallela su molteplici profondità candidate.

Concetti Chiave

• Earliest Available Depth (EAD): Gli autori definiscono EAD come il layer più superficiale in cui la previsione di un token corrisponde all'output a profondità finale e rimane stabile attraverso i layer successivi.
• Insieme di Esplorazione ($X$): Invece di scegliere un'uscita, DEX definisce un insieme di profondità candidate $\{d_0, d_1, \dots, d_{N-1}\}$, dove $d_{N-1}$ è la profondità finale.
• Risoluzione: L'overhead di DEX è limitato dalla "risoluzione" di questo insieme (il gap massimo tra un layer e il successivo livello esplorato). Man mano che l'insieme diventa più fine, l'overhead tende a zero.

L'Algoritmo: Expand–Commit–Collapse

DEX opera su un reticolo profondità-posizione (depth–position lattice) gestito da una Macchina a Stati Finiti (FSM). Il processo consiste in tre fasi:

1. Espansione del Reticolo (Esplorazione Parallela):

   • Il modello è partizionato in $N$ stadi di profondità, ciascuno gestito da un separato "esploratore di profondità" (potenzialmente su diverse unità hardware).
   • Ad ogni posizione di commit, il sistema espande i rami candidati in parallelo. Un token di esplorazione può subire un'espansione di profondità (passando il proprio stato nascosto al successivo esploratore del layer) o un'espansione di posizione (usando l'attuale candidato come prefisso per generare la successiva posizione del token).
   • Ciò crea un albero di potenziali futuri, dove diversi rami rappresentano previsioni effettuate a diverse profondità.

2. Commit (Validazione):

   • Una volta che il "percorso ancora" (l'anchor path, ovvero il ramo che attraversa tutti i layer) raggiunge la profondità finale, produce il token di riferimento a profondità finale.
   • DEX valida tutti i token candidati non finali del percorso ancora rispetto a questo riferimento.
   • Seleziona il ramo più superficiale il cui token corrisponde al riferimento a profondità finale. Se non viene trovato alcun match intermedio, si torna al ramo a profondità finale.
   • Fondamentalmente, il token aggiunto all'output è sempre quello prodotto dal riferimento a profondità finale, garantendo l'equivalenza lossless alla decodifica AR standard.

3. Collapse (Potatura dello Stato):

   • Dopo il commit di un token, il sistema collassa il reticolo. I rami incoerenti con il nuovo prefisso impegnato vengono scartati.
   • Solo gli stati (cache KV, stati nascosti) del ramo selezionato vengono preservati per il passaggio successivo.
   • Ciò assicura che, sebbene il sistema esplori più profondità, mantenga l'esatto stato richiesto per il successivo passo AR.

Ottimizzazioni

• Depth-Coupled Sampling (DECS): Per evitare che diverse profondità nello stesso slot chain propongano lo stesso token (il che sprecherebbe l'esplorazione), DEX maschera i token già proposti dalle profondità più superficiali durante la decodifica dei candidati più profondi. Questo incoraggia la diversità nelle proposte intermedie senza alterare la distribuzione di campionamento finale.
• Inducing Adapter: Per gli LLM standard dove i layer intermedi non sono naturalmente allineati con l'output head, gli autori propongono di attaccare piccoli adapter residuali alle profondità intermedie. Questi vengono addestrati tramite auto-distillazione per allineare gli stati nascosti intermedi con l'LM head finale, migliorando la "prontezza" dei token a profondità più superficiali.

3. Contributi Chiave

1. Formalizzazione della Readiness del Token: Il paper definisce la Earliest Available Depth (EAD) per quantificare quando un token diventa prevedibile. Utilizza questo concetto per dimostrare matematicamente il "collo di bottiglia della selezione" nei metodi esistenti, mostrando come il loro overhead sia guidato dal rischio di una selezione errata piuttosto che solo dal costo di computazione.
2. L'Algoritmo DEX: Un algoritmo di decodifica lossless che istanzia l'esplorazione della profondità tramite una procedura expand–commit–collapse. Preserva l'esatta equivalenza alla decodifica AR standard riducendo al contempo il costo del commit dei token.
3. Validazione Empirica: Gli autori dimostrano che DEX supera le baseline di selezione della profondità (LayerSkip, AdaDecode, DEL) e ottiene una throughput end-to-end competitiva rispetto ai metodi di speculazione e decodifica distribuita (Lookahead, PEARL, EAGLE).

4. Risultati Sperimentali

Gli autori hanno valutato DEX sia su modelli addestrati per l'early-exit (LayerSkip-trained CodeLlama-34B e Llama-2-70B) sia su LLM standard (CodeLlama-34B, Llama-2-70B, Qwen3-32B) attraverso task quali ragionamento matematico (GSM8K), generazione di codice (HumanEval) e riassunto (XSum).

• Guadagni di Throughput: DEX ha costantemente ottenuto una velocità superiore in termini di token al secondo (TPS) rispetto alle baseline di selezione della profondità. Sui modelli LayerSkip, DEX ha realizzato una frazione maggiore dell'headroom teorico dell'oracolo rispetto ai metodi di selezione.
• Scalabilità con la Risoluzione: La throughput è migliorata man mano che la risoluzione dell'esplorazione diventava più fine (ovvero, aggiungendo più esploratori). Ad esempio, su Llama-2-70B, aumentare gli esploratori da 3 a 8 ha ridotto significativamente il divario rispetto allo speedup dell'oracolo.
• Competitività: Nei confronti end-to-end, DEX (con 3–4 esploratori) è competitivo con il parallelismo tensoriale multi-GPU (TP) e la decodifica speculativa distribuita (PEARL, EAGLE) su risorse hardware equivalenti.
• Ablation: Rimuovere i rami intermedi (variante a uscita singola) ha mostrato che il guadagno di performance deriva specificamente dall'esplorazione di molteplici profondità, non solo dal parallelizzare i layer.

5. Significato e Claim

Il paper sostiene che l'esplorazione parallela della profondità fornisce un modo scalabile per sfruttare l'asse della profondità sottoutilizzato nella decodifica degli LLM.

• Cambio di Paradigma: A differenza dei metodi precedenti che scommettono su un singolo punto di uscita, DEX tratta l'adattamento della profondità come un problema di risoluzione controllata. Esplorando un insieme di profondità, converte l'incertezza del "quando uscire" in un overhead limitato determinato dalla spaziatura dell'insieme di esplorazione.
• Efficienza Lossless: DEX ottiene questi guadagni senza compromettere la correttezza dell'output (lossless) e senza richiedere modifiche all'architettura del modello (eccetto per i piccoli adapter opzionali nei modelli standard).
• Complementarità: Gli autori notano che DEX opera sull' "asse della profondità" (quanto presto un token viene determinato), il che è complementare ai metodi di accelerazione dell' "asse del token" (quanti token futuri vengono predetti). Suggeriscono che, sebbene la combinazione sia possibile, essa richiede una coordinazione complessa.

In conclusione, il paper stabilisce che sostituendo la selezione a profondità singola con l'esplorazione parallela, è possibile ridurre significativamente la latenza di decodifica mantenendo l'esatto comportamento dei modelli autoregressivi standard.