SwiGLU Adaptativo por Confiança para Mistura de Especialistas

Resumo

O SwiGLU tornou-se uma ativação com portão padrão nos MLPs modernos de Transformers, mas sua nitidez de portão — a suavidade e seletividade da função de portão — é tipicamente fixa durante o treinamento. Neste trabalho, propomos o SwiGLU Ciente da Confiança (κ-SwiGLU), uma variante do SwiGLU para modelos de Mistura de Especialistas (MoE) que ajusta a nitidez do portão do especialista de acordo com a confiança do roteamento em nível de token. Especificamente, o κ-SwiGLU parametriza o coeficiente de nitidez do portão SiLU como uma função aprendível do logit do roteador, permitindo que cada unidade de portão do especialista interpole entre um portão suave e amplamente ativo e um portão nítido e seletivo. Avaliamos o κ-SwiGLU no conjunto de dados FineWeb-Edu em modelos Transformer MoE variando de 8 a 28 camadas. Nessas configurações, o κ-SwiGLU melhora o desempenho médio do CORE, adicionando parâmetros desprezíveis e incorrendo apenas em uma pequena sobrecarga computacional, demonstrando que a nitidez de portão ciente da confiança é um mecanismo promissor para melhorar MLPs MoE. O código está disponível em 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.