Anchor-MoE: A Mean-Anchored Mixture of Experts For Probabilistic Regression

TL;DR

Anchor-MoE introduces a flexible mixture of experts model for probabilistic regression, combining an anchor mean with local expert dispatching to achieve minimax-optimal risk rates and state-of-the-art empirical performance.

Contribution

The paper proposes Anchor-MoE, a novel mixture of experts framework with theoretical risk guarantees and superior empirical results in probabilistic regression.

Findings

Achieves minimax-optimal $L^2$ risk rate $O(N^{-2\alpha/(2\alpha+d)})$.

Scales favorably with number of experts and latent dimensions in CRPS and NLL.

Outperforms strong baselines on standard UCI regression benchmarks.

Abstract

Regression under uncertainty is fundamental across science and engineering. We present an Anchored Mixture of Experts (Anchor-MoE), a model that handles both probabilistic and point regression. For simplicity, we use a tuned gradient-boosting model to furnish the anchor mean; however, any off-the-shelf point regressor can serve as the anchor. The anchor prediction is projected into a latent space, where a learnable metric-window kernel scores locality and a soft router dispatches each sample to a small set of mixture-density-network experts; the experts produce a heteroscedastic correction and predictive variance. We train by minimizing negative log-likelihood, and on a disjoint calibration split fit a post-hoc linear map on predicted means to improve point accuracy. On the theory side, assuming a H\"older smooth regression function of order~$\alpha$ and fixed Lipschitz partition-of-unity weights with bounded overlap, we show that Anchor-MoE attains the minimax-optimal $L^2$ risk rate $O\!\big(N^{-2\alpha/(2\alpha+d)}\big)$. In addition, the CRPS test generalization gap scales as $\widetilde{O}\!\Big(\sqrt{(\log(Mh)+P+K)/N}\Big)$; it is logarithmic in $Mh$ and scales as the square root in $P$ and $K$. Under bounded-overlap routing, $K$ can be replaced by $k$, and any dependence on a latent dimension is absorbed into $P$. Under uniformly bounded means and variances, an analogous $\widetilde{O}\!\big(\sqrt{(\log(Mh)+P+K)/N}\big)$ scaling holds for the test NLL up to constants. Empirically, across standard UCI regressions, Anchor-MoE consistently matches or surpasses the strong NGBoost baseline in RMSE and NLL; on several datasets it achieves new state-of-the-art probabilistic regression results on our benchmark suite. Code is available at https://github.com/BaozhuoSU/Probabilistic_Regression.