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:

  1. Canonicalization: Layer-wise centering and scalar normalization, followed by GPA to align all layers into a common coordinate system.

  2. Shared dynamics fitting: A pooled linear transition is fitted via ridge regression across all layer transitions, compared against layer-specific models and persistence baselines.

  3. 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.

  4. Causal transport: Replacing native routing states at selected layers with transported states to measure functional impact via ∆NLL.

  5. Matched-rank readout comparison: Comparing router-control states against matched-rank PCA residual subspaces with equal parameter budgets.

Empirical Validation / Results

Key findings:

  1. 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.

  2. 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.

  3. 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.

  4. 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.

  5. 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