引言
2017 年，谷歌提出的 Transformer 架构彻底改变了自然语言处理领域。然而，仅仅一年后，谷歌又发布了 Universal Transformer (UT)，它通过跨层共享参数来实现深度循环，在逻辑推理等组合问题上展现出优于 Transformer 的性能。但 UT 的计算效率远低于标准 Transformer，使其在以参数为王的语言建模任务中难以施展拳脚。近日，包括 LSTM 之父 Jürgen Schmidhuber 和斯坦福大学教授 Christopher Manning 在内的研究者，提出了一种名为 Mixture-of-Experts Universal Transformers (MoEUT) 的新架构，它利用 MoE 的优势，有效地解决了 UT 的计算效率问题，并在多个语言建模任务上超越了标准 Transformer。
MoEUT：高效升级 Universal Transformer
MoEUT 是一种混合专家 (MoE) 架构，它允许 UT 以计算和内存高效的方式扩展。研究者利用了 MoE 在前馈和自注意力层方面的最新进展，并结合了以下两项创新：
- 层分组 (Layer Grouping)：循环堆叠 MoE 层组，以提高模型的效率。
- Peri-layernorm 方案：将层范数应用于紧接 sigmoid 或 softmax 激活之前的层，专门为共享层 MoE 架构设计，并得到实证验证。
这些创新使得 MoEUT 能够构建参数和资源高效的 UT 语言模型，在性能上超越标准 Transformer，同时在计算力和内存需求方面也更具优势。
MoEUT 架构概览
MoEUT 架构的核心是利用 MoE 来解决 UT 的参数计算比问题。研究者使用 MoE 前馈网络块 (FFN) 和 MoE 自注意力层，并引入了层分组和Peri-layernorm 方案，以提高模型的效率和性能。
- MoE 前馈块：将前馈块划分为多个专家，每个专家都有两组权重。MoE 层仅选择得分最高的 K 个专家来产生层输出，并通过平衡损失来避免训练期间的损失激增。
- MoE 自注意力层：使用 SwitchHead 方法将 MoE 扩展到注意力层，每个头都有一个与之关联的查询和键投影以及 N_A 值和输出投影，它们针对每个输入进行动态选择。
实验结果
研究者在多个语言建模数据集和下游任务上测试了 MoEUT 的性能，结果表明：
- MoEUT 在 C4、SlimPajama 和 peS2o 语言建模数据集上超越了标准 Transformer，并且在所有规模（最高 10 亿参数）上对算力和内存的需求更低。
- MoEUT 在 BLiMP、儿童图书测试、Lambada、HellaSwag、PIQA 和 ARC-E 等下游任务上展现出良好的零样本性能。
结论
MoEUT 架构的出现为 UT语言模型的应用打开了新的局面。它有效地解决了 UT 的计算效率问题，使其在以参数为王的语言建模任务中展现出更强的竞争力。MoEUT 的成功也证明了 MoE 在 Transformer 架构中的巨大潜力，未来将会在更多领域得到应用。
参考文献
- MoEUT: Mixture-of-Experts Universal Transformers
- 论文地址：https://arxiv.org/pdf/2405.16039
- 代码地址：https://github.com/robertcsordas/moeut
未来展望
MoEUT 的出现为 UT 语言模型的应用打开了新的局面，未来将会在以下几个方面继续探索：
- 探索 MoEUT 在其他领域，例如图像识别和语音识别中的应用。
- 进一步优化 MoEUT 架构，以提高其效率和性能。
- 研究 MoEUT 的可解释性，以更好地理解其工作机制。
MoEUT 的出现标志着 Transformer 架构的又一次重大突破，它将推动自然语言处理领域的发展，并为人工智能技术的应用带来更多可能性。
Views: 0