Vertrouwensadaptieve SwiGLU voor Mixture-of-Experts

Samenvatting

SwiGLU is een standaard gated-activatiefunctie geworden in moderne Transformer MLP's, maar de scherpte van de poortfunctie – de soepelheid en selectiviteit van de poortfunctie – blijft doorgaans vast tijdens de training. In dit werk stellen we Confidence-Aware SwiGLU (κ-SwiGLU) voor, een variant van SwiGLU voor Mixture-of-Experts (MoE)-modellen die de scherpte van de expertpoort aanpast op basis van het token-level routeringsvertrouwen. Specifiek parametriseert κ-SwiGLU de SiLU-poortscherptecoëfficiënt als een leerbare functie van de router-logit, waardoor elke expertpoorteenheid kan interpoleren tussen soepele, breed actieve poortwerking en scherpe, selectieve poortwerking. We evalueren κ-SwiGLU op de FineWeb-Edu-dataset over MoE Transformer-modellen variërend van 8 tot 28 lagen. In deze settings verbetert κ-SwiGLU de gemiddelde CORE-prestaties, terwijl het verwaarloosbare parameters toevoegt en slechts een kleine rekenkundige overhead met zich meebrengt, wat aantoont dat vertrouwensbewuste poortscherpte een veelbelovend mechanisme is voor het verbeteren van MoE MLP's. De code is beschikbaar op 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.