← All questions
LLM FoundationsHard· Mixture of Experts· Loss Functions

Explain the Switch Transformer auxiliary loss and why i...

The Switch Transformer auxiliary loss is defined as αN ∑ i fiPi, where fi is the hard token fraction (non-differentiable) and Pi is the soft router probability (differentiable). The gradient ∂L/∂Pi = αN fi flows through Pi to the router embeddings, using fi as a load-proportional multiplier. Because no gradient passes through the discrete top-K decision itself, the loss remains differentiable.

Code playground (Python)