Part IV · Findings and OutlookSection 21 of 24

Section 21: Merging Routed and Shared Expert Activations

  • 6 min read
A broad teal dry-brush stroke glides down from the upper left and a narrow pink stroke rises from below; they join into one stroke running right, the pink becoming a vein inside it that runs dry first.

Selected router probabilities have mean \(1/K\). Multiplying them by \(K\) makes the mean routed coefficient one, matching the coefficient of each block in an additively partitioned dense MLP. The shared branch then enters the same sum with unit gain.

  1. Speed-Read

    the rule that merges routed and shared outputs changes the initialization scale (Section 21).

Section 20 established that several ungated shared-expert partitions collapse into one wider shared MLP. After choosing that shared branch, the model recipe must specify how to weight it against the token-dependent routed branches.

An MoE layer may contain both routed experts, selected separately for each token, and a shared expert, evaluated for every token. The routed and shared branches are projected back to the same model width before they are added, but that does not determine how strongly each branch should enter the sum. A normalization that appears innocuous can shrink the MoE output, change its Jacobian, and rescale the gradients of routed and shared parameters by different amounts.

For token \(x\), let \(\mathcal{S}(x)\) contain the \(K\) selected routed experts, let \(f_i(x)\in\mathbb{R}^{d_{\mathrm{model}}}\) be routed expert \(i\)’s output, and let \(g(x)\in\mathbb{R}^{d_{\mathrm{model}}}\) be the shared-expert output. We write the merge as \[ y_{\mathrm{MoE}}(x) = \sum_{i\in\mathcal{S}(x)} r_i(x) f_i(x) + s(x)g(x), \qquad |\mathcal{S}(x)|=K. \tag{80} \] The coefficients \(r_i(x)\) determine the routed contribution, while \(s(x)\) sets the gain of the shared branch (Figure 69). As Appendix E shows, public MoE implementations make substantially different choices for both quantities (DeepSeek-AI et al. 2024; Muennighoff et al. 2024; Yan et al. 2026). We therefore treat the merge rule as part of the model and initialization recipe.

Figure 69Interactive · conceptual coefficients, no measured quantities

Reweight one token’s routed experts: under the Olmo rule the gains still average one

Merge rule
Router weights
012gainwidth×1.12Expert 02560p₀ = 0.28×0.88Expert 12560p₁ = 0.22×1.20Expert 22560p₂ = 0.30×0.80Expert 32560p₃ = 0.20×1.00Shared1280āwRouter · selected probabilities pi, Σpi = 1One dense SwiGLU MLP · dact = 4 × 2560 + 1280 = 11,520token xyMoE(x)

This token’s merge · Equation 80

yMoE = 1.12f0 + 0.88f1 + 1.20f2 + 0.80f3 + 1.00g

Routed gains
Σri = 4.00
mean 1.00 per selected expert
Shared gain
s = 1.00
every token, no router
Width-weighted mean
āw = 1.000
equals a dense block’s 1 · Equation 84
Variance ratio
Rvar ≈ 1.024
Σpi2 = 0.2568 · Equation 85

Drag a routed gain bar, or focus it and use the arrow keys (or the − / + buttons); the other routed weights rescale so that Σpi stays 1.

  • Routed expert gain ri
  • Shared expert gain s
  • 1 = a dense block’s coefficient
  • Width-weighted mean āw

Box widths are proportional to intermediate width (dr = 2560, ds = 1280, the widths of the report’s study); every branch returns a model-width tensor before the addition. Bar heights are merge coefficients, so a bar’s area (width × coefficient) is its term in the width-weighted mean āw (Equation 84); the areas fill the boxes’ total width only when āw = 1. Gains are the figure’s example or weights you set; the variance ratio is Equation 85 evaluated for them, not a measurement.

Figure 69. Merging four routed-expert outputs with one shared-expert output. In the Olmo rule, selected routed coefficients vary by token but sum to \(K=4\), so their mean is one; the shared-expert coefficient is also one. Routed and shared experts may have different intermediate widths because every branch returns a model-width tensor before the addition. Conceptual illustration with no measured quantities.

21.1 A Dense MLP Provides the Scale Reference

A dense SwiGLU MLP provides a useful reference because its intermediate dimension can be partitioned exactly. Split its up, gate, and down-projection weights into \(K\) blocks of width \(d_r\), the routed-expert intermediate width, and one block of width \(d_s\), the total shared width (\(m_s\) in Equation 79). Matrix multiplication distributes across this partition, giving \[ f_{\mathrm{dense}}(x) = \sum_{i=1}^{K} f_i(x) + g(x), \qquad d_{\mathrm{act}}=K d_r+d_s. \tag{81} \] Unit coefficients exactly recover a dense MLP formed by concatenating the active intermediate blocks. For the initialization-scale experiment below, \(K=4\), \(d_r=2560\), and \(d_s=1280\), so the matched dense width is \[ d_{\mathrm{act}}=4\times2560+1280=11{,}520. \tag{82} \]

Routing replaces the four unit routed coefficients with data-dependent route weights. Let \(p_i(x)\) denote the selected router probabilities after L1 normalization, so \(\sum_i p_i(x)=1\). Olmo’s production recipes configure the Olmo-core router to use \[ r_i(x)=Kp_i(x), \qquad s(x)=1. \tag{83} \] The combination weight \(p_{t,e}\) of Equation 1 corresponds to \(r_i\) here, not to \(p_i\). If the selected probabilities are uniform, \(p_i=1/K\) and every active branch again has coefficient one. The routed mixture therefore approaches the exact dense decomposition continuously as the selected probabilities approach a uniform distribution.

This choice also accounts for different expert widths without an additional width multiplier. Define the width-weighted mean coefficient \[ \bar a_w = \frac{d_r\sum_i r_i+d_s s}{K d_r+d_s}. \tag{84} \] Equation 83 gives \(\bar a_w=1\) for any \(d_r\) and \(d_s\). Since the width-weighted mean is already one, a further common factor would rescale the whole MoE output rather than compensate for unequal widths. The narrower shared expert already contributes fewer independently initialized intermediate units, so multiplying its output by a second width fraction would count the same width difference twice under this initialization.

Matching the Coefficients of a Partitioned Dense MLP

Selected router probabilities have mean \(1/K\). Multiplying them by \(K\) makes the mean routed coefficient one, matching the coefficient of each block in an additively partitioned dense MLP. The shared branch then enters the same sum with unit gain.

21.2 Predicting Variance at Initialization

The dense decomposition also yields a simple scale prediction. Assume the randomly initialized branch outputs are centered and approximately independent, and that output variance is proportional to intermediate width under the common initialization. Conditional on the router coefficients, the MoE-to-dense output-variance ratio is \[ R_{\mathrm{var}} \approx \frac{d_r\sum_i r_i^2+d_s s^2}{K d_r+d_s}. \tag{85} \] For the Olmo rule, \[ R_{\mathrm{var}}^{\mathrm{Olmo}} \approx \frac{d_r K^2\sum_i p_i^2+d_s}{K d_r+d_s}. \tag{86} \] Uniform routing gives \(\sum_i p_i^2=1/K\) and hence \(R_{\mathrm{var}}^{\mathrm{Olmo}}=1\). Router concentration increases the ratio according to the concentration term \(\sum_i p_i^2\).

Table 37 contrasts this rule with the three alternatives used in our controlled experiment. All four define valid model parameterizations, but they do not begin at the same forward or backward scale.

Convention Coefficients Immediate effect
Olmo \(r_i=Kp_i,\quad s=1\) Unit-average routed coefficients and a unit-gain shared branch.
Routed unit sum \(r_i=p_i,\quad s=1\) The routed mixture receives one unit in total, while the shared branch independently receives unit gain.
Equal branch mass \(r_i=\tfrac{1}{2}p_i,\quad s=\tfrac{1}{2}\) Routed and shared contributions split one unit of coefficient mass.
Width-proportional unit sum \(r_i=\dfrac{K d_r}{d_{\mathrm{act}}}p_i,\quad s=\dfrac{d_s}{d_{\mathrm{act}}}\) A single unit of coefficient mass is divided in proportion to active intermediate width.
Table 37. Merge conventions evaluated in the initialization-scale study. The selected probabilities satisfy \(\sum_i p_i=1\), and \(d_{\mathrm{act}}=K d_r+d_s\).

21.3 The Measured Forward and Backward Scales Agree

We test the prediction in a controlled random-initialization experiment summarized in Figure 70. The forward study uses the layer shape of an earlier Olmo development recipe, not the production recipe: \(d_{\mathrm{model}}=2560\), \(d_r=2560\), \(d_s=1280\), and \(K=4\) experts selected from \(N=48\) router candidates. It instantiates only the four selected experts and evaluates 512 random tokens for each of eight seeds. The dense reference concatenates the same selected and shared weights into an exact width-11,520 SwiGLU MLP, so differences arise from merge coefficients rather than different sampled parameters.

The backward study instantiates all 48 routed experts and gathers each token’s four. To make repeated gradient measurements practical, it scales every dimension down by a factor of ten and raises the initialization standard deviation by \(\sqrt{10}\), which preserves the leading-order projection variances. It uses 1024 tokens for each of twelve seeds and applies the same random upstream gradient to every normalization. We report input gradients with router coefficients both fixed and differentiable, separating the expert-path Jacobian from the additional gradient through the router.

Figure 70Interactive · measured, at random initialization · forward full size, backward scaled surrogate

Which merge rule starts at the dense scale, forward and backward?

Convention
Panel (a)

(a) Forward output RMS relative to matched dense

0.00.20.40.60.81.01.2RMS ratio1.016Olmori = Kpi, s = 10.411Routedunit sum0.205Equalbranch mass0.216Width-proportionalunit summatched dense = 1

(b) Input-gradient RMS relative to matched dense

  • router coefficients fixed
  • including router gradient
0.00.20.40.60.81.01.2RMS ratio1.0161.054Olmori = Kpi, s = 10.4100.417Routedunit sum0.2050.208Equalbranch mass0.2160.225Width-proportionalunit summatched dense = 1

(c) Weight-gradient RMS relative to the Olmo rule

  • routed down weight
  • shared down weight
0.00.20.40.60.81.01.2Gradient RMS ratio1.0001.000Olmori = Kpi, s = 10.2501.000Routedunit sum0.1250.500Equalbranch mass0.2220.111Width-proportionalunit sumOlmo rule = 1

(d) Cosine similarity to matched dense output

0.00.20.40.60.81.0Cosine similarity0.984Olmori = Kpi, s = 10.812Routedunit sum0.812Equalbranch mass0.970Width-proportionalunit sum

Olmori = Kpi, s = 1

(a) Output RMS
1.016× dense
per-seed 1.014–1.017variance 1.032×; Equation 85 predicts 1.032×
(b) Input-gradient RMS
1.016× · 1.054×
per-seed 1.013–1.018 · 1.052–1.057router fixed · including router gradient
(c) Down-weight gradient RMS
1.000× · 1.000×
identical in all 12 seedsrouted · shared, relative to Olmo
(d) Cosine to dense
0.984
per-seed 0.983–0.986direction of the merged output

Ranges are the per-seed minimum and maximum: forward panels (a) and (d) over 8 seeds × 512 tokens, backward panels (b) and (c) over 12 seeds × 1,024 tokens.

Plotted from the study’s raw CSV summaries, the values of the original figure. Forward: the full-size initialization configuration (dmodel = 2560, dr = 2560, ds = 1280, N = 48, K = 4), 512 tokens × 8 seeds. Backward: the scaled surrogate (dmodel = 256, dr = 256, ds = 128, all 48 routed experts), 1,024 tokens × 12 seeds. Pick a convention, or hover, tap or focus a bar, to read its values and per-seed range. Random initialization only: no trained-model quality is measured.

Figure 70. Olmo’s merge rule preserves dense-like scale at random initialization. The four conventions are those of Table 37. (a) The Olmo output RMS is \(1.016\times\) the matched dense reference, corresponding to a \(1.032\times\) variance ratio. (b) Its fixed-router input-gradient RMS is also \(1.016\times\) dense; differentiating through the router raises the ratio to \(1.054\times\). (c) Alternatives that constrain the routed or total coefficient mass to one reduce the routed down-projection gradient to a quarter or less of the Olmo value while the shared gradient stays at one or falls with it. (d) Olmo’s output also stays closest in direction to the dense output. Forward measurements use the full-size initialization configuration; backward measurements use the dimensionally scaled surrogate described in the text.

The selected probabilities have \(\mathbb{E}[\sum_i p_i^2]=0.2589\), only slightly above the uniform top-4 value of 0.25. Substituting this measurement into Equation 86 predicts a variance ratio of 1.032, matching the measured value to within 0.03 percentage points. The shared-to-routed branch RMS ratio is \(0.7082\pm0.0028\), close to \(\sqrt{1280/2560}=0.7071\). This confirms that intermediate width already controls branch scale under the tested initialization.

The backward results expose a second advantage of the Olmo parameterization. With router coefficients fixed, its input-Jacobian scale remains close to the dense reference. Its routed and shared down-projection gradients also retain the unit-coefficient scale: the dense reference’s routed and shared down-projection gradient RMS values are \(0.98\times\) and \(1.00\times\) the Olmo values (Figure 70(c) is normalized to Olmo, so it does not plot these ratios). By contrast, routed-unit-sum normalization reduces the routed down-projection gradient RMS to \(0.25\times\) the Olmo value for \(K=4\), while leaving the shared down-projection gradient unchanged. Changing the merge rule is therefore also a parameter-family-specific learning-rate change unless the initialization and optimizer are adjusted with it.

21.4 Scope of the Initialization Result

Olmo’s \(r_i=Kp_i,\ s=1\) rule preserves dense-like forward and expert-path backward scale under the stated architecture, initialization, and approximately uniform selected routing. It has a direct additive interpretation, handles unequal expert widths without a second width penalty, and makes its forward departure from dense scale predictable through \(\sum_i p_i^2\).

The result does not establish a universal MoE normalization. Public models use retained routed mass below one, unit routed mass, learned shared gates, and model-specific routed scales; Appendix E provides a source-pinned comparison. These recipes also differ in initialization, upcycling, router form, and expert width, so their coefficients cannot be ranked in isolation.

Fixed coefficient changes can often be absorbed into expert down projections, but that reparameterization does not make the training recipes identical. It changes initial parameter scale, optimizer states, weight decay, gradient clipping, and effective learning rates. The measurements here characterize random initialization and local backward scale; they do not measure trained language-model quality.

Search the full report

Search across sections, appendices, and references.

    Enter a term to find passages. Esc closes search.