Confidence-Adaptive SwiGLU für Mixture-of-Experts

Zusammenfassung

SwiGLU hat sich als Standard-Gating-Aktivierung in modernen Transformer-MLPs etabliert, doch seine Gate-Schärfe – die Glätte und Selektivität der Gating-Funktion – bleibt während des Trainings typischerweise festgelegt. In dieser Arbeit schlagen wir Confidence-Aware SwiGLU (κ-SwiGLU) vor, eine Variante von SwiGLU für Mixture-of-Experts (MoE)-Modelle, die die Expert-Gate-Schärfe gemäß der tokenweisen Routing-Konfidenz anpasst. Konkret parametrisiert κ-SwiGLU den Schärfekoeffizienten des SiLU-Gates als eine lernbare Funktion des Router-Logits, sodass jede Expert-Gate-Einheit zwischen glattem, breit aktivem Gating und scharfem, selektivem Gating interpolieren kann. Wir evaluieren κ-SwiGLU auf dem FineWeb-Edu-Datensatz mit MoE-Transformer-Modellen, deren Größe von 8 bis 28 Schichten reicht. In diesen Einstellungen verbessert κ-SwiGLU die durchschnittliche CORE-Leistung, fügt dabei vernachlässigbare Parameter hinzu und verursacht nur einen geringen Rechenaufwand, was zeigt, dass konfidenzbewusste Gate-Schärfe ein vielversprechender Mechanismus zur Verbesserung von MoE-MLPs ist. Der Code ist verfügbar unter https://github.com/askerlee/kappa-swiglu.

English

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.