MoE Top-K Routing & Auxiliary Load Balancing Loss

• Ensure proper multi-dimensional tensor broadcasting and keepdims retention.
• Enforce numerical guards (eps clamping and overflow thresholds) during exponentiation and division.
• Verify train versus eval mode behavioral distinctions (e.g. frozen running statistics and dropout bypass).

💻 Industrial Code Implementation

import numpy as np

def moe_topk_routing_and_aux_loss(
    x: np.ndarray,          # (N, D) 输入 Token 特征
    W_g: np.ndarray,        # (D, num_experts) 门控路由权重
    top_k: int = 2,
    alpha: float = 0.01
):
    N, D = x.shape
    num_experts = W_g.shape[1]
    
    # 1. 计算所有专家的原始门控 Logits: (N, E)
    router_logits = x @ W_g
    
    # 2. 计算完整 Softmax 概率矩阵 P: (N, E) (用于辅助损失计算)
    l_max = np.max(router_logits, axis=-1, keepdims=True)
    exp_l = np.exp(router_logits - l_max)
    P_matrix = exp_l / np.sum(exp_l, axis=-1, keepdims=True)
    
    # 3. 选取每个 Token 的 Top-K 专家索引与分数
    topk_indices = np.argsort(-router_logits, axis=-1)[:, :top_k] # (N, k)
    
    # 提取 top_k 的 logits 并做局部 Softmax 归一化权重
    topk_logits = np.take_along_axis(router_logits, topk_indices, axis=-1)
    exp_topk = np.exp(topk_logits - np.max(topk_logits, axis=-1, keepdims=True))
    topk_weights = exp_topk / np.sum(exp_topk, axis=-1, keepdims=True) # (N, k)
    
    # 4. 计算负载均衡辅助损失 (Auxiliary Loss)
    # P_i: 全局平均分配概率 (E,)
    P_i = np.mean(P_matrix, axis=0)
    
    # f_i: 实际被分派给每个专家的 Token 比例 (E,)
    # 将 topk_indices 展平统计频次
    counts = np.bincount(topk_indices.reshape(-1), minlength=num_experts)
    f_i = counts / (N * top_k)
    
    aux_loss = alpha * num_experts * np.sum(f_i * P_i)
    
    return topk_indices, topk_weights, float(aux_loss)

import numpy as np
N, D, E = 10, 8, 4
x = np.random.randn(N, D)
Wg = np.random.randn(D, E)
idx, weights, loss = moe_topk_routing_and_aux_loss(x, Wg, top_k=2)
assert idx.shape == (N, 2)
assert weights.shape == (N, 2)
assert np.allclose(weights.sum(axis=-1), 1.0), "局部权重之和必须为 1"
assert loss > 0, "辅助损失必须为正"
print("✓ MoE Top-K 门控与辅助损失自测通过")