DAG-MoE: From Simple Mixture to Structural Aggregation in Mixture-of-Experts
Abstract
Mixture-of-Experts (MoE) models have become a leading approach for decoupling parameter count from computational cost in large language models, yet effectively scaling MoE performance remains a challenge. Prior work shows that fine-grained experts enlarge the space of expert combinations and improve flexibility, but they also impose substantial routing overhead, creating a new scalability bottleneck. In this paper, we explore a complementary axis for scaling—how expert outputs are aggregated. We theoretically show that replacing the standard weighted-summation aggregation with structural aggregation expands the expert-combination space without altering the experts or router, and enables possible multi-step reasoning within a single MoE layer. To this end, we propose DAG-MoE, a sparse MoE framework that employs a lightweight module to automatically learn the optimal aggregation structure among the selected experts. Extensive experiments under standard language modeling settings show that DAG-MoE consistently improves performance in both pretraining and fine-tuning, surpassing traditional MoE baselines.
Lay Summary
Today's largest AI language models often use a mixture-of-experts design: instead of one giant network handling every input, the model keeps many smaller specialist networks and picks a few to answer each piece of text. Current systems combine these specialists' answers by simple averaging, like averaging votes from a panel of judges—which throws away any information about how they could build on each other's reasoning. We propose DAG-MoE, which instead arranges the chosen experts as a small flowchart, where each expert can pass intermediate results to the next and refine the answer step by step. The model learns this flowchart on its own, separately for every word it processes. We show mathematically that this is strictly more flexible than the averaging used today, and it can mimic step-by-step reasoning—all without using more compute. Across language modeling and downstream tests, DAG-MoE consistently outperforms standard mixture-of-experts models, with the largest gains on tasks requiring multi-step thinking. This suggests that how experts cooperate matters as much as which experts are chosen, pointing to a new way of improving AI language models without simply making them bigger.