Beyond the Previous Layer: Residual Predictive Structure in Sparse MoE Routing

2026-09-15

AI Summary

研究通过扩展历史选择层数，使用OLMoE和JetMoE模型提高了稀疏混合专家模型中路由预测的准确性。

Abstract

Sparse mixture-of-experts models route each token through a sequence of expert selections. We ask whether the immediately preceding selection adequately summarizes this trajectory for predicting the next router. Using frozen OLMoE and JetMoE models, we measure the held-out predictive gain from earlier expert selections while retaining the most recent selection as a common baseline. In OLMoE, extending the history from one to eleven layers raises router-logit $R^2$ from 0.59879 to 0.66544. A preregistered JetMoE replication yields four-layer gains of 0.14275 and 0.20528 at two target depths, with paired bootstrap intervals above zero. These gains survive nonlinear decoding: adding history to a small multilayer perceptron improves $R^2$ by 0.17137 and 0.21861, whereas nonlinear decoding of the recent state alone adds 0.00139 and 0.00936 over a linear probe. Parameter-matched controls preserve the advantage, and cross-fitted history residuals predict target residuals with $R^2$ of 0.20549 and 0.23556. These findings identify residual predictive structure in expert-selection trajectories beyond adjacent-layer persistence.

Problem

• Sparse Mixture-of-Experts
• Router Prediction
• Expert Selection Trajectory

Innovation

• sparse mixture-of-experts
• router prediction
• historical selection impact
• nonlinear decoding
• residual predictive structure