Summary
This paper investigates the cross-layer geometric and dynamical structure of routers in sparse Mixture-of-Experts (MoE) models. The authors propose a method to align router-control states across layers using orthogonal Procrustes analysis and generalized Procrustes analysis (GPA), revealing a shared, low-dimensional dynamical core that is functionally relevant for routing decisions.
Introduction and Theoretical Foundation
Sparse MoE models activate only a subset of experts per token, with a router determining expert selection. While prior work has explored cross-layer predictability of routing decisions, the underlying geometric structure has remained unclear. This paper addresses a fundamental question: Is there consistent, reusable dynamical structure in router-control states across layers? The authors introduce a theoretical framework based on:
- Gauge invariance: Router-control coordinates are not unique—any orthogonal transformation preserves routing outputs, creating an equivalence class of representations.
- Orthogonal Procrustes alignment: A method to align layer-specific coordinate systems by finding the optimal orthogonal transformation, enabling meaningful cross-layer comparison.
- Generalized Procrustes Analysis (GPA): Extends pairwise alignment to find a globally consistent canonical gauge across all layers.
The central dynamical hypothesis is that after alignment, a single low-capacity linear transition can approximate the layer-specific dynamics: Qₗ₊₁ᵀFₗQₗ ≈ A for many layers ℓ.
Methodology
The authors evaluate four sparse MoE architectures: Granite (24 layers, 32 experts), OLMoE-SFT (16 layers, 64 experts), Phi-tiny (32 layers, 16 experts), and IBM Shared (40 layers, 62 experts with explicit shared experts).
Key methodological components:
-
Canonicalization: Layer-wise centering and scalar normalization, followed by GPA to align all layers into a common coordinate system.
-
Shared dynamics fitting: A pooled linear transition is fitted via ridge regression across all layer transitions, compared against layer-specific models and persistence baselines.
-
Routing stability analysis: A theoretical result (Proposition 2) connecting state-space prediction error to routing decisions via the top-k margin and decoder operator norm.
-
Causal transport: Replacing native routing states at selected layers with transported states to measure functional impact via ∆NLL.
-
Matched-rank readout comparison: Comparing router-control states against matched-rank PCA residual subspaces with equal parameter budgets.
Empirical Validation / Results
Key findings:
-
Gauge alignment exposes reusable dynamics: A shared transition reaches 79–90% of the R² obtained by layer-specific models across all architectures. The improvement is specific to orthogonal Procrustes alignment—random gauges and PCA-basis alignments do not produce the same effect.
-
Routing specificity vs. residual smoothness: Residual PCA states are easier to predict across layers (R² = 0.788–0.861) than router-control states (R² = 0.371–0.542), but router states achieve much higher top-k routing recall (0.989–0.999 vs. 0.405–0.798). This demonstrates that temporal predictability and routing relevance are distinct properties.
-
Low-dimensional shared core: For OLMoE, very low-rank states (rank 2) are highly predictable (R² = 0.695) but preserve routing poorly (recall = 0.316), while full-rank states (rank 63) are less predictable (R² = 0.487) but preserve routing much better (recall = 0.699). A small shared component captures cross-layer dynamics, while additional dimensions preserve finer routing information.
-
Causal functional relevance: Canonical coordinates dramatically reduce transport degradation (∆NLL = 0.0070) compared to raw coordinates (0.1120) or random subspaces (0.1495). Learned dynamics outperform persistence at longer horizons for OLMoE and Phi-tiny.
-
Limits of the claim: Direct weight sharing degrades model quality; the IBM shared-expert model shows weaker universality; the claim is limited to local functional transport and reusable dynamics, not exact dynamical equivalence.
Theoretical and Practical Implications
-
Theoretical: The work formalizes a geometric framework for understanding router structure, introducing the concept of gauge invariance in MoE routing and demonstrating that cross-layer dynamics are approximately equivalent up to layer-specific orthogonal transformations.
-
Practical: The findings complement cross-layer prefetching systems by providing a coordinate system where simple shared predictors become effective. The results suggest potential for router-skipping or state-transport mechanisms, though no wall-clock speedup is claimed.
-
Architecture insights: The explicit shared-expert architecture (IBM) exhibits weaker universality, suggesting that architectural choices affect the degree of reusable cross-layer structure. The matched-readout result provides a methodology for evaluating whether any representation captures routing-relevant information.
Conclusion
Across four sparse MoE architectures, router-control states exhibit consistent cross-layer structure that becomes visible after orthogonal gauge alignment. A single linear transition captures most of the predictive power of layer-specific models with far fewer parameters. This effect is not explained by generic hidden-state smoothness—router-control states uniquely preserve expert choices. Causal interventions confirm functional relevance, with learned evolution outperforming persistence at longer horizons. The results support reusable cross-layer structure in MoE routers, though not exact dynamical equivalence, and establish a theoretical framework for future work on routing geometry and dynamics.
Related papers
- What Attention Recalls and Recurrence Controls in Hybrid Language Models
In hybrid language models, attention's KV cache enables exact item retrieval while the recurrent state controls output language and persona, revealing a functional dissociation.
- FlexComp: One Model for Every Ratio in Context Compression
FlexComp enables one model to compress contexts at any ratio via Matryoshka training, with per-input budget selection preserving accuracy at up to 266x compression and 47% higher throughput.
- Cacheable by Design? Training MoE Routers for Locality Against the Edge Memory-Bandwidth Wall: A Pre-Registered Negative Result
Training routers for cache locality cuts MoE cache misses by up to 60% but always costs perplexity, with no weight satisfying strict quality gates.