L'Allineamento del Manifold di Instradamento Migliora la Generalizzazione dei Modelli Linguistici Mixture-of-Experts

Routing Manifold Alignment Improves Generalization of Mixture-of-Experts LLMs

November 10, 2025

Autori: Zhongyang Li, Ziyue Li, Tianyi Zhou

cs.AI

Abstract

I modelli Sparse Mixture-of-Experts (MoE) sono stati ampiamente adottati nei recenti modelli linguistici di grandi dimensioni poiché consentono di aumentare efficientemente le capacità del modello senza incrementare i costi di inferenza. Tuttavia, le valutazioni su un ampio spettro di task downstream rivelano una sistematica sub-ottimalità dei router negli attuali LLM MoE, che si traduce in un divario prestazionale significativo (ad esempio, del 10-20% in accuratezza) rispetto a un routing ottimale. In questo articolo, dimostriamo che allineare la varietà (manifold) dei pesi di routing con quella dell'embedding del task può ridurre efficacemente questo divario e migliorare le prestazioni di generalizzazione dei modelli MoE. Il nostro metodo, "Routing Manifold Alignment (RoMA)", introduce un termine aggiuntivo di regolarizzazione sulla varietà nella funzione obiettivo del post-training e richiede solo un fine-tuning leggero dei router (mantenendo congelati gli altri parametri). Nello specifico, la regolarizzazione incoraggia i pesi di routing di ciascun campione ad avvicinarsi a quelli dei suoi vicini di successo (i cui pesi di routing portano a risposte corrette) in uno spazio di embedding del task. Di conseguenza, campioni relativi a task simili condivideranno scelte di esperti simili attraverso i vari layer. Costruire tali legami tra task ed esperti su campioni diversi è essenziale per ottenere una migliore generalizzazione. Inoltre, RoMA dimostra il vantaggio di unificare la comprensione del task (tramite modelli di embedding) con la generazione di soluzioni (tramite LLM MoE). Negli esperimenti, abbiamo applicato il fine-tuning ai router di OLMoE, DeepSeekMoE e Qwen3-MoE utilizzando RoMA. Le valutazioni su benchmark diversificati e ampi confronti con i baseline mostrano il sostanziale miglioramento apportato da RoMA.

English

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.