Speculative Streaming — La Variante Più Efficiente del Decoding Speculativo

Cos'è: Speculative Streaming è una tecnica di accelerazione dell'inferenza per LLM proposta da Nikhil Bhendawade, Irina Belousova, Qichen Fu, Henry Mason, Mohammad Rastegari e Mahyar Najibi di Apple Machine Learning Research nel febbraio 2024. È una variante dello speculative decoding classico che elimina la necessità di un modello draft separato: lo stesso modello target produce contemporaneamente la sua predizione canonica e diverse predizioni speculative per i token successivi, tramite multi-stream prediction heads aggiunte al modello durante un fine-tuning leggero. Lo speedup misurato è 1.8-3.1x rispetto al decoding autoregressivo standard, con qualità identica all'output greedy o ai sampling stocastici.

Speculative decoding classico: il problema del draft model

Lo speculative decoding originale (Leviathan, Kalman, Matias, Google Research, novembre 2022) parte da un'osservazione semplice: l'inferenza autoregressiva genera un token alla volta, costando un forward pass completo del modello target per ogni token prodotto. Su un modello da 70B parametri ogni forward pass è dispendioso, anche se la maggior parte dei token prodotti sono prevedibili (articoli, congiunzioni, pattern frequenti). L'idea è usare un secondo modello molto più piccolo — il "draft" — per generare velocemente N token candidati, e poi verificare tutti gli N in un singolo forward pass del target.

La verifica funziona perché un forward pass del target su un prefisso di N token già noti produce, per costruzione, le probabilità di tutti i token possibili in ogni posizione. Confrontando le probabilità del target con quelle del draft si decide quanti token accettare (procedura rejection sampling, formalmente equivalente a samplare dal target). Se in media si accettano 3 token per ogni forward pass del target, lo speedup è circa 3x.

Il problema operativo è il draft model. Deve essere abbastanza grande da indovinare correttamente i token "facili" del target (altrimenti l'acceptance rate è basso e lo speedup svanisce), ma abbastanza piccolo da essere quasi gratuito da eseguire. Tipicamente si usa un modello della stessa famiglia ma 10-20 volte più piccolo (Llama 7B come draft per Llama 70B). Mantenere due modelli sincronizzati, fare-tunarli insieme, deployarli accanto: tutto questo aggiunge complessità infrastrutturale che alcuni team preferiscono evitare.

L'idea di Speculative Streaming: teste multi-stream sul target

La proposta di Apple è elegante. Invece di un modello draft esterno, si aggiunge al modello target una serie di "stream heads": piccole teste di predizione lineari che condividono il backbone del transformer e producono, per ogni posizione, predizioni per i token futuri t+1, t+2, ..., t+N. La testa "stream 0" è la testa canonica originale (predice t+1, lo stesso obiettivo del training next-token standard). Le teste "stream 1", "stream 2", ..., "stream N-1" sono nuove e predicono token più avanti.

Durante l'inferenza, ad ogni forward pass del target si ottengono contemporaneamente N predizioni speculative. Queste vengono usate esattamente come quelle che produrrebbe un draft model esterno: il forward pass successivo del target verifica tutte le N in batch e accetta il prefisso più lungo coerente. La differenza fondamentale è che non c'è overhead di esecuzione di un secondo modello: tutto avviene nello stesso forward pass del target, sfruttando il fatto che le teste aggiuntive sono leggere (poche centinaia di milioni di parametri totali, contro i miliardi del backbone).

L'addestramento delle teste è un fine-tuning leggero: si parte da un modello pre-addestrato standard, si congelano i pesi del backbone, si addestrano solo le N teste su un piccolo corpus (qualche miliardo di token) con una loss multi-task che minimizza la cross-entropy di ogni testa rispetto al token alla sua posizione target. Costo computazionale: percentuali singole rispetto al pre-training originale.

Risultati: speedup 1.8-3.1x, qualità invariata

Il paper riporta misurazioni su modelli pubblici nella famiglia OPT (1.3B, 6.7B, 13B) e su Llama 2 7B/13B. Lo speedup end-to-end è 1.8x su task di summarization, 2.6x su QA estrattivo, fino a 3.1x su task di generazione strutturata (codice, JSON). Il pattern è coerente: più i token futuri sono prevedibili dal contesto, più alta è l'acceptance rate delle teste speculative, più alto lo speedup.

Crucialmente, l'output è identico per costruzione a quello del decoding autoregressivo standard (con greedy decoding o con sampling stocastico, a seconda della modalità configurata). Non è un'approssimazione: la procedura di verifica garantisce che ogni token finale accettato sia un campione valido dalla distribuzione del target. Questo è un punto centrale: a differenza di tecniche come quantizzazione o distillazione che cambiano le risposte del modello, Speculative Streaming è un'ottimizzazione invisibile dal punto di vista dell'utente.

Il consumo di memoria aggiuntivo è limitato. Le teste totali aggiungono tipicamente 5-10% di parametri rispetto al modello target. Su Apple Silicon — dove la memoria unificata CPU/GPU è il vincolo principale per inferenza on-device di LLM — questo è un trade-off interessante: si spende 5-10% di memoria extra per ottenere 2-3x in throughput.

Implementazione: MLX, PyTorch, e il contesto Apple Intelligence

L'implementazione di riferimento è in PyTorch, e il paper include un porting per MLX, il framework di machine learning open-source di Apple ottimizzato per Apple Silicon (M1, M2, M3, M4). MLX è progettato per sfruttare la memoria unificata: tensori vivono in uno spazio di memoria condiviso tra CPU e GPU, eliminando i trasferimenti espliciti. Speculative Streaming si integra naturalmente in questo modello perché tutte le teste possono accedere allo stesso pool di memoria del backbone senza overhead di copia.

Il contesto strategico è chiaro: Apple Intelligence, annunciato a WWDC 2024 e disponibile su iOS 18/macOS Sequoia, esegue modelli on-device (un modello da circa 3B parametri per task locali, con fallback a modelli più grandi via Private Cloud Compute). L'inferenza on-device su iPhone è il caso d'uso primario: vincoli stretti su latenza (l'utente deve percepire risposte fluide) e consumo energetico (batteria). Speculative Streaming è ottimale per questo scenario perché elimina il modello draft (che raddoppierebbe l'occupazione di RAM, già scarsa su mobile) e raddoppia o triplica il throughput sullo stesso hardware.

C'è anche una connessione concettuale con il paper Multi-Token Prediction di Meta (Gloeckle et al., aprile 2024, due mesi dopo Speculative Streaming). Entrambi i lavori arrivano indipendentemente alla stessa idea — usare teste multiple per predire token futuri — ma con motivazioni diverse: Meta lo usa principalmente come obiettivo di training (migliora la qualità), Apple come obiettivo di inferenza (migliora la velocità). Le due tecniche sono complementari e c'è ricerca attiva sulla loro combinazione.

---