条件计算原理

随着大型神经网络，特别是 Transformer 模型，参数数量持续增加，训练和推理相关的计算成本成为一个重要瓶颈。标准的密集架构要求每个参数都参与每个输入 token 的计算。模型大小（参数数量）与计算负载（$FLOPs$）之间的这种关联限制了实际的可扩展性。如果一个密集层的参数数量翻倍，通常通过它的每个输入的计算量也会翻倍。

条件计算提供了一种不同的方法。其基本思想是根据输入本身，只为特定输入激活模型参数的一部分。网络不再通过一个庞大的计算块处理信息，而是动态地选择专门的计算路径。

原理：专业化与效率

设想一个包含众多专门子网络或“专家”的完整网络。对于特定的输入 token（或序列），也许只有其中几个专家拥有相关的知识或功能来有效处理它。条件计算允许模型识别并只调用这些相关的专家，而其他专家则保持非活跃状态。

这种方法受生物系统神经元稀疏激活的启发。在计算方面，主要优点是将模型容量与每输入计算成本解耦。我们可以显著增加模型的参数总量（添加更多专家），而无需按比例增加处理单个 token 所需的 $FLOPs$。

考虑一个简化比较：

• 密集层： 输入 $x$ 由权重矩阵 $W$ 处理。计算涉及 $xW$。如果 $W$ 的大小翻倍（例如，隐藏维度翻倍），则 $FLOPs$ 大致翻倍。
• 条件计算： 输入 $x$ 首先由路由机制分析。基于 $x$，路由器从更大的池 $\{W_1, W_2, ..., W_N\}$ 中选择参数块（专家）的一个小集合 $\{W_{i_1}, W_{i_2}, ..., W_{i_k}\}$，其中 $k \ll N$。计算只涉及选定的专家，例如 $\sum_{j=1}^{k} \alpha_j (x W_{i_j})$。即使所有 $N$ 个专家的参数总量很大，每个输入 $x$ 的计算主要取决于激活专家的规模和数量 ($k$)。

密集层与条件计算设置中的数据流比较。在条件路径中，路由器对于给定输入选择性地只激活某些专家（此处显示专家 1 和 K）。

扩展优势

主要优势源于这种关注点分离：

1. 容量增加： 模型可以包含显著更多的参数，潜在地学习更多样和专门的功能，而不会导致每 token 计算成本的过高增长。
2. 计算量恒定（相对）： 尽管总模型大小增长，通过固定激活专家数量 ($k$)，处理单个 token 的计算成本 ($FLOPs$) 可以大致保持不变（或增长慢得多）。

这种理论上的分离在下方可视化，比较了参数数量和计算成本可能如何扩展。

密集与条件计算模型中参数数量和每 token 计算成本如何扩展的比较（假设在条件情况下激活专家数量 $k$ 是固定的）。请注意 Y 轴的对数刻度。条件计算允许参数数量增长，而每 token 计算量不按比例增加。

与专家混合模型的关联

专家混合模型（MoE）是深度学习架构，尤其是 Transformer 中条件计算原理的一种直接有效的实现。在 MoE 层中，“专家”通常是前馈网络（FFN），而路由机制是一个小型可训练神经网络，称为“门控网络”或“路由器”。

路由器 $G(x)$ 检查输入表示 $x$（通常是 Transformer 中自注意力层的输出），并生成概率或权重，指示哪些专家应该处理此输入。选定专家 $E_i(x)$ 的输出随后被组合，通常由路由器的分数加权，正如章节概述中所述：

$y = \sum_{i=1}^{N} G(x)_i E_i(x)$

在这里，$G(x)_i$ 体现了条件性。理想情况下，对于给定的 $x$，大多数 $G(x)_i$ 值是零（或接近零），从而确保只有稀疏的专家子集 $E_i$ 对输出 $y$ 有显著贡献。

尽管原理上简洁，但通过 MoE 实现条件计算的优势也带来了实际挑战，包括路由决策、专家间的负载均衡以及分布式环境中的高效实现。这些复杂性是后续章节的重点。然而，理解条件计算的核心原理，为理解现代 MoE 模型中采用的设计选择和优化策略提供了必要的基础。