摘要

本文介绍了一种名为PEER（Parameter Efficient Expert Retrieval）的新型层设计，它利用产品键技术从大量微型专家（超过一百万）中进行稀疏检索，以解决标准transformer架构中前馈（FFW）层随着隐藏层宽度增加而导致的计算成本和激活内存线性增长的问题。PEER层在语言建模任务中展示了优于密集FFW层和粗粒度MoE（Mixture-of-Experts）的性能-计算权衡。通过实现大量专家的高效利用，PEER解锁了transformer模型进一步扩展的潜力，同时保持了计算效率。

原理

PEER层是一种混合专家架构，使用产品键（product keys）在路由器中，并以单神经元多层感知器（MLP）作为专家。其工作原理包括三个部分：一个包含N个专家的池E，每个专家ei与f具有相同的签名；一组对应的N个产品键K；以及一个将输入向量x映射到查询向量q(x)的查询网络q。给定输入x，首先检索与查询q(x)内积最高的k个专家对应的产品键子集，然后应用非线性激活函数（如softmax或sigmoid）到这些top k专家的查询-键内积上以获得路由器分数，最后通过路由器分数加权线性组合专家输出以计算输出。

流程

PEER层的工作流程如下：首先，输入x通过查询网络q映射到查询向量q(x)，然后与产品键集K中的所有键进行内积计算，选择内积最高的k个专家。接着，对这些top k专家的查询-键内积应用非线性激活函数以获得路由器分数，最后，通过路由器分数加权线性组合这些专家的输出，得到最终的输出。具体实现中，产品键通过将两个独立集合的子键连接形成，从而允许高效地找到top k专家。

应用

PEER层的设计不仅适用于语言建模任务，还具有广泛的潜在应用范围，包括但不限于终身学习、知识密集型任务和高效模型扩展。通过解耦计算成本与参数计数，PEER为大规模模型的扩展提供了新的可能性，特别是在需要处理大量数据流和持续学习的场景中。