Transformer——Q97 Switch Transformer的专家负载均衡损失公式推导

该问题归类到Transformer架构问题集——架构变体——稀疏/混合专家。请参考LLM数学推导——Transformer架构问题集。

1. 问题背景：当专家 “忙闲不均” 时，MoE 如何破局？

在 Switch Transformer 构建的混合专家世界里，每个输入样本都会触发一场 “专家选秀”：门控网络像星探，根据输入特征为每个样本挑选 1 个或多个 “最合适” 的专家。理想状态下，1000 个专家应像交响乐团的乐手，各自在合适的时刻奏响乐章。但现实却像流量明星效应 —— 少数专家被高频翻牌，承担 70% 以上的负载，而大量 “冷门” 专家长期坐冷板凳，参数更新频率不足热门专家的 1/10。

这种 “马太效应” 带来三重困境：

• 参数浪费：未激活专家的万亿参数沦为摆设，违背 MoE “用稀疏激活驾驭大规模参数” 的设计初衷

• 训练失衡：过载专家因输入单一过拟合，闲置专家因缺乏 “锻炼” 导致参数退化，模型在低频任务上准确率暴跌 20%

• 硬件低效：GPU 资源向少数专家倾斜，设备利用率从 75% 降至 40%，电费账单飙升却换不来性能提升

负载均衡损失函数正是为打破这种失衡而生，它像一位严格的班主任，强制要求 “优生” 分享机会，“差生” 获得关注，让整个专家团队在协作中迸发最大能量。

2. 技术原理：从数学推导看均衡策略如何 “劫富济贫”

2.1 专家负载的量化 “标尺”

要实现均衡，先得看清失衡。假设共有m个专家，门控网络输出概率分布，每个样本激活k个专家（通过 one-hot 掩码M表示）：

• 期望负载：，理想值为，如同专家的 “月度 KPI”，过低或过高都需警惕

• 负载熵：，单样本选择的 “分散度” 指标，值越高说明选择越均匀，避免 “所有样本都扎堆选少数专家” 的假均衡

• 标准差：，全局失衡的 “体温计”，意味着系统已 “发烧”

2.2 损失函数的三层 “均衡魔法”

2.2.1 第一层：全局均衡的 “长期规划”（KL 散度约束）<