Beyond Gaussian Mixtures

Standard mixture models—whether Gaussian, Student-t, or otherwise—model the marginal distribution $p(\mathbf{y})$ of observations. Each component has a fixed mixing proportion $\pi_k$ that applies uniformly across the entire data space. But what if the "right" component to use varies depending on context?

Consider predicting house prices. In urban areas, prices might follow one distribution (high mean, moderate variance). In suburban areas, another pattern emerges (lower mean, higher variance). A standard mixture model would blend these patterns into a single, averaged prediction. But intuitively, if we know whether a house is urban or suburban, we should use different models.

Mixture of Experts (MoE) models formalize this intuition. Instead of fixed mixing proportions $\pi_k$, they use input-dependent gating functions $g_k(\mathbf{x})$ that determine which "expert" (component) is most relevant for each input. This transforms mixture models from unconditional density estimators into powerful conditional models $p(\mathbf{y} | \mathbf{x})$.

MoE models are foundational to modern machine learning: they connect classical mixture models to neural networks, form the basis of Gated Mixture of Experts used in large language models, and provide a principled framework for combining specialized models.

Before diving into Mixture of Experts, let's understand precisely what problem it solves by examining the limitations of standard approaches.

The Regression Problem

In supervised learning, we want to model the conditional distribution $p(y | \mathbf{x})$ or, at minimum, predict $\mathbb{E}[y | \mathbf{x}]$. Standard approaches:

Linear regression: Assumes $p(y | \mathbf{x}) = \mathcal{N}(y | \mathbf{w}^T\mathbf{x}, \sigma^2)$. This is a single, global model.

Gaussian Mixture Model: Models the marginal $p(y) = \sum_k \pi_k \mathcal{N}(y | \mu_k, \sigma_k^2)$. This captures multimodality in $y$ but ignores $\mathbf{x}$ entirely.

Neither approach directly addresses input-dependent multimodality—situations where:

• The relationship between $\mathbf{x}$ and $y$ changes across the input space
• Different mechanisms or regimes govern different regions
• A single global model cannot capture local structure

Example: Inverse Kinematics

Consider a robot arm: given a target position $(x, y)$ in Cartesian coordinates, find joint angles $\theta$ that reach that position. For most targets, multiple solutions exist (the arm can reach the same point in different configurations). The mapping from position to angles is one-to-many.

A standard regression model averages over solutions, predicting joint angles that might not correspond to any valid configuration. What we need is a model that:

1. Recognizes there are multiple modes
2. Uses context (the input $\mathbf{x}$) to weight which modes are likely
3. Outputs a full conditional distribution, not just a mean

The Mixture of Experts model defines the conditional distribution as a mixture where both the mixing proportions and the component densities depend on the input.

Basic Formulation

$$p(y | \mathbf{x}, \boldsymbol{\Theta}) = \sum_{k=1}^K g_k(\mathbf{x} | \boldsymbol{\theta}_g) \cdot p(y | \mathbf{x}, \boldsymbol{\theta}_k)$$

Components:

1. Gating network $g_k(\mathbf{x} | \boldsymbol{\theta}_g)$: Computes the input-dependent probability of selecting expert $k$. Must satisfy:

   • $g_k(\mathbf{x}) \geq 0$ for all $k$, $\mathbf{x}$
   • $\sum_{k=1}^K g_k(\mathbf{x}) = 1$ for all $\mathbf{x}$

2. Expert networks $p(y | \mathbf{x}, \boldsymbol{\theta}_k)$: Each expert provides a conditional distribution over outputs given inputs. Common choices:

   • Linear experts: $p(y | \mathbf{x}) = \mathcal{N}(y | \mathbf{w}_k^T\mathbf{x}, \sigma_k^2)$
   • Neural network experts: Arbitrary neural networks

3. Latent variable interpretation: Let $z \in {1, \ldots, K}$ indicate which expert generated the output:

   • $p(z = k | \mathbf{x}) = g_k(\mathbf{x})$
   • $p(y | z = k, \mathbf{x}) = p(y | \mathbf{x}, \boldsymbol{\theta}_k)$
   • Marginalizing: $p(y | \mathbf{x}) = \sum_k p(z = k | \mathbf{x}) p(y | z = k, \mathbf{x})$

Softmax Gating

The most common gating function uses a softmax over linear combinations:

$$g_k(\mathbf{x}) = \frac{\exp(\mathbf{v}k^T \mathbf{x})}{\sum{j=1}^K \exp(\mathbf{v}_j^T \mathbf{x})} = \text{softmax}(\mathbf{V}^T \mathbf{x})_k$$

where $\mathbf{V} = [\mathbf{v}_1, \ldots, \mathbf{v}_K]$ are the gating parameters.

Properties of softmax gating:

• Smooth transitions: Gating weights change continuously with $\mathbf{x}$
• Interpretable: Weights can be visualized to understand expert selection
• Gradient-friendly: Softmax is differentiable, enabling end-to-end training

Linear Experts Example

With Gaussian linear experts:

$$p(y | \mathbf{x}) = \sum_{k=1}^K g_k(\mathbf{x}) \cdot \mathcal{N}(y | \mathbf{w}_k^T \mathbf{x}, \sigma_k^2)$$

This is a piecewise linear model with soft boundaries. In regions where $g_k(\mathbf{x}) \approx 1$, the model behaves like expert $k$'s linear regression. Transition regions blend multiple experts.

Parameter count: For $D$-dimensional inputs and $K$ experts:

• Gating: $K \times D$ parameters
• Experts: $K \times (D + 1)$ parameters (weights + variance)
• Total: $K \times (2D + 1)$ parameters

Training Mixture of Experts models follows the EM framework, treating the expert assignment as latent. However, because the gating function is parameterized, the M-step typically requires gradient-based optimization.

Problem Setup

Given data ${(\mathbf{x}n, y_n)}{n=1}^N$, we want to maximize:

$$\mathcal{L}(\boldsymbol{\Theta}) = \sum_{n=1}^N \log p(y_n | \mathbf{x}n, \boldsymbol{\Theta}) = \sum{n=1}^N \log \sum_{k=1}^K g_k(\mathbf{x}_n) \cdot p(y_n | \mathbf{x}_n, \boldsymbol{\theta}_k)$$

E-Step: Compute Responsibilities

For each observation, compute the posterior probability that expert $k$ generated $(\mathbf{x}_n, y_n)$:

$$h_{nk} = p(z_n = k | y_n, \mathbf{x}_n) = \frac{g_k(\mathbf{x}_n) \cdot p(y_n | \mathbf{x}_n, \boldsymbol{\theta}k)}{\sum{j=1}^K g_j(\mathbf{x}_n) \cdot p(y_n | \mathbf{x}_n, \boldsymbol{\theta}_j)}$$

Note the key difference from standard GMM: the responsibilities $h_{nk}$ depend on both $\mathbf{x}_n$ (through gating) and $y_n$ (through expert likelihood).

M-Step: Update Parameters

The expected complete-data log-likelihood is:

$$Q(\boldsymbol{\Theta} | \boldsymbol{\Theta}^{(t)}) = \sum_{n=1}^N \sum_{k=1}^K h_{nk} \left[ \log g_k(\mathbf{x}_n) + \log p(y_n | \mathbf{x}_n, \boldsymbol{\theta}_k) \right]$$

This decomposes into two independent optimization problems:

Gating parameters: Maximize $$Q_g = \sum_{n=1}^N \sum_{k=1}^K h_{nk} \log g_k(\mathbf{x}_n)$$

For softmax gating, this is equivalent to weighted multinomial logistic regression with soft labels $h_{nk}$. Solved via gradient descent:

$$\frac{\partial Q_g}{\partial \mathbf{v}k} = \sum{n=1}^N (h_{nk} - g_k(\mathbf{x}_n)) \mathbf{x}_n$$

Expert parameters: For each expert $k$, maximize $$Q_k = \sum_{n=1}^N h_{nk} \log p(y_n | \mathbf{x}_n, \boldsymbol{\theta}_k)$$

For Gaussian linear experts, this is weighted least squares:

$$\mathbf{w}k^{\text{new}} = \left( \sum{n=1}^N h_{nk} \mathbf{x}n \mathbf{x}n^T \right)^{-1} \sum{n=1}^N h{nk} \mathbf{x}_n y_n$$

$$\sigma_k^{2, \text{new}} = \frac{\sum_{n=1}^N h_{nk} (y_n - \mathbf{w}k^T \mathbf{x}n)^2}{\sum{n=1}^N h{nk}}$$

A single level of gating may not capture complex input-output relationships. Hierarchical Mixture of Experts (HME) extends the framework by organizing experts in a tree structure, with gating functions at each internal node.

Tree Structure

In a two-level HME:

• Level 1 (top): A gating network partitions inputs among $M$ groups
• Level 2: Within each group $m$, another gating network selects among $K_m$ experts

The conditional distribution becomes:

$$p(y | \mathbf{x}) = \sum_{m=1}^M g_m^{(1)}(\mathbf{x}) \sum_{k=1}^{K_m} g_{mk}^{(2)}(\mathbf{x}) \cdot p(y | \mathbf{x}, \boldsymbol{\theta}_{mk})$$

Advantages of hierarchy:

1. Divide and conquer: Top-level gating performs coarse partitioning; lower levels refine
2. Soft decision trees: HME can be viewed as a probabilistic, soft-boundary decision tree
3. Modularity: Different branches can specialize in different aspects of the problem
4. Reduced parameters: With shared structure, HME can be more parameter-efficient than a flat model with the same effective capacity

EM for Hierarchical MoE

The E-step computes responsibilities at all levels:

$$h_{nmk} = p(z^{(1)} = m, z^{(2)} = k | y_n, \mathbf{x}_n) = \frac{g_m^{(1)}(\mathbf{x}n) \cdot g{mk}^{(2)}(\mathbf{x}n) \cdot p(y_n | \mathbf{x}n, \boldsymbol{\theta}{mk})}{\sum{m', k'} \cdots}$$

We can also compute marginal responsibilities:

• $h_{nm} = \sum_k h_{nmk}$ (responsibility for group $m$)
• $h_{nk|m} = h_{nmk} / h_{nm}$ (responsibility for expert $k$ within group $m$)

The M-step updates each gating network and expert independently, using the appropriate responsibilities.

Connection to Decision Trees

HME provides a probabilistic generalization of decision trees:

This connection makes HME attractive for interpretable modeling—the tree structure provides some insight into the decision process.

The Mixture of Experts framework has experienced a remarkable resurgence in deep learning, particularly in large language models (LLMs). Modern implementations differ from classical MoE in several important ways.

Sparse Gating

In classical MoE, all experts contribute to every prediction (soft gating). This becomes prohibitive with hundreds or thousands of experts. Sparse MoE addresses this:

$$p(y | \mathbf{x}) = \sum_{k \in \text{Top-}K} g_k(\mathbf{x}) \cdot p(y | \mathbf{x}, \boldsymbol{\theta}_k)$$

where Top-$K$ selects only the $K$ experts with highest gating weights. Common choice: $K = 1$ or $K = 2$.

Benefits:

• Computational sparsity: Only $K$ experts are evaluated per input
• Capacity scaling: Total parameters can be huge, but per-example compute is fixed
• Specialization: Experts can develop distinct competencies

Challenges:

• Load balancing: Without intervention, some experts become over-used while others are ignored
• Training instability: Hard selection creates gradient discontinuities
• Routing collapse: All inputs may route to the same experts

Key Innovations in Modern MoE

1. Noisy Top-K Gating (Shazeer et al., 2017)

Add noise before taking top-K to encourage exploration: $$g(\mathbf{x}) = \text{softmax}(\mathbf{x}^T W_g + \epsilon \cdot \text{Softplus}(\mathbf{x}^T W_{\text{noise}}))$$

2. Load Balancing Losses

Add auxiliary losses that penalize uneven expert usage: $$\mathcal{L}_{\text{balance}} = N \cdot \sum_k f_k \cdot P_k$$

where $f_k$ is the fraction of tokens routed to expert $k$, and $P_k$ is the average gating probability for expert $k$.

3. Expert Capacity Limits

Limit the number of tokens each expert can process: $$\text{capacity} = \frac{\text{tokens per batch}}{\text{number of experts}} \times \text{capacity factor}$$

Tokens exceeding capacity are dropped or processed by a residual connection.

4. Switch Transformer (Fedus et al., 2021)

Simplified MoE with $K = 1$ routing and careful initialization. Demonstrated that sparse MoE can scale to trillion-parameter models while maintaining training efficiency.

Beyond linear models, experts can be arbitrary neural networks. This dramatically increases expressiveness but requires end-to-end gradient training rather than EM.

Neural MoE Architecture

Standard form: Replace each expert $p(y | \mathbf{x}, \boldsymbol{\theta}_k)$ with a neural network $f_k(\mathbf{x})$:

$$\hat{y} = \sum_{k=1}^K g_k(\mathbf{x}) \cdot f_k(\mathbf{x})$$

where:

• $g_k(\mathbf{x}) = \text{softmax}(\text{GatingNetwork}(\mathbf{x}))_k$
• $f_k(\mathbf{x})$ is the output of expert network $k$

In transformer layers: MoE typically replaces the feedforward network (FFN) sublayer:

$$\text{FFN}{\text{MoE}}(\mathbf{x}) = \sum{k \in \text{Top-}K} g_k(\mathbf{x}) \cdot \text{FFN}_k(\mathbf{x})$$

Each expert is a separate FFN with distinct parameters, enabling specialization.

Training Considerations

End-to-end backpropagation: With soft gating, gradients flow through all experts. With hard (Top-K) gating, gradients only flow through selected experts.

Straight-through estimator: For hard gating, use soft gating in the backward pass to enable gradient flow to non-selected experts (for the gating network).

Gradient scaling: Experts receiving fewer tokens may have unstable gradients. Some implementations scale gradients by routing frequency.

Mixture of Experts models find application wherever different subsets of the input space benefit from specialized treatment.

Natural Language Processing

Large Language Models: MoE enables scaling to trillion+ parameters while maintaining tractable inference. Different experts may specialize in:

• Different languages (multilingual models)
• Different domains (code, legal, medical)
• Different reasoning types (factual recall vs. arithmetic)

Machine Translation: Language pairs with different linguistic properties benefit from expert specialization. One expert might handle Romance languages, another Germanic, etc.

Multimodal Learning

Vision-Language Models: Experts can specialize in:

• Visual reasoning (spatial relationships, counting)
• Textual understanding (semantics, syntax)
• Cross-modal alignment

Recommender Systems

User Modeling: Different user segments (casual browsers, power users, bargain hunters) exhibit distinct behavior patterns. MoE naturally segments users and applies appropriate models.

This page has provided a comprehensive treatment of Mixture of Experts, from classical formulations to modern neural network implementations.

What's Next: Hidden Markov Models

Mixture of Experts addresses input-dependent gating but assumes observations are independent. Many real-world problems involve sequential data where past observations inform future ones. The next page introduces Hidden Markov Models (HMMs), which extend mixture models to sequences by introducing temporal dependencies between latent states. This leads to elegant dynamic programming algorithms for inference and learning.