Source-linked AI summary
Confidence-Adaptive SwiGLU for Mixture-of-Experts
Shaohua Li, Xiuchao Sui, Xiaobing Sun, Yuhang Wu, Liangli Zhen, Yong Liu, Rick Siow Mong Goh
TL;DR
Modern Transformer MLPs typically keep SwiGLU gate sharpness fixed, leaving open how MoE experts could adapt gating to token-level routing confidence. κ-SwiGLU makes sharpness a learnable function of router signals, and experiments report improved mean CORE performance with negligible parameter growth and small computational overhead. The paper therefore presents confidence-aware gate sharpness as a promising mechanism for MoE MLPs within the evaluated scope.
Problem
SwiGLU gate sharpness is typically fixed, motivating investigation of whether MoE expert gates can adapt their smoothness and selectivity to routing confidence.
Method
κ-SwiGLU explicitly couples router logits to each expert’s SiLU gate through a token-dependent multiplicative sharpness coefficient.
Results
κ-SwiGLU improves mean CORE performance across FineWeb-Edu MoE architectures and model depths while adding negligible parameters and small computational overhead.
Takeaways & Limitations
Confidence-aware gate sharpness provides a more flexible and expressive gating mechanism for the evaluated MoE models.
Takeaways & Limitations
The evaluation uses relatively small-scale MoE language models and does not establish whether the trends hold for substantially larger models, longer training, or larger corpora.
Abstract
from arXiv · showhide
SwiGLU has become a standard gated activation in modern Transformer MLPs, yet its gate sharpness -- the smoothness and selectivity of the gating function -- is typically fixed throughout training. In this work, we propose Confidence-Aware SwiGLU ($κ$-SwiGLU), a variant of SwiGLU for Mixture-of-Experts (MoE) models that adjusts expert gate sharpness according to token-level routing confidence. Specifically, $κ$-SwiGLU parameterizes the SiLU gate sharpness coefficient as a learnable function of the router logit, enabling each expert gate unit to interpolate between smooth, broadly active gating and sharp, selective gating. We evaluate $κ$-SwiGLU on the FineWeb-Edu dataset across MoE Transformer models ranging from 8 to 28 layers. Across these settings, $κ$-SwiGLU improves mean CORE performance while adding negligible parameters and incurring only a small computational overhead, demonstrating that confidence-aware gate sharpness is a promising mechanism for improving MoE MLPs. The code is available at https://github.com/askerlee/kappa-swiglu.
1 Introduction
SwiGLU uses a SiLU gate to selectively modulate activations, but κ-SwiGLU makes gate sharpness responsive to routing confidence in MoE experts. It builds on observed router–gate coupling and improves benchmark performance across evaluated model settings.
- SwiGLU’s SiLU gate selectively suppresses or amplifies intermediate features while retaining low computational cost.
- κ controls the transition between inactive and active gate states, with larger values producing sharper and more selective gating.The standard SiLU gate corresponds to κ = 1 in the supplied method description.
- MoE router scores provide a token-level confidence signal that can control gate activation inside selected experts.Higher router scores indicate greater confidence in a token–expert assignment.
- Training induces router–gate alignment, shifting tokens with different router affinities across regions of an expert’s SiLU transition curve.This coupling can systematically amplify or suppress subsets of expert activations.
- κ-SwiGLU explicitly modulates the SiLU curve’s sharpness with a token-dependent multiplicative coefficient rather than additively shifting its input.Each expert learns a confidence–sharpness mapping that can produce smoother or sharper gating.
- κ-SwiGLU shows a consistent positive trend in model quality across FineWeb-Edu MoE models spanning 8 to 28 layers.The evaluation reports stronger pretraining benchmark performance across the evaluated settings.
2 Related Work
Related work addresses MoE routing balance, routing geometry, representation collapse, and gated activation design. κ-SwiGLU differs by using the selected expert’s router signal to modulate gate sharpness inside the expert MLP.
- MoE Load balancing and routing stability: MoE research commonly uses load-balancing losses and router regularization to prevent routing collapse and encourage more uniform expert utilization.
- MoE Load balancing and routing stability: Recent methods study batch-statistics sensitivity, auxiliary-loss-free balancing, and alternative formulations such as ϕ-balancing.
- Geometry of MoE routing: Geometry-aware routing methods encourage expert assignment to reflect token-representation geometry, while other work addresses representational collapse and router–expert coupling.
- Gated activations in Transformer MLPs: Transformer MLPs have shifted from ReLU- and GELU-based networks toward gated activations including GLU, GeGLU, and SwiGLU.Recent work continues to explore alternative activation ranges and adaptive mixtures of activation functions.
3 Method
The method identifies router–gate coupling as an implicit confidence-induced bias in MoE expert gates, then introduces κ-SwiGLU to adapt gate sharpness directly from routing confidence. The resulting gates can become sharper or smoother per token and expert while bounded regularization controls modulation.
- Confidence-adaptive SwiGLU: κ-SwiGLU uses the router logit as a token-level signal to learn a multiplicative sharpness coefficient for each expert gate.This contrasts with emergent router–gate coupling, which affects gates through an additive input shift.
- Router–gate alignment: Router–gate alignment emerges during training, with gate projections rapidly reaching cosine similarities of 0.2–0.4 before changing across layers and time.The coupling remains non-negligible across seven independently trained 8-layer MoEs.
- Confidence-induced gate bias: Router-aligned gate components create confidence-induced input shifts, systematically producing positive or negative biases for different subsets of expert gates.The top and bottom 5% of induced bias values remain substantially positive and negative throughout training.
- Confidence-adaptive SwiGLU: The sharpness-adjusted gate SiLUκ(z) = z · σ(κz) makes larger κ values sharper and more selective, while smaller values produce smoother, more broadly active gating.The standard SiLU gate is the fixed-sharpness case κ = 1.
- Regularization on κ parameters: A bounded exponential mapping constrains κ to (1/U, U), and L2 regularization limits deviations from standard SiLU behavior.When the conditioned signal is zero, the mapping yields κ = 1 and recovers the standard SwiGLU gate.
4 Experiments
Experiments on FineWeb-Edu evaluate κ-SwiGLU across standard and sandwiched MoE architectures from 8 to 28 layers. κ-SwiGLU generally improves centered CORE performance with negligible parameter growth and small computational overhead.
- Experimental Setup: Models use FineWeb-Edu and centered CORE accuracy, averaging 22 benchmarks relative to a fixed-answer baseline across three random-seed runs.The evaluated architectures vary MoE layers, dense layers, candidate experts, and memory-saving configurations.
- Main Results: 0.6–1.0 percentage points: κ-SwiGLU improves centered CORE in standard MoE models, with slightly larger gains at shallower depths.The passage attributes the possible depth trend to differences in expert-pool size under the memory budget.
- Main Results: 0.4–0.6 percentage points: κ-SwiGLU improves centered CORE in sandwiched MoE models deeper than 16 layers.The gains are slightly larger at greater depths.
- Analysis: κ-SwiGLU improves or matches the baseline on most CORE tasks, while learned κ values become persistently sharper or smoother than the fixed value.Removing the router-logit-dependent scale term consistently hurts performance, whereas removing the offset has a smaller impact.
- Efficiency: κ-SwiGLU adds negligible parameters and small computational overhead because its learned scalars require no additional matrix multiplications.On MoE-14L, active parameters increase by 0.02%, while inference throughput differs from standard SwiGLU by 4.0%.
5 Conclusion
The paper proposes κ-SwiGLU, which adapts expert gate sharpness to token-level routing confidence. On FineWeb-Edu, it improves mean CORE performance across MoE architectures and depths with negligible additional parameters and small computational overhead.
- Conclusion: κ-SwiGLU dynamically adjusts each expert’s gate sharpness using token-level routing confidence.The method explicitly couples router logits with expert gate sharpness to adapt activation selectivity.
- Conclusion: The method introduces negligible additional parameters and small computational overhead.The paper identifies alternative confidence-modulation parameterizations and non-language-modeling MoE applications as future work.
Limitations
The evaluation is limited to relatively small FineWeb-Edu MoE models and pretrained CORE performance, while the sharpness parameterization explores only a simple design.
- The experiments use relatively small MoE language models trained on FineWeb-Edu, so scaling to larger models, longer schedules, and larger corpora remains unverified.The study spans multiple depths and standard and sandwiched MoE architectures, but its largest models remain below frontier scale.
- The evaluation primarily measures pretrained performance with CORE, which does not fully capture instruction-tuned, long-context, reasoning-heavy, or deployment-oriented behavior.Broader evaluation is needed to characterize where confidence-aware gate sharpness is most beneficial.
- κ-SwiGLU adds approximately 4–7% computational overhead from extra elementwise operations despite introducing only a small number of parameters.Kernel-level optimization may reduce this overhead.
- The method uses a simple affine transformation of router logits followed by a bounded mapping, leaving alternative signals, parameterizations, initializations, and regularization schemes unexplored.The authors also leave applications beyond language modeling for future work.
A Breakdown of CORE Results
The CORE analysis examines performance across 22 datasets and tests whether aggregate gains persist after removing BoolQ. κ-SwiGLU generally retains gains without BoolQ, though they become smaller.
- CORE comprises 22 datasets spanning textbook knowledge, commonsense reasoning, and language modeling, with Tables 5 and 6 providing per-task results.Centered CORE is supplemented by per-task performance to show behavior beyond the aggregate score.
- κ-SwiGLU still improves CORE without BoolQ in most settings, although gains are smaller than in the full CORE average.The comparison is designed to test whether aggregate improvements are driven by one benchmark.
B Empirical Analysis of κ-SwiGLU
Training analysis shows that κ-SwiGLU explores token-dependent gate sharpness, with κ values separating into sharper and smoother groups before settling toward moderate modulation. Router-linked parameters follow similar dynamics.
- During warm-up, κ remains fixed at 1, after which the highest- and lowest-κ token groups rapidly diverge toward sharper and smoother gates.The dynamics are shown for the top and bottom 5% of learned κ values in a representative 12-layer MoE layer.
- The positive means of α and b in the 9th layer follow trends similar to κ, while their negative subsets are approximately symmetric.Figure 10 reports the positive subsets because the negative values follow the same trend up to a sign flip.
- The top 5% of κ reaches around 2.5 while the bottom 5% drops to around 0.4 shortly after unfreezing, then both groups move back toward 1.By the end of training, both groups remain substantially separated from 1.
- The learned κ dynamics indicate input-dependent gate selectivity rather than uniformly sharper or smoother gates.
- The confidence-dependent scale contribution is approximately 0.335, or 1.675× the bias contribution b = 0.2, indicating that the scale term dominates the learned sharpness modulation.
C Scaling of Computational Budgets
The computational-budget analysis estimates effective MoE capacity using a token-to-parameter ratio and a square-root scaling rule that discounts expert parameters while preserving their sublinear contribution.
- All models use a token-to-parameter ratio of 5 to estimate effective parameter count because sparse routing does not activate all MoE parameters per token.The estimate follows a square-root scaling rule.
- The scaling rule uses top-k routing with k = 2 and E total experts to adjust the expert-pool contribution.
- The rule discounts expert MLP parameters by a sublinear factor, producing greater effective capacity than the active-parameter count alone.
- Table 6 reports per-task centered accuracy percentages and Centered CORE for Sandwiched-MoE configurations, including a no-BoolQ aggregate.