🤖 machine learning

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

This paper proposes that standard affinity-based routing in Mixture-of-Experts models fails at domain transitions, and demonstrates that integrating Free Energy Principle-inspired mechanisms—specifically temporal memory, precision-weighted gating, and anticipatory routing—can dramatically improve expert selection accuracy and predictive performance by enabling the model to detect and prepare for distributional shifts before they occur.

Original authors: Man Yung Wong (Russell)

Published 2026-05-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Man Yung Wong (Russell)

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are navigating a massive, shifting landscape with a team of four specialized guides (the "Experts"). Your job is to pick the right guide for every step of your journey.

In current AI systems (called Mixture-of-Experts or MoE), the person picking the guides is like a tourist with no map, no compass, and no memory. They only look at the ground directly under their feet right now to decide who to hire.

The paper argues this is a fatal flaw. If you are walking through a forest and suddenly the terrain is about to turn into a desert, a guide who only looks at the current tree under your foot will keep hiring the "forest guide." They won't realize the switch is coming until it's too late.

The author, Russell Wong, suggests we need to upgrade this "hiring manager" using three simple tricks inspired by how the brain works (specifically, the Free Energy Principle). Here is how the paper explains these upgrades using everyday analogies:

1. The Problem: "The Amnesiac Tourist"

Current AI routing is stateless. It forgets everything that happened in the previous 100 steps.

The Result: When the AI switches from talking about "cats" to "cars," it often fails to switch experts immediately. In the paper's experiments, at the exact moment of the switch, the AI assigned a 0.6% chance to the correct expert. It was almost completely wrong.
The Analogy: It's like a chef who only tastes the spoonful of soup in their mouth right now, ignoring the fact that they just poured in a whole bucket of salt. They can't predict the soup is going to be too salty until it's already ruined.

2. The Solution: Three New Tools for the Hiring Manager

The paper proposes three lightweight changes to fix this.

A. Temporal Memory (The "Backpack")

What it is: Instead of looking only at the current step, the system carries a "backpack" (called a membrane potential) that holds a weighted memory of the last few steps.
The Analogy: Imagine the hiring manager now carries a backpack. If they've been walking through a forest for 30 minutes, the backpack gets heavy with "forest dust." Even if the ground under their foot looks like a forest, the heavy backpack tells them, "We've been here a long time; we are likely about to leave."
The Result: The AI learns to "forget" old information at just the right speed. If the terrain changes fast, the backpack empties quickly. If the terrain is stable, it holds on longer. This alone improved the AI's ability to catch the switch from 0.6% to 30%.

B. Precision-Weighted Gating (The "Trust Score")

What it is: The system tracks how reliable each expert has been recently. If an expert keeps making mistakes, the system lowers their "trust score" (precision) and stops listening to them, even if they look like a good fit for the current word.
The Analogy: Imagine you have four guides. Guide A is great at forests but terrible at deserts. Guide B is the opposite. The current system hires Guide A just because they are standing in a forest. The new system says, "Wait, Guide A made a huge mistake 5 minutes ago when we were near a desert. I'm lowering their trust score. I'll only hire them if they are 100% sure."
The Result: This helps the AI adapt when experts become unreliable. It creates a 31x difference in how much the system trusts a reliable expert versus an unreliable one.

C. Anticipatory Routing (The "Crystal Ball")

What it is: The system tries to predict what the next step will look like and hires the expert for that step, not the current one.
The Analogy: A normal tourist hires a guide for the path they are on now. An anticipatory tourist looks at the path ahead and hires the guide for the next turn.
The Catch: The paper found a secret ingredient. A "crystal ball" doesn't work if the tourist has no memory. If you try to predict the future without remembering the past, you can't tell when a switch is coming.
The Magic Combo: When you combine the Backpack (Memory) with the Crystal Ball (Prediction), the system becomes superhuman. The backpack fills up with "forest dust," and the crystal ball reads that saturation to say, "We are full of forest dust; the desert is coming next!"
The Result: This combination is super-additive.
- Backpack alone: +30% improvement.
- Crystal Ball alone: +0% improvement (it was blind).
- Backpack + Crystal Ball: +74% improvement.
- Together, they closed 75% of the gap between the AI and a perfect "Oracle" (a perfect guide).

3. Why This Matters (The "Navigation" Metaphor)

The paper concludes that current AI is just reacting to the present. It's a machine that says, "This word looks like a cat, so I'll use the cat expert."

The new system navigates. It says:

Where have I been? (Backpack/Memory: "I've been in the forest for a long time.")
Who can I trust? (Trust Score: "The desert guide is reliable; the forest guide is tired.")
Where am I going? (Crystal Ball: "The terrain is shifting to desert next.")

The Bottom Line

The paper claims that by adding these three simple, low-cost mechanisms (Memory, Trust, and Prediction), we can stop the AI from getting "lost" at the moments when it matters most: when the topic changes.

In their tests, this upgrade reduced the number of experts needed to get the job done from an impossible number (hundreds) to a tiny, manageable number (3 or 4). It turns a confused tourist into a skilled navigator.

Note: The paper does not claim this fixes all AI problems or applies to medical diagnosis. It strictly focuses on making the "routing" (the decision of which part of the AI to use) smarter during language generation.

1. Problem Statement

Sparse Mixture-of-Experts (MoE) models (e.g., DeepSeek-V3, Mixtral) currently rely on stateless affinity-based routing. The routing gate selects experts based solely on the current token embedding ( $x_t$ ), ignoring temporal context, expert reliability history, and future requirements.

The paper identifies a critical failure mode: Domain Transitions. When a sequence shifts from one distribution (Domain A) to another (Domain B), the token at the transition step ( $t$ ) still belongs to Domain A, but the optimal expert for the next step ( $t+1$ ) belongs to Domain B.

The Issue: Standard routers, seeing only $x_t$ , assign near-zero probability to the correct expert at the transition step.
Empirical Evidence: In controlled experiments, standard routing assigns only 0.006 ± 0.001 probability mass to the correct expert at transition steps.
Consequence: To achieve 99% routing coverage with such low confidence, a model would require an infeasibly large number of experts ( $K \approx 765$ ), negating the efficiency benefits of MoE.

2. Methodology: The Free Energy Principle (FEP) Framework

The author proposes that MoE routing should be reframed through Friston's Free Energy Principle, which dictates that an optimal system under uncertainty must:

Update state recursively (not from scratch).
Weight signals by precision (inverse variance of error).
Minimize expected future free energy (act on predictions, not just current observations).

To implement this, the paper introduces three lightweight modifications to the standard routing gate, grounded in Spiking Neural Network (SNN) dynamics (specifically Leaky Integrate-and-Fire, or LIF):

A. Temporal Routing Memory ( $\beta$ )

Mechanism: Replaces the stateless input $x_t$ with a hidden state $h_t$ updated via LIF dynamics:
$h_t = \beta \odot h_{t-1} + x_t$
where $\beta$ is a learnable decay vector (initialized at 0.9).
Function: Accumulates context across tokens. It allows the router to "remember" that the system has been in Domain A for a long time, creating a "saturation" signal that indicates an imminent switch, even before the switch token appears.
Timescale: Per-token (fast).

B. Precision-Weighted Gating ( $\Pi$ )

Mechanism: Maintains an online, non-parametric tracker of the Mean Squared Error (MSE) for each expert. The routing gate is modulated by the inverse variance (precision) of these errors:
$g_t = \text{softmax}(W \cdot x_t \odot \Pi_t)$
where $\Pi_i = 1 / (\hat{\sigma}^2_i + \epsilon)$ .
Function: Dynamically down-weights experts that are currently unreliable (high error) and up-weights reliable ones. Unlike standard load balancing (which tracks frequency), this tracks accuracy.
Timescale: Per-batch (medium).

C. Anticipatory Routing

Mechanism: Uses a predictor network $f(x_t, h_t)$ to estimate the next token embedding $\hat{x}_{t+1}$ . The routing decision is conditioned on this prediction:
$g_t = \text{softmax}(W \cdot h_t + W_{pred} \cdot \hat{x}_{t+1})$
Function: Routes based on the predicted future state rather than the current state. This allows the router to switch to the new domain's expert before the new domain tokens appear in the input.
Crucial Insight: A stateless predictor ( $f(x_t)$ ) fails because pre-transition tokens are distributionally identical. The predictor must be conditioned on the $\beta$ -accumulated state $h_t$ to detect the saturation signal indicating a transition.

3. Key Contributions

Theoretical Unification: Connects MoE routing to the Free Energy Principle, demonstrating that standard MoE lacks recursive state estimation, precision weighting, and future-state minimization.
Three Lightweight Mechanisms: Proposes $\beta$ (memory), $\Pi$ (precision), and Anticipation (prediction) as necessary components, implementable with minimal parameter overhead (~200 lines of code total).
Discovery of Super-Additivity: The paper's most significant finding is the $\beta \times \text{Ant}$ interaction.
- $\beta$ alone: +0.295 gain.
- Anticipation alone: +0.000 gain (fails without memory).
- Combined: +0.741 gain.
- Interaction Effect: The combined gain exceeds the sum of individual parts by +0.446, proving that memory is a prerequisite for anticipation.
Structural Explanation for DeepSeek-V3: Explains why DeepSeek-V3 discards its Multi-Token Prediction (MTP) module at inference: a stateless MTP head cannot detect domain transitions. The paper argues MTP should be retained if conditioned on temporal memory ( $\beta$ ).

4. Experimental Results

Experiments were conducted on controlled domain-switch tasks and a character-level MoE language model (5 seeds).

Metric	Standard MoE	$\beta$ -MoE	$\beta$ + Ant MoE	Oracle
Transition Accuracy	0.006	0.301	0.748	0.994
BPC at Transition	6.56	4.01	4.19	N/A
Experts for 99% Coverage ( $K$ )	~765 (Infeasible)	~6.0	2.4	N/A

Transition Accuracy: The combined mechanism raises the probability of selecting the correct expert at the transition step from 0.6% to 74.8%, closing 75% of the gap to the oracle.
Efficiency: The required number of experts ( $K$ ) to achieve 99% coverage drops from an infeasible ~765 to a small constant (2.4) at transition steps.
Language Model Validation: In a character-level LM, $\beta$ -routing reduced Bits Per Character (BPC) at the transition step from 6.56 to 4.01. Adding anticipation stabilized the routing decision (reduced variance) but did not further reduce BPC at this small scale, suggesting the primary benefit is routing stability and latency reduction via prefetching.

5. Significance and Implications

Navigation vs. Classification: The paper reframes routing from a static classification problem ("which expert fits this token?") to a dynamic navigation problem ("where is the system going, and which path is reliable?").
Latency Reduction: By anticipating domain switches one step ahead, the system can prefetch expert weights from memory (HBM) to compute units (SRAM) while the current step is executing, eliminating memory-latency stalls at critical transition points.
Re-evaluating MTP: The work suggests that Multi-Token Prediction modules in models like DeepSeek-V3 are underutilized. If conditioned on temporal state, they could serve as powerful anticipatory routing mechanisms rather than just training aids.
Biological Plausibility: The mechanisms mirror biological processes (LIF neurons, precision weighting in predictive coding, active inference), suggesting that current MoE architectures are missing fundamental principles of intelligent resource allocation.

In conclusion, the paper demonstrates that affinity is insufficient for robust MoE routing. By integrating temporal memory, precision weighting, and anticipation, MoE models can achieve near-oracle routing performance at domain transitions, drastically reducing the computational cost of maintaining high accuracy.