Sintesi Tecnica: Dual-Flow Transformers

Definizione del Problema

Il costo dell'erogazione dei Large Language Models (LLM) è sempre più dominato dall'inferenza piuttosto che dall'addestramento. Tuttavia, le due fasi principali dell'inferenza — prefill (elaborazione del prompt) e decode autoregressivo (generazione di token) — presentano colli di bottiglia hardware fondamentalmente diversi:

• Prefill: L'elaborazione parallela dei token del prompt è tipicamente compute-bound (limitata dal calcolo), vincolata al throughput aritmetico.
• Decode: La generazione sequenziale è tipicamente memory-bandwidth-bound (limitata dalla larghezza di banda della memoria), vincolata alla necessità di leggere ripetutamente la cache Key-Value (KV) in crescita e i pesi del modello per ogni nuovo token.

Le strategie di scaling convenzionali (aumento della larghezza o della profondità del modello) aumentano i costi per entrambe le fasi simultaneamente, poiché ogni strato aggiunto viene valutato sia durante il prefill che durante il decode. Ciò crea una rigidità in cui migliorare la qualità della generazione (che spesso beneficia di calcolo aggiuntivo) aumenta inevitabilmente il costo dell'elaborazione del prompt. Il paper si chiede se sia possibile allocare calcolo appreso aggiuntivo specificamente alla predizione della continuazione (decode) preservando il calcolo su tutto il prompt e una singola cache KV persistente.

Metodologia: Il Dual-Flow Transformer

Gli autori propongono il Dual-Flow Transformer, un'architettura che scinde strutturalmente il percorso primario di elaborazione del prompt dal percorso ausiliario di predizione della continuazione.

Architettura Core

Il modello consiste in due traiettorie di stato residuo che condividono la maggior parte dei parametri del modello ma differiscono nella gestione dello stato e nell'ambito delle attivazioni:

1. Flusso Primario (Primary Flow): Un percorso Transformer causale completo e standard. Elabora l'intero prompt, scrive la cache KV persistente e determina lo stato primario. È autosufficiente e non dipende dal flusso ausiliario.
2. Flusso Ausiliario (Auxiliary Flow): Una traiettoria residua che legge lo stato primario e la cache, ma non scrive mai nella cache KV persistente e non alimenta mai di ritorno il flusso primario.
   • Attivazione: Il flusso ausiliario è omesso durante l'elaborazione iniziale del prompt (prefill). Viene attivato solo all'ultima posizione del prompt (per predire il primo token di continuazione) e in ogni successiva posizione generata.
   • Condivisione dei Pesi: Entrambi i flussi condividono tutte le principali matrici di attenzione, MLP e di output.
   • Separazione: Utilizzano embedding di token separati e vettori di accoppiamento leggeri e appresi.

Meccanismi Chiave

• Attenzione Asimmetrica con KV Condiviso (Asymmetric Shared-KV Attention): Il flusso primario esegue la classica auto-attenzione causale. Il flusso ausiliario genera una nuova query ($Q_2$), che è accoppiata alla query primaria ($Q_1$) tramite un vettore appreso ($a_\ell$). Fondamentalmente, il flusso ausiliario presta attenzione alle chiavi e ai valori (Keys e Values) del flusso primario ($K_1, V_1$). Ciò consente al flusso ausiliario di sfruttare lo stato in cache senza scrivere il proprio, consentendo l'esecuzione raggruppata dove pesi e regioni KV vengono caricati una sola volta per servire entrambi i flussi.
• Accoppiamento Leggero (Lightweight Coupling): L'informazione fluisce unidirezionalmente dal flusso primario a quello ausiliario tramite vettori appresi ($a_\ell, b_\ell, c_\ell$) ai livelli di query, intermedio MLP e output MLP. Questo arricchisce la rappresentazione ausiliaria senza creare dipendenze che altererebbero il percorso primario.
• Obiettivo di Miscela (Mixture Objective): La predizione finale è una miscela delle distribuzioni primaria e ausiliaria: $p_{mix} = \alpha_t p_{primary} + (1-\alpha_t) p_{aux}$. Il peso di miscelazione $\alpha_t$ è regolato dinamicamente in base alla confidenza (probabilità di collisione) della distribuzione primaria, garantendo che il flusso primario rimanga dominante (almeno il 50% del peso) durante l'inferenza.
• Router Replay (per MoE): Nelle configurazioni Mixture-of-Experts (MoE), il flusso ausiliario riutilizza lo stesso set di esperti selezionati come il flusso primario (determinato dal router primario) ma li applica con i propri pesi di miscela. Ciò preserva la logica di selezione degli esperti del flusso primario pur aumentando l'aritmetica eseguita sugli esperti.

Costi di Inferenza vs. Addestramento

• Inferenza: La fase di prefill non incorre in costi aggiuntivi rispetto a un Transformer standard; il flusso ausiliario viene saltato completamente. La fase di decode incorre in un calcolo aggiuntivo (circa il doppio dell'aritmetica per token) ma non aumenta la dimensione della cache KV o l'insieme di pesi unici letti, poiché i due flussi sono raggruppati.
• Addestramento: Entrambi i flussi vengono valutati attraverso l'intera sequenza, risultando in circa 2x i FLOP di addestramento del modello primario. Tuttavia, l'aumento del numero di parametri è minimo (solo embedding separati e vettori di accoppiamento).

Contributi Chiave

1. Architettura: Introduzione di un design asimmetrico con KV condiviso che scinde il calcolo su tutto il prompt dal calcolo della continuazione. Il flusso primario gestisce la persistenza dello stato, mentre il flusso ausiliario fornisce una traiettoria "residua" per una migliore predizione senza alterare lo stato primario.
2. Qualità Empirica: Dimostrazione di riduzioni costanti della loss di validazione su tre assi sperimentali:
   • Scaling dei Dati (NanoGPT): Dual-Flow supera i Transformer standard a parità di budget di token, con un limite inferiore della loss (fitted loss floor) più basso.
   • Scaling della Dimensione del Modello (Dense LLaMA): I guadagni persistono all'aumentare della dimensione del modello.
   • MoE Sparsa: Validato in configurazioni MoE in stile LLaMA con router replay.
3. Allocazione MoE Specifica per Fase: A differenza dei precedenti design a flussi paralleli che applicano un'allocazione uniforme del calcolo, Dual-Flow permette di configurare indipendentemente i fan-out degli esperti primari e ausiliari. Ciò abilita uno spazio di trade-off tra costo di prefill, costo di decode e qualità predittiva, permettendo ai sistemi di ottimizzare per carichi di lavoro specifici (ad esempio, fissando il costo del prefill mentre si aumenta il calcolo del decode).

Risultati

• Loss di Validazione: In esperimenti su budget di token corrispondenti in NanoGPT, Dual-Flow raggiunge costantemente una loss di validazione inferiore rispetto ai Transformer standard. Il limite asintotico della loss (fitted asymptotic loss floor) è sceso da 2.9416 (Transformer) a 2.9013 (Dual-Flow).
• Efficienza Computazionale: In un confronto a parità di calcolo di addestramento (dove un modello Dual-Flow con metà dei passi di dati esegue circa gli stessi FLOP dello scheletro condiviso di un Transformer con il doppio dei passi), Dual-Flow ha comunque ottenuto una loss di validazione inferiore, suggerendo che allocare il calcolo tra flussi interagenti è più efficace rispetto al semplice aumento dei dati per un singolo flusso.
• Ablazioni: L'analisi dei componenti ha mostrato che la costruzione minima di Dual (embedding separati, KV condiviso) supera già l'arte precedente più vicina (PHD-2). L'aggiunta dell'accoppiamento tra i flussi e della miscela di probabilità ha ulteriormente migliorato le prestazioni.

Significato e Rivendicazioni

Il paper sostiene che Dual-Flow stabilisce un'opportunità strutturale per l'allocazione specifica di esperti per fase. Decoupling il percorso del prompt dal percorso di continuazione, l'architettura consente ai sistemi di serving di trattare il calcolo durante la generazione come una risorsa scalabile senza aumentare proporzionalmente il costo dell'elaborazione del prompt.

Nello specifico, gli autori evidenziano che nei modelli MoE, i flussi primario e ausiliario possono avere fan-out di esperti indipendenti. Ciò crea un controllo a tre vie:

1. Costo di Prefill: Determinato dall'uso degli esperti del flusso primario.
2. Costo di Decode: Determinato dalla somma dell'uso degli esperti primari e ausiliari.
3. Qualità Predittiva: Migliorata aumentando il calcolo totale allocato alla fase di continuazione.

Ciò consente una flessibilità specifica per l'implementazione, dove un sistema può preservare i costi di calcolo del prompt mantenendo il budget di decode disponibile per più esperti ausiliari, o viceversa, senza cambiare il pool di parametri condivisi o la struttura della singola cache primaria. Il lavoro supporta la premessa che il calcolo del next-token è una risorsa di scaling utile che può essere allocata internamente al modello piuttosto che esternamente tramite tracce di ragionamento più lunghe o campionamenti multipli.