Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Auxiliary Balancing Loss › Implementation

Michael BrenndoerferNovember 17, 20258 min read

Language AI HandbookAll chapters

Part 33 of 67Mixture of Experts

Chapter 244Auxiliary Balancing Loss › Implementation

Reading section 8 of 13 · Implementation · 8 minSection 8 of 13 · 8 min

Section 8 of 13 · 8 min read

Auxiliary Balancing Loss › ImplementationOne section · Sections in this chapter8 / 13 ↓

Sections in this chapter13

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

Let's implement the auxiliary balancing loss step by step. We'll create a module that can be integrated into any MoE training loop.

In[10]:

Code

import torch
import torch.nn.functional as F


def compute_load_balancing_loss(
    router_probs: torch.Tensor, expert_indices: torch.Tensor, num_experts: int
) -> torch.Tensor:
    """
    Compute the auxiliary load balancing loss.

    Args:
        router_probs: Gate probabilities, shape (batch_size, seq_len, num_experts)
        expert_indices: Selected expert indices, shape (batch_size, seq_len)
        num_experts: Number of expert networks

    Returns:
        Scalar auxiliary loss
    """
    # Flatten batch and sequence dimensions
    batch_size, seq_len, _ = router_probs.shape
    num_tokens = batch_size * seq_len

    router_probs_flat = router_probs.view(-1, num_experts)  # (N, E)
    expert_indices_flat = expert_indices.view(-1)  # (N,)

    # Compute fraction of tokens routed to each expert (f_i)
    # One-hot encode the expert selections and average
    expert_mask = F.one_hot(expert_indices_flat, num_experts).float()  # (N, E)
    tokens_per_expert = expert_mask.sum(dim=0)  # (E,)
    f = tokens_per_expert / num_tokens  # (E,)

    # Compute fraction of router probability per expert (P_i)
    P = router_probs_flat.mean(dim=0)  # (E,)

    # Load balancing loss: scaled dot product
    load_balancing_loss = num_experts * (f * P).sum()

    return load_balancing_loss

This function computes the exact formulation from the Switch Transformer paper. Let's verify it works correctly with a simple example.

In[11]:

Code

import torch
import torch.nn.functional as F

# Create example data: 4 tokens, 3 experts
batch_size, seq_len, num_experts = 2, 2, 3
num_tokens = batch_size * seq_len

# Simulate router probabilities (softmax output)
torch.manual_seed(42)
router_logits = torch.randn(batch_size, seq_len, num_experts)
router_probs = F.softmax(router_logits, dim=-1)

# Select top expert for each token
expert_indices = router_probs.argmax(dim=-1)

# Compute loss and breakdown for inspection
loss = compute_load_balancing_loss(router_probs, expert_indices, num_experts)

router_probs_flat = router_probs.view(-1, num_experts)
expert_indices_flat = expert_indices.view(-1)
expert_mask = F.one_hot(expert_indices_flat, num_experts).float()
f = expert_mask.sum(dim=0) / num_tokens
P = router_probs_flat.mean(dim=0)
f_P_product = f * P

Out[12]:

Console

Router probabilities:
tensor([[0.3683, 0.2992, 0.3325],
        [0.5215, 0.1348, 0.3438],
        [0.8114, 0.0471, 0.1415],
        [0.2484, 0.3246, 0.4271]])

Selected experts: [0, 0, 0, 2]

Load balancing loss: 1.3300

Fraction of tokens per expert (f): ['0.7500', '0.0000', '0.2500']
Average probability per expert (P): ['0.4874', '0.2014', '0.3112']
f * P products: ['0.3655', '0.0000', '0.0778']

The loss is above 1.0 because expert 2 receives most tokens (high $f_2$ ) and also has high average probability (high $P_2$ ). With perfect balance, we'd expect $f = P = [0.333, 0.333, 0.333]$ and loss $= 1.0$ .

Handling Top-K RoutingLink Copied

When using top-2 or top-k routing, each token contributes to multiple experts' load. We need to adjust the computation to account for the fact that each token sends computation to $k$ experts rather than just one. The normalization factor changes accordingly, since the total number of expert assignments is $N \times k$ rather than just $N$ .

In[13]:

Code

import torch


def compute_load_balancing_loss_topk(
    router_probs: torch.Tensor, expert_indices: torch.Tensor, num_experts: int
) -> torch.Tensor:
    """
    Compute auxiliary loss for top-k routing.

    Args:
        router_probs: Gate probabilities, shape (batch_size, seq_len, num_experts)
        expert_indices: Selected expert indices, shape (batch_size, seq_len, k)
        num_experts: Number of expert networks

    Returns:
        Scalar auxiliary loss
    """
    batch_size, seq_len, k = expert_indices.shape
    num_tokens = batch_size * seq_len

    router_probs_flat = router_probs.view(-1, num_experts)  # (N, E)
    expert_indices_flat = expert_indices.view(-1, k)  # (N, k)

    # Count how many times each expert is selected across all tokens
    # Each token contributes k selections total
    tokens_per_expert = torch.zeros(num_experts, device=router_probs.device)
    for i in range(k):
        expert_counts = torch.bincount(
            expert_indices_flat[:, i], minlength=num_experts
        ).float()
        tokens_per_expert += expert_counts

    # Normalize by total selections (N * k)
    f = tokens_per_expert / (num_tokens * k)

    # Average router probability per expert
    P = router_probs_flat.mean(dim=0)

    return num_experts * (f * P).sum()

In[14]:

Code

# Example with top-2 routing
k = 2
_, topk_indices = router_probs.topk(k, dim=-1)

loss_topk = compute_load_balancing_loss_topk(
    router_probs, topk_indices, num_experts
)

Out[15]:

Console

Top-2 expert indices:
tensor([[[0, 2],
         [0, 2]],

        [[0, 2],
         [2, 1]]])

Top-k load balancing loss: 1.0907

The top-k loss is typically lower than top-1 because selecting two experts per token naturally spreads load more evenly.

Integrating with TrainingLink Copied

The following training loop incorporates the auxiliary loss:

In[16]:

Code

import torch
import torch.nn as nn
import torch.nn.functional as F


class MoELayer(nn.Module):
    """Simplified MoE layer with auxiliary loss tracking."""

    def __init__(self, hidden_size: int, num_experts: int, expert_size: int):
        super().__init__()
        self.num_experts = num_experts

        # Router network
        self.router = nn.Linear(hidden_size, num_experts)

        # Expert networks (simplified as linear layers)
        self.experts = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(hidden_size, expert_size),
                    nn.GELU(),
                    nn.Linear(expert_size, hidden_size),
                )
                for _ in range(num_experts)
            ]
        )

        # Store auxiliary loss
        self.aux_loss = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, hidden_size = x.shape

        # Compute router probabilities
        router_logits = self.router(x)  # (B, S, E)
        router_probs = F.softmax(router_logits, dim=-1)

        # Select top-1 expert
        expert_indices = router_probs.argmax(dim=-1)  # (B, S)

        # Compute auxiliary loss
        self.aux_loss = compute_load_balancing_loss(
            router_probs, expert_indices, self.num_experts
        )

        # Route tokens to experts (simplified dense computation for clarity)
        output = torch.zeros_like(x)
        x_flat = x.view(-1, hidden_size)
        output_flat = output.view(-1, hidden_size)
        indices_flat = expert_indices.view(-1)

        for i in range(self.num_experts):
            mask = indices_flat == i
            if mask.any():
                output_flat[mask] = self.experts[i](x_flat[mask])

        return output.view(batch_size, seq_len, hidden_size)

In[17]:

Code

import torch
import torch.nn.functional as F


def train_step(model, x, y, optimizer, aux_loss_coef=0.01):
    """Single training step with auxiliary loss."""
    optimizer.zero_grad()

    # Forward pass
    output = model(x)

    # Task loss (e.g., cross-entropy for language modeling)
    task_loss = F.mse_loss(output, y)  # Simplified example

    # Collect auxiliary losses from all MoE layers
    total_aux_loss = 0
    for module in model.modules():
        if isinstance(module, MoELayer) and module.aux_loss is not None:
            total_aux_loss += module.aux_loss

    # Combined loss
    total_loss = task_loss + aux_loss_coef * total_aux_loss

    # Backward and update
    total_loss.backward()
    optimizer.step()

    return {
        "task_loss": task_loss.item(),
        "aux_loss": total_aux_loss.item()
        if isinstance(total_aux_loss, torch.Tensor)
        else total_aux_loss,
        "total_loss": total_loss.item(),
    }

In[18]:

Code

import torch

# Create model and run training step
hidden_size, num_experts, expert_size = 64, 4, 128
model = MoELayer(hidden_size, num_experts, expert_size)

# Random input and target
x = torch.randn(2, 8, hidden_size)
y = torch.randn(2, 8, hidden_size)

optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
metrics = train_step(model, x, y, optimizer, aux_loss_coef=0.01)

Out[19]:

Console

Training metrics:
  task_loss: 1.0804
  aux_loss: 1.0224
  total_loss: 1.0906

The total loss is the sum of task loss and weighted auxiliary loss. The auxiliary loss is initially high (>1.0). This reflects the random initialization of the router.

Monitoring Expert UtilizationLink Copied

Tracking expert utilization during training helps diagnose load balancing issues. The key metrics to watch are the auxiliary loss value and the actual distribution of tokens across experts. The auxiliary loss is a differentiable proxy for this distribution, but the actual distribution is what matters for inference efficiency and model capacity utilization.

In[20]:

Code

import torch


def compute_expert_utilization(
    router_probs: torch.Tensor, expert_indices: torch.Tensor, num_experts: int
) -> dict:
    """Compute expert utilization statistics."""
    num_tokens = expert_indices.numel()
    expert_indices_flat = expert_indices.view(-1)

    # Count tokens per expert
    counts = torch.bincount(expert_indices_flat, minlength=num_experts).float()
    fractions = counts / num_tokens

    # Compute statistics
    mean_fraction = 1.0 / num_experts  # Expected under perfect balance
    max_fraction = fractions.max().item()
    min_fraction = fractions.min().item()

    # Imbalance ratio: max/min (1.0 = perfect balance)
    imbalance_ratio = max_fraction / max(min_fraction, 1e-8)

    return {
        "tokens_per_expert": counts.tolist(),
        "fraction_per_expert": fractions.tolist(),
        "max_fraction": max_fraction,
        "min_fraction": min_fraction,
        "imbalance_ratio": imbalance_ratio,
    }

In[21]:

Code

import torch
import torch.nn.functional as F

# Generate more data to see utilization patterns
torch.manual_seed(123)
batch_size, seq_len, num_experts = 8, 32, 4
router_logits = torch.randn(batch_size, seq_len, num_experts)
router_probs = F.softmax(router_logits, dim=-1)
expert_indices = router_probs.argmax(dim=-1)

utilization = compute_expert_utilization(
    router_probs, expert_indices, num_experts
)

Out[22]:

Console

Expert utilization statistics:
  Tokens per expert: [77.0, 65.0, 62.0, 52.0]
  Fraction per expert: ['0.301', '0.254', '0.242', '0.203']
  Max/min fractions: 0.301 / 0.203
  Imbalance ratio: 1.48x

With random initialization, some imbalance naturally occurs. The auxiliary loss should push this ratio closer to 1.0 during training.

Implementation NotesLink Copied

Several practical issues arise when integrating auxiliary losses into real training pipelines. First, when a model has multiple MoE layers, each layer computes its own auxiliary loss, and the total auxiliary contribution should sum across all layers. The pattern shown in the train_step function above handles this correctly by iterating over all modules and accumulating. In a transformer with MoE layers at every other block, a 24-layer model might have 12 MoE layers, each contributing an auxiliary loss. The coefficient $\alpha$ therefore applies to the sum of all these contributions, which means the effective balancing pressure scales with the number of MoE layers. Some implementations divide by the number of MoE layers to normalize this effect.

Second, the auxiliary loss should not be included in validation metrics. When you compute validation perplexity to assess model quality, you should report only the task loss, not the total loss including balancing. This is because the auxiliary loss measures a training artifact (routing distribution), not the quality of the model's predictions. Including it would obscure real performance trends during model development.

Third, gradient clipping can interact unexpectedly with auxiliary losses. If you clip gradients by global norm, the auxiliary loss gradients are included in the norm computation. A large auxiliary loss during early training can dominate the gradient norm, causing excessive clipping of task loss gradients. Some practitioners compute gradient norms separately for the task and auxiliary components, or apply the auxiliary loss only after a warmup period when the model has begun to stabilize.