Transformer——Q99 推导稀疏门控（Sparse Gating）的Top-k选择梯度近似

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

1. 问题背景：当离散选择遇到连续梯度 ——Top-k 门控的梯度困境

在混合专家模型（MoE）的稀疏门控机制中，Top-k 选择（如每个样本激活概率最高的 2-4 个专家）是实现 “稀疏激活” 的核心操作。然而，Top-k 本质是离散的硬选择（选择概率最高的 k 个专家，其余置零），这种 “非黑即白” 的决策在反向传播时会遇到梯度消失问题 —— 未选中专家的梯度为零，选中专家的梯度依赖于不连续的指示函数，导致优化器难以有效更新门控网络参数。

举个直观例子：假设门控输出概率为 [0.3, 0.3, 0.2, 0.2]，Top-2 选择前两个专家（掩码 [1,1,0,0]）。但在反向传播时，梯度无法告知门控网络 “第三个专家的概率需要提高多少才能进入 Top-2”，因为离散选择切断了概率与掩码之间的连续映射。梯度近似方法正是为解决这一困境而生，通过平滑的连续函数近似离散选择，让梯度能够流经门控网络，实现端到端训练。

2. 技术原理：从离散选择到连续近似的梯度桥梁

2.1 Top-k 选择的前向传播与梯度困境

假设门控网络输出未归一化得分向量，Top-k 操作选择得分最高的 k 个专家，生成 one-hot 掩码：

梯度困境：

• 离散掩码对的导数几乎处处为零（仅当是 top-k 得分边界时可能非零，但概率为零）

• 传统反向传播无法更新未选中专家的得分，导致门控网络优化失效

2.2 Gumbel-Softmax：用噪声实现梯度桥接

2.2.1 松弛 Top-k 选择

引入 Gumbel 噪声，对得分添加扰动：

选择前 k 个带噪声的得分对应的专家，生成软掩码，其中