[论文解读] Scaling Expert Language Models with Unsupervised Domain Discovery

Suchin Gururangan, Margaret Li+5 more|arXiv (Cornell University)|Mar 24, 2023

Topic Modeling被引用 7

一句话总结

本文提出 c-BTM，一种通过无监督聚类将文本语料库划分为特定领域的组，从而实现通信高效、无需同步的并行训练方法，以训练稀疏专家语言模型。通过让专家专门化于学习到的聚类，该方法在显著降低 FLOPs 的同时，性能优于稠密模型，且随着聚类数量和数据规模的增加，性能增益也持续提升。

ABSTRACT

Large language models are typically trained densely: all parameters are updated with respect to all inputs. This requires synchronization of billions of parameters across thousands of GPUs. We introduce a simple but effective method to asynchronously train large, sparse language models on arbitrary text corpora. Our method clusters a corpus into sets of related documents, trains a separate expert language model on each cluster, and combines them in a sparse ensemble for inference. This approach generalizes embarrassingly parallel training by automatically discovering the domains for each expert, and eliminates nearly all the communication overhead of existing sparse language models. Our technique outperforms dense baselines on multiple corpora and few-shot tasks, and our analysis shows that specializing experts to meaningful clusters is key to these gains. Performance also improves with the number of experts and size of training data, suggesting this is a highly efficient and accessible approach to training large language models.

研究动机与目标

• 为解决在数千个 GPU 上训练大规模稠密语言模型时产生的高通信和同步开销问题。
• 实现在无需依赖元数据定义领域的情况下，对稀疏专家语言模型进行可扩展、通信高效的训练。
• 探究无监督聚类是否能发现有意义的领域，从而在性能上优于稠密基线模型。
• 探索在自动发现的聚类上训练的专家模型的扩展特性，特别是专家数量和数据规模的关系。

提出的方法

• 对文本语料库应用 k-means 聚类，无需使用元数据即可发现类似领域的组别。
• 从预训练的基础模型（如 OPT-1.3B）初始化专家语言模型（ELM），并独立地在分配给其的聚类上微调每个 ELM。
• 使用稀疏集成推理机制，根据输入上下文嵌入与每个专家聚类中心之间的余弦相似度，选择 top-k 专家。
• 以完全并行的方式训练专家，通过避免参数同步更新，最大限度减少 GPU 间通信。
• 通过离线聚类平衡聚类，确保专家使用均衡，避免路由瓶颈。
• 在推理阶段通过加权平均组合 top-k 专家的输出，实现高效的稀疏计算。

实验结果

研究问题

• RQ1 与稠密训练相比，无监督聚类是否能有效发现有意义的领域，从而提升语言模型性能？
• RQ2 聚类数量（即专家数量）如何影响模型性能和训练效率？
• RQ3 当扩展到更大数据集时，c-BTM 在困惑度和 FLOP 效率方面是否优于稠密基线模型？
• RQ4 与具有学习路由机制的现有稀疏 MoE 模型相比，c-BTM 的性能如何？
• RQ5 与随机或不平衡聚类相比，平衡聚类在多大程度上影响模型性能？

主要发现

• c-BTM 在验证困惑度上优于参数量相当的稠密语言模型，且随着聚类数量的增加，性能增益持续提升。
• 在仅使用 8 个 GPU 同时训练 128 个专家语言模型（总参数量 168B）并处理 168B 标记数据时，展示了高效的并行化和可扩展性。
• 使用 128 个专家并采用 top-4 推理时，性能与使用全部专家相当，甚至 top-1 推理也优于稠密模型。
• c-BTM 使用 1.3B 参数专家和稀疏推理，达到与 6.7B 参数稠密模型相当的困惑度，但训练 FLOPs 仅为其 29%。
• 平衡聚类的性能与使用真实元数据进行专家分配相当，且显著优于随机或不平衡聚类。
• 最优聚类数量随训练数据规模增加而增加，且性能随聚类数量和数据规模的增加而单调提升。