1 Introduction↩︎

SwiGLU MLPs [1], [2] have become a standard component of modern Transformer architectures, including dense language models [3], [4] and Mixture-of-Experts (MoE) models [4]–[10]. In a SwiGLU MLP, a SiLU gate modulates intermediate activations conditioned on the input, selectively suppressing or amplifying features and improving expressivity at low computational cost.

The SwiGLU layer is commonly defined as: \[\begin{align} \mathrm{SwiGLU}(x) &= \mathrm{SiLU}(W_g x) \odot (W_u x), \\ \mathrm{SiLU}(z) &= z \cdot \sigma(z) = \frac{z}{1 + e^{-z}} . \end{align}\] SiLU can be viewed as a fixed-sharpness instance of Swish. We denote the corresponding sharpness-adjusted SiLU gate as \[\begin{align} \mathrm{SiLU}_{\kappa}(z) &= z \cdot \sigma(\kappa z), \end{align}\] where \(\kappa\) is a sharpness coefficient controlling the transition between inactive and active gate states. The standard SiLU gate corresponds to \(\kappa=1\)¹.

Larger \(\kappa\) yields a sharper, more selective gate, while smaller \(\kappa\) produces smoother, more broadly active gating. Allowing this sharpness to adapt could provide a more expressive mechanism for regulating feature activation. This is especially relevant in MoE models, where the router dynamically assigns each token to a small subset of experts based on router logits. These scores provide a natural signal of routing certainty: when a token receives a high score for an expert, the router is more confident in that token–expert assignment; when the score is lower, the assignment is more uncertain. Beyond expert selection, such scores could therefore serve as a control signal for modulating gate activations inside the selected experts.

We identify a previously overlooked factor in training MoEs with SwiGLU experts: a hidden co-evolution between the router and the expert gate. Specifically, we observe that gate projection directions within an expert rapidly become aligned or anti-aligned with the corresponding router weight vector during training. This alignment makes the expert gate implicitly sensitive to routing certainty: tokens with different router affinities are shifted to different regions of the SiLU gate’s transition curve, causing subsets of expert activations to be systematically amplified or suppressed, as illustrated in the left panel of Figure 2. This emergent coupling further motivates using routing confidence to regulate how each selected expert processes a token.

Motivated by this implicit router–gate coupling, we propose \(\kappa\)-SwiGLU, which explicitly uses routing confidence to modulate the sharpness of each expert’s SiLU gate. Unlike the emergent alignment effect, which influences gate behavior by additively shifting the gate input, \(\kappa\)-SwiGLU directly adjusts the gate’s smoothness and selectivity through a token-dependent multiplicative sharpness coefficient. Each expert learns its own confidence–sharpness mapping, allowing routing confidence to induce either sharper, more selective gates or smoother, more broadly active gates depending on the learned parameters. This mechanism is illustrated in Figure 1 and in the right panel of Figure 2, where it is contrasted with the implicit additive-shift effect shown in the left panel.

We evaluate \(\kappa\)-SwiGLU by training SwiGLU-based MoE models on FineWeb-Edu, scaling from 8 to 28 layers. Across model settings, \(\kappa\)-SwiGLU shows a consistent positive trend in model quality, as reflected in stronger pretraining benchmark performance. Our analysis suggests that confidence-aware gate sharpness enables more expressive gating patterns and promotes synergistic interactions between the router and expert gates.

2 Related Work↩︎

2.0.0.1 MoE Load balancing and routing stability.

One of the central challenges in MoE training is preventing routing collapse, where only a small subset of experts receives most tokens. Prior work commonly addresses this issue using auxiliary load-balancing losses that encourage more uniform expert utilization [12]–[16]. Switch Transformer [14] also introduced the router \(z\)-loss to regularize router logits. Subsequent work has studied how load-balancing losses are sensitive to implementation details such as batch statistics [17], proposed auxiliary-loss-free balancing strategies [18], and alternative balancing formulations such as \(\phi\)-balancing [19]. These approaches primarily regulate how tokens are assigned to experts. In contrast, \(\kappa\)-SwiGLU uses the selected expert’s router signal to modulate gate sharpness inside the expert MLP.

2.0.0.2 Geometry of MoE routing.

Recent work has explored geometry-aware routing mechanisms, including Routing Manifold Alignment and kNN-augmented routing, which encourage expert assignment to better reflect token-representation geometry [20], [21]. Other studies address representational collapse in MoE training [22], geometric regularization of expert weights and activations [23], and auxiliary losses that couple routers and experts [24]. Rather than adding geometric constraints, our work uses the router logit directly as a token-level signal for adaptive gate sharpness.

2.0.0.3 Gated activations in Transformer MLPs.

Early Transformer models commonly used ReLU- or GELU-based feed-forward networks [25], [26], while modern LLMs have largely converged toward gated variants such as GLU, GeGLU, and SwiGLU [1], [2]. SwiGLU has become a common choice due to its strong empirical performance and modest computational cost. Recent work has revisited this activation design space, including ReLU\(^2\) [27], expanded gating ranges in xGELU and xSiLU [28], and adaptive mixtures of activation functions [29]. In practice, SwiGLU gates can require stabilization techniques such as clamping to avoid excessively large activations during training [8], [30]; relatedly, Power Linear Unit (PowLU) controls activation magnitude using a bounded rational power function [31]. However, the interaction between expert routing and gated activations remains relatively underexplored. Our work studies this interaction by using MoE routing confidence to adapt the sharpness of the SwiGLU gate.

3 Method↩︎

In MoEs, routed tokens tend to be biased toward the corresponding router direction. Meanwhile, expert gate projection vectors naturally become aligned or anti-aligned with the same router direction during training. These two effects induce an implicit, confidence-induced gate bias in the expert gates: tokens with different router logits are shifted to different regions of the SiLU gate.

To complement this implicit bias mechanism, we propose \(\kappa\)-SwiGLU, which uses router logits to produce a token-dependent sharpness coefficient for each expert gate. Unlike the emergent implicit gate bias which shifts the gate input additively, \(\kappa\)-SwiGLU modulates the gate multiplicatively, allowing each expert to adapt its activation selectivity according to routing confidence.

3.1 Router Logits as Routing Confidence↩︎

Mixture-of-Experts (MoE) models typically route each token to a small subset of experts according to router logits, computed from the inner product between an expert router vector \(r_e\) and the input token representation \(x\). Geometrically, the tokens routed to expert \(e\) occupy a narrow region of representation space biased toward the router direction \(r_e\)², as illustrated in Figure 3. Tokens that are more closely aligned with \(r_e\) receive higher router scores and can therefore be interpreted as higher-confidence assignments to expert \(e\). Based on this observation, we use the router logit as a token-level routing confidence signal³ \[\begin{align} s_e(x) = r_e^\top x \label{router-logits} \end{align}\tag{1}\]

3.2 Confidence-Induced Gate Bias↩︎

We first empirically verify the existence of router–gate alignment by tracking the cosine similarity between each expert’s router vector and its gate projection vectors during training. Figure 4 shows the router–gate alignment dynamics for two representative layers across 7 independently trained 8-layer MoE models. Expert gate projections rapidly become aligned with the corresponding router vector within the first few hundred training steps, reaching peak cosine similarities of \(0.2\)–\(0.4\). Although the coupling strength changes over training and varies across layers, it remains non-negligible throughout. This suggests that router–gate coupling emerges naturally during training while exhibiting diverse, layer-dependent signed dynamics.

To better understand the impact of router–gate alignment, let \(w_{e,j}\) denote the \(j\)-th gate projection vector of expert \(e\), and let \(\hat{r}_e = r_e / \|r_e\|_2\) be the unit-normalized router vector. We decompose \(w_{e,j}\) into components parallel and orthogonal to the router direction: \[w_{e,j} = (w_{e,j}^{\top}\hat{r}_e)\hat{r}_e + w_{e,j}^{\perp}.\] Similarly, we write the input representation as \(x = a\hat{r}_e + x^\perp\), where \(a = \hat{r}_e^\top x\). The gate input can then be written as \[w_{e,j}^{\top}x = (w_{e,j}^{\top}\hat{r}_e)a + (w_{e,j}^{\perp})^\top x^\perp.\] The first term, \((w_{e,j}^{\top}\hat{r}_e)a\), acts as a confidence-induced gate bias: tokens with larger projection onto the router direction induce a systematic shift in the corresponding gate input. Thus, router–gate alignment implicitly changes the bias of expert gates, as illustrated in the left panel of Figure 2.

Figure 5 empirically confirms this effect: the top and bottom \(5\%\) of induced bias values remain substantially positive and negative throughout training, indicating that router–gate alignment produces non-negligible bidirectional shifts in expert gates.

3.3 Confidence-Adaptive SwiGLU↩︎

Motivated by the naturally emerging confidence-induced gate bias, we introduce an explicit confidence-aware sharpness parameterization to make the router–gate coupling more expressive and controllable. While the emergent alignment effect modulates expert gates implicitly through confidence-induced gate bias shifts, \(\kappa\)-SwiGLU directly controls the smoothness and selectivity of the SiLU gate using a token-dependent sharpness coefficient. This provides a more flexible mechanism for routing confidence to influence expert computation, allowing each expert to learn whether higher confidence should induce sharper, more selective gating or smoother, more broadly active gating.

We first define a sharpness-adjusted SiLU gate: \[\begin{align} \mathrm{SiLU}_{\kappa}(z) &= z \cdot \sigma(\kappa z), \end{align}\] where \(\kappa\) is a positive sharpness coefficient.

The standard SiLU gate corresponds to the fixed-sharpness case \(\kappa=1\). Larger values of \(\kappa\) produce sharper and more selective gating, while smaller values produce smoother and more broadly active gating. To intuitively understand the effect of \(\kappa\), we plot \(\mathrm{SiLU}_{\kappa}(z)\) and its gradient under different \(\kappa\) values in the left and right panels of Figure 6, respectively. As \(\kappa\) increases, the transition between inactive and active states becomes sharper, with a more pronounced change around zero. Conversely, as \(\kappa\) decreases, the transition becomes smoother, with a more gradual change across a wider range of inputs. These differences are most pronounced in the transition region around zero, where the gate output changes rapidly with respect to the input, while different \(\kappa\) values produce more similar behavior in the saturated regions far from zero. In the gradient panel of Figure 6, the corresponding gradients show even greater sensitivity to \(\kappa\): larger \(\kappa\) yields steeper gradient changes around the transition region, indicating that the gate becomes more responsive to small variations in near-zero inputs.

We then parameterize the sharpness coefficient \(\kappa\) as a learnable function of the routing confidence: \[\begin{align} \kappa_{e,j}(x) &= \phi\left(\alpha_{e,j} \cdot s_e(x) + b_{e,j}\right), \label{eq:kappa} \end{align}\tag{2}\] where \(s_e(x)\) is the router logit defined in Eq. 1 , \(\alpha_{e,j}\) and \(b_{e,j}\) are learnable parameters, and \(\phi\) is a monotonically increasing function that maps the potentially unbounded confidence-conditioned signal to a positive sharpness coefficient.

To avoid extremely large or small sharpness values that could destabilize training, we use the following bounded exponential mapping for \(\phi\): \[\begin{align} \phi(z) &= U^{\tanh(z)}, \end{align}\] where \(U > 1\) is a hyperparameter controlling the range of the sharpness coefficient.

This choice of \(\phi\) has two useful properties. First, since \(\tanh(z) \in (-1,1)\), it maps the confidence-conditioned signal to a bounded sharpness coefficient, \(\kappa_{e,j}(x) \in (1/U, U)\), preventing extreme \(\kappa\) values. Second, when \(z=0\), we have \(\tanh(z)=0\) and \(\phi(z)=1\), so \(\kappa\)-SwiGLU reduces to the standard SwiGLU gate.

Compared with the original SwiGLU, \[\begin{align} &\quad\quad \mathrm{SwiGLU}_e(x) = \mathrm{SiLU}(W_{g,e}x) \odot (W_{u,e}x) \nonumber \\ &= \left[ (W_{g,e}x) \odot \underbrace{ \sigma\left(W_{g,e}x\right) }_{\text{fixed sharpness}} \right] \odot (W_{u,e}x), \end{align}\] \(\kappa\)-SwiGLU replaces the fixed-sharpness SiLU gate with the confidence-conditioned gate \(\mathrm{SiLU}_{\kappa_e(x)}\): \[\begin{align} & \mathrm{\kappa\text{-}SwiGLU}_e(x) = \mathrm{SiLU}_{\kappa_e(x)}(W_{g,e}x) \odot (W_{u,e}x) \nonumber \\ &= \left[ (W_{g,e}x) \odot \underbrace{ \sigma\left(\kappa_e(x) W_{g,e}x\right) }_{\text{confidence-aware sharpness}} \right] \odot (W_{u,e}x). \end{align}\] For notational simplicity, we omit the gate-unit index \(j\) here. Thus, \(\kappa\)-SwiGLU generalizes the fixed-sharpness gate in standard SwiGLU by allowing each expert to adapt its activation selectivity according to token-level routing confidence. When \(\kappa_{e,j}(x)>1\), the gate becomes sharper and more selective than the standard SiLU gate; when \(\kappa_{e,j}(x)<1\), it becomes smoother and more broadly active, as illustrated in Figure 6.

3.4 Regularization on \(\kappa\) parameters↩︎

We apply L2 regularization to \(\alpha_{e,j}\) and \(b_{e,j}\) to prevent the sharpness modulation from deviating too aggressively from the standard SiLU: \[\begin{align} \mathcal{L}_{\mathrm{reg}} = \lambda_{\alpha} \sum_{e,j} \alpha_{e,j}^2 + \lambda_b \sum_{e,j} b_{e,j}^2. \end{align}\]

4 Experiments↩︎

4.1 Experimental Details↩︎

We implement \(\kappa\)-SwiGLU in a standard MoE Transformer architecture with SwiGLU MLPs. Our training pipeline is based on the widely used Nanochat codebase⁴, with modifications to incorporate a standard token-choice router and the proposed \(\kappa\)-SwiGLU activation. We train all models on the FineWeb-Edu dataset [32].

For matrix parameters, we use Aurora [33], an emerging state-of-the-art optimizer for large language model training, with a learning rate of \(0.01\) and weight decay of \(0.05\). For non-matrix parameters, we use AdamW with a learning rate of \(0.3\) and betas \((0.8, 0.95)\). These hyperparameters are inherited from the default Nanochat settings for dense models and slightly adapted to better accommodate the MoE architecture. An auxiliary load-balancing loss [14] is applied to the router logits to encourage balanced expert utilization, with weight \(10^{-3}\). A router \(z\)-loss [14] is also applied to suppress excessively large router logits, with weight \(10^{-5}\).

All models are trained using four H200 GPUs, each equipped with 141GB of memory. To fit within the available GPU memory, we vary the number of candidate experts and number of MoE layers across settings, as detailed in the next subsection. Table 1 summarizes the model configurations and training budgets for all settings.

We evaluate pretraining performance using the average score across 22 CORE benchmarks [34], which cover a diverse set of tasks including textbook knowledge, commonsense reasoning, and language modeling. We report Centered CORE accuracy, computed as the average benchmark score relative to a fixed-answer baseline. This metric mitigates potential biases arising from differences in answer format, class imbalance, and task difficulty across benchmarks.

Since models trained with the same settings but different random seeds can show noticeable variation on the CORE benchmark [35], we train three independent runs for each setting using random seeds 24, 26, and 28. We report the mean and standard deviation across these runs to ensure that our results are robust to random variation.

4.2 Model Settings and Training Budgets↩︎

To fit within the available GPU memory, we adopt three memory-saving strategies. First, we use at most 10 MoE layers and implement any remaining layers as dense layers. Second, as model depth increases, we gradually reduce the number of candidate experts from 64 to 16, reducing the memory cost of MoE layers at the expense of a smaller expert pool. Third, for deeper models with 16–28 layers, we adopt a sandwiched MoE architecture, in which only the middle two layers are MoE layers and all remaining layers are dense. Despite using only two MoE layers in this setting, \(\kappa\)-SwiGLU yields gains in the deeper sandwiched models, suggesting that confidence-aware gate sharpness can also benefit mixed MoE-dense architectures.

As summarized in Table 1, we train standard MoE models with 8, 10, 12, and 14 layers, consisting of 2–4 dense layers and 6–10 MoE layers, as well as sandwiched MoE models with 16, 20, 24, and 28 layers, consisting of 14–26 dense layers and 2 MoE layers. We set the MoE routing top-\(k\) to \(2\) for all models.

Since MoE models contain substantially more total parameters than dense models, we use a token-to-parameter ratio of 5 for all models. Following the common observation that not all MoE parameters are activated for each token, we estimate the effective parameter count using a square-root scaling rule, detailed in the appendix 8.

4.3 Optimization of \(\kappa\) Parameters↩︎

Optimizing \(\kappa\)-SwiGLU requires learning the gate-wise parameters \(\alpha_{e,j}\) and \(b_{e,j}\) for each expert. Since \(\alpha_{e,j}\) and \(b_{e,j}\) are scalars, they introduce only a negligible increase in parameter count compared with the original SwiGLU. Moreover, the router logits \(s_e(x)\) are already available during the forward pass, so computing \(\kappa_{e,j}(x)\) requires no additional matrix multiplications and only a few elementwise operations, resulting in a small computational overhead.

The range hyperparameter \(U\) is set to \(3\), constraining \(\kappa\) to the interval \((1/3, 3)\). For L2 regularization of \(\alpha_{e,j}\) and \(b_{e,j}\), we set \(\lambda_{\alpha}=2\times10^{-2}\) and \(\lambda_b=10^{-2}\), which we find effective in preventing overfitting while preserving sufficient flexibility for learning sharpness modulation.

To ensure stable training, we initialize \(\alpha_{e,j}=b_{e,j}=0\), so that \(\phi(\alpha_{e,j}s_e(x)+b_{e,j})=\phi(0)=1\) for all tokens. This initializes \(\kappa\)-SwiGLU to the standard SiLU gate. During the first \(1/10\) of training iterations, we keep \(\alpha_{e,j}\) and \(b_{e,j}\) frozen at 0, allowing the model to establish stable initial routing behavior and expert representations before introducing confidence-aware sharpness modulation. After this initial phase, we unfreeze \(\alpha_{e,j}\) and \(b_{e,j}\) and update them by backpropagation together with the rest of the model parameters. We use a learning rate schedule of linear warmup followed by linear decay: the learning rate is warmed up to \(0.12\) during the first 1000 iterations and then linearly decayed to \(0.06\) by the end of training.

4.4 Main Results↩︎

Figures 7 and 8 compare the performance of \(\kappa\)-SwiGLU and standard SwiGLU across model depths for standard and sandwiched MoE architectures, respectively. Across the evaluated settings, \(\kappa\)-SwiGLU improves the mean centered CORE score in 7 out of 8 configurations.

In the standard MoE setting, \(\kappa\)-SwiGLU improves the centered CORE score by approximately 0.6–1.0 percentage points, with slightly larger gains at shallower depths. One possible explanation is that shallower models use larger expert pools, whereas deeper standard MoE models reduce the number of candidate experts from 64 to 16 to fit within the memory budget. This makes the deeper models effectively denser and may reduce the relative benefit of confidence-aware expert gating.

In the sandwiched MoE setting, \(\kappa\)-SwiGLU yields gains of approximately 0.4–0.6 percentage points for models deeper than 16 layers, with slightly larger gains at greater depths.

Although the improvement in each individual setting is modest relative to run-to-run variation, \(\kappa\)-SwiGLU improves the mean centered CORE score in 7 out of 8 settings. This cross-setting consistency suggests that confidence-aware gate sharpness provides a robust positive trend across MoE architectures and model depths, while introducing only negligible additional parameters and small computational overhead.

A full breakdown over the 22 CORE benchmarks is provided in Appendix 6. Across the 22 benchmarks, \(\kappa\)-SwiGLU improves or matches the baseline on the majority of tasks, suggesting that the gains are not driven by a single benchmark.

In addition, Appendix 7 analyzes the learned \(\kappa\) parameters in detail. We show that after the warm-up phase, \(\kappa\) rapidly diverges from the standard fixed value of 1, with some gate units becoming substantially sharper and others substantially smoother. Over training, these values gradually return toward a more moderate range while remaining separated from 1, indicating that \(\kappa\)-SwiGLU learns a nontrivial and persistent modulation of gate sharpness. We further analyze the learned scale and bias parameters, \(\alpha\) and \(b\), and find that the router-logit-dependent scale term \(\alpha_{e,j}s_e(x)\) contributes more strongly than the offset term \(b_{e,j}\), supporting the importance of routing confidence in driving the learned sharpness modulation.

4.5 Ablation Studies↩︎

We perform ablation studies to understand the contributions of the two components in the \(\kappa\) parameterization in Eq. 2 : the router-logit-dependent scale term \(\alpha_{e,j} \cdot s_e(x)\) and the offset term \(b_{e,j}\). We compare three variants: \(\kappa\)-SwiGLU\(_{-\alpha}\), which removes the scale term by setting \(\alpha_{e,j}=0\); \(\kappa\)-SwiGLU\(_{-b}\), which removes the offset term by setting \(b_{e,j}=0\); and the full \(\kappa\)-SwiGLU method.

As shown in Table 3, removing the scale term \(\alpha_{e,j} \cdot s_e(x)\) consistently hurts performance. In contrast, removing the offset term \(b_{e,j}\) has a smaller impact. This suggests that the confidence-dependent scale term accounts for most of the benefit, while the offset mainly provides additional flexibility.

4.6 Computational Overhead of \(\kappa\)-SwiGLU↩︎

Table 4 compares the computational overhead of \(\kappa\)-SwiGLU with standard SwiGLU on the MoE-14L model, measured by active parameter count and tokens per second (TPS) during training and inference. \(\kappa\)-SwiGLU introduces only \(0.02\%\) additional active parameters. Its inference throughput differs from the standard SwiGLU baseline by only \(4.0\%\), indicating that \(\kappa\)-SwiGLU achieves performance gains with a small computational overhead.

5 Conclusion↩︎

In this work, we propose \(\kappa\)-SwiGLU, a confidence-aware variant of SwiGLU for Mixture-of-Experts (MoE) models that dynamically adjusts expert gate sharpness based on token-level routing confidence. By explicitly coupling router logits with expert gate sharpness, \(\kappa\)-SwiGLU allows each expert to adapt its activation selectivity according to routing confidence, providing a more flexible and expressive gating mechanism. Experiments on FineWeb-Edu show that \(\kappa\)-SwiGLU improves mean CORE performance across a range of MoE architectures and model depths, while introducing only negligible additional parameters and a small computational overhead. Future work could explore alternative parameterizations of confidence-aware gate modulation, as well as applications of \(\kappa\)-SwiGLU to MoE models beyond language modeling.

Limitations↩︎

This work evaluates \(\kappa\)-SwiGLU on relatively small-scale MoE language models trained on FineWeb-Edu. Although our experiments cover multiple model depths and both standard and sandwiched MoE architectures, the largest models remain much smaller than frontier-scale MoE systems due to limited computational resources. It remains to be verified whether the same trends hold at substantially larger parameter counts, longer training schedules, and larger-scale pretraining corpora.

Our evaluation is primarily based on pretrained model performance measured by CORE. While CORE covers a diverse set of benchmarks, it does not fully capture downstream behavior after instruction tuning, long-context use, reasoning-heavy evaluation, or deployment-oriented metrics. Broader evaluation is needed to better characterize where confidence-aware gate sharpness is most beneficial.

The proposed method introduces only a small number of additional parameters, but it incurs a small computational overhead of approximately \(4\)–\(7\%\) due to the extra elementwise operations needed to compute token-dependent sharpness coefficients. Further kernel-level optimization may reduce this overhead to a negligible level.

Finally, our method parameterizes the sharpness coefficient using a simple affine transformation of router logits followed by a bounded mapping. Other confidence signals, parameterizations, initialization strategies, or regularization schemes may lead to different trade-offs between stability, expressivity, and performance. We leave a more systematic exploration of these design choices, as well as applications beyond language modeling, to future work.

6 Breakdown of CORE Results↩︎

The CORE benchmark for pretrained model evaluation consists of 22 datasets, spanning a diverse set of tasks including textbook knowledge, commonsense reasoning, and language modeling. To provide a more fine-grained view of model behavior beyond the aggregate CORE score, Tables 5 and 6 report per-task performance across all evaluated configurations.

We note that BoolQ contributes noticeably to the improvement in several settings, especially for sandwiched MoE models. To verify that the observed gains are not solely driven by this benchmark, we also report Centered CORE without BoolQ. As shown in Tables 5 and 6, \(\kappa\)-SwiGLU still improves CORE (no BoolQ) in most settings, although the gains are smaller than in the full CORE average. This suggests that BoolQ amplifies the aggregate improvement, but the benefit of confidence-aware gate sharpness is not entirely explained by a single benchmark.

7 Empirical Analysis of \(\kappa\)-SwiGLU↩︎

Figure 9 shows the training dynamics of the top and bottom \(5\%\) of the learned \(\kappa\) values in a representative layer of a 12-layer MoE, with similar behavior observed across all MoE layers.

During the initial warm-up phase, the \(\kappa\) values remain fixed at 1, reducing \(\kappa\)-SwiGLU to the standard SiLU gate. Once unfrozen, the two groups rapidly diverge: tokens with the largest learned \(\kappa\) values move toward much sharper gates, while those with the smallest learned \(\kappa\) values move toward smoother gates. Interestingly, this separation is strongest shortly after unfreezing, with the top \(5\%\) reaching around \(2.5\) and the bottom \(5\%\) dropping to around \(0.4\). As training proceeds, both groups gradually move back toward 1, suggesting that the model initially explores a wide range of gate sharpness but later settles on a more moderate modulation. However, by the end of training, both groups remain substantially separated from 1, indicating that the learned \(\kappa\) values continue to have a non-negligible effect on the gating behavior.

This behavior indicates that \(\kappa\)-SwiGLU introduces a flexible, input-dependent adjustment of gate selectivity rather than simply making all gates uniformly sharper or smoother.

Since the \(\kappa\) parameterization in Eq. 2 contains two learnable parameters, \(\alpha\) and \(b\), we further track their training dynamics to understand their relative contributions to the learned \(\kappa\) values. We observe that the positive and negative values are approximately symmetric; therefore, we report only the mean of the positive values, with the negative values following the same trend up to a sign flip. As shown in Figure 10, both parameters exhibit trends similar to the resulting \(\kappa\): they increase sharply during the early phase of training, then gradually decay toward zero while remaining around \(0.05\) by the end of training. Because \(\alpha\) is multiplied by the router logit \(s_e(x)\), which empirically usually falls in the range \([2,4]\), the term \(\alpha \cdot s_e(x)\) contributes substantially more to \(\kappa\) than the bias term \(b_{e,j}\). For example, using a typical router logit value of \(s_e(x)=2.5\), the scale term in the middle of training contributes approximately \(0.134 \times 2.5 = 0.335\) to the affine input of \(\phi\). This is about \(1.675\times\) the bias contribution \(b=0.2\), indicating that the confidence-dependent scale term dominates the learned sharpness modulation.

8 Scaling of Computational Budgets↩︎

Since MoE models contain substantially more total parameters than dense models, we use a token-to-parameter ratio of 5 for all models. Following the common observation that not all MoE parameters are activated for each token, we estimate the effective parameter count using a square-root scaling rule \[\begin{align} N_{\mathrm{eff}} = N_{\mathrm{shared}} + N_{\mathrm{attn}}^{\mathrm{MoE}} + N_{\mathrm{MLP}}^{\mathrm{all\;experts}} \sqrt{\frac{k}{E}}. \end{align}\] where \(k=2\) is the top-\(k\) routing value, and \(E\) is the total number of experts. This rule discounts the expert MLP parameters by \(\sqrt{k/E}\) rather than the raw active fraction \(k/E\), assigning a larger effective parameter count to the expert pool. This is consistent with the intuition that sparse routing makes the effective capacity of MoE layers grow sublinearly with the total expert parameter count, while still exceeding the capacity implied by the active parameter count alone.

References↩︎

---

1. This sharpness-adjusted SiLU gate is equivalent to the Swish activation, with \(\kappa\) corresponding to the Swish sharpness coefficient. Although Swish allows this coefficient to be learned in principle [11], modern Transformer MLPs typically use the fixed SiLU gate of \(\kappa=1\).↩︎

2. For a random unit vector in 512-D, the volume fraction satisfying \(\cos(x,r_e)\ge 0.15\) is approximately \(0.03\%\).↩︎

3. We do not use the routing probability because it depends on the set of competing experts selected for the token and on the normalization over their logits, which introduces additional variation unrelated to the token–expert affinity itself.↩︎

4. https://github.com/karpathy/nanochat/↩︎