Mixture of Experts (MoE)

The expert networks are the specialized, independent sub-models within an MoE layer. Each expert is typically a feed-forward neural network (FFN) with its own set of parameters, trained to handle specific types or features of the input data. In a standard Transformer MoE layer, these replace the single, monolithic FFN. The system's total capacity scales with the number of experts (e.g., thousands), but only a small subset (e.g., 2-8) is activated per token, keeping computational cost manageable. Experts can be homogeneous or heterogeneous in architecture.

The router or gating network is a lightweight, trainable component that performs conditional computation. For each input token or sequence, it outputs a routing distribution—a set of scores or probabilities over all available experts. Its primary functions are:

• Selection: Dynamically chooses the top-k experts (e.g., top-2) for each input.
• Load Balancing: Incorporates auxiliary losses (like load balancing loss) to ensure all experts receive a roughly equal amount of training data, preventing expert collapse where only a few experts are ever used.
• The router is often a simple linear layer followed by a softmax, making its computational overhead negligible compared to the experts.

Sparse activation is the core mechanism that enables MoE efficiency. Instead of passing every input through every expert (dense computation), the router selects only the top-k experts (where k is a small integer like 1, 2, or 4). The input is then processed solely by this small, active subset. This creates a sparsely activated model where the total floating-point operations (FLOPs) scale with k * (expert size), not (number of experts) * (expert size). This allows for models with trillions of parameters (like Google's Switch Transformer) while requiring far less computation per forward pass than a dense model of equivalent size.

The auxiliary loss is a critical regularization term added to the main task loss (e.g., language modeling loss) to ensure stable MoE training. Its primary purpose is load balancing—distributing training examples evenly across experts. Without it, a "rich-get-richer" dynamic can occur, where a few experts are selected repeatedly, becoming specialized, while others are underutilized and fail to learn (expert collapse). A common formulation, as used in the Switch Transformer, encourages uniformity in both the router's probability distribution and the proportion of tokens assigned to each expert over a batch.

Noise is added to the router's logits during training to encourage exploration across all experts. This is implemented as tunable Gaussian or Gumbel noise. It serves two key purposes:

• Breaking Symmetry: Prevents identical experts from receiving identical routing signals at initialization.
• Promoting Discovery: Helps the router explore the expert space early in training, preventing premature convergence to a suboptimal routing strategy. The noise magnitude is often annealed over time. This technique is analogous to exploration strategies in reinforcement learning and is essential for training large, stable MoE layers.

The capacity factor is a hyperparameter that defines a soft limit on the number of tokens an expert can process in a given forward pass. Since token-to-expert assignment is dynamic, an expert might be the top choice for more tokens than it can physically compute in parallel. The capacity factor sets a buffer: expert_capacity = (tokens_per_batch / num_experts) * capacity_factor. Tokens exceeding an expert's capacity are dropped (skipped) or auxiliary-routed. A factor >1.0 (e.g., 1.25-2.0) provides slack to handle assignment fluctuations but increases memory and computation. Tuning this is crucial for balancing efficiency and model quality.

Adaptive computation is a family of techniques where a neural network dynamically adjusts its computational cost per input based on its perceived complexity or difficulty. Like MoE, it is a form of conditional computation.

• Core Principle: Avoids applying a uniform, maximal computational budget to all inputs.
• Common Techniques:
  • Early Exiting: Allows intermediate layers to produce a final output if a confidence threshold is met.
  • Conditional Computation: Dynamically activates different sub-networks or pathways (as in MoE).
  • Adaptive Attention Span: Adjusts the context length considered by attention mechanisms.
• Goal: Reduce average latency and power consumption during inference.

Structured sparsity is a model compression paradigm where weights are pruned in contiguous blocks or predefined patterns, rather than removing individual weights arbitrarily. This is crucial for achieving actual speedups on standard hardware.

• Contrast with Unstructured Sparsity: Unstructured sparsity removes individual weights but requires specialized hardware to exploit. Structured sparsity targets formats hardware can accelerate.
• Common Patterns:
  • Channel/Filter Pruning: Removes entire channels or filters from convolutional layers.
  • Block Sparsity: Prunes blocks of weights (e.g., 4x4).
  • N:M Sparsity: Enforces that for every block of M weights, at most N are non-zero (e.g., 2:4 sparsity).
• Hardware Alignment: Patterns like N:M sparsity are directly supported by modern NVIDIA Ampere/Ada/Hopper architecture Tensor Cores.

Model parallelism is a distributed training strategy that partitions a single model's layers or tensors across multiple devices (GPUs/TPUs) to train models that are too large to fit on one device.

• Contrast with Data Parallelism: Data parallelism replicates the entire model on each device and splits the batch. Model parallelism splits the model itself.
• Common Strategies:
  • Tensor Parallelism: Splits individual weight matrices and the associated computation across devices (e.g., splitting the feed-forward or attention layers in a transformer).
  • Pipeline Parallelism: Places different groups of model layers on different devices. A micro-batch moves through this pipeline sequentially, using techniques like GPipe or PipeDream to maintain high hardware utilization.
• Use Case: Essential for training modern large language models (LLMs) like GPT, Llama, and Mixtral (which itself uses expert parallelism, a form of model parallelism for MoE).