路由流形对齐提升混合专家大语言模型的泛化能力
Routing Manifold Alignment Improves Generalization of Mixture-of-Experts LLMs
AI总结:
针对现有MoE大语言模型路由器次优导致与最优路由存在10-20%准确率差距的问题,提出RoMA方法通过对齐路由权重与任务嵌入流形、轻量微调路由器,在OLMoE等模型上实现泛化性能显著提升。
摘要:
稀疏混合专家(MoE)因能在不增加推理成本的情况下高效提升模型能力,已被广泛应用于近年的大语言模型中。然而,在广泛下游任务上的评估显示,现有MoE大语言模型的路由器存在持续的次优性,导致其与最优路由间存在严重性能差距(例如准确率低10-20%)。本文表明,将路由权重的流形与任务嵌入的流形对齐,可有效缩小该差距并提升MoE大语言模型的泛化性能。我们提出的“路由流形对齐(RoMA)”方法,在训练后目标函数中引入额外的流形正则项,且仅需对路由器进行轻量级微调(其他参数冻结)。具体而言,该正则项促使每个样本的路由权重在任务嵌入空间中接近其成功邻居(路由权重能得到正确答案的样本)的路由权重。因此,针对相似任务的样本在各层将共享相似的专家选择。在不同样本上建立任务与专家间的此类绑定,对实现更好泛化至关重要。此外,RoMA展示了将任务理解(通过嵌入模型)与解决方案生成(通过MoE大语言模型)相统一的优势。实验中,我们使用RoMA对OLMoE、DeepSeekMoE和Qwen3-MoE的路由器进行微调。在多样基准上的评估及与基线的广泛比较表明,RoMA带来了显著性能提升。
英文摘要:
Sparse Mixture-of-Experts (MoE) have been widely adopted in recent large language models since it can efficiently scale up the model capability without increasing the inference cost. However, evaluations on broad downstream tasks reveal a consistent suboptimality of the routers in existing MoE LLMs, which results in a severe performance gap (e.g., 10-20% in accuracy) to the optimal routing. In this paper, we show that aligning the manifold of routing weights with that of task embedding can effectively reduce the gap and improve MoE LLMs' generalization performance. Our method, "Routing Manifold Alignment (RoMA)", introduces an additional manifold regularization term in the post-training objective and only requires lightweight finetuning of routers (with other parameters frozen). Specifically, the regularization encourages the routing weights of each sample to be close to those of its successful neighbors (whose routing weights lead to correct answers) in a task embedding space. Consequently, samples targeting similar tasks will share similar expert choices across layers. Building such bindings between tasks and experts over different samples is essential to achieve better generalization. Moreover, RoMA demonstrates the advantage of unifying the task understanding (by embedding models) with solution generation (by MoE LLMs). In experiments, we finetune routers in OLMoE, DeepSeekMoE, and Qwen3-MoE using RoMA. Evaluations on diverse benchmarks and extensive comparisons with baselines show the substantial improvement brought by RoMA.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。