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

1. Probleemstelling

Sparse Mixture-of-Experts (MoE) modellen (bijv. DeepSeek-V3, Mixtral) vertrouwen momenteel op stateless affiniteitsgebaseerde routing. De routing-gate selecteert experts uitsluitend op basis van de huidige token-embedding ($x_t$), waarbij temporale context, de geschiedenis van expert-betrouwbaarheid en toekomstige vereisten worden genegeerd.

Het paper identificeert een kritieke faalmodus: Domeinovergangen. Wanneer een sequentie verschuift van de ene verdeling (Domein A) naar de andere (Domein B), behoort de token op het overgangstijdstip ($t$) nog steeds tot Domein A, maar behoort de optimale expert voor de volgende stap ($t+1$) tot Domein B.

• Het Probleem: Standaard routers, die alleen $x_t$ zien, wijzen een kans van bijna nul toe aan de juiste expert op het overgangstijdstip.
• Empirisch Bewijs: In gecontroleerde experimenten wijst standaard routing slechts 0,006 ± 0,001 waarschijnlijkheidsmassa toe aan de juiste expert op overgangstijdstippen.
• Gevolg: Om 99% routingdekking te bereiken met zo'n lage zekerheid, zou een model een onhaalbaar groot aantal experts nodig hebben ($K \approx 765$), waardoor de efficiëntievoordelen van MoE teniet worden gedaan.

2. Methodologie: Het Kader van het Vrije-Energieprincipe (FEP)

De auteur stelt dat MoE-routing moet worden herformuleerd via Friston's Vrije-Energieprincipe, dat voorschrijft dat een optimaal systeem onder onzekerheid moet:

1. Toestand recursief updaten (niet vanaf nul).
2. Signalen wegen op precisie (inverse variantie van fouten).
3. Verwachte toekomstige vrije energie minimaliseren (handelen op basis van voorspellingen, niet alleen huidige waarnemingen).

Om dit te implementeren, introduceert het paper drie lichtgewicht modificaties aan de standaard routing-gate, gebaseerd op Spiking Neural Network (SNN) dynamiek (specifiek Leaky Integrate-and-Fire, of LIF):

A. Temporale Routinggeheugen ($\beta$)

• Mechanisme: Vervangt de stateless input $x_t$ door een verborgen toestand $h_t$ die wordt bijgewerkt via LIF-dynamiek:
  $$h_t = \beta \odot h_{t-1} + x_t$$
  waarbij $\beta$ een leerbare vervaldvector is (geïnitieerd op 0,9).
• Functie: Accumuleert context over tokens. Het stelt de router in staat om te "onthouden" dat het systeem al geruime tijd in Domein A heeft verbleven, waardoor een "verzadigings"-signaal ontstaat dat een naderende schakeling aangeeft, zelfs voordat de schakeltoken verschijnt.
• Tijdschaal: Per-token (snel).

B. Precisie-gewogen Gating ($\Pi$)

• Mechanisme: Houdt een online, niet-parametrische tracker bij van de Mean Squared Error (MSE) voor elke expert. De routing-gate wordt gemoduleerd door de inverse variantie (precisie) van deze fouten:
  $$g_t = \text{softmax}(W \cdot x_t \odot \Pi_t)$$
  waarbij $\Pi_i = 1 / (\hat{\sigma}^2_i + \epsilon)$.
• Functie: Weegt experts die momenteel onbetrouwbaar zijn (hoge fout) dynamisch af en weegt betrouwbare experts op. In tegenstelling tot standaard load balancing (dat frequentie bijhoudt), houdt dit nauwkeurigheid bij.
• Tijdschaal: Per-batch (gemiddeld).

C. Anticiperende Routing

• Mechanisme: Gebruikt een predictor-netwerk $f(x_t, h_t)$ om de volgende token-embedding $\hat{x}_{t+1}$ te schatten. De routing-beslissing is geconditioneerd op deze voorspelling:
  $$g_t = \text{softmax}(W \cdot h_t + W_{pred} \cdot \hat{x}_{t+1})$$
• Functie: Routet op basis van de voorspelde toekomstige toestand in plaats van de huidige toestand. Dit stelt de router in staat om voordat de tokens van het nieuwe domein in de input verschijnen, over te schakelen naar de expert van het nieuwe domein.
• Cruciaal Inzicht: Een stateless predictor ($f(x_t)$) faalt omdat pre-overgang tokens distributieel identiek zijn. De predictor moet geconditioneerd zijn op de door $\beta$ geaccumuleerde toestand $h_t$ om het verzadigings-signaal te detecteren dat een overgang aangeeft.

3. Belangrijkste Bijdragen

1. Theoretische Unificatie: Verbindt MoE-routing met het Vrije-Energieprincipe, en toont aan dat standaard MoE ontbreekt aan recursieve toestandsschatting, precisieweging en minimalisatie van toekomstige toestanden.
2. Drie Lichtgewicht Mechanismen: Stelt $\beta$ (geheugen), $\Pi$ (precisie) en Anticipatie (voorspelling) voor als noodzakelijke componenten, implementeerbaar met minimale parameter-overhead (totaal ~200 regels code).
3. Ontdekking van Super-additiviteit: De meest significante bevinding van het paper is de $\beta \times \text{Ant}$ interactie.
   • $\beta$ alleen: +0,295 winst.
   • Anticipatie alleen: +0,000 winst (faalt zonder geheugen).
   • Gecombineerd: +0,741 winst.
   • Interactie-effect: De gecombineerde winst overtreft de som van de individuele delen met +0,446, wat bewijst dat geheugen een voorwaarde is voor anticipatie.
4. Structurele Verklaring voor DeepSeek-V3: Verklaart waarom DeepSeek-V3 zijn Multi-Token Prediction (MTP) module tijdens inferentie verwijdert: een stateless MTP-head kan domeinovergangen niet detecteren. Het paper betoogt dat MTP behouden moet blijven als het geconditioneerd is op temporair geheugen ($\beta$).

4. Experimentele Resultaten

Experimenten werden uitgevoerd op gecontroleerde domein-switch taken en een character-level MoE taalmodel (5 zaden).

• Overgangsnauwkeurigheid: Het gecombineerde mechanisme verhoogt de kans op het selecteren van de juiste expert op het overgangstijdstip van 0,6% naar 74,8%, waarmee 75% van de kloof met de oracle wordt gedicht.
• Efficiëntie: Het benodigde aantal experts ($K$) om 99% dekking te bereiken, daalt van een onhaalbaar ~765 naar een kleine constante (2,4) op overgangstijdstippen.
• Validatie Taalmodel: In een character-level LM verlaagde $\beta$-routing de Bits Per Character (BPC) op het overgangstijdstip van 6,56 naar 4,01. Het toevoegen van anticipatie stabiliseerde de routingbeslissing (verlaagde variantie) maar verlaagde de BPC niet verder op deze kleine schaal, wat suggereert dat het primaire voordeel ligt in routingstabiliteit en latency-reductie via prefetching.

5. Betekenis en Implicaties

• Navigatie versus Classificatie: Het paper herformuleert routing van een statisch classificatieprobleem ("welke expert past bij deze token?") naar een dynamisch navigatieprobleem ("waar gaat het systeem heen, en welk pad is betrouwbaar?").
• Latency-reductie: Door domeinwissels één stap vooruit te anticiperen, kan het systeem expert-weights vooraf ophalen uit het geheugen (HBM) naar de rekeneenheden (SRAM) terwijl de huidige stap wordt uitgevoerd, waardoor geheugen-latency-stalls op kritieke overgangspunten worden geëlimineerd.
• Herwaardering van MTP: Het werk suggereert dat Multi-Token Prediction modules in modellen zoals DeepSeek-V3 onderbenut zijn. Als ze geconditioneerd zijn op temporale toestand, kunnen ze dienen als krachtige anticiperende routingmechanismen in plaats van slechts als hulpmiddelen voor training.
• Biologische Plausibiliteit: De mechanismen spiegelen biologische processen na (LIF-neuronen, precisieweging in predictieve codering, actieve inferentie), wat suggereert dat huidige MoE-architecturen fundamentele principes van intelligente resource-allocatie missen.

Concluderend toont het paper aan dat affiniteit onvoldoende is voor robuuste MoE-routing. Door temporair geheugen, precisieweging en anticipatie te integreren, kunnen MoE-modellen bijna-oracle routingprestaties bereiken bij domeinovergangen, waardoor de rekenkosten voor het behouden van hoge nauwkeurigheid drastisch worden verlaagd.