Affinity Is Not Enough: Recovering the Free Energy Principle in Mixture-of-Experts

1. Enunciato del Problema

I modelli Sparse Mixture-of-Experts (MoE) (ad esempio DeepSeek-V3, Mixtral) attualmente si affidano a un routing basato sull'affinità senza stato. Il gate di routing seleziona gli esperti basandosi esclusivamente sull'embedding del token corrente ($x_t$), ignorando il contesto temporale, la storia dell'affidabilità degli esperti e i requisiti futuri.

Il paper identifica una modalità di fallimento critica: Transizioni di Dominio. Quando una sequenza passa da una distribuzione (Dominio A) a un'altra (Dominio B), il token al passo di transizione ($t$) appartiene ancora al Dominio A, ma l'esperto ottimale per il passo successivo ($t+1$) appartiene al Dominio B.

• Il Problema: I router standard, vedendo solo $x_t$, assegnano una probabilità prossima allo zero all'esperto corretto al passo di transizione.
• Evidenza Empirica: In esperimenti controllati, il routing standard assegna solo una massa di probabilità di 0,006 ± 0,001 all'esperto corretto ai passi di transizione.
• Conseguenza: Per ottenere una copertura di routing del 99% con una tale bassa confidenza, un modello richiederebbe un numero di esperti inattuabilmente grande ($K \approx 765$), annullando i benefici di efficienza dell'MoE.

2. Metodologia: Il Framework del Principio dell'Energia Libera (FEP)

L'autore propone che il routing MoE debba essere riformulato attraverso il Principio dell'Energia Libera di Friston, che stabilisce che un sistema ottimale in condizioni di incertezza deve:

1. Aggiornare lo stato in modo ricorsivo (non da zero).
2. Pesare i segnali in base alla precisione (varianza inversa dell'errore).
3. Minimizzare l'energia libera futura attesa (agire sulle previsioni, non solo sulle osservazioni attuali).

Per implementare ciò, il paper introduce tre modifiche leggere al gate di routing standard, basate sulla dinamica delle Reti Neurali a Spike (SNN) (specificamente Leaky Integrate-and-Fire, o LIF):

A. Memoria di Routing Temporale ($\beta$)

• Meccanismo: Sostituisce l'input senza stato $x_t$ con uno stato nascosto $h_t$ aggiornato tramite dinamiche LIF:
  $$h_t = \beta \odot h_{t-1} + x_t$$
  dove $\beta$ è un vettore di decadimento apprendibile (inizializzato a 0,9).
• Funzione: Accumula il contesto attraverso i token. Permette al router di "ricordare" che il sistema è stato nel Dominio A per lungo tempo, creando un segnale di "saturazione" che indica un imminente cambio, anche prima che appaia il token di transizione.
• Scala Temporale: Per-token (veloce).

B. Gate Pesato per Precisione ($\Pi$)

• Meccanismo: Mantiene un tracciatore online e non parametrico dell'Errore Quadratico Medio (MSE) per ogni esperto. Il gate di routing è modulato dalla varianza inversa (precisione) di questi errori:
  $$g_t = \text{softmax}(W \cdot x_t \odot \Pi_t)$$
  dove $\Pi_i = 1 / (\hat{\sigma}^2_i + \epsilon)$.
• Funzione: Soppesa dinamicamente gli esperti che sono attualmente inaffidabili (alto errore) e potenzia quelli affidabili. A differenza del bilanciamento del carico standard (che traccia la frequenza), questo traccia l'accuratezza.
• Scala Temporale: Per-batch (media).

C. Routing Anticipatorio

• Meccanismo: Utilizza una rete predittiva $f(x_t, h_t)$ per stimare l'embedding del token successivo $\hat{x}_{t+1}$. La decisione di routing è condizionata a questa previsione:
  $$g_t = \text{softmax}(W \cdot h_t + W_{pred} \cdot \hat{x}_{t+1})$$
• Funzione: Instrada in base allo stato futuro previsto piuttosto che allo stato corrente. Ciò permette al router di passare all'esperto del nuovo dominio prima che i token del nuovo dominio appaiano nell'input.
• Insight Cruciale: Un predittore senza stato ($f(x_t)$) fallisce perché i token pre-transizione sono distribuzionalmente identici. Il predittore deve essere condizionato sullo stato accumulato da $\beta$, $h_t$, per rilevare il segnale di saturazione che indica una transizione.

3. Contributi Chiave

1. Unificazione Teorica: Collega il routing MoE al Principio dell'Energia Libera, dimostrando che l'MoE standard manca di stima dello stato ricorsiva, pesatura per precisione e minimizzazione dello stato futuro.
2. Tre Meccanismi Leggeri: Propone $\beta$ (memoria), $\Pi$ (precisione) e Anticipazione (previsione) come componenti necessari, implementabili con un sovraccarico di parametri minimo (~200 righe di codice in totale).
3. Scoperta della Super-additività: La scoperta più significativa del paper è l'interazione $\beta \times \text{Ant}$.
   • $\beta$ da solo: +0,295 di guadagno.
   • Anticipazione da sola: +0,000 di guadagno (fallisce senza memoria).
   • Combinati: +0,741 di guadagno.
   • Effetto di Interazione: Il guadagno combinato supera la somma delle singole parti di +0,446, dimostrando che la memoria è un prerequisito per l'anticipazione.
4. Spiegazione Strutturale per DeepSeek-V3: Spiega perché DeepSeek-V3 scarta il suo modulo Multi-Token Prediction (MTP) durante l'inferenza: un capo MTP senza stato non può rilevare le transizioni di dominio. Il paper sostiene che l'MTP dovrebbe essere mantenuto se condizionato sulla memoria temporale ($\beta$).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su compiti controllati di cambio di dominio e su un modello linguistico MoE a livello di caratteri (5 semi).

• Accuratezza di Transizione: Il meccanismo combinato alza la probabilità di selezionare l'esperto corretto al passo di transizione da 0,6% al 74,8%, colmando il 75% del divario rispetto all'oracle.
• Efficienza: Il numero richiesto di esperti ($K$) per ottenere una copertura del 99% scende da un inattuabile ~765 a una piccola costante (2,4) ai passi di transizione.
• Validazione del Modello Linguistico: In un LM a livello di caratteri, il routing $\beta$ ha ridotto i Bit per Carattere (BPC) al passo di transizione da 6,56 a 4,01. L'aggiunta dell'anticipazione ha stabilizzato la decisione di routing (ridotto la varianza) ma non ha ulteriormente ridotto il BPC a questa piccola scala, suggerendo che il beneficio principale è la stabilità del routing e la riduzione della latenza tramite il prefetching.

5. Significato e Implicazioni

• Navigazione vs Classificazione: Il paper riformula il routing da un problema di classificazione statica ("quale esperto si adatta a questo token?") a un problema di navigazione dinamica ("dove sta andando il sistema e quale percorso è affidabile?").
• Riduzione della Latenza: Anticipando i cambi di dominio un passo avanti, il sistema può precaricare i pesi degli esperti dalla memoria (HBM) alle unità di calcolo (SRAM) mentre il passo corrente è in esecuzione, eliminando i blocchi di latenza di memoria ai punti critici di transizione.
• Rivalutazione dell'MTP: Il lavoro suggerisce che i moduli Multi-Token Prediction in modelli come DeepSeek-V3 sono sottoutilizzati. Se condizionati sullo stato temporale, potrebbero servire come potenti meccanismi di routing anticipatorio invece che come semplici aiuti all'addestramento.
• Plausibilità Biologica: I meccanismi rispecchiano processi biologici (neuroni LIF, pesatura per precisione nella codifica predittiva, inferenza attiva), suggerendo che le attuali architetture MoE stanno perdendo principi fondamentali dell'allocazione intelligente delle risorse.

In conclusione, il paper dimostra che l'affinità non è sufficiente per un routing MoE robusto. Integrando memoria temporale, pesatura per precisione e anticipazione, i modelli MoE possono raggiungere prestazioni di routing vicine all'oracle alle transizioni di dominio, riducendo drasticamente il costo computazionale del mantenimento di un'alta accuratezza.