Mixture-of-experts models

Overview

scikit-learn's sklearn.mixture fits unsupervised Gaussian mixtures to a sample of points. A mixture of experts (MoE) is the supervised analogue: it models the conditional density of a target \(y\) given inputs \(x\) as an input-dependent, gated mixture of local (expert) regressions. This package mirrors the structure of sklearn/mixture and adds two estimators, GaussianMixtureOfExperts and BayesianGaussianMixtureOfExperts, that port the estimation algorithms from two peer-reviewed papers.

The model

For inputs \(x\in\mathbb{R}^p\) and a scalar target \(y\in\mathbb{R}\), write \(\tilde{x}=(1,x)\). A mixture of \(K\) linear-Gaussian experts models

$$ p(y\mid x)=\sum_{k=1}^{K} g_k(x)\,\mathcal{N}\!\big(y \mid a_k^\top \tilde{x},\ \sigma_k^2\big), $$

where each expert \(k\) is a local linear regression with coefficients \(a_k\) and noise variance \(\sigma_k^2\), and the gating functions \(g_k(x)\ge 0,\ \sum_k g_k(x)=1\) softly partition the input space. Two gating families are provided.

Softmax gating

A multinomial-logistic (softmax) gate, with the last expert as the zero reference for identifiability:

$$ g_k(x)=\frac{\exp(\beta_k^\top \tilde{x})}{\sum_{j=1}^{K}\exp(\beta_j^\top \tilde{x})},\qquad \beta_K=0. $$

This is the model studied in Nguyen, Nguyen and Ho (NeurIPS 2023, Spotlight). The gate M-step maximises the responsibility-weighted gate log-likelihood by a Bohning majorise-minimise update: the softmax Hessian is bounded by \(1/4\), giving the working response \(z_k=\tilde{x}\beta_k+4(r_k-P_k)\) solved as a weighted least squares, with a backtracking step that keeps the objective monotone.

Gaussian gating (GLLiM)

A generative gate places a Gaussian on the inputs of each component, so the joint model is

$$ p(x,y)=\sum_{k=1}^{K}\pi_k\,\mathcal{N}(x\mid c_k,\Gamma_k)\,\mathcal{N}(y\mid A_k x+b_k,\ \sigma_k^2),\qquad g_k(x)=\frac{\pi_k\mathcal{N}(x\mid c_k,\Gamma_k)}{\sum_j \pi_j\mathcal{N}(x\mid c_j,\Gamma_j)}. $$

This is the Gaussian locally-linear mapping (GLLiM; Deleforge, Forbes and Horaud, 2015), the estimator whose parameter-estimation rates are studied in Nguyen, Nguyen, Nguyen and Ho (AISTATS 2024). Its EM has closed-form M-steps for \(\pi_k, c_k, \Gamma_k\) (weighted mean and full covariance of \(x\)) and weighted least squares for the experts. In both gates the experts may be multi-output, and covariance_type in {full, tied, diag, spherical} sets the structure of the expert covariance \(\Sigma_k\) (as in sklearn.mixture.GaussianMixture); the Gaussian gate itself always uses a full \(\Gamma_k\).

Estimation by EM

Both gates are fitted by Expectation-Maximisation. The E-step forms responsibilities \(r_{ik}\propto g_k(x_i)\,\mathcal{N}(y_i\mid \text{expert}_k)\) (softmax, conditional objective) or the joint responsibilities (Gaussian, joint objective); the M-step updates experts and gate as above. The estimators expose the scikit-learn mixture API: fit, predict, predict_proba, score, score_samples, bic, aic, and the attributes converged_, n_iter_, lower_bound_.

Quick start

import numpy as np
from mixture import GaussianMixtureOfExperts

rng = np.random.default_rng(0)
x = rng.uniform(-3, 3, size=(600, 1))
y = np.where(x[:, 0] < 0, 2*x[:, 0] + 1, -x[:, 0] - 1) + 0.3*rng.standard_normal(600)

moe = GaussianMixtureOfExperts(n_components=2, gate="softmax", random_state=0).fit(x, y)
moe.predict(x)         # E[y | x]
moe.predict_proba(x)   # gating weights g_k(x), rows sum to 1
moe.bic(x, y)          # model selection over the number of experts

The estimators compose with Pipeline and GridSearchCV. Install with pip install -e . from the package directory.

Scope and honesty

• GaussianMixtureOfExperts is a frequentist maximum-likelihood EM and a faithful port of the published estimators. It supports scalar or multivariate ``y`` and all expert covariance types {full, tied, diag, spherical}.
• BayesianGaussianMixtureOfExperts is a full mean-field variational Bayes (posterior distributions, an evidence lower bound, and automatic expert pruning). For the Gaussian gate it is fully conjugate: covariance_type selects a diagonal (Normal-Gamma) or full-covariance posterior, where full gives a Normal-Inverse-Wishart gate and matching Matrix-Normal-Wishart experts (so the expert covariances are full too). The weight prior (weight_prior_type) is Dirichlet, a Bayesian nonparametric Dirichlet process, or the full Pitman-Yor process (discount \(\sigma>0\), with \(\sigma\) and the concentration \(\alpha\) estimated by importance sampling); the last two select the number of experts automatically, Pitman-Yor with heavier tails. All of this follows the BNP-GLLiM model (Nguyen, Forbes, Arbel and Nguyen, 2024). For the non-conjugate softmax gate, the experts are treated with full variational Bayes while the gate is fitted by the Bohning update (a regularised point estimate), a variational-EM.
• Verified on synthetic data: the Dirichlet and Dirichlet-process lower bounds are monotone (validating the Wishart and stick-breaking KL terms), full-covariance experts recover a correlated \(\Sigma_k\), the Pitman-Yor sampler estimates a genuine discount \(\sigma\in(0,1)\), parameters are recovered, and redundant experts are pruned. The Pitman-Yor free energy is stochastic (its \((\alpha,\sigma)\) are refreshed by importance sampling each iteration), so it converges to a noise floor rather than being exactly monotone.

Choosing the number of experts (live)

How many experts does the data support? This demo fits a mixture of experts to one dataset and compares seven selectors at once: the two sweep-free merging procedures, the dendrogram of mixing measures (Hai et al., AISTATS 2026) and Merge-Truncate-Merge (Nguyen et al., 2024), the slope heuristic / dimension jump (Nguyen et al., 2022), the variational ELBO pruning, and the classical AIC, BIC and ICL. Switch the gate between Gaussian- and softmax-gated, set the true number of experts, and read each criterion's choice off the bars; pick which criterion drives the clustering on the left, and compare the two merging procedures on the right. These are the browser counterparts of the model-selection papers and the compare_selectors function in the package.

Variational inference versus sampling (MCMC)

The BayesianGaussianMixtureOfExperts here uses variational inference (VI): it replaces the intractable posterior with the closest member of a tractable family and maximises an evidence lower bound. VI is fast, deterministic and scales well, which is why it is the default for large models (Blei, Kucukelbir and McAuliffe, 2017). The price is that the mean-field family tends to underestimate posterior uncertainty and is mode-seeking: it locks onto one region of the posterior and can miss others entirely.

The exact-inference alternative is Markov chain Monte Carlo (MCMC): a Gibbs sampler for the conjugate mixture, or gradient-based Hamiltonian Monte Carlo and the No-U-Turn Sampler (NUTS; Hoffman and Gelman, 2014) for the continuous parameters. MCMC is asymptotically exact and recovers the full posterior, including multimodal uncertainty, but it is slower and, for mixtures of experts, meets two structural difficulties:

• Label switching. The likelihood is invariant to permuting the experts, so the posterior has at least \(K!\) symmetric copies of every solution; a chain that mixes across them makes the marginal expert parameters meaningless unless the labels are resolved after the fact (Stephens, 2000; Jasra, Holmes and Stephens, 2005).
• Mode trapping. The genuinely distinct solutions form well-separated modes with low-probability valleys between them, so a single local sampler, NUTS included, settles into one mode and never sees the rest.

A mixture of experts is therefore a canonical multimodal test case. A recent review, Latuszynski, Moores and Stumpf-Fetizon (2025), surveys the MCMC machinery built for exactly this setting; the leading ideas are:

• Parallel tempering (replica exchange): run several chains at increasing "temperatures" that flatten the posterior valleys, and swap states between them so a cold, exact chain can hop between modes it could never cross on its own (Geyer, 1991; Earl and Deem, 2005). Using NUTS for the cold chain is the practical gold standard for scientific, low-to-moderate-dimensional Bayesian MoE.
• Mode jumping and the Wang-Landau algorithm: propose long-range moves between known modes, or adaptively learn the density of states so the sampler spends equal time in every mode (Wang and Landau, 2001; Bornn et al., 2013).
• Sparse finite mixture MCMC: place a sparse Dirichlet prior on the weights so the sampler itself decides how many experts are active, and remove label switching with a relabeling step, so the model order is chosen during sampling rather than by a separate sweep (Malsiner-Walli, Fruhwirth-Schnatter and Grun, 2016). This is the sampling analogue of the Dirichlet-process and Pitman-Yor pruning in this package.
• Nonparametric posterior bootstrap: replace the Markov chain with an embarrassingly parallel set of randomized-objective optimisations, resolving massively multimodal posteriors without mixing at all (Fong, Lyddon and Holmes, 2019).

The trade-offs, with a mixture of experts as the running example. Every row is now runnable from this package: the samplers are implemented in mixture._samplers (conjugate Gibbs, a hand-coded No-U-Turn Sampler with finite-difference-checked gradients, parallel tempering with a NUTS kernel per temperature, sparse finite mixture MCMC, and the posterior bootstrap), and compare_inference() reproduces this table empirically on any dataset.

Verified empirically by compare_inference on the two-expert crossing-slope benchmark (the posterior has symmetric well-separated modes): the finite-difference check puts the NUTS gradients at machine precision; Gibbs, NUTS, parallel tempering and the bootstrap all recover the sorted expert slopes; mean-field VI underestimates the slope posterior spread relative to Gibbs; a single NUTS chain records zero mode hops while parallel tempering hops between the modes; the sparse-mixture sampler's posterior of the number of active experts concentrates on the truth from an overfitted start; and the bootstrap replicates land in both modes with no Markov chain at all. The samplers implement the conjugate Gaussian-gated model with univariate input and scalar response (the setting of the demos) as a reference implementation of the paradigms.

Watch the samplers run

The animation below runs each inference method live on the two-expert crossing-slope benchmark (a bimodal posterior: the two label assignments are symmetric, well-separated modes). Pick a method and press Run: the left panel shows the sampled expert fits (a fading trail of draws; for the bootstrap, the accumulating replicate cloud), the right panel shows the method's signature diagnostic. Watch variational Bayes ascend deterministically into one mode; Gibbs and a single HMC chain wander within a mode (mode hops stay near zero); parallel tempering hop between the modes as heated chains relay states down; the sparse-mixture sampler build its posterior over the number of active experts; and the posterior bootstrap fill both modes with independent parallel replicates.

When does each method win?

Run the animation across the five scenarios and the pattern is consistent. The scorecard summarises when each method is the right tool, with the numbers measured by compare_inference on the crossing-slope benchmark: mean-field VB underestimates the slope spread (sd ratio ~0.88 versus Gibbs), a single NUTS chain records 0 mode hops where parallel tempering records dozens, the sparse-mixture sampler concentrates its K+ posterior on the truth (about 0.99 at K=2), and the bootstrap replicates fill both modes.

The one-line rule: use variational Bayes to decide how many experts and to get there fast; switch to (or warm-start) parallel tempering or the sparse-mixture sampler the moment mode-complete, calibrated uncertainty is the deliverable.

In practice the two families are complementary. This package's fast variational fit, with Dirichlet-process or Pitman-Yor pruning to choose the number of experts, is an excellent initialiser and proposal for a parallel-tempered NUTS or a sparse-finite-mixture sampler when gold-standard, mode-complete uncertainty is needed, for example in a scientific inverse problem where the number and identity of the experts is itself the question. This pipeline is implemented as vb_warm_start(): the variational fit supplies the mode, the pruned K and the starting states (hot chains start at label permutations of the fit, seeding the symmetric modes), so the tempered sampler needs a far shorter warmup; the Hybrid option in the animation above shows it in action. For very high-dimensional deep mixtures of experts, where exact MCMC is infeasible, the current direction is deep variational surrogates instead: distilling an expressive generative model into a mixture-of-experts head, or a mixture-of-products-of-experts variational autoencoder for multimodal fusion (Sutter, Daunhawer and Vogt, 2021).

Connections to the research

This package is the practical, runnable face of a body of work on the theory of mixtures of experts. The three questions it quietly answers, what can these models represent, how large a model should you fit, and how do you use them when the likelihood is intractable, are exactly the themes of the publications below. Each links to a full-text reading page and to an interactive demo.

References

1. F. Pedregosa et al. Scikit-learn: Machine Learning in Python. JMLR 12:2825-2830, 2011. (template and framework for this extension).
2. H. Nguyen, TrungTin Nguyen and N. Ho. Demystifying Softmax Gating Function in Gaussian Mixture of Experts. NeurIPS 2023 (Spotlight).
3. H. Nguyen, TrungTin Nguyen, K. Nguyen and N. Ho. Towards Convergence Rates for Parameter Estimation in Gaussian-gated Mixture of Experts. AISTATS 2024.
4. A. Deleforge, F. Forbes and R. Horaud. High-dimensional regression with Gaussian mixtures and partially-latent response variables. Statistics and Computing, 2015.
5. TrungTin Nguyen, F. Forbes, J. Arbel and H. D. Nguyen. Bayesian nonparametric mixture of experts for inverse problems. Journal of Nonparametric Statistics, 2024 (the full-covariance Normal-Inverse-Wishart gate and Dirichlet-Process / Pitman-Yor weight prior are ported from this BNP-GLLiM work). details.
6. R. A. Jacobs, M. I. Jordan, S. J. Nowlan and G. E. Hinton. Adaptive mixtures of local experts. Neural Computation 3(1):79-87, 1991.
7. D. M. Blei, A. Kucukelbir and J. D. McAuliffe. Variational Inference: A Review for Statisticians. Journal of the American Statistical Association 112(518):859-877, 2017.
8. M. D. Hoffman and A. Gelman. The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo. JMLR 15:1593-1623, 2014.
9. A. Jasra, C. C. Holmes and D. A. Stephens. Markov Chain Monte Carlo Methods and the Label Switching Problem in Bayesian Mixture Modeling. Statistical Science 20(1):50-67, 2005.
10. G. Malsiner-Walli, S. Fruhwirth-Schnatter and B. Grun. Model-based clustering based on sparse finite Gaussian mixtures. Statistics and Computing 26:303-324, 2016.
11. E. Fong, S. Lyddon and C. Holmes. Scalable Nonparametric Sampling from Multimodal Posteriors with the Posterior Bootstrap. ICML 2019.
12. K. Latuszynski, M. T. Moores and T. Stumpf-Fetizon. MCMC for multi-modal distributions. arXiv:2501.05908, 2025. arXiv.
13. C. J. Geyer. Markov chain Monte Carlo maximum likelihood. Computing Science and Statistics: Proc. 23rd Symposium on the Interface, 156-163, 1991 (parallel tempering).
14. D. J. Earl and M. W. Deem. Parallel tempering: Theory, applications, and new perspectives. Physical Chemistry Chemical Physics 7:3910-3916, 2005.
15. F. Wang and D. P. Landau. Efficient, multiple-range random walk algorithm to calculate the density of states. Physical Review Letters 86(10):2050-2053, 2001.
16. M. Stephens. Dealing with label switching in mixture models. Journal of the Royal Statistical Society B 62(4):795-809, 2000.
17. T. M. Sutter, I. Daunhawer and J. E. Vogt. Generalized Multimodal ELBO. ICLR 2021 (the mixture-of-products-of-experts VAE).