Expert-Choice-Routing ermöglicht adaptive Berechnung in Diffusionssprachmodellen
Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models
April 2, 2026
Autoren: Shuibai Zhang, Caspian Zhuang, Chihan Cui, Zhihan Yang, Fred Zhangzhi Peng, Yanxin Zhang, Haoyue Bai, Zack Jia, Yang Zhou, Guanhua Chen, Ming Liu
cs.AI
Zusammenfassung
Diffusions-Sprachmodelle (DLMs) ermöglichen parallele, nicht-autoregressive Texterzeugung, doch bestehende DLM-Mixture-of-Experts (MoE)-Modelle übernehmen Token-Choice (TC)-Routing von autoregressiven Systemen, was zu Lastenungleichgewicht und starrer Berechnungszuweisung führt. Wir zeigen, dass Expert-Choice (EC)-Routing besser für DLMs geeignet ist: Es bietet durch sein Design deterministischen Lastenausgleich, was einen höheren Durchsatz und schnellere Konvergenz als TC ermöglicht. Aufbauend auf der Eigenschaft, dass die EC-Kapazität extern steuerbar ist, führen wir eine zeitstufenabhängige Expertenkapazität ein, die die Expertenzuweisung entsprechend des Denoising-Schritts variiert. Wir stellen fest, dass die Zuweisung von mehr Kapazität zu Schritten mit niedrigem Maskierungsanteil unter gleichen FLOPs durchweg die beste Leistung erzielt, und liefern eine mechanistische Erklärung: Tokens in Kontexten mit niedrigem Maskierungsanteil weisen eine um eine Größenordnung höhere Lerneffizienz auf, sodass die Konzentration von Rechenleistung auf diese Schritte den größten Grenznutzen bringt. Schließlich zeigen wir, dass bestehende vortrainierte TC-DLMs durch Austausch nur des Routers auf EC umgerüstet werden können, was schnellere Konvergenz und verbesserte Genauigkeit über verschiedene Downstream-Aufgaben hinweg erreicht. Zusammengenommen etablieren diese Ergebnisse EC-Routing als überlegenes Paradigma für DLM-MoE-Modelle und demonstrieren, dass Berechnung in DLMs als adaptive Strategie statt als feste architektonische Konstante behandelt werden kann. Code ist verfügbar unter https://github.com/zhangshuibai/EC-DLM.
English
Diffusion language models (DLMs) enable parallel, non-autoregressive text generation, yet existing DLM mixture-of-experts (MoE) models inherit token-choice (TC) routing from autoregressive systems, leading to load imbalance and rigid computation allocation. We show that expert-choice (EC) routing is a better fit for DLMs: it provides deterministic load balancing by design, yielding higher throughput and faster convergence than TC. Building on the property that EC capacity is externally controllable, we introduce timestep-dependent expert capacity, which varies expert allocation according to the denoising step. We find that allocating more capacity to low-mask-ratio steps consistently achieves the best performance under matched FLOPs, and provide a mechanistic explanation: tokens in low-mask-ratio contexts exhibit an order-of-magnitude higher learning efficiency, so concentrating compute on these steps yields the largest marginal return. Finally, we show that existing pretrained TC DLMs can be retrofitted to EC by replacing only the router, achieving faster convergence and improved accuracy across diverse downstream tasks. Together, these results establish EC routing as a superior paradigm for DLM MoE models and demonstrate that computation in DLMs can be treated as an adaptive policy rather than a fixed architectural constant. Code is available at https://github.com/zhangshuibai/EC-DLM.