Switch Transformer › Complete Switch Layer Implementation

class SwitchLayer(nn.Module):
    """
    Switch Transformer layer: routes each token to one expert.
    """

    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_experts: int,
        capacity_factor: float = 1.25,
        dropout: float = 0.1,
        balance_loss_weight: float = 0.01,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor
        self.balance_loss_weight = balance_loss_weight

        # Router
        self.router = SwitchRouter(d_model, num_experts)

        # Expert networks
        self.experts = nn.ModuleList(
            [SwitchExpert(d_model, d_ff, dropout) for _ in range(num_experts)]
        )

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        batch_size, seq_len, d_model = x.shape
        total_tokens = batch_size * seq_len

        # Get routing decisions
        expert_indices, expert_weights, router_logits = self.router(x)

        # Flatten for easier processing
        x_flat = x.view(total_tokens, d_model)
        expert_indices_flat = expert_indices.view(total_tokens)
        expert_weights_flat = expert_weights.view(total_tokens)

        # Calculate capacity
        capacity = compute_capacity(
            batch_size, seq_len, self.num_experts, self.capacity_factor
        )

        # Initialize output
        output = torch.zeros_like(x_flat)

        # Track which tokens are processed for load balancing
        tokens_per_expert = torch.zeros(self.num_experts, device=x.device)

        # Process each expert
        for expert_idx in range(self.num_experts):
            # Find tokens assigned to this expert
            mask = expert_indices_flat == expert_idx
            token_indices = mask.nonzero(as_tuple=True)[0]

            # Apply capacity limit
            num_tokens = min(len(token_indices), capacity)
            token_indices = token_indices[:num_tokens]
            tokens_per_expert[expert_idx] = num_tokens

            if num_tokens > 0:
                # Get token representations
                expert_input = x_flat[token_indices]

                # Process through expert
                expert_output = self.experts[expert_idx](expert_input)

                # Weight by router probability
                weights = expert_weights_flat[token_indices].unsqueeze(-1)
                output[token_indices] = weights * expert_output

        # Reshape output
        output = output.view(batch_size, seq_len, d_model)

        # Add residual connection (necessary for dropped tokens!)
        output = output + x

        # Compute load balancing loss
        balance_loss = self._compute_balance_loss(
            router_logits.view(total_tokens, -1), expert_indices_flat
        )

        return output, balance_loss * self.balance_loss_weight

    def _compute_balance_loss(
        self, router_logits: torch.Tensor, expert_indices: torch.Tensor
    ) -> torch.Tensor:
        """Auxiliary loss to encourage balanced routing."""
        # Fraction of tokens routed to each expert
        num_tokens = router_logits.shape[0]

        # f_i: fraction of tokens to expert i
        expert_counts = torch.zeros(
            self.num_experts, device=router_logits.device
        )
        for i in range(self.num_experts):
            expert_counts[i] = (expert_indices == i).float().sum()
        f = expert_counts / num_tokens

        # P_i: mean router probability for expert i
        router_probs = F.softmax(router_logits, dim=-1)
        P = router_probs.mean(dim=0)

        # Balance loss: sum of f_i * P_i, scaled by num_experts
        balance_loss = self.num_experts * torch.sum(f * P)

        return balance_loss

# Setup parameters
B, S, D = 4, 64, 256
num_experts = 8

# Create a Switch layer
switch_layer = SwitchLayer(
    d_model=D, d_ff=1024, num_experts=num_experts, capacity_factor=1.25
)

# Test input
test_input = torch.randn(B, S, D)
output, aux_loss = switch_layer(test_input)

# Analyze routing distribution
with torch.no_grad():
    expert_indices, _, _ = switch_layer.router(test_input)
    expert_counts = torch.bincount(
        expert_indices.flatten(), minlength=num_experts
    )