Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Switch Transformer › Capacity Factor: Managing Overflow

Michael BrenndoerferUpdated January 5, 20269 min read

Language AI HandbookAll chapters

Part 33 of 67Mixture of Experts

Chapter 247Switch Transformer › Capacity Factor: Managing Overflow

Reading section 4 of 14 · Capacity Factor: Managing Overflow · 9 minSection 4 of 14 · 9 min

Section 4 of 14 · 9 min read

Switch Transformer › Capacity Factor: Managing OverflowOne section · Sections in this chapter4 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Even with perfect load balancing losses pushing the router toward uniform distribution, routing can never be perfectly uniform in practice. The discrete nature of expert selection means that random fluctuations will always cause some experts to receive more tokens than others in any given batch. This creates a practical problem that cannot be ignored: how many tokens should each expert be prepared to handle? If we allocate too few slots, important tokens get dropped. If we allocate too many, we waste memory that could be used for more experts or larger expert networks.

The capacity factor ( $C$ ) is Switch Transformer's solution to this challenge. It defines how much buffer capacity each expert has beyond the perfectly balanced allocation. This gives a mechanism to handle the inevitable routing imbalances without catastrophic failure. The capacity factor is one of the most practically important hyperparameters in a Switch Transformer, because it directly governs the tradeoff between memory usage and the frequency of token dropping.

Think of the capacity factor like reserving seats at a restaurant. If you know on average 10 people will want to sit at each of your 8 tables, but arrivals are random, you do not set every table for exactly 10 guests. You set each table with a few extra chairs to handle groups that show up together. The capacity factor tells each expert how many extra chairs to set out. Setting $C = 1.25$ is like setting 12.5 chairs at each table: you can accommodate bursts of extra demand without anyone being turned away, at the cost of a modest amount of wasted seating capacity when demand is below peak.

Capacity CalculationLink Copied

If a batch contains $n$ total tokens and we have $E$ experts, perfect balance would assign $n/E$ tokens to each expert. In an ideal world, the router would reach this exact distribution, and we could set the capacity to exactly $n/E$ without any waste. But in reality, some experts will be more popular than others for any given batch, driven by the content distribution of that particular batch. The capacity factor scales the allocation to accommodate this variance:

\text{capacity} = \left\lfloor \frac{n \cdot C}{E} \right\rfloor

where:

$n$ : the total number of tokens in the batch, equal to batch size times sequence length
$C$ : the capacity factor (a scalar multiplier greater than 1.0 for typical use)
$E$ : the number of experts available
$\lfloor \cdot \rfloor$ : the floor function, rounding down to ensure an integer number of token slots

Understanding this formula requires thinking about what happens at different values of $C$ . With $C = 1.0$ , each expert has exactly enough slots for perfect balance. Any deviation from uniform routing causes some tokens to be dropped. With $C = 1.25$ , each expert can handle 25% more tokens than the perfectly balanced allocation, so routing variance must be quite substantial before any tokens are dropped. With $C = 2.0$ , experts can absorb twice the expected load, in effect eliminating dropping at the cost of using twice the memory per expert.

Why does this formula make sense? Notice that $n/E$ is the perfectly balanced allocation per expert. Multiplying by $C$ scales this up proportionally. If $C = 1.5$ and each expert would receive 100 tokens under perfect balance, each expert instead gets capacity for 150 tokens, leaving a 50-token buffer for routing imbalance. The floor function ensures integer capacity, since you cannot process a fractional token.

A capacity factor of $C = 1.0$ means experts have exactly enough slots for perfect balance. In practice, routing is never perfect, so some experts overflow. The Switch solution is simple and pragmatic: tokens that exceed an expert's capacity skip the expert entirely and pass through via the residual connection.

This design creates a spectrum of tradeoffs that you can select based on your constraints:

$C = 1.0$ : Minimal memory usage, but many tokens may be dropped (skipped)
$C = 1.25$ : 25% buffer handles typical load variance well
$C = 2.0$ : Generous buffer, fewer drops, but double the memory per expert

In[8]:

Code

def compute_capacity(
    batch_size: int, seq_len: int, num_experts: int, capacity_factor: float
) -> int:
    """Calculate expert capacity given batch parameters."""
    total_tokens = batch_size * seq_len
    # Perfect balance allocation with capacity factor buffer
    capacity = int((total_tokens * capacity_factor) / num_experts)
    return max(capacity, 1)  # At least 1 token per expert

In[9]:

Code

# Example setup: 32 batch size, 512 sequence length, 16 experts
batch_size = 32
seq_len = 512
num_experts = 16
total_tokens = batch_size * seq_len

# Calculate capacities for different factors
capacity_results = []
for cf in [1.0, 1.25, 1.5, 2.0]:
    cap = compute_capacity(batch_size, seq_len, num_experts, cf)
    perfect_balance = total_tokens // num_experts
    buffer = cap - perfect_balance
    capacity_results.append((cf, cap, perfect_balance, buffer))

Out[10]:

Console

C=1.00: capacity=1024 tokens/expert (perfect balance=1024, buffer=0)
C=1.25: capacity=1280 tokens/expert (perfect balance=1024, buffer=256)
C=1.50: capacity=1536 tokens/expert (perfect balance=1024, buffer=512)
C=2.00: capacity=2048 tokens/expert (perfect balance=1024, buffer=1024)

With a capacity factor of 1.25, each expert can handle 25% more tokens than perfect balance would require. This buffer accommodates the natural variance in routing decisions without requiring excessive memory. In practice, the authors found $C = 1.25$ gives a good tradeoff between memory efficiency and token coverage, though this value may need adjustment for different batch sizes and expert counts. Smaller batches have higher relative routing variance and may benefit from a larger capacity factor, while very large batches average out the variance and can use smaller factors.

Out[11]:

Visualization

Line plot showing tokens per expert increasing linearly as capacity factor grows from 0.8 to 2.5. A horizontal red dashed line marks the perfect balance level and a vertical green dotted line marks the recommended C=1.25 value. The region above perfect balance is shaded blue to indicate buffer capacity. — Expert capacity scaling with capacity factor. Capacity increases linearly with the capacity factor, giving a buffer against routing imbalances. The recommended value of C=1.25 balances memory efficiency with token coverage.

Line plot with red circles showing simulated token drop rate (percentage) decreasing steeply as capacity factor increases from 0.8 to 2.5. A vertical green dotted line marks the recommended C=1.25 value, and the area under the curve is shaded red. — Token dropping rates for different capacity factors. Higher capacity factors reduce token dropping rates, with diminishing returns observed beyond C=1.5, making C=1.25 a common practical choice.

Token Dropping BehaviorLink Copied

When an expert reaches capacity, additional tokens routed to it are "dropped": they bypass the expert entirely and only the residual connection preserves their information. This might seem problematic at first glance. After all, if the router determined that a token should go to a particular expert, doesn't skipping that expert harm the model's performance? The answer is fine-grained: token dropping does have a cost, but several design choices mitigate this cost enough that moderate dropping is acceptable and the overall system works well.

Understanding why token dropping is tolerable requires considering what information is lost when a token gets dropped. The token's representation entering the Switch layer is not destroyed; it continues unchanged through the residual connection and remains available for all subsequent transformer layers. What the token misses is the specific expert processing in this one layer. Given that a large Switch Transformer might have 12 or 24 such layers, missing one expert processing in one layer is a relatively small fraction of the total computation applied to that token.

Several factors mitigate the impact of token dropping, making this design more reliable than it initially appears:

Load balancing reduces drops. The auxiliary balancing loss (covered in Chapter 6 of this part) incentivizes the router to distribute tokens evenly across experts. A well-trained router rarely overloads any single expert severely, minimizing the frequency of capacity overflow. During training, the model learns routing strategies that naturally keep expert loads within bounds, so capacity overflow becomes increasingly rare as training progresses.
Residual connections preserve information. Dropped tokens still retain their original representation through the residual connection; they simply do not receive the expert's specialized processing for that layer. The token's information is not lost, just unchanged by that particular expert. This is a graceful degradation rather than a catastrophic failure.
Drops are distributed across tokens. With proper load balancing, drops rarely concentrate on specific tokens across multiple layers. A token might skip one expert in layer 5 but be processed normally in layers 1 through 4 and 6 through 12. The impact of a single dropped expert processing is diluted across the many layers of the network, preventing any single token from being systematically disadvantaged.
Training adapts to token dropping. The model learns to function with some level of token dropping, developing representations that do not critically depend on every expert processing at every layer. This adaptation means that occasional capacity overflow does not catastrophically harm performance. The model has seen token dropping during training and has learned to work around it.