Fast Model Selection and Stable Optimization for Softmax-Gated Multinomial-Logistic Mixture of Experts Models

TL;DR

This paper introduces a stable, efficient training algorithm for softmax-gated Mixture-of-Experts models, along with a principled method for selecting the number of experts, validated on biological data.

Contribution

It develops a globally convergent MM algorithm for softmax-gated MoE and proposes a new model selection method based on dendrograms, improving stability and interpretability.

Findings

Guaranteed monotone ascent and convergence of the training algorithm.

Near-parametric optimal rates for model selection and parameter recovery.

Improved accuracy and calibration on biological protein interaction data.

Abstract

Mixture-of-Experts (MoE) architectures combine specialized predictors through a learned gate and are effective across regression and classification, but for classification with softmax multinomial-logistic gating, rigorous guarantees for stable maximum-likelihood training and principled model selection remain limited. We address both issues in the full-data (batch) regime. First, we derive a batch minorization-maximization (MM) algorithm for softmax-gated multinomial-logistic MoE using an explicit quadratic minorizer, yielding coordinate-wise closed-form updates that guarantee monotone ascent of the objective and global convergence to a stationary point (in the standard MM sense), avoiding approximate M-steps common in EM-type implementations. Second, we prove finite-sample rates for conditional density estimation and parameter recovery, and we adapt dendrograms of mixing measures to the classification setting to obtain a sweep-free selector of the number of experts that achieves near-parametric optimal rates after merging redundant fitted atoms. Experiments on biological protein--protein interaction prediction validate the full pipeline, delivering improved accuracy and better-calibrated probabilities than strong statistical and machine-learning baselines.