Gating Networks › Code Implementation

import numpy as np
import torch

# Set random seed for reproducibility
torch.manual_seed(42)
np.random.seed(42)

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


class Router(nn.Module):
    """
    Gating network for Mixture of Experts.
    Routes tokens to top-k experts based on learned routing scores.
    """

    def __init__(self, d_model: int, num_experts: int, top_k: int = 2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k

        # Single linear layer: projects from d_model to num_experts
        # No bias term, common in production MoE implementations
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x: torch.Tensor):
        """
        Compute routing weights for input tokens.

        Args:
            x: Input tensor of shape (batch_size, seq_len, d_model)

        Returns:
            routing_weights: Weights for selected experts (batch, seq, top_k)
            selected_experts: Indices of selected experts (batch, seq, top_k)
            router_probs: Full probability distribution (batch, seq, num_experts)
        """
        # Compute routing logits
        logits = self.gate(x)  # (batch, seq, num_experts)

        # Convert to probabilities via softmax
        router_probs = F.softmax(logits, dim=-1)

        # Select top-k experts
        routing_weights, selected_experts = torch.topk(
            router_probs, self.top_k, dim=-1
        )

        # Renormalize weights to sum to 1
        routing_weights = routing_weights / routing_weights.sum(
            dim=-1, keepdim=True
        )

        return routing_weights, selected_experts, router_probs

# Create a router with 8 experts, selecting top-2
d_model = 256
num_experts = 8
router = Router(d_model, num_experts, top_k=2)

# Simulate a batch of token representations
batch_size = 2
seq_len = 4
x = torch.randn(batch_size, seq_len, d_model)

# Get routing decisions
routing_weights, selected_experts, router_probs = router(x)

def analyze_routing_distribution(router, num_samples=1000, d_model=256):
    """
    Analyze how the router distributes tokens across experts.
    """
    # Generate random token representations
    x = torch.randn(num_samples, 1, d_model)

    with torch.no_grad():
        _, selected_experts, router_probs = router(x)

    # Count how often each expert is selected
    expert_counts = torch.zeros(router.num_experts)
    for i in range(router.num_experts):
        expert_counts[i] = (selected_experts == i).sum().item()

    # Compute mean probability assigned to each expert
    mean_probs = router_probs.squeeze(1).mean(dim=0)

    return expert_counts, mean_probs

# Analyze an untrained router
counts, mean_probs = analyze_routing_distribution(router)

def compute_routing_entropy(router_probs):
    """
    Compute entropy of routing distributions.
    Higher entropy = more uniform, lower = more concentrated.
    """
    # Avoid log(0)
    probs = router_probs.clamp(min=1e-10)
    entropy = -torch.sum(probs * torch.log(probs), dim=-1)
    return entropy


# Compute entropy for our sample
_, _, router_probs = router(x)
entropies = compute_routing_entropy(router_probs)

# Maximum possible entropy (uniform distribution over 8 experts)
max_entropy = np.log(num_experts)

# Create synthetic domain representations
# Imagine these are averaged representations from different document types
num_domains = 4
prototype_scale = 1.0
noise_std = 0.8
batch_size_per_domain = 64

# NOTE: If you crank up prototype_scale and reduce noise_std, this becomes
# an almost trivial supervised classification problem and the loss can drop
# to ~0 in just a few optimization steps (softmax saturation).
domain_prototypes = torch.randn(num_domains, d_model) * prototype_scale

# Create a router with one expert per domain
domain_router = Router(d_model, num_domains, top_k=1)

# Simulate "labeled" training: each sample comes from a known domain
# We'll train the router to route domain samples to corresponding experts

z_loss_coef = 1e-2
optimizer = torch.optim.Adam(
    domain_router.parameters(), lr=0.005, weight_decay=1e-4
)

losses = []
for epoch in range(200):
    epoch_loss = 0
    for domain_idx in range(num_domains):
        # Generate samples from this domain (prototype + noise)
        samples = (
            domain_prototypes[domain_idx]
            + torch.randn(batch_size_per_domain, d_model) * noise_std
        )
        samples = samples.unsqueeze(1)  # Add seq dimension

        # Train on logits directly for numerical stability
        logits = domain_router.gate(samples).squeeze(1)  # (batch, num_domains)
        labels = torch.full(
            (logits.shape[0],),
            domain_idx,
            dtype=torch.long,
            device=logits.device,
        )

        # Cross-entropy: encourage routing to the correct expert
        nll = F.cross_entropy(logits, labels)

        # Router z-loss: discourages runaway logits (keeps softmax from saturating)
        z = torch.logsumexp(logits, dim=-1)
        z_loss = (z**2).mean()

        loss = nll + z_loss_coef * z_loss

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()

    losses.append(epoch_loss / num_domains)

# Test routing accuracy on new samples
correct = 0
total = 0
domain_accuracies = []

with torch.no_grad():
    for domain_idx in range(num_domains):
        samples = (
            domain_prototypes[domain_idx] + torch.randn(100, d_model) * 0.3
        )
        samples = samples.unsqueeze(1)

        _, selected, _ = domain_router(samples)
        selected = selected.squeeze()

        # Calculate accuracy for this domain
        domain_acc = (selected == domain_idx).float().mean().item()
        domain_accuracies.append(domain_acc)

        correct += (selected == domain_idx).sum().item()
        total += 100

# Extract router weights for visualization
weights = domain_router.gate.weight.detach().numpy()