On Minimax Estimation of Parameters in Softmax-Contaminated Mixture of Experts

TL;DR

This work analyzes minimax estimation of gating and prompt parameters in a softmax-contaminated MoE, where a fixed pre-trained expert is augmented by a trainable prompt. A distinguishability condition is introduced to distinguish the prompt from the pre-trained model, enabling a clear dichotomy: under distinguishability, the MLE converges at near-parametric rates $\widetilde{\mathcal{O}}(n^{-1/2})$ with minimax lower bounds matching; under non-distinguishability, rates degrade depending on how fast the prompt parameters approach the pre-trained ones. Theoretical results are complemented by empirical experiments that validate the rates and show that softmax gating improves sample efficiency relative to input-free gating. Practically, the results guide how to fine-tune prompts (ensuring distinct expertise) and favor softmax gating to achieve faster and more reliable parameter estimation in downstream tasks.

Abstract

The softmax-contaminated mixture of experts (MoE) model is deployed when a large-scale pre-trained model, which plays the role of a fixed expert, is fine-tuned for learning downstream tasks by including a new contamination part, or prompt, functioning as a new, trainable expert. Despite its popularity and relevance, the theoretical properties of the softmax-contaminated MoE have remained unexplored in the literature. In the paper, we study the convergence rates of the maximum likelihood estimator of gating and prompt parameters in order to gain insights into the statistical properties and potential challenges of fine-tuning with a new prompt. We find that the estimability of these parameters is compromised when the prompt acquires overlapping knowledge with the pre-trained model, in the sense that we make precise by formulating a novel analytic notion of distinguishability. Under distinguishability of the pre-trained and prompt models, we derive minimax optimal estimation rates for all the gating and prompt parameters. By contrast, when the distinguishability condition is violated, these estimation rates become significantly slower due to their dependence on the prompt convergence rate to the pre-trained model. Finally, we empirically corroborate our theoretical findings through several numerical experiments.

Figures (2)

• Figure 1: (Distinguishable Setting:$f_0$ is the density of a Laplace distribution.) Log-log graphs depicting the empirical convergence rates of the MLE $(\widehat{\beta}_n,\widehat{\tau}_n,\widehat{\eta}_n,\widehat{\nu}_n)$ to the ground-truth values $(\beta^*,\tau^*,\eta^*,\nu^*)$. The blue lines display the parameter estimation errors, while the orange dashed dotted lines are the fitted lines, highlighting the empirical MLE convergence rates.
• Figure 2: (Non-distinguishable Setting:$f_0$ is a Gaussian density.) Log-log graphs depicting the empirical convergence rates of the MLE $(\widehat{\beta}_n,\widehat{\tau}_n,\widehat{\eta}_n,\widehat{\nu}_n)$ to the ground-truth values $(\beta^*,\tau^*,\eta^*,\nu^*)$. The blue lines display the parameter estimation errors, while the orange dashed dotted lines are the fitted lines, highlighting the empirical MLE convergence rates. Figure \ref{['fig:non_dis_eta']} and Figure \ref{['fig:non_dis_nu']} illustrates results for Case (i) and Case (ii), respectively.

Theorems & Definitions (27)

• Definition 1: Distinguishability
• Proposition 1
• Proposition 2: Identifiability
• Proposition 3: Model Convergence
• Theorem 1
• Theorem 2
• Definition 2: Strong Identifiability
• Theorem 3
• Theorem 4
• proof : Proof of Theorem \ref{['thm:not_equal']}