Mixture of Depths — DeepMind Insegna ai Transformer ad Allocare il Compute Dinamicamente

Cos'è: "Mixture-of-Depths: Dynamically Allocating Compute in Transformer-Based Language Models" (Raposo et al., Google DeepMind, aprile 2024) introduce un'architettura in cui ogni token decide, layer per layer, se "spendere" il compute completo del transformer block o saltarlo via residual. Un router top-k seleziona quali token processare in ogni layer; gli altri passano oltre senza essere toccati. Risultato: stessa accuracy con fino al 50% in meno di FLOPs, e un complemento naturale a Mixture-of-Experts (MoE) — MoE distribuisce orizzontalmente (quale esperto), MoD distribuisce verticalmente (quale profondità).

Il problema: il transformer spende troppo compute su token banali

I transformer standard hanno una proprietà semplice e costosa: ogni token attraversa tutti i layer della rete, sempre. Un articolo determinativo ("il", "the") riceve esattamente lo stesso budget di compute di una parola tecnica ambigua che richiede ragionamento multi-step. Per la maggior parte dei token, gran parte dei layer non sta facendo lavoro utile — sta semplicemente passando l'informazione avanti.

Questa uniformità di compute è una semplificazione di implementazione, non una necessità computazionale. Idealmente, un modello dovrebbe spendere molto compute sui token difficili e poco su quelli facili. Tentativi precedenti — adaptive computation time (Graves, 2016), CALM (Schuster et al., 2022), early exit dinamico — avevano provato a implementare questa idea ma con limitazioni: aggiungevano complessità al training, rompevano le forme tensoriali statiche richieste da TPU/GPU per il batching efficiente, e raramente offrivano gains netti in produzione.

L'idea di MoD: top-k routing per layer, decisione binaria per token

Il paper risolve il problema con una mossa elegante. Per ogni layer del transformer e per ogni token nella sequenza, una piccola rete di routing (essenzialmente un singolo vettore di pesi linear seguito da uno scalare per token) produce un punteggio. Si ordina la sequenza per punteggio e si selezionano i top-k token — dove k è un iperparametro fissato, ad esempio il 12.5% della sequenza.

Solo i token selezionati attraversano il blocco transformer completo (self-attention + MLP). Gli altri saltano il blocco e procedono via connessione residual al layer successivo. Il punto cruciale: k è fisso e noto a priori, quindi le forme dei tensori restano statiche e l'esecuzione parallela su hardware moderno resta efficiente. Non c'è il problema di "code dinamiche" che affligge i metodi early-exit.

Per token, la decisione di attraversare o saltare un layer è binaria — ma diventa interessante quando si compone su tutti i layer del modello. Un token può attraversare il blocco 1 ma saltare il 2 e il 3, attraversare il 4, e così via. Il "depth" effettivo che ogni token esperisce è dinamico e dipende dalla sua difficoltà, dal contesto e dalla posizione nella sequenza.

I numeri: riduzione FLOPs fino al 50%, accuracy invariata

I risultati sperimentali del paper sono convincenti. Su modelli da 220M a 1.4B parametri addestrati su dataset standard di language modeling, le configurazioni MoD con capacity k = 12.5% (un token su otto attraversa effettivamente ogni layer) mostrano:

• FLOPs di training ridotti del 30-50% rispetto al transformer denso equivalente, con la stessa loss finale.
• A pari FLOPs di training, i modelli MoD ottengono perplexity migliore — sono "compute-optimal" in modo più aggressivo.
• Velocità di inferenza più alta del 50% con qualità invariata, perché molti token saltano la maggior parte dei blocchi.

Un dettaglio importante: l'allocazione di compute appresa dal modello non è random. Analizzando quali token vengono selezionati nei vari layer, gli autori trovano pattern interpretabili — token rari, token in posizioni semanticamente cariche, token che iniziano costrutti complessi ricevono più compute. Token funzionali frequenti ne ricevono meno. Il routing impara una forma di "attention sulla difficoltà".

Complementare a MoE: orizzontale e verticale

Mixture-of-Experts (MoE), usato da modelli come Mixtral, GPT-4 (secondo le ricostruzioni della community) e Gemini 1.5, distribuisce il compute orizzontalmente: in ogni layer, ogni token sceglie uno o più "esperti" — sotto-reti MLP specializzate — tra molti disponibili. Il modello totale ha più parametri di quelli attivati per token, ma il compute per token resta proporzionale al numero di esperti attivati.

MoD distribuisce il compute verticalmente: in ogni layer, ogni token decide se essere processato o no. Le due tecniche sono ortogonali e componibili. Il paper Mixture-of-Depths-and-Experts (MoDE), discusso nello stesso lavoro, combina le due: per ogni layer ogni token decide (a) se attraversare il layer e (b) quale esperto usare se lo attraversa. Il risultato è un'allocazione di compute molto più granulare e adattiva di qualsiasi architettura transformer dense.

Concettualmente, MoE permette specializzazione (esperti diversi per domini diversi), MoD permette modulazione di sforzo (compute differenziato per difficoltà). Insieme, suggeriscono che il transformer denso uniforme — la cosa più ovvia da fare nel 2017 — è probabilmente la peggiore allocazione di compute possibile per un budget dato.

Perché potrebbe diventare standard nei modelli frontier 2025+

MoD è un'architettura nuova e non ancora dispiegata su scala frontier al momento del paper, ma ha tutte le caratteristiche per diventarlo. Tre ragioni principali. Primo: è hardware-friendly. A differenza di alcune tecniche di adaptive compute, MoD mantiene forme tensoriali statiche, rendendola compatibile con TPU, GPU e i compilatori esistenti senza modifiche architetturali fondamentali.

Secondo: si compone con tutto il resto. Funziona insieme a MoE, a quantization (4-bit/8-bit), a Flash Attention, a context window estesi. Non è una variante esotica del transformer: è un'aggiunta modulare. Questo facilita l'adozione progressiva nei laboratori frontier.

Terzo: la pressione economica spinge in questa direzione. Il costo di servire un LLM in produzione è dominato dall'inferenza, e l'inferenza è dominata dal numero di FLOPs per token. Una tecnica che dimezza i FLOPs con qualità invariata genera margini significativi. Gemini 2, Claude 4 e i futuri modelli OpenAI hanno tutti incentivi forti a integrare qualcosa di simile a MoD, anche se non sempre i dettagli architetturali sono pubblici. Il paper DeepMind ha probabilmente delineato una delle direzioni dominanti dell'architettura LLM per il 2025-2026.

---