Switch Transformer › Training Stability Strategies

Switch Transformers, like other MoE models, can be unstable during training. The sparse routing mechanism introduces discontinuities that dense models do not face, and the discrete expert selection can create optimization challenges that compound over the course of a long training run. The authors identified several strategies necessary for stable training at scale, each addressing a specific source of instability. These strategies are not just engineering tricks; they stand for real insights about the pathological training dynamics that arise in sparse models and how to counteract them.

One counterintuitive finding was that using lower precision (bfloat16 or float16) in most of the model while keeping the router in float32 materially improved stability. This selective precision strategy shows a deep understanding of where numerical precision matters most in the network.

The router's softmax operation is particularly sensitive to numerical precision because small differences in logits get amplified exponentially. When computing $e^{z_i}$ for logits $z_i$, even small rounding errors can cause some experts to receive disproportionately high or low probabilities. In lower precision, these errors accumulate and can cause the router to make inconsistent or degenerate routing decisions. A rounding error of 0.001 in a logit becomes a rounding error of $e^{0.001} \approx 1.001$ in the exponential, but if the logit magnitude itself is large (say 20), then small absolute errors correspond to large relative errors in the resulting probabilities.

By keeping the router computation in float32 while letting the rest of the model to use memory-efficient bfloat16, Switch Transformer reaches both numerical stability and computational efficiency. The router is a tiny fraction of total computation (a single linear layer followed by a softmax), so the precision overhead is minimal in terms of memory and compute. The expert networks, which dominate the parameter count and computation, can safely use lower precision because the relative smoothness of their FFN computations makes them less sensitive to rounding errors.

As we discussed in Chapter 7, the router z-loss penalizes large router logits, preventing the softmax from becoming too peaked. This is particularly important for Switch Transformer because top-1 routing already creates hard decisions. Extremely peaked softmax distributions worsen the discontinuity, making tiny input changes flip routing decisions sharply. The z-loss is a direct intervention against this failure mode.

The z-loss gives a gentle pressure against this behavior. To understand what it penalizes, we first need to understand the log-sum-exp function, which is closely related to the softmax denominator. For a batch of $n$ tokens with logit vectors $z^{(i)}$ each containing $E$ elements, the z-loss is:

$$L_z = \frac{1}{n} \sum_{i=1}^{n} \left( \log \sum_{j=1}^{E} e^{z_j^{(i)}} \right)^2$$

where:

• $L_z$: the auxiliary router z-loss, added to the total training objective with a small coefficient
• $n$: the total number of tokens in the batch
• $E$: the total number of experts
• $z_j^{(i)}$: the logit (pre-softmax score) for expert $j$ given token $i$
• $\sum_{j=1}^{E} e^{z_j^{(i)}}$: the softmax partition function (denominator), which captures the total magnitude of the logit vector

This formula penalizes large logits to maintain numerical stability. Each component serves a specific purpose:

1. Sum of Exponentials: The inner sum $\sum e^{z_j}$ captures the aggregate magnitude of the logits. When logits are large, this sum grows exponentially. This gives a sensitive detector for logit magnitude.
2. Logarithm: The $\log$ scales this value back to the linear domain, in effect computing something close to $\max_j z_j$ plus a small correction term. This makes the penalty proportional to the maximum logit rather than exponential in it.
3. Square: Squaring the result penalizes large positive values strongly while being smooth near zero, creating a strong incentive to keep logits bounded without overly constraining small logit values.

The log-sum-exp is tightly related to the maximum logit value. For a vector with one dominant entry $z_{\max}$ and all others much smaller, $\log \sum_j e^{z_j} \approx z_{\max}$. Minimizing $(\log \sum e^{z_j})^2$ is therefore approximately equivalent to minimizing $z_{\max}^2$, directly penalizing large maximum logits. Squaring the log-sum-exp creates a smooth, differentiable penalty that applies strong pressure when logits are large.

Without this loss, the model could increase logits indefinitely. Due to softmax's translation invariance, logits of [5, 3, 2] and [500, 498, 497] produce identical probability distributions. But the second case risks floating-point overflow and creates extremely peaked distributions that harm optimization. The z-loss prevents this degenerate behavior.

The Switch paper recommended several initialization and regularization tweaks that collectively improve training stability. These recommendations emerged from extensive empirical experimentation and address specific failure modes that appeared during early Switch Transformer training runs.

Smaller router initialization. The router weights should be initialized with a smaller standard deviation than the default, typically by using a scale factor of 0.1 relative to the default initialization. This prevents the router from making extremely confident routing decisions early in training, before it has seen enough data to form reliable preferences. With large initial weights, the router commits strongly to certain routing patterns from the very first batches, which can create path dependency: early routing decisions determine which tokens experts see, which determines how experts specialize, which in turn reinforces the original routing decisions. Starting with small weights keeps the router more agnostic early on, letting more exploratory routing before specialization solidifies.

Increased dropout in experts. Higher dropout rates of 0.1 to 0.2 inside expert networks help prevent overfitting to specific routing patterns. Because each expert only sees the subset of tokens routed to it, the effective training sample for each expert is smaller than the full batch. Higher dropout compensates for this reduced sample diversity and prevents experts from memorizing the specific tokens they receive rather than learning general patterns.

Expert dropout as structured regularization. Occasionally dropping entire experts during training encourages the model to develop redundancy and prevents over-reliance on any single expert. When an expert is dropped during a training step, its tokens must be handled by the residual connection alone, which trains the rest of the network to function reliably without that expert's contribution. This technique is analogous to dropout at the expert level rather than the neuron level. This gives structured regularization that improves generalization.

These techniques work together to prevent premature specialization, where the model commits too strongly to certain routing patterns before it has seen enough data to make informed decisions. The goal is to keep the training dynamics exploratory during the early phase, letting the router and experts to co-evolve gradually toward a stable and beneficial division of labor rather than locking into a suboptimal pattern early and never recovering.