理解 PEER：并行专家的力量
扩展大型语言模型（LLM）的一大挑战在于 Transformer 模块中前馈层的计算与内存需求。MoE 通过以稀疏激活的专家模块替换这些层来应对，每个专家专长于任务的不同方面。该方法只激活与给定输入相关的专家，从而降低计算开销并提升效率。
当前 MoE 方法存在局限，例如固定路由器在新增专家时需要重新调整。因此，引入了一种新的路由方法，以可学习的索引替代固定路由器。
参数高效专家检索（PEER）减少了 MoE 层中的活跃参数数量，影响了预训练和推理期间的计算与激活内存消耗。
PEER 证明，通过采用合适的检索与路由机制，MoE 可扩展至数百万专家，从而降低训练和服务超大语言模型的成本与复杂度。
在上图中，输入查询首先经过乘积键检索以确定前k名专家。随后，这些入选专家基于其专业知识对输入进行处理，并在推理阶段对其输出进行聚合，生成模型的最终输出。
百万专家混合体：技术细节
PEER 的核心创新在于使用乘积键检索。其目标与传统 MoE 相同：找到最适合当前任务的前 k 名专家。然而，当专家数量极大（可能超过一百万）时，既有技术会变得计算昂贵或效率低下。
设想一个包含N名专家的场景，每个专家由一个d维向量表示。直接计算前k名专家，需要计算输入查询与全部N个专家键的相似度，其时间复杂度为 O(Nd)。当N非常大（例如N ≥ 10^6）时，代价将高得难以承受。
PEER 采用了一种巧妙策略：不再使用N个相互独立的d维专家键，而是将每个键拆分为两个相互独立的子集，每个子集维度为d/2。同样地，将查询向量分为两个子查询。随后，对这些子查询与子键之间的内积执行 top-k 操作。
这种笛卡尔积结构的键将计算复杂度从O(Nd)大幅降低到O((N^.5+ k2)d)，即便专家数量庞大，也能高效找出前 k 名专家。
PEER 层架构
参数高效专家检索（PEER）层是一种 MoE 架构，在路由器中使用乘积键，并以单神经元 MLP 作为专家。
一个 PEER 层由三个组件组成：
- 包含N名专家的池E
- 对应的一组N 个乘积键K
- 查询网络q
其工作流程如下：
- 给定输入查询x，检索一组前k名专家，这些专家的对应乘积键与查询q(x)的内积最大。
- 对这些检索到的前k名专家应用如 sigmoid 或 softmax 等激活函数，获得路由分数。
- 将专家输出按路由分数线性加权组合，得到最终输出。
本质上，PEER 层能高效识别与输入最相关的专家，使得在计算可控的前提下有效利用海量专家。这一创新是将 MoE 扩展至百万专家的关键支撑，为更强大、更高效的 LLM 铺平道路。
PEER 方法的优势
PEER 与 MoME 架构结合，相较传统 MoE 带来多项显著优势，推动 LLM 能力边界：
- 效率提升：通过利用海量专用专家，PEER 能更准确高效地检索相关信息，从而提升整体模型表现。
- 终身学习：PEER 促进持续学习与模型增量扩展。可以逐步添加新专家，在无需完全重新训练的情况下扩充模型知识库，从而保留既有知识。
- 可扩展性：PEER 解决了传统 MoE 架构在可扩展性上的限制，使构建如 GPT-MoE-1.8T 等万亿参数模型成为可能，为打造更强大、更能干的 LLM 打开空间。
MoME 的潜在应用
专家混合已经是业界广泛采用的模型范式，诸如 YouTube 等公司已将其用于推荐系统。MoME 的未来发展与此相似。在英伟达 GTC 2024 上也有所暗示，他们提到了 GPT 的 1.8 万亿参数模型。
百万专家混合体（MoME）及其 PEER 架构，尤其适用于需要广博知识与快速检索响应的复杂 NLP 任务。它解决了训练与服务超大语言模型所固有的可扩展性挑战，为其在计算机视觉、内容生成、推荐系统与智能计算等领域的应用开辟新可能。
挑战与局限
将模型扩展到百万专家在效率层面颇具吸引力，但管理如此庞大的网络也带来不少挑战。我们来看其中一些：
- 计算复杂度与效率：
- 专家检索成本：从百万规模的专家池中计算前k名专家可能计算开销巨大。即便采用乘积键检索等优化，如何高效管理检索流程仍然具有挑战。
- 训练开销：训练拥有海量专家的超大模型需要巨大的计算资源，并非所有研究机构或公司都能承担。
- 内存约束：
- 激活内存：随着专家数量增长，训练与推理过程中存储激活和中间结果所需的内存可能变得难以承受。
- 参数存储：存放一百万名专家需要巨大的内存容量，这会限制在资源受限设备上的部署。
- 正则化与稳定性：
- 过拟合：参数量如此之大，容易对训练数据过拟合。有效的正则化技术至关重要，但设计与实现并不容易。
- 稳定性：训练稳定性可能成为问题，对百万专家进行梯度与参数更新的管理会引发数值不稳定与收敛困难。
结语
本文探讨了百万专家混合体（MoME）技术——一种可扩展的大语言模型方法。
MoME 通过专用专家网络与 PEER 路由机制来提升效率与性能。
我们讨论了其核心组件、优势与潜在应用。若需更深入的技术细节与基准结果，请参阅相关研究论文。
更多信息请参阅这篇研究论文，了解技术细节与基准结果。
多年来，扩展 AI 模型主要意味着增加参数数量和训练数据。
这种方法虽然能提升性能，但也会显著提高计算成本。专家混合（MoE）被视为应对这一挑战的有前景方案，用稀疏激活的专家模块替代传统的稠密前馈层。
MoE 的工作方式是根据主题专长将任务分配给不同专家。每个专家都在特定数据集上进行强化训练以服务某一目的，另一个组件——门控网络——负责完成这些任务的分配。
虽然 MoE 模型优于采用前馈层的传统模型，但随着模型规模增大，在使用固定数量训练 token 的限制下，其效率可能会出现平台期。
解决这一问题需要高粒度架构并配备大量专家。然而，唯一支持超过一万名专家的现有架构——词专家混合体（MoWE）——具有语言相关性，并依赖固定的路由方案。
本文介绍的百万专家混合体（MoME），在这篇论文中提出，通过引入参数高效专家检索（PEER）架构，采用乘积键检索（product key retrieval）实现对海量专家的高效路由，从而应对上述挑战。