Адаптивный по уверенности SwiGLU для смеси экспертов

Аннотация

SwiGLU стал стандартной активацией с гейтированием в современных трансформерных MLP, однако его резкость гейта — плавность и избирательность функции гейтирования — обычно фиксирована на протяжении обучения. В данной работе мы предлагаем Confidence-Aware SwiGLU (κ-SwiGLU), вариант SwiGLU для моделей смеси экспертов (MoE), который регулирует резкость гейта экспертов в зависимости от уверенности маршрутизации на уровне токенов. В частности, κ-SwiGLU параметризует коэффициент резкости гейта SiLU как обучаемую функцию логита маршрутизатора, позволяя каждому элементу гейта эксперта интерполировать между плавным, широко активирующим гейтированием и резким, селективным гейтированием. Мы оцениваем κ-SwiGLU на наборе данных FineWeb-Edu на моделях трансформеров MoE с 8 по 28 слоев. Во всех этих конфигурациях κ-SwiGLU улучшает среднюю производительность CORE, добавляя незначительное количество параметров и внося лишь небольшой вычислительный оверхед, что демонстрирует перспективность механизма резкости гейта, зависящего от уверенности, для улучшения MoE MLP. Код доступен по адресу 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.