Switch Transformer › Key Parameters

The Switch Layer has four hyperparameters that you will need to configure for any practical deployment, and the tradeoffs they govern are nontrivial. Understanding what each parameter controls helps you make principled choices when adapting Switch Transformer to a new scale or training environment.

The key parameters for the Switch Layer are:

• num_experts: Number of expert networks. More experts increase total parameter capacity without increasing compute per token, since only one expert activates. However, more experts also means more communication overhead in distributed settings and more routing complexity. Typical values in the Switch paper range from 128 to 2048, depending on the number of devices available.
• capacity_factor: Multiplier defining the expert buffer size ($C$). Controls the tradeoff between memory and token dropping. The recommended default of 1.25 works well for most training configurations, but you should increase it to 1.5 or 2.0 if you observe high drop rates during early training.
• balance_loss_weight: Coefficient for the auxiliary load balancing loss to prevent routing collapse. If set too high, it forces uniform routing at the cost of specialization. If set too low, routing collapse can occur. The paper found 0.01 to be a reliable default across a range of configurations.
• d_ff: Hidden dimension of the expert feed-forward networks. This should match the FFN dimension of the dense layers in your baseline transformer to ensure a fair parameter count comparison for each individual expert.