SwiGLU adaptativo por confianza para mezcla de expertos

Resumen

SwiGLU se ha convertido en una activación con compuerta estándar en los MLP de los Transformers modernos; sin embargo, su nitidez de compuerta —la suavidad y selectividad de la función de activación— suele ser fija durante todo el entrenamiento. En este trabajo, proponemos Confidence-Aware SwiGLU (κ-SwiGLU), una variante de SwiGLU para modelos de Mixtura de Expertos (MoE) que ajusta la nitidez de la compuerta del experto según la confianza de enrutamiento a nivel de token. Específicamente, κ-SwiGLU parametriza el coeficiente de nitidez de la compuerta SiLU como una función aprendible del logit del enrutador, lo que permite que cada unidad de compuerta del experto interpole entre una activación suave y de amplio espectro, y una activación nítida y selectiva. Evaluamos κ-SwiGLU en el conjunto de datos FineWeb-Edu con modelos Transformer MoE de 8 a 28 capas. En estas configuraciones, κ-SwiGLU mejora el rendimiento CORE medio, a la vez que añade parámetros insignificantes e incurre solo en una pequeña sobrecarga computacional, lo que demuestra que la nitidez de compuerta consciente de la confianza es un mecanismo prometedor para mejorar los MLP de MoE. El código está disponible en 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.