[论文解读] On the Representation Collapse of Sparse Mixture of Experts

Zewen Chi, Dong Li+10 more|arXiv (Cornell University)|Apr 20, 2022

Expert finding and Q&A systems被引用 16

一句话总结

本文识别出稀疏专家混合模型（SMoE）中的表征坍缩问题，即标记表征聚集在专家中心附近，导致性能下降。本文提出X-MoE方法，将隐藏表征投影到低维超球面，并使用L2归一化的路由分数与可学习温度门控机制，实现在七个多语言基准上的稳定性能提升，同时改善了路由稳定性。

ABSTRACT

Sparse mixture of experts provides larger model capacity while requiring a constant computational overhead. It employs the routing mechanism to distribute input tokens to the best-matched experts according to their hidden representations. However, learning such a routing mechanism encourages token clustering around expert centroids, implying a trend toward representation collapse. In this work, we propose to estimate the routing scores between tokens and experts on a low-dimensional hypersphere. We conduct extensive experiments on cross-lingual language model pre-training and fine-tuning on downstream tasks. Experimental results across seven multilingual benchmarks show that our method achieves consistent gains. We also present a comprehensive analysis on the representation and routing behaviors of our models. Our method alleviates the representation collapse issue and achieves more consistent routing than the baseline mixture-of-experts methods.

研究动机与目标

• 识别并分析稀疏专家混合（SMoE）模型中尚未被充分探索的表征坍缩问题。
• 通过重新定义路由评分机制，提升SMoE模型中的路由稳定性和表征多样性。
• 通过更鲁棒的路由算法，增强多语言预训练和微调性能。
• 在多次训练运行中实现一致的路由行为，降低模型行为的方差。

提出的方法

• 在计算路由分数前，将标记隐藏表征和专家嵌入投影到低维空间。
• 对标记表征和专家嵌入均应用L2归一化，以在超球面上测量路由分数。
• 在软专家门控中使用可学习温度参数，以控制专家激活，提升路由灵活性。
• 将路由公式化为归一化表征在超球面上的点积相似度，替代标准的未归一化点积。
• 将新型路由机制集成到基于Transformer的SMoE模型中，采用top-1路由及softmax和sigmoid门控函数。
• 在多语言预训练和下游微调任务上训练并评估模型，以评估性能和路由一致性。

实验结果

研究问题

• RQ1 在预训练和微调过程中，表征坍缩在稀疏专家混合模型中在多大程度上表现出来？
• RQ2 将表征投影到低维超球面是否能缓解表征坍缩并提升模型性能？
• RQ3 所提出的路由机制在微调过程中对不同随机种子的路由一致性有何影响？
• RQ4 使用L2归一化的路由分数与可学习温度门控是否能带来更稳定和高效的专家分配？
• RQ5 与基线SMoE模型相比，该方法在多语言基准性能和表征多样性方面表现如何？

主要发现

• 所提出的X-MoE方法在预训练和微调阶段均实现了在七个多语言基准上的稳定性能提升。
• 在SMoE模型中实证观察到表征坍缩，即标记表征紧密聚集在专家中心附近。
• X-MoE通过超球面路由和L2归一化显著减少了表征坍缩，增强了表征多样性。
• X-MoE中的路由行为表现出高度的运行间一致性，无论随机种子如何，均收敛到相似的路由模式，而SMoE基线则不然。
• 该方法同时提升了语言建模和下游任务性能，增益归因于更稳定和多样的专家路由。
• 在软门控中使用可学习温度参数增强了路由适应性，对性能和稳定性提升有显著贡献。