Switch Transformer

The defining mechanism of the Switch Transformer. Instead of routing each token to the top-k experts (e.g., k=2), it uses top-1 gating, sending each token to exactly one expert. This simplification reduces router computation and communication overhead while maintaining model capacity. The router is a simple linear layer followed by a softmax over experts.

A critical auxiliary loss function designed to ensure all experts are utilized evenly. Without it, the router can collapse, always selecting the same few experts. The Switch Transformer uses an expert-level auxiliary loss that encourages a uniform distribution of routing probabilities across the batch. This is essential for training stability and efficient hardware utilization.

The strategy for distributing experts across multiple devices (e.g., GPUs/TPUs).

• Each expert is placed on a different device.
• After routing, an all-to-all communication operation scatters tokens to the devices hosting their assigned expert.
• Experts process their assigned tokens in parallel.
• A second all-to-all gathers the processed tokens back. This pattern is the backbone of scalable MoE training and inference.

A hyperparameter that defines a buffer for expert processing.

• Expert Capacity = (tokens_per_batch / num_experts) * Capacity Factor.
• A factor >1.0 (e.g., 1.25) provides slack, preventing dropped tokens—tokens that cannot be processed because an expert's buffer is full.
• Dropped tokens are typically passed to the next layer unchanged via a residual connection. Tuning this factor balances compute efficiency against token loss.

The core efficiency proposition. For a given input, only the parameters of the single selected expert per layer are activated and computed. This enables:

• Massive parameter counts (e.g., 1.6 trillion parameters) with a constant computational cost proportional to the active, not total, parameters.
• A sparse computational graph that differs for each input, contrasting with the dense, fixed graph of standard transformers.

The primary architectural goal. By decoupling model size (number of experts) from computational cost (active experts per token), the Switch Transformer demonstrated predictable scaling laws.

• Achieved a 1.6 trillion parameter model with comparable training time to a dense T5-Base model.
• Showed improved sample efficiency and downstream task performance compared to dense models of equivalent computational budget, validating the MoE scaling hypothesis.

A neural network architecture where the model is composed of multiple sub-networks called experts. A gating network (router) dynamically selects a sparse subset of these experts to process each input. This enables models with trillions of parameters while maintaining a manageable computational cost per token, as only a fraction of the total parameters are activated for any given input.

The fundamental property of conditional computation models like MoE. Instead of a dense model where all parameters are used for every input, sparse activation means only a small, dynamically chosen subset of the model's total parameters is computed. This is the mechanism that allows massive models like the Switch Transformer to be feasible, trading parameter count for activation cost.

A small, trainable neural network component within a MoE model. For each input token or hidden state, the router outputs a set of scores or logits. These scores determine expert routing—which expert(s) should process that input. The Switch Transformer uses a simplified top-1 gating strategy, routing each token to exactly one expert.

A critical training challenge in MoE. Without intervention, the router can collapse, always selecting the same few experts. Load balancing techniques prevent this.

• Auxiliary Loss (Load Loss): An additional term added to the training objective that penalizes unbalanced routing, encouraging uniform expert utilization.
• Noise Top-k Gating: A common method that adds tunable Gaussian noise to router logits before selection to encourage exploration.

A model parallelism strategy designed for MoE models deployed across multiple devices (e.g., GPUs). Different experts are placed on different devices. After the routing decision, tokens must be communicated to their assigned expert's device.

• All-to-All Communication: The collective operation where tokens are scattered from all devices to their target expert devices, and outputs are gathered back. This communication overhead is a major bottleneck in distributed MoE inference.

A key implementation hyperparameter that defines a soft limit on the number of tokens a single expert can process in a forward pass. It's calculated as: (batch_size * sequence_length * k) / num_experts * capacity_factor

• A factor >1.0 (e.g., 1.25) provides a buffer to handle uneven routing and minimize dropped tokens.
• A lower factor saves memory and compute but increases drop risk. Balancing this is essential for efficient inference.