PSGaBloME: Penalized Model Selection for Softmax-Gated Mixture of Experts
A softmax-gated mixture of experts (MoE) needs its hyper-parameters chosen from the data: how many experts, how complex, and which gating weights are worth keeping. This demo runs the paper's slope-heuristic mechanism live. Fit a family of MoE models by EM, then let a penalized maximum-likelihood criterion pick the complexity: the empirical contrast (negative log-likelihood) plus a penalty κ·D/n proportional to model dimension, with an ℓ1 term that sparsifies the gating so the selected model keeps a data-driven number of active experts. Watch the dimension jump: as κ grows, the selected complexity falls in a staircase, and the slope-heuristic rule κ̂ = 2κmin lands the selected model near the oracle.
The data are drawn from a softmax-gated mixture of experts: for input x, a softmax gate gk(x) (here a normalized Gaussian gate, which is exactly a softmax of a quadratic in x) picks one of K₀ linear experts y = akx + bk + noise. We fit the whole family of candidate MoE models, K = 1..8 experts, by EM and record each fitted model's empirical contrast ɣn(K) = −(1/n)·loglik and its dimension D(K), the number of free parameters. Model selection minimizes the penalized criterion ɣn(K) + κ·D(K)/n. Because ɣn keeps falling as the fit gets richer while D grows linearly, the selected complexity K̂ is a decreasing staircase in κ: a large dimension jump marks κmin, the smallest penalty that stops the model from exploding, and Birge-Massart's slope heuristic recommends κ̂ = 2κmin. The paper proves this choice satisfies a non-asymptotic oracle inequality: the selected model's risk is within a constant of the best (oracle) model in the family. The ℓ1 slider λ adds a sparsity penalty on the gating: experts whose average usage falls below λ are pruned and the gate renormalized, so the selected model reports a data-driven number of active experts (a proxy for the paper's joint rank-and-variable selection and block-diagonal covariance sparsity). The oracle line is the complexity minimizing the true risk, estimated on a large fresh test sample. (Honest caveats: this is a 1D regression toy fit in-browser with a Gaussian gate and linear experts, not the paper's high-dimensional multivariate polynomial SGaBloME with block-diagonal covariances; the oracle inequality and non-asymptotic risk bound are the asymptotic-in-spirit theory, while the demo shows the selection + sparsification mechanism: EM fits, the dimension penalty, the jump, and κ̂ = 2κmin.)
Run the experiments
Every animation runs live in your browser. Click a button to run that experiment on the demo (it scrolls up and starts); drag any control to take over. Nothing is downloaded, it is generated on the fly.
Sweep the penalty constant
Drag the star κ across the dimension-jump staircase: small κ keeps the largest model, and past the big jump the selected complexity collapses toward the oracle.
Sparsify the gating
Raise the ℓ1 constant λ to prune low-usage gating weights, so the selected model reports a shrinking, data-driven number of active experts.
Change the truth
Vary the true number of experts generating the data and watch the oracle complexity and the selected model track it.
Grow the sample
Increase the sample size; with more data the non-asymptotic penalty matters less and the selection sharpens onto the oracle.
Two faces of the slope heuristic
Toggle between the dimension-jump staircase (selected complexity vs κ) and the slope-estimation view (empirical contrast vs dimension, whose linear part gives κmin).
Fresh samples
Draw new samples from the true softmax-gated MoE and refit the whole family; the jump and the selected model move with the data.
The idea in three steps
Penalization trades fit against complexity. The slope heuristic makes that trade-off data-driven, with a non-asymptotic guarantee rather than an asymptotic criterion like BIC.
A family of MoE models
Fit softmax-gated mixtures of experts with K = 1..8 by EM. Each fit gives an empirical contrast (its negative log-likelihood) that keeps falling as K grows, plus a dimension D(K) counting its free parameters.
The dimension jump
Minimize contrast + κ·D/n. As κ increases, the selected complexity collapses in a staircase. The largest jump marks κmin, and the slope heuristic doubles it: κ̂ = 2κmin.
Active experts and the oracle
An ℓ1 penalty prunes low-usage gating weights, leaving a data-driven number of active experts. The oracle inequality guarantees the selected model's risk stays within a constant of the best model in the family.
For the polynomial softmax-gated block-diagonal MoE (SGaBloME) model, the joint rank and variable selection, the block-diagonal covariance structures, and the full non-asymptotic oracle inequality, see A non-asymptotic risk bound for model selection in a high-dimensional mixture of experts via joint rank and variable selection (TrungTin Nguyen, Dung Ngoc Nguyen, Hien Duy Nguyen & Chamroukhi, AJCAI 2023).