---
title: Stage-Aware Mixture of Experts Overview
url: https://www.emergentmind.com/topics/stage-aware-mixture-of-experts-moe
type: topic
---

# Stage-Aware Mixture of Experts Overview

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" [2602.17038], the time-dependent gating of mechanistic, graph-diffusion, and neural-reaction experts in neurodegenerative progression modelling [2508.07032], the varying-coefficient statistical MoE of Zhao et al. for longitudinal and indexed data [2601.01699], and the two-stage dense-to-sparse training framework of EvoMoE [2112.14397].

## 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 [2602.17038]. In disease modelling, the motivating problem is that traditional models assume fixed mechanisms throughout disease progression despite stage-dependent pathological dynamics [2508.07032]. 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 [2601.01699]. In large-scale Transformer MoE training, EvoMoE identifies immature experts and unstable sparse gate learning when both experts and gate are random at initialization [2112.14397].

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 [2602.17038]. IGND-MoE uses time-dependent weights $\beta_j(t)$ over three continuous-time expert vector fields [2508.07032]. VCMoE lets both $\pi_j(X,t)$ and expert parameters vary smoothly in the indexing variable $t$ [2601.01699]. 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 [2112.14397].

| 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 $s_t \approx (o_t,\mathrm{history},h_t,\mathrm{goal}\ g)$, uses a short LSTM over the past $L=5$ $(o,a)$ pairs, and selects $k_t = \arg\max p_t$ deterministically with a straight-through estimator. The policy at step $t$ is then $\pi(a_t \mid s_t;\theta_{base}, B_{k^\*}, A_{k^\*})$, where the backbone $\theta_{base}$ remains frozen and each expert is a LoRA adapter on the query/value projections in each Transformer layer [2602.17038].

IGND-MoE uses a different mechanism: the gate depends only on time. The three expert vector fields
$$
f_M(c(t),t;\theta_M), \quad f_S(c(t),A,t;\theta_S), \quad f_L(c(t),t;\theta_L)
$$
are combined by stage-dependent weights $\beta_j(t)$ satisfying $\beta_j(t)\ge 0$ and $\sum_{j=1}^3 \beta_j(t)=1$. The weights are produced by a small temporal-attention network $g_\phi:\mathbb R \to \mathbb R^3$ with
$$
\beta_j(t)=\frac{\exp(\lambda_j(t))}{\sum_{k=1}^3 \exp(\lambda_k(t))}, \qquad \lambda(t)=g_\phi(t).
$$
In practice $g_\phi$ is a small MLP in the single scalar input $t$, yielding smooth stage-wise expert contributions [2508.07032].

VCMoE generalizes this principle statistically. For latent experts $j=1,\dots,J$, the gating probabilities are
$$
\pi_j(X,t)=\frac{\exp\{X^\top \beta_j(t)\}}{\sum_{k=1}^J \exp\{X^\top \beta_k(t)\}},
$$
while each expert density $f_j(y\mid Z,t;\alpha_j(\cdot),\delta_j(\cdot))$ also depends on coefficient curves that vary with $t$. The model therefore makes both mixture weights and expert responses stage-dependent, rather than only modulating a fixed expert bank [2601.01699].

EvoMoE uses Gumbel-Softmax routing scores during gate-sparsify:
$$
g'(x_s)=\mathrm{softmax}\!\Bigl((x_sW_g+\zeta)/\tau\Bigr).
$$
A content-based threshold $c$ prunes low-score experts,
$$
g(x_s)_i=
\begin{cases}
g'(x_s)_i, & g'(x_s)_i>c,\\
0, & \text{otherwise,}
\end{cases}
$$
and after a designated number of iterations $T_D$, 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 [2112.14397].

## 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
$$
z_t = \arg\max_k g_\phi(s_t)_k
$$
and adds a switching penalty
$$
\mathcal L_{switch} = (\lambda_s/(T-1))\sum_{t=1}^{T-1} \mathbf 1[z_t \neq z_{t+1}],
$$
with the backward pass using the soft surrogate $1-\sum_{k=1}^K p_t^k p_{t+1}^k$. The full loss is
$$
\mathcal L = \mathcal L_{RL} + \alpha \mathcal L_{div} + \beta \mathcal L_{bal} + \gamma \mathcal L_{switch},
$$
where $\mathcal L_{div} = \sum_{i\neq j} \max(0,\tau_{div} - \mathrm{KL}(\pi_{exp}^i\|\pi_{exp}^j))$ and $\mathcal L_{bal}=\sum_{k=1}^K (f_k-1/K)^2$. No explicit phase labels or semi-MDP definitions are given; the router and experts are trained jointly under the single RL objective, and $\tau$ is annealed from $2.0 \to 0.5$ [2602.17038].

IGND-MoE is trained by dual iterative optimization. Because observations are irregular and infrequent, each subject $i$ is assigned a time shift $t_0^i$. Temporal alignment solves
$$
t_0^i \leftarrow \arg\min_{\tau \in [0,T_{\max}]} \sum_{s=1}^{S_i} \left\|c_{\rm obs}^i(s)-c(\tau+t_s^i)\right\|^2,
$$
and, with shifts fixed, trajectory construction minimizes
$$
\mathcal L_{\rm traj}(\Theta)=\sum_{i=1}^N \sum_{s=1}^{S_i} \left\|c(t_0^i+t_s^i;\Theta)-c_{\rm obs}^i(s)\right\|^2.
$$
The total loss is
$$
\mathcal L_{\rm total}=\mathcal L_{\rm traj}+\lambda_1 \mathcal L_{\rm norm}+\lambda_2 \mathcal L_{\rm ortho},
$$
where $\mathcal L_{\rm norm}$ constrains learned experts relative to the mechanistic model and $\mathcal L_{\rm ortho}$ encourages diversity between experts. All ODE parameters are updated by gradient-based ODE backprop, while $\{t_0^i\}$ are updated in a separate one-dimensional minimization [2508.07032].

VCMoE is estimated by a label-consistent EM algorithm built on local likelihood. At a target stage $t$, each coefficient curve is approximated by a local linear expansion, and the local weighted log-likelihood is
$$
\ell_n(t)=\frac1n \sum_{i=1}^n \log\Bigl\{\sum_{j=1}^J \pi_j(X_i,t)\,f_j(Y_i\mid Z_i,t)\Bigr\}K\!\Bigl(\frac{t_i-t}{h}\Bigr).
$$
The E-step computes responsibilities
$$
\gamma_{ij}^{(m)}=\frac{\pi_j(X_i,t_i)\,f_j(Y_i\mid Z_i,t_i)}{\sum_{k=1}^J \pi_k(X_i,t_i)\,f_k(Y_i\mid Z_i,t_i)},
$$
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 [2601.01699].

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_S$ and $T_D$, with temperature annealed from $\tau_{max}$ to $\tau_{min}$ and an optional balance loss for expert loads [2112.14397].

## 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 $r$ 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 [2602.17038].

IGND-MoE is more heterogeneous. The overall continuous dynamics satisfy
$$
\frac{d\,c(t)}{dt}=\sum_{j=1}^3 \beta_j(t;\phi)\, f_j(c(t),t;\theta_j),
$$
with $f_1=f_M$, $f_2=f_S$, and $f_3=f_L$. The inhomogeneous graph neural diffusion expert computes a low-rank latent representation
$$
Z(t)=\mathrm{Enc}_{\theta_S}(c(t),A)\in \mathbb R^{n\times r},
$$
builds a refined adjacency
$$
\hat A(t)=\sigma(Z(t)Z(t)^\top),
$$
and outputs
$$
f_S(c(t),A,t;\theta_S)= (\hat A(t)-I)c(t).
$$
The localized neural reaction expert is
$$
f_L(c(t),t;\theta_L)=\mathrm{MLP}_{\theta_L}(c(t)) \in \mathbb R^n.
$$
This composition preserves a mechanistic component while allowing stage-dependent graph-refined diffusion and residual local dynamics [2508.07032].

VCMoE formalizes stage dependence at the level of statistical function classes. In the Gaussian expert example,
$$
f_j(y\mid Z,t)=\phi\bigl(y;\mu_j(Z,t)=Z^\top \alpha_j(t),\sigma_j^2(t)=\delta_j(t)\bigr),
$$
so both regression effects and dispersion vary smoothly with $t$. 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 [2601.01699].

EvoMoE retains the standard expert aggregation rule
$$
y_s=\sum_{i=1}^N g(x_s)_i\, e_i(x_s),
$$
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 [2112.14397].

## 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 $50\%$ of total. Under a single GiGPO policy, simple pick-and-place dominates $75\%$ of updates, while complex manipulation tasks Heat/Cool/Clean account for only $5\%$. This is associated with high success on Pick ($\sim98\%$) but poorer results on Heat/Cool ($\sim80\%/60\%$). With $K=4$ experts, each expert sees $\sim30\%$ of batches, overall success on ALFWorld improves from $86.1\%$ to $93.8\%$ $(+7.7$ pp$)$, and WebShop improves from $67.4\%$ to $82.3\%$ $(+14.9$ pp$)$. Ablations show $K=0$ gives $88.3\%$, $K=2$ gives $91.4\%$, $K=4$ gives $93.8\%$, and $K=6$ drops to $85.9\%$ because of data fragmentation. Token-level MoE yields $45$ step-level switches per episode and $85.7\%$ success, trajectory-level routing yields only $3$ switches but $88.5\%$ success, and phase-level routing yields approximately $8$ switches and $93.8\%$ success. Post-hoc labeling aligns experts with exploration, manipulation, navigation, and recovery, with entropies of approximately $3.5$, $0.5$, $1.5$, and $2.0$ bits, and router assignments overlap by approximately $87\%$ step-wise with human-annotated phase boundaries. The paper also reports that single-policy gradients from different phases form $>90^\circ$ conflicts and that a single policy converges to average entropy $\sim2.3$ bits, whereas PA-MoE experts match phase-optimal entropy [2602.17038].

IGND-MoE is evaluated on tau-PET data from ADNI with $N=216$ and $378$ scans. It achieves test SSE $= 13.97 \pm 0.30$ and mean Pearson $R = 0.717 \pm 0.014$, whereas the pure mechanistic model has SSE $\approx 17.8$ and $R \approx 0.66$. The learned gate weights show a clear stage dependence: at early stages $(t \approx 0)$, graph diffusion $\beta_2(t)$ dominates and the mechanistic component is moderate; in middle stages, both graph and mechanistic components decline; at late stages $(t \gg 0)$, the neural reaction weight $\beta_3(t)$ rises to $>50\%$. Error maps over time show correction of systematic under-prediction in medial temporal and late-stage neocortical regions [2508.07032].

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 $n$, 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 $p \approx 0.03$ under the null of constancy. The *Ywhaz* coefficient remains essentially zero in both experts with GLRT $p \approx 0.73$, and the fitted gating probabilities achieve AUC $\approx 0.885$ in separating the two cell types [2601.01699].

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 $220$B, GLUE average score is $89.75$ for a standard Transformer, $90.31$ for Switch MoE, $90.65$ for BASE Layer, and $90.88$ for EvoMoE. Language modeling perplexity is $15.14$ for a standard Transformer, $13.12$ for Switch MoE, $12.45$ for BASE Layer, and $12.24$ for EvoMoE. For En$\to$De translation, Transformer-Base obtains $28.1$ BLEU, Switch $28.4$, and EvoMoE $29.6$. Convergence curves show GPT-DTS reaches a given validation PPL in approximately half the iterations of Switch, corresponding to a $2\times$ speed-up and $1.42\times$ FLOPs-efficiency. Removing stage 1 costs approximately $5\%$ more FLOPs with little quality loss, while removing stage 2 degrades quality by approximately $1$ BLEU [2112.14397].

## 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 [2602.17038]. 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 [2508.07032]. 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 [2601.01699]. 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 [2112.14397].

Interpretability appears in several distinct forms. PA-MoE produces emergent expert specialization that can be post-hoc labeled as exploration, manipulation, navigation, and recovery [2602.17038]. IGND-MoE yields stage-wise trajectories of $\beta_j(t)$ that are interpretable as changing contributions of mechanistic, graph, and local processes [2508.07032]. VCMoE supports simultaneous confidence bands and generalized likelihood ratio testing for whether specific coefficients genuinely vary with stage [2601.01699]. 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 [2112.14397].

Several open issues are visible from these works. One is the distinction between **known** versus **latent** stage variables: VCMoE assumes a known index $t$, 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 [2602.17038]. A further issue is label consistency and identifiability, which VCMoE treats formally but which are less explicitly resolved in neural MoE settings [2601.01699]. 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.

Source: https://www.emergentmind.com/topics/stage-aware-mixture-of-experts-moe