模型架構與推論
UniF-MoE Extracts Shared Blocks Before Routing Residual Computation, Cutting Inference Time by 45.2% Versus Top-2 GMoE
Rather than deciding shared experts, FFN width, and the number of activated experts separately, UniF-MoE lets each token first select shared channels and then accumulates routing probabilities according to its remaining requirements. In the authors’ DeiT and BERT experiments, it improves accuracy and MoE inference cost, but it has not yet been validated on generative large language models or distributed expert parallelism.
Traditional sparse MoE systems use a top-k router to send each token to a fixed number of full FFN experts, while shared experts, intra-expert pruning, and dynamic k are typically handled by separate controllers. The newly proposed UniF-MoE argues that these three decisions are sequentially dependent: after reusable computation is extracted, the residual content, optimal experts, and required number of experts may all change. The system must therefore “share first, then route the residual.”
The authors first study experts created by copying the same dense FFN and then sparsely upcycling it, splitting each FFN into aligned key-value channel blocks. Experts that are frequently activated together show greater overlap in their value positions. After these shared positions are removed, only 5.7% of tokens retain their original top-2 expert combination. Among expert pairs with high sharing ratios, 56.3% can reconstruct the original output using only two residual experts, compared with 11.8% for the low-sharing group. This supports rerouting instead of reusing the original routing decision.
UniF-MoE calculates a shared-demand value α for each token and uses it to determine both the number of shared blocks and the shared-path weight. The mean of the up-projection keys in each block serves as a prototype for selecting the actual shared content. The residual demand, β = 1 − α, then becomes the cumulative routing-probability threshold: the system activates the minimum number of residual experts needed for their summed probability to cover β. Gram regularization constrains routing vectors to remain close to orthogonal, preventing multiple experts from repeatedly performing the same work.
Across five DomainBed datasets, UniF-MoE achieves an average accuracy of 69.5%, outperforming the listed baselines for which complete comparisons are available. Its average across five GLUE tasks is 82.76%, also exceeding the 81.95% obtained by selecting a fixed k separately for each task. On VLCS, compared with top-2 GMoE, it uses 9.1% fewer activated parameters and 16.1% fewer FLOPs, reduces inference step time from 0.31 seconds to 0.17 seconds, and lowers memory usage from 0.55 GiB to 0.26 GiB. However, dense DeiT still requires only 0.04 seconds and 0.29 GiB, showing that this is an improvement within MoE rather than evidence that it comprehensively outperforms dense models.
The Apache 2.0-licensed code includes vision and language implementations, tests, and reproduction scripts, but provides neither training checkpoints nor result arrays. The next question is whether this routing approach can scale to decoder-only LLMs, long-sequence batching, cross-GPU expert communication, and heterogeneous experts that are no longer initialized from the same FFN.