Label-Efficient Residual Correction for Mixture-of-Experts via Their Routed Readouts

Pretrained mixture-of-experts (MoE) predictors are increasingly reused on heterogeneous deployment data. In such settings, the model's error, whether caused by distribution shift or by limited training, can differ across the input regions served by different experts, yet retraining is often infeasible due to restricted access to data or model parameters. Existing post-hoc adaptation methods typically treat the model as a black-box predictor, ignoring the routing weights and expert outputs that form each prediction. We propose ATF-kNN (Affine Teacher Features with kNN residual imputation), a post-hoc adaptation method for a frozen MoE teacher, which leverages the routing probabilities and expert-specific outputs exposed at inference time. Our approach fits one offset and one gain per expert in one step, gated by the model's own routing probabilities, and imputes residuals for unlabeled deployment data via k-nearest neighbors in this readout space. We characterize the population gain over scalar teacher correction and prove a finite-sample risk certificate under regularity conditions. Experiments on real and synthetic data show consistent improvements over the frozen teacher under both limited training and distribution shift, robustness to a misspecified number of experts and over a range of imputation hyperparameters, and, given enough labels, gains over scalar correction when the teacher's residual depends on its routing.