[论文解读] BASE Layers: Simplifying Training of Large, Sparse Models

Mike Lewis, Shruti Bhosale+3 more|arXiv (Cornell University)|Mar 30, 2021

Topic Modeling参考文献 29被引用 63

一句话总结

BASE 层通过线性分配问题引入一个最优、平衡的 token-to-expert 分配，实现每个 token 由单个 expert 路由的稀疏化，且无额外损失或超参数，从而提高计算效率。

ABSTRACT

We introduce a new balanced assignment of experts (BASE) layer for large language models that greatly simplifies existing high capacity sparse layers. Sparse layers can dramatically improve the efficiency of training and inference by routing each token to specialized expert modules that contain only a small fraction of the model parameters. However, it can be difficult to learn balanced routing functions that make full use of the available experts; existing approaches typically use routing heuristics or auxiliary expert-balancing loss functions. In contrast, we formulate token-to-expert allocation as a linear assignment problem, allowing an optimal assignment in which each expert receives an equal number of tokens. This optimal assignment scheme improves efficiency by guaranteeing balanced compute loads, and also simplifies training by not requiring any new hyperparameters or auxiliary losses. Code is publicly released at https://github.com/pytorch/fairseq/

研究动机与目标

• 激励使用稀疏专家模型以扩展语言模型规模，同时降低训练成本。
• 提出一个平衡、线性分配的 BASE 层，消除平衡损失和容量因子。
• 证明单个 token 对应一个 expert 的路由配合平衡分配，在大模型中能实现强计算效率。

提出的方法

• 将 token-to-expert 分配建模为线性分配问题，以确保每个 expert 处理等量的 token。
• 使用拍卖算法在跨工作者并行求解分配，并通过 token 置换处理鲁棒性边界情况。
• 实现一个 BASE 层，使每个 token 路由到单个 expert，然后通过残差连接对 expert 输出进行软混合。
• 使用平衡的 token 分配进行训练，以最大化吞吐量并避免额外的平衡损失，同时在推理阶段使用贪婪的 expert 选择进行测试。
• 将计算效率以在固定 GPU 运行时间下的 perplexity 来衡量，并与密集和先前的稀疏方法进行比较。

实验结果

研究问题

• RQ1 基于线性分配的路由方案能否在不增加额外损失项的情况下实现专家使用的平衡？
• RQ2 单个 expert 的 BASE 层在大型模型中是否能达到或超过现有密集和稀疏方法的效率？
• RQ3 BASE 的放置和大小如何影响性能和计算效率？
• RQ4 平衡路由对训练过程中的专家专门化和负载平衡有何影响？

主要发现

• BASE 层通过线性分配实现均衡的 token-to-expert 路由，无需 Balancing losses 或 capacity factors。
• 单个 BASE 层在更高的计算预算下可以显著优于密集数据并行和模型并行基线。
• BASE 与 Sparsely Gated MoE 和 Switch transformers 在计算效率方面相比具有优势，有时在类似预算下达到或超过它们的性能。
• 交错多个 BASE 层在总参数量大致保持不变的情况下带来越来越多的性能提升。
• BASE 层在不同放置和 BASE 子层数量下显示出稳健的性能，表明架构鲁棒性。