Source-linked AI summary
Evidence for Shared Routing Geometry and Dynamics in Sparse Mixture-of-Experts
Kirill Labzin, Stepan Kulibaba, Artem Dzhalilov, Artem Gorokhov
TL;DR
The paper investigates whether predictable routing across MoE depth reflects shared geometry rather than unrelated layer-specific coordinates. It aligns router-control subspaces with GPA and models their evolution using a shared linear transition, finding reusable dynamics that preserve routing information while improving longer-horizon transport over persistence on selected models.
Problem
The study asks what geometric structure underlies predictable routing across depth and whether routing-relevant states can be aligned into a common representation.
Method
The paper isolates each router’s control subspace, aligns layer-wise spaces with generalized orthogonal Procrustes analysis, and models canonical-state evolution with a shared linear transition.
Results
Across four architectures, one shared transition captures 79–90% of layer-specific dynamics, while router-control states preserve expert choices more faithfully than more predictable residual representations.
Takeaways & Limitations
The results support reusable cross-layer routing structure and partially reusable dynamics rather than exact dynamical equivalence across all layers.
Takeaways & Limitations
The study is empirical, uses one common text corpus, and finds that explicit shared-expert architectures show weaker universality; recursive open-loop skipping accumulates error.
Abstract
from arXiv · showhide
Sparse mixture-of-experts (MoE) models use an independently parameterized router at each sparse layer to select experts for every token. Prior work has shown that routing decisions across depth can often be predicted from earlier routing signals, suggesting that routing is not fully independent across layers. However, the structure behind this predictability remains unclear. In this work, we provide evidence that routing-relevant states across layers share a common geometric structure that is obscured by layer-specific coordinate systems. We isolate the control subspace of each router and align these spaces into a shared canonical representation using generalized orthogonal Procrustes analysis. After alignment, a single linear transition reaches $R^2=0.39$--$0.71$ and retains 79--90\% of the predictive power of separately fitted layer-specific dynamics, indicating that much of routing-state evolution follows a reusable process across depth. We then ask whether this shared dynamics is specific to routing or simply reflects the smooth evolution of hidden representations. A matched-rank comparison shows that residual representations are often easier to predict across layers, while router-control states preserve the model's expert choices much more faithfully. This separates generic cross-layer predictability from routing-specific information. Finally, we test whether the predicted canonical states remain meaningful when used in place of native routing states. The transported states preserve local routing behavior, while learned state evolution reduces $Δ\mathrm{NLL}$ relative to simple persistence by 15.7\% on OLMoE and 6.2\% over a 10-router horizon on Phi.
1 Introduction
This section asks what geometric structure explains predictable routing across depth and introduces router-control alignment as a test of reusable dynamics. The paper argues that aligned states support shared modeling while retaining routing-specific information.
- Sparse MoE routers determine expert selection, computation, load distribution, and inference-time expert availability.
- Cross-layer routing predictability and rotating routing-visible residual components motivate searching for shared geometry across depth.
- The study isolates each router’s functionally visible control subspace and aligns layer-wise spaces with generalized orthogonal Procrustes analysis.
- A single shared linear transition captures 79–90% of the R2 of layer-specific dynamics across four MoE architectures after gauge fixing.
- Matched-rank probes separate generic residual smoothness from information specifically useful for expert selection.
2 Related Work
Related work establishes that routing signals and routing-visible representations exhibit cross-layer structure, but does not determine whether that structure reflects shared dynamics or layer-specific coordinates. This paper addresses that distinction through aligned router-control states.
- Prior MoE research studies routing schemes, load balancing, failures, and expert specialization, whereas this paper examines routing representations across depth.
- Earlier work finds routing-visible control components that rotate across layers and shows that future expert choices can be predicted from earlier routing signals.
- These findings reveal cross-layer routing structure without resolving whether it arises from shared dynamics or layer-specific coordinates.
- Representation-alignment methods such as CKA, Procrustes alignment, and GPA provide tools for comparing neural representations beyond native coordinates.
- Dynamical-systems and Koopman-style analyses motivate treating depth as discrete time while modeling router-control rather than full hidden states.
3 Method
The method extracts router-control states, resolves their orthogonal coordinate ambiguity with GPA, and models their cross-layer evolution using a shared transition decoded into target-layer routing decisions. A top-k stability analysis links state prediction error to routing outcomes.
- The pipeline extracts low-dimensional router-control states, aligns them with GPA, evolves them using one linear transition, and decodes them with the target layer’s readout.
- Centering removes the common expert-logit shift because softmax probabilities and top-k selection are shift-invariant.
- The control basis spans exactly the residual-stream directions that can change relative router logits.
- At full numerical rank, the factorization reconstructs centered logits precisely, preserves native top-k sets exactly, and leaves language-model loss unchanged.
- 3.2 Exact orthogonal gauge symmetry: Orthogonal basis changes leave router outputs unchanged, so cross-layer comparison must account for this exact gauge symmetry.
- 3.2 Exact orthogonal gauge symmetry: Procrustes alignment resolves basis ambiguity, while GPA fits a single canonical coordinate system across depth using training data.
- The shared dynamical hypothesis posits a low-capacity operator fitted by pooled ridge regression, with layer-specific affine models as an upper bound.
- 3.5 Routing stability under transport: Canonical-state error induces centered-logit error through the layer decoder, and top-k stability depends on the routing margin.
4 Experimental Setup
The experiments compare four sparse MoE architectures using shared corpus splits, held-out tokens, and metrics covering geometry, dynamics, routing fidelity, and causal intervention effects.
- The evaluation spans four MoE architectures varying depth, expert count, routing sparsity, and explicit shared experts.
- Table 1 reports evaluated architectures and gives rank as the numerical rank of the centered router.
- Data and splits: WikiText-2 supplies common natural-language probe data, with independent sequence-level splits for geometry, alignment, and dynamics analyses.
- Data and splits: Final matched-readout validation and targeted causal replications use 50,176 held-out tokens per seed.
- Metrics: Metrics assess geometry, dynamics, routing fidelity, and causal interventions through ΔNLL.
5 Results
After gauge fixing, a single transition captures much of the predictable evolution of router-control states across layers. Matched-rank and causal tests show that routing fidelity requires information beyond the most predictable low-dimensional structure, with benefits varying by architecture and horizon.
- 5.1 Gauge alignment exposes reusable dynamics: Shared prediction strengthens after orthogonal Procrustes alignment, whereas raw, PCA-basis, random-gauge, and CCA coordinates remain weak.Shuffling token correspondences removes most of the alignment effect, supporting meaningful cross-layer alignment rather than an arbitrary orthogonal transformation.
- 5.1 Gauge alignment exposes reusable dynamics: 79–90% of layer-specific R2 is retained by one shared transition across four architectures, using far fewer parameters.The comparison is against separately fitted transitions for each adjacent layer.
- 5.2 Residual smoothness is not the same as routing specificity: Residual representations are easier to predict across layers, while router-control states recover expert choices more accurately under matched-rank, equal-budget readouts.The result separates generic cross-layer smoothness from information directly used for expert selection.
- 5.3 A low-dimensional shared core coexists with routing detail: Higher ranks improve routing recovery and causal transport even though very low-rank states can be highly predictable, especially in OLMoE.The rank sweep therefore distinguishes predictable shared structure from finer routing information.
- 5.3 A low-dimensional shared core coexists with routing detail: The IBM shared-expert model is less compressible: both predictability and routing fidelity continue improving as rank increases.This contrasts with the clearer low-rank trade-off observed for OLMoE.
- 5.4 Causal transport tests functional relevance: Canonical transport reduces causal degradation relative to raw, shuffled-gauge, and random matched-rank controls, while learned dynamics outperform persistence across replicated OLMoE and Phi tests.The benefit is architecture- and horizon-dependent, and repeated router skipping accumulates errors quickly.
6 Discussion
The paper distinguishes generic hidden-state smoothness from routing-specific structure by aligning router-control states across layers. The aligned representation supports reusable dynamics and preserves expert-selection information, but shared dynamics are approximate rather than universal.
- Orthogonal gauge fixing maps layer-specific control coordinates into canonical states where a low-capacity transition becomes reusable across depth.
- Residual PCA predicts future residual states more accurately, while router-control coordinates preserve expert choices more faithfully.This separates temporal predictability from routing relevance.
- A shared transition captures most of the predictive power of independently fitted layer-specific transitions with fewer parameters, but does not make all layers dynamically identical.The shared structure is best understood as approximate reuse after layer-specific gauge changes.
- The shared predictor provides a geometric coordinate system for operationalizing future-expert predictability in cross-layer prefetching.
7 Limitations
The study is an empirical, controlled comparison with several scope constraints. Its causal interventions assess routing-state transport rather than broad language-modeling or deployment benefits.
- The study uses one common text corpus for controlled cross-model comparison and remains empirical.
- Phi-tiny-MoE is instruction-tuned and distribution-mismatched to WikiText perplexity, so its causal results compare interventions rather than improved language modeling.
- The canonical state is identifiable only up to a common orthogonal transform because GPA selects one convenient global gauge.
- Recursive open-loop router skipping accumulates error, and the study makes no wall-clock speedup claim.
8 Conclusion
Across four sparse MoE architectures, the paper finds that orthogonal alignment reveals consistent router-control structure and partially reusable cross-layer dynamics. These states remain more routing-relevant than generic residual representations, and learned evolution improves longer-horizon transport on Phi and OLMoE.
- Across four sparse MoE architectures, router-control states show consistent cross-layer structure after orthogonal gauge alignment.
- A single linear transition captures most of the predictive power of layer-specific models while using far fewer parameters.
- Residual states are often easier to predict, but router-control states preserve expert choices more accurately.
- Learned state evolution outperforms persistence at longer horizons on Phi and OLMoE, while the evidence supports partial rather than exact dynamical equivalence.
A Mathematical Details
The mathematical details formalize router-control gauge invariance, multi-layer alignment, routing-stability bounds, and shared linear dynamics. Together, these constructions connect coordinate alignment to predictive structure and preservation of expert choices.
- Under an orthogonal coordinate transform, centered logits, softmax probabilities, and top-k decisions remain identical.
- The orthogonal Procrustes solution minimizes representation mismatch using the SVD of X^⊤Y.
- GPA minimizes average pairwise disagreement among aligned layer representations.
- If the logit perturbation satisfies ||δg||∞ < γ_k/2, the selected top-k expert set is unchanged.
- The shared operator is fitted by pooled least squares, with shared dynamics determined by cross-layer covariance in the canonical gauge.
B Additional Experimental Details
Additional experiments examine how causal transport changes with horizon, verify reconstruction and alignment procedures, and compare routing-control representations with matched-rank residual representations. The analyses also include seed-level replication controls and rank-sweep checks.
- Causal transport: Absolute ΔNLL generally increases with transport horizon, while learned shared transitions become more competitive with identity persistence at longer horizons, especially for OLMoE and Phi.Figure 4 reports median ΔNLL with interquartile ranges and compares shared dynamics against identity persistence.
- Validation checks: Centered router reconstruction matches native top-k routing exactly across all layers and models before every causal run.Evaluation scripts also verify model-weight checksums before and after interventions.
- Alignment controls: Alignment baselines compare independently sampled orthogonal gauges and PCA coordinate conventions against token-wise Procrustes alignment using mean shared-dynamics R2 across three split seeds.These controls are reported in Table 7 and used for Figure 2.
- Matched readout: Matched-readout comparisons use learned linear logit decoders with identical per-layer parameterization for router-control and residual-PCA representations.The parameterization is (r + 1)E, with random matched-rank residual subspaces used as a negative control.
- Seed-level replications: Targeted causal replication blocks are fixed before evaluation on additional split seeds rather than re-selected separately for each seed.Table 8 expands Table 6 and uses lower ΔNLL as the better outcome; Shared−Id. below zero favors learned dynamics.
- Rank sweeps: Rank truncation precedes gauge alignment, and the rank sweep is treated as structural evidence rather than a final benchmark.Representative OLMoE points appear in Table 4, with full cross-model curves in Figure 3.