Stage-Aware Mixture of Experts Overview

• Stage-Aware Mixture of Experts (MoE) are models that condition expert allocation on continuous or discrete stage variables, adapting to temporal, developmental, or pathological regimes.
• They integrate stage-conditioned inference with stage-structured optimization to address issues like simplicity bias in RL and evolving dynamics in disease modelling.
• Empirical results show improved expert specialization, interpretability, and performance across applications in reinforcement learning, neurodegenerative progression, and language tasks.

Stage-aware Mixture of Experts (MoE) denotes a class of MoE formulations in which expert selection, expert weighting, or the training regime is explicitly conditioned on a stage-like variable such as a latent decision phase, disease stage, known indexing variable, or training stage. In contrast to standard MoE designs that rely on token-level routing or fixed-coefficient gating, these models make expert allocation depend on temporally or structurally coherent regimes. Recent formulations include the phase-aware policy architecture for agentic reinforcement learning in "Phase-Aware Mixture of Experts for Agentic Reinforcement Learning" (Yang et al., 19 Feb 2026), the time-dependent gating of mechanistic, graph-diffusion, and neural-reaction experts in neurodegenerative progression modelling (He et al., 9 Aug 2025), the varying-coefficient statistical MoE of Zhao et al. for longitudinal and indexed data (Zhao et al., 5 Jan 2026), and the two-stage dense-to-sparse training framework of EvoMoE (Nie et al., 2021).

1. Conceptual scope and design space

The central idea is that heterogeneous data often exhibit regime structure that is not well captured by a single shared policy or by a standard token-level gate. In agentic RL, the motivating problem is simplicity bias: a single policy network causes simple tasks to occupy most parameters and dominate gradient updates, leaving insufficient capacity for complex tasks (Yang et al., 19 Feb 2026). In disease modelling, the motivating problem is that traditional models assume fixed mechanisms throughout disease progression despite stage-dependent pathological dynamics (He et al., 9 Aug 2025). In statistical MoE, the limitation is that constant coefficients in gating and expert models can be inadequate when covariate influences and latent subpopulation structure evolve across a known dimension (Zhao et al., 5 Jan 2026). In large-scale Transformer MoE training, EvoMoE identifies immature experts and unstable sparse gate learning when both experts and gate are random at initialization (Nie et al., 2021).

A useful organizing distinction is between stage-conditioned inference and stage-structured optimization. Stage-conditioned inference appears in models where the gate or the expert coefficients are explicit functions of a phase or stage variable. PA-MoE learns latent phase boundaries directly from the RL objective and routes each environment step to exactly one expert (Yang et al., 19 Feb 2026). IGND-MoE uses time-dependent weights $\beta_j(t)$ over three continuous-time expert vector fields (He et al., 9 Aug 2025). VCMoE lets both $\pi_j(X,t)$ and expert parameters vary smoothly in the indexing variable $t$ (Zhao et al., 5 Jan 2026). By contrast, EvoMoE is stage-aware in the sense of training procedure: expert-diversify precedes gate-sparsify, and the gate evolves from dense to sparse (Nie et al., 2021).

Formulation	Stage signal	Mechanism
PA-MoE	Latent phase from RL state/history/goal	Top-1 expert routing with a lightweight phase router
IGND-MoE	Global pseudo-time $t$	Softmax-gated mixture of three ODE experts
VCMoE	Known index $t$	Gating and expert coefficients vary smoothly in $t$
EvoMoE	Training stage	Expert-diversify followed by dense-to-sparse gate learning

This landscape suggests that “stage awareness” is not a single mechanism but a family of design choices for imposing regime sensitivity on MoE models.

2. Routing and weighting mechanisms

PA-MoE implements stage awareness through a lightweight phase router $g_\phi(s_t) \equiv p_t \in \Delta^K$, where

$$p_t = \mathrm{softmax}\!\left(\mathrm{MLP}\left([o_t^{align};h_t^{enc}]\right)/\tau\right).$$

Here $$o_t^{align} = \mathrm{CrossAttn}(\mathrm{Enc}(o_t),\mathrm{Enc}(g),\mathrm{Enc}(g))$$ and $$h_t^{enc} = \mathrm{LSTM}(\mathrm{Embed}(a_{t-L:t-1}),\mathrm{Embed}(o_{t-L:t-1}))$$. The router takes as input the current state $\pi_j(X,t)$0, uses a short LSTM over the past $\pi_j(X,t)$1 $\pi_j(X,t)$2 pairs, and selects $\pi_j(X,t)$3 deterministically with a straight-through estimator. The policy at step $\pi_j(X,t)$4 is then $\pi_j(X,t)$5, where the backbone $\pi_j(X,t)$6 remains frozen and each expert is a LoRA adapter on the query/value projections in each Transformer layer (Yang et al., 19 Feb 2026).

IGND-MoE uses a different mechanism: the gate depends only on time. The three expert vector fields

$\pi_j(X,t)$7

are combined by stage-dependent weights $\pi_j(X,t)$8 satisfying $\pi_j(X,t)$9 and $t$0. The weights are produced by a small temporal-attention network $t$1 with

$t$2

In practice $t$3 is a small MLP in the single scalar input $t$4, yielding smooth stage-wise expert contributions (He et al., 9 Aug 2025).

VCMoE generalizes this principle statistically. For latent experts $t$5, the gating probabilities are

$t$6

while each expert density $t$7 also depends on coefficient curves that vary with $t$8. The model therefore makes both mixture weights and expert responses stage-dependent, rather than only modulating a fixed expert bank (Zhao et al., 5 Jan 2026).

EvoMoE uses Gumbel-Softmax routing scores during gate-sparsify:

$t$9

A content-based threshold $t$0 prunes low-score experts,

$t$1

and after a designated number of iterations $t$2, routing switches to Top-1. This is not stage-conditioned by the input domain in the same sense as PA-MoE, IGND-MoE, or VCMoE, but it is explicitly stage-aware as an optimization schedule (Nie et al., 2021).

3. Learning objectives and estimation procedures

PA-MoE couples routing and expert learning through a joint RL objective. To encourage temporally contiguous assignments, it defines the discrete phase index

$t$3

and adds a switching penalty

$t$4

with the backward pass using the soft surrogate $t$5. The full loss is

$t$6

where $t$7 and $t$8. No explicit phase labels or semi-MDP definitions are given; the router and experts are trained jointly under the single RL objective, and $t$9 is annealed from $t$0 (Yang et al., 19 Feb 2026).

IGND-MoE is trained by dual iterative optimization. Because observations are irregular and infrequent, each subject $t$1 is assigned a time shift $t$2. Temporal alignment solves

$t$3

and, with shifts fixed, trajectory construction minimizes

$t$4

The total loss is

$t$5

where $t$6 constrains learned experts relative to the mechanistic model and $t$7 encourages diversity between experts. All ODE parameters are updated by gradient-based ODE backprop, while $t$8 are updated in a separate one-dimensional minimization (He et al., 9 Aug 2025).

VCMoE is estimated by a label-consistent EM algorithm built on local likelihood. At a target stage $t$9, each coefficient curve is approximated by a local linear expansion, and the local weighted log-likelihood is

$t$0

The E-step computes responsibilities

$t$1

and the M-step maximizes the weighted local objective on a grid. The paper further establishes identifiability up to label-swapping, pointwise asymptotic normality, sup-norm approximations, simultaneous confidence bands, and a generalized likelihood ratio test for testing whether a coefficient is genuinely varying across the index variable (Zhao et al., 5 Jan 2026).

EvoMoE decomposes training into two phases. Stage 1, expert-diversify, warms up a single shared expert and then spawns multiple diverse experts. Stage 2, gate-sparsify, trains the gate with DTS-Gate, starting from a dense gate and gradually becoming sparse while routing tokens to fewer experts. The stage transition is controlled by $t$2 and $t$3, with temperature annealed from $t$4 to $t$5 and an optional balance loss for expert loads (Nie et al., 2021).

4. Mathematical structures of the experts

A distinctive feature of stage-aware MoE is that the experts are often structurally heterogeneous rather than merely replicated feed-forward blocks. PA-MoE uses sparse experts implemented as LoRA adapters of rank $t$6 on the query/value projections in each Transformer layer. The backbone is frozen, and only one expert’s LoRA weights receive gradients at each step because routing is top-1. This yields hard parameter isolation (Yang et al., 19 Feb 2026).

IGND-MoE is more heterogeneous. The overall continuous dynamics satisfy

$t$7

with $t$8, $t$9, and $g_\phi(s_t) \equiv p_t \in \Delta^K$0. The inhomogeneous graph neural diffusion expert computes a low-rank latent representation

$g_\phi(s_t) \equiv p_t \in \Delta^K$1

builds a refined adjacency

$g_\phi(s_t) \equiv p_t \in \Delta^K$2

and outputs

$g_\phi(s_t) \equiv p_t \in \Delta^K$3

The localized neural reaction expert is

$g_\phi(s_t) \equiv p_t \in \Delta^K$4

This composition preserves a mechanistic component while allowing stage-dependent graph-refined diffusion and residual local dynamics (He et al., 9 Aug 2025).

VCMoE formalizes stage dependence at the level of statistical function classes. In the Gaussian expert example,

$g_\phi(s_t) \equiv p_t \in \Delta^K$5

so both regression effects and dispersion vary smoothly with $g_\phi(s_t) \equiv p_t \in \Delta^K$6. The identifiability theorem requires differentiability and separation conditions involving both the coefficient functions and their derivatives. This places stage-aware MoE within a nonparametric or semiparametric framework rather than only a neural-routing framework (Zhao et al., 5 Jan 2026).

EvoMoE retains the standard expert aggregation rule

$g_\phi(s_t) \equiv p_t \in \Delta^K$7

but its contribution lies in the path by which the experts and gate are trained. The model begins from a single shared expert, then diversifies experts, and only afterward learns sparse routing. This suggests that expert maturity and gate maturity can be treated as distinct optimization problems (Nie et al., 2021).

5. Empirical evidence and observed specialization

In agentic RL, PA-MoE is evaluated on ALFWorld and WebShop with PPO, RLOO, GRPO, and GiGPO on Qwen2.5-1.5B and 7B. The paper defines a “parameter occupancy” metric as the fraction of batches where a task’s loss is greater than $g_\phi(s_t) \equiv p_t \in \Delta^K$8 of total. Under a single GiGPO policy, simple pick-and-place dominates $g_\phi(s_t) \equiv p_t \in \Delta^K$9 of updates, while complex manipulation tasks Heat/Cool/Clean account for only $$p_t = \mathrm{softmax}\!\left(\mathrm{MLP}\left([o_t^{align};h_t^{enc}]\right)/\tau\right).$$0. This is associated with high success on Pick ($$p_t = \mathrm{softmax}\!\left(\mathrm{MLP}\left([o_t^{align};h_t^{enc}]\right)/\tau\right).$$1) but poorer results on Heat/Cool ($$p_t = \mathrm{softmax}\!\left(\mathrm{MLP}\left([o_t^{align};h_t^{enc}]\right)/\tau\right).$$2). With $$p_t = \mathrm{softmax}\!\left(\mathrm{MLP}\left([o_t^{align};h_t^{enc}]\right)/\tau\right).$$3 experts, each expert sees $$p_t = \mathrm{softmax}\!\left(\mathrm{MLP}\left([o_t^{align};h_t^{enc}]\right)/\tau\right).$$4 of batches, overall success on ALFWorld improves from $$p_t = \mathrm{softmax}\!\left(\mathrm{MLP}\left([o_t^{align};h_t^{enc}]\right)/\tau\right).$$5 to $$p_t = \mathrm{softmax}\!\left(\mathrm{MLP}\left([o_t^{align};h_t^{enc}]\right)/\tau\right).$$6 $$p_t = \mathrm{softmax}\!\left(\mathrm{MLP}\left([o_t^{align};h_t^{enc}]\right)/\tau\right).$$7 pp$$p_t = \mathrm{softmax}\!\left(\mathrm{MLP}\left([o_t^{align};h_t^{enc}]\right)/\tau\right).$$8, and WebShop improves from $$p_t = \mathrm{softmax}\!\left(\mathrm{MLP}\left([o_t^{align};h_t^{enc}]\right)/\tau\right).$$9 to $$o_t^{align} = \mathrm{CrossAttn}(\mathrm{Enc}(o_t),\mathrm{Enc}(g),\mathrm{Enc}(g))$$0 $$o_t^{align} = \mathrm{CrossAttn}(\mathrm{Enc}(o_t),\mathrm{Enc}(g),\mathrm{Enc}(g))$$1 pp$$o_t^{align} = \mathrm{CrossAttn}(\mathrm{Enc}(o_t),\mathrm{Enc}(g),\mathrm{Enc}(g))$$2. Ablations show $$o_t^{align} = \mathrm{CrossAttn}(\mathrm{Enc}(o_t),\mathrm{Enc}(g),\mathrm{Enc}(g))$$3 gives $$o_t^{align} = \mathrm{CrossAttn}(\mathrm{Enc}(o_t),\mathrm{Enc}(g),\mathrm{Enc}(g))$$4, $$o_t^{align} = \mathrm{CrossAttn}(\mathrm{Enc}(o_t),\mathrm{Enc}(g),\mathrm{Enc}(g))$$5 gives $$o_t^{align} = \mathrm{CrossAttn}(\mathrm{Enc}(o_t),\mathrm{Enc}(g),\mathrm{Enc}(g))$$6, $$o_t^{align} = \mathrm{CrossAttn}(\mathrm{Enc}(o_t),\mathrm{Enc}(g),\mathrm{Enc}(g))$$7 gives $$o_t^{align} = \mathrm{CrossAttn}(\mathrm{Enc}(o_t),\mathrm{Enc}(g),\mathrm{Enc}(g))$$8, and $$o_t^{align} = \mathrm{CrossAttn}(\mathrm{Enc}(o_t),\mathrm{Enc}(g),\mathrm{Enc}(g))$$9 drops to $$h_t^{enc} = \mathrm{LSTM}(\mathrm{Embed}(a_{t-L:t-1}),\mathrm{Embed}(o_{t-L:t-1}))$$0 because of data fragmentation. Token-level MoE yields $$h_t^{enc} = \mathrm{LSTM}(\mathrm{Embed}(a_{t-L:t-1}),\mathrm{Embed}(o_{t-L:t-1}))$$1 step-level switches per episode and $$h_t^{enc} = \mathrm{LSTM}(\mathrm{Embed}(a_{t-L:t-1}),\mathrm{Embed}(o_{t-L:t-1}))$$2 success, trajectory-level routing yields only $$h_t^{enc} = \mathrm{LSTM}(\mathrm{Embed}(a_{t-L:t-1}),\mathrm{Embed}(o_{t-L:t-1}))$$3 switches but $$h_t^{enc} = \mathrm{LSTM}(\mathrm{Embed}(a_{t-L:t-1}),\mathrm{Embed}(o_{t-L:t-1}))$$4 success, and phase-level routing yields approximately $$h_t^{enc} = \mathrm{LSTM}(\mathrm{Embed}(a_{t-L:t-1}),\mathrm{Embed}(o_{t-L:t-1}))$$5 switches and $$h_t^{enc} = \mathrm{LSTM}(\mathrm{Embed}(a_{t-L:t-1}),\mathrm{Embed}(o_{t-L:t-1}))$$6 success. Post-hoc labeling aligns experts with exploration, manipulation, navigation, and recovery, with entropies of approximately $$h_t^{enc} = \mathrm{LSTM}(\mathrm{Embed}(a_{t-L:t-1}),\mathrm{Embed}(o_{t-L:t-1}))$$7, $$h_t^{enc} = \mathrm{LSTM}(\mathrm{Embed}(a_{t-L:t-1}),\mathrm{Embed}(o_{t-L:t-1}))$$8, $$h_t^{enc} = \mathrm{LSTM}(\mathrm{Embed}(a_{t-L:t-1}),\mathrm{Embed}(o_{t-L:t-1}))$$9, and $\pi_j(X,t)$00 bits, and router assignments overlap by approximately $\pi_j(X,t)$01 step-wise with human-annotated phase boundaries. The paper also reports that single-policy gradients from different phases form $\pi_j(X,t)$02 conflicts and that a single policy converges to average entropy $\pi_j(X,t)$03 bits, whereas PA-MoE experts match phase-optimal entropy (Yang et al., 19 Feb 2026).

IGND-MoE is evaluated on tau-PET data from ADNI with $\pi_j(X,t)$04 and $\pi_j(X,t)$05 scans. It achieves test SSE $\pi_j(X,t)$06 and mean Pearson $\pi_j(X,t)$07, whereas the pure mechanistic model has SSE $\pi_j(X,t)$08 and $\pi_j(X,t)$09. The learned gate weights show a clear stage dependence: at early stages $\pi_j(X,t)$10, graph diffusion $\pi_j(X,t)$11 dominates and the mechanistic component is moderate; in middle stages, both graph and mechanistic components decline; at late stages $\pi_j(X,t)$12, the neural reaction weight $\pi_j(X,t)$13 rises to $\pi_j(X,t)$14. Error maps over time show correction of systematic under-prediction in medial temporal and late-stage neocortical regions (He et al., 9 Aug 2025).

VCMoE reports simulation studies under two Gaussian experts, two binomial experts, and three Gaussian experts, with small RASE, bootstrap simultaneous confidence bands achieving nominal coverage even for moderate $\pi_j(X,t)$15, and GLRT statistics exhibiting the predicted Wilks phenomenon. In the mouse embryonic snRNA-seq application, the deep-layer neuron expert shows the Satb2 coefficient shifting from weakly positive at early stages to negative later, with GLRT $\pi_j(X,t)$16 under the null of constancy. The Ywhaz coefficient remains essentially zero in both experts with GLRT $\pi_j(X,t)$17, and the fitted gating probabilities achieve AUC $\pi_j(X,t)$18 in separating the two cell types (Zhao et al., 5 Jan 2026).

EvoMoE evaluates RoBERTa-24L for masked language modeling, GPT-24L for causal language modeling, and Transformer-12L for machine translation. At matched inference FLOPs of approximately $\pi_j(X,t)$19B, GLUE average score is $\pi_j(X,t)$20 for a standard Transformer, $\pi_j(X,t)$21 for Switch MoE, $\pi_j(X,t)$22 for BASE Layer, and $\pi_j(X,t)$23 for EvoMoE. Language modeling perplexity is $\pi_j(X,t)$24 for a standard Transformer, $\pi_j(X,t)$25 for Switch MoE, $\pi_j(X,t)$26 for BASE Layer, and $\pi_j(X,t)$27 for EvoMoE. For En$\pi_j(X,t)$28De translation, Transformer-Base obtains $\pi_j(X,t)$29 BLEU, Switch $\pi_j(X,t)$30, and EvoMoE $\pi_j(X,t)$31. Convergence curves show GPT-DTS reaches a given validation PPL in approximately half the iterations of Switch, corresponding to a $\pi_j(X,t)$32 speed-up and $\pi_j(X,t)$33 FLOPs-efficiency. Removing stage 1 costs approximately $\pi_j(X,t)$34 more FLOPs with little quality loss, while removing stage 2 degrades quality by approximately $\pi_j(X,t)$35 BLEU (Nie et al., 2021).

6. Relation to standard MoE, interpretability, and open issues

A recurring theme is that standard MoE assumptions can be mismatched to regime-structured problems. Token-level routing fragments contiguous behavioral structure in agentic RL; PA-MoE explicitly argues that phase-consistent patterns become scattered expert assignments under traditional MoE and that this undermines expert specialization (Yang et al., 19 Feb 2026). Fixed-mechanism disease models obscure changing pathological contributions across progression; IGND-MoE instead exposes stage-wise weights that align with literature, with graph-related processes more influential early and other unknown physical processes dominant later (He et al., 9 Aug 2025). Constant-coefficient MoE can be statistically inadequate in settings where covariate influences and latent subpopulation structure evolve along time, space, or another index variable; VCMoE addresses this by making both gating and expert coefficients varying functions and by proving identifiability and consistency (Zhao et al., 5 Jan 2026). Sparse gating from scratch can produce unstable optimization because experts and routing are simultaneously immature; EvoMoE addresses this with dense-to-sparse gate evolution and expert warm-up (Nie et al., 2021).

Interpretability appears in several distinct forms. PA-MoE produces emergent expert specialization that can be post-hoc labeled as exploration, manipulation, navigation, and recovery (Yang et al., 19 Feb 2026). IGND-MoE yields stage-wise trajectories of $\pi_j(X,t)$36 that are interpretable as changing contributions of mechanistic, graph, and local processes (He et al., 9 Aug 2025). VCMoE supports simultaneous confidence bands and generalized likelihood ratio testing for whether specific coefficients genuinely vary with stage (Zhao et al., 5 Jan 2026). EvoMoE is less directly interpretable at the level of domain mechanisms, but it makes the training dynamics of expert formation and sparse routing more explicit (Nie et al., 2021).

Several open issues are visible from these works. One is the distinction between known versus latent stage variables: VCMoE assumes a known index $\pi_j(X,t)$37, IGND-MoE estimates temporal alignment through subject-specific shifts, and PA-MoE discovers latent phase boundaries directly from reward signals. Another is the trade-off between specialization and fragmentation: PA-MoE shows that too fine a granularity harms performance, while too coarse a granularity also underperforms (Yang et al., 19 Feb 2026). A further issue is label consistency and identifiability, which VCMoE treats formally but which are less explicitly resolved in neural MoE settings (Zhao et al., 5 Jan 2026). Taken together, these results suggest that stage-aware MoE is most appropriate when the latent heterogeneity is structured along contiguous temporal, developmental, pathological, or optimization regimes rather than being exchangeable across isolated tokens or samples.