Mathematical Prerequisites for Mixture of Experts
Building the math foundations you need for Mixture of Experts — expected value, Gaussian densities, likelihood, Bayes' theorem, softmax, mixture models, conditional probability, multinomial distributions, and the Sherman-Morrison-Woodbury formula — all derived step by step with one consistent example.
Before diving into Mixture of Experts, you need a handful of mathematical ideas. This post builds every one of them inside a single, consistent example so nothing feels abstract or disconnected. By the end, you will have all the tools required to derive the MoE framework — cooperative and competitive errors, posterior responsibilities, the mixture-of-Gaussians interpretation, hierarchical mixtures, and the EM algorithm — from scratch in Part 1.
The Running Example
Imagine two weather forecasters predicting tomorrow’s temperature. Forecaster A predicts 20°C. Forecaster B predicts 26°C. The actual temperature turns out to be 22°C.
We will use this setup — two forecasters, one observed outcome — for every concept in this post. The forecasters will become “experts” in the MoE framework, the actual temperature will become the “target,” and the question of how to combine or select between forecasters will motivate every mathematical tool we build.
Throughout, we will use the following concrete numbers:
| Forecaster | Prediction |
|---|---|
| A | |
| B |
Observed temperature: .
1. Squared Error
We need a way to measure how wrong a prediction is. The most natural choice is the squared error: take the difference between the prediction and the actual value, then square it.
Why square? Two reasons. First, squaring makes all errors positive — a prediction that is 2 degrees too high is just as “wrong” as one that is 2 degrees too low. Second, squaring penalises large errors more heavily than small ones: an error of 4 is penalised 16 times, not 4 times.
When dealing with vector-valued outputs (predictions with multiple components), we use the squared norm:
This sums the squared errors across all components. For our scalar example, reduces to .
Numerical check
For forecaster A: .
For forecaster B: .
Forecaster A has a squared error of 4, forecaster B has a squared error of 16. By this measure, A is four times better than B on this particular observation. We will use these two numbers — 4 and 16 — throughout the rest of the post.
2. Expected Value
Suppose we do not know which forecaster will be selected — a gating mechanism picks forecaster A with probability and forecaster B with probability . What is the average squared error we would see across many such selections?
The expected value answers this question. It multiplies each possible outcome by its probability, then sums:
For the expected squared error when the gating mechanism assigns to A and to B:
Numerical check
On average, if we randomly pick a forecaster according to these probabilities and measure the squared error, we get 8.8. This formula is exactly the competitive error function in MoE — each expert is measured against the full target independently, and the errors are averaged using the gating probabilities as weights. We will see this in Part 1 as equation (1.2) from Jacobs et al. (1991).
3. The Softmax Function
In MoE, a gating network must produce probabilities that are all positive and sum to 1. Given arbitrary real-valued scores (which could be any real numbers — positive, negative, or zero), how do we convert them into valid probabilities?
The softmax function does exactly this:
Each score is exponentiated (making it positive, since for all ), then divided by the sum of all exponentiated scores (ensuring the result sums to 1).
Let us verify both requirements. Positivity: since and , we have for all . Summation: .
An important property. Softmax preserves the ranking: if , then (because the exponential is monotonically increasing: , so ). The forecaster with the higher score gets the higher probability.
Numerical check
Suppose the gating network produces scores and for our two forecasters. Then:
Check: . Both positive, sum to 1. Forecaster A, which had the higher score (), gets the higher probability (). In Part 1, the gating network computes these scores as linear functions of the input: .
4. The Gaussian Density
We now arrive at the concept that connects squared error to probability. The Gaussian (or normal) probability density function describes how likely a value is, given that it was drawn from a bell-shaped distribution centred at with spread controlled by the variance :
Let us unpack each piece:
- is the squared error between the observation and the centre — the same squared error from Section 1.
- Dividing by scales the error by the variance. Larger variance means the distribution is wider and more tolerant of errors.
- The exponential converts the scaled squared error into a positive number (since for all ). Smaller errors give larger density values.
- The factor is the normalising constant — it ensures the density integrates to 1 over all possible values.
A common simplification. When comparing two forecasters on the same observation, the normalising constant is the same for both (since is the same). In MoE, Jacobs et al. (1991) use unit variance () for simplicity. With , the Gaussian density simplifies to:
The normalising constant is now the same for all experts and does not affect which expert is “better.” This is why the MoE papers often write and drop the constant.
Numerical check
Using unit variance ():
Forecaster A (, ):
Forecaster B (, ):
Forecaster A’s density is , roughly 400 times larger than forecaster B’s density of . The Gaussian density is telling us: if the true temperature-generating process is centred at 20 (forecaster A’s prediction) with unit variance, the observed value of 22 is plausible. If the process is centred at 26, the value 22 is extremely unlikely. This ratio — 400 to 1 — is exactly the kind of comparison that will drive the posterior responsibilities in MoE.
Why the densities are small
Both values — and — are small numbers. That is normal. A continuous density is not a probability. It measures how densely the probability is concentrated around . The actual probability of observing exactly 22.000… is zero for any continuous distribution. What matters is the relative comparison between densities, not their absolute size. Forecaster A’s density is 400 times larger than forecaster B’s — that is the relevant information.
5. Likelihood
We now flip the perspective. In Section 4, we fixed a forecaster () and asked how probable the observation is. Now we fix the observation () and ask: which forecaster makes the observation more plausible?
The likelihood function uses the same formula as the probability density, but views it as a function of (the forecaster’s prediction) rather than (the observation):
Same formula. Different question.
Probability asks: given forecaster A’s prediction , how likely is the observation ?
Likelihood asks: given the observation , how plausible is it that the generating process was centred at ?
The answer is the same number — — but the interpretation is different. Probability varies while holding fixed. Likelihood varies while holding fixed.
Numerical check
The likelihood of forecaster A is 400 times larger than the likelihood of forecaster B. If we had to pick one forecaster based solely on this observation, we would pick A — it assigns far more plausibility to the data we actually saw.
This is exactly what happens inside MoE. Each expert’s output defines a Gaussian centre, and the likelihood measures how well that expert explains the observed target. Experts with higher likelihood get more responsibility. We will formalise this in the Bayes’ theorem section.
6. Log-Likelihood and the Logarithm
Working directly with likelihoods leads to numerical problems — when you multiply many small numbers together (one per data point), the product quickly underflows to zero. The logarithm fixes this by converting products into sums.
The natural logarithm is the inverse of the exponential: if , then . We need three properties.
Property 1: Monotonically increasing. If , then . Maximising is equivalent to maximising — the logarithm preserves the location of maxima and minima.
Property 2: Log of a product is a sum of logs. This is the logarithm product rule:
More generally, . This property is what makes the EM algorithm tractable — it will allow us to bring the logarithm inside the summation signs and separate the parameters of different experts.
Property 3: Log of an exponential cancels. . This converts the Gaussian density into a simple quadratic.
The log-likelihood of a Gaussian
Applying property 3 to our Gaussian likelihood:
By property 2 (log of a product):
By property 3 (log of exponential cancels):
The first term is a constant (it does not depend on ). So maximising the log-likelihood is equivalent to minimising — the squared error from Section 1.
This is a deep result: maximum likelihood estimation with a Gaussian model is equivalent to minimising squared error. When you see squared error in MoE, you are implicitly doing maximum likelihood under a Gaussian assumption.
Numerical check
Forecaster A:
Verify: . Correct.
Forecaster B:
Verify: . Correct.
The log-likelihood difference is . We can verify: . The logarithm has turned the ratio of 400:1 in likelihoods into an additive difference of 6.0 in log-likelihoods — much easier to work with numerically.
7. Negative Log-Likelihood as a Loss Function
In optimisation, we typically minimise a loss function. But likelihood and log-likelihood are things we want to maximise — higher likelihood means the model fits the data better. The standard trick is to negate the log-likelihood, turning maximisation into minimisation:
This is the negative log-likelihood (NLL). Minimising the NLL is mathematically identical to maximising the likelihood.
For a single Gaussian:
The constant term does not affect the optimisation, so the NLL is proportional to the squared error.
In MoE, the error function from Jacobs et al. (1991) is:
This is the negative log-likelihood of a mixture model — we will build up to it in Section 10. The negative sign converts “maximise likelihood” into “minimise error.”
Numerical check
Forecaster A: .
Forecaster B: .
Minimising the NLL means preferring forecaster A (loss 2.919) over forecaster B (loss 8.919). This is the same conclusion as before — just expressed as a minimisation problem.
8. Bayes’ Theorem
We now have all the ingredients for the most important result in this post. Suppose we have two forecasters, and before seeing the actual temperature, a gating network assigns prior probabilities and to each. After seeing the actual temperature , we want to update these probabilities to reflect how well each forecaster did. The updated probabilities are called posterior probabilities.
Bayes’ theorem tells us exactly how to update:
Let us read this formula piece by piece:
- is the prior probability — the gating network’s belief about forecaster before seeing the data. In MoE, these are the softmax outputs from Section 3.
- is the likelihood — how well forecaster ‘s prediction explains the observed data. This is the Gaussian density from Section 4.
- is the joint probability — the probability of choosing forecaster AND observing from it.
- is the total probability of observing under the full mixture. We will call this the marginal likelihood. It sums over all possible forecasters.
- The ratio gives the posterior probability — the updated belief about forecaster after seeing the data.
Derivation
Bayes’ theorem follows directly from the definition of conditional probability. The conditional probability of event given event is:
Let = “forecaster was selected” and = “we observed .” Then:
- . This is the probability of picking forecaster (probability ) and then observing from it (density ).
- . This is the total probability of observing , obtained by summing over all possible forecasters. This step uses the law of total probability: the probability of is the sum of the probabilities of through each possible path.
Substituting:
That is Bayes’ theorem. Nothing more.
Numerical check
Let the gating network assign prior probabilities , (same as Section 2). Using unit-variance Gaussian likelihoods (dropping the normalising constant, since it cancels in the ratio):
Joint probabilities:
Marginal likelihood (total probability):
Posterior probabilities:
Check: . The posteriors sum to 1, as they must.
What happened. Forecaster A started with a prior of and ended with a posterior of . Its responsibility increased dramatically because it fits the data much better (squared error 4 vs. 16). Forecaster B started with a prior of and ended with a posterior of — nearly zero.
This is exactly the mechanism that drives specialisation in MoE. The posteriors replace the priors in the gradient computation, ensuring that experts who fit the data well get larger gradients and learn faster on the cases they handle best. In Part 1, these posteriors are called in equation (1.5).
Prior vs. posterior: the key distinction
This distinction is easy to gloss over, but the entire MoE framework depends on it. Let us pin it down precisely.
Prior depends on the input only. It is computed before seeing the target . It answers: “Based on the input alone, which expert should handle this case?”
Posterior depends on both the input and the target . It is computed after seeing the target. It answers: “Given that we now see the target, which expert actually explains it best?”
The prior is the gating network’s best guess. The posterior is reality’s correction. In MoE, using posteriors instead of priors in the gradient is what makes the system discover the right task decomposition — it is not just a mathematical nicety, it is the mechanism that makes experts specialise.
9. Mixture Models
We now combine all the tools we have built. A mixture model says that the data was generated by one of component distributions, but we do not know which one. The overall probability of observing is a weighted sum of the component densities:
Each term has two parts:
- is the mixing weight (or prior probability) for component . These are non-negative and sum to 1 — exactly the output of a softmax (Section 3).
- is the component density — for Gaussians, this is the formula from Section 4.
This is the law of total probability applied to continuous densities. The observation could have come from any component, so we sum over all possible sources, weighting each by its probability.
The mixture model as a generative story
A mixture model describes a two-step process for generating data:
- Select a component with probability . (Roll a weighted die.)
- Generate the observation from component ‘s density . (Draw from the selected distribution.)
In MoE, step 1 is the gating network selecting an expert, and step 2 is the selected expert generating a prediction (with Gaussian noise). The entire MoE framework is a mixture model where the mixing weights and component means are input-dependent.
Mixture of Gaussians
When every component density is a Gaussian (Section 4), the mixture model is called a mixture of Gaussians (also known as a Gaussian mixture model or GMM). The total density becomes:
Each component is a bell curve centred at with width , and the mixture is a weighted sum of these bell curves. The result is a density that can have multiple peaks — one per component — allowing the model to capture data that clusters in several regions.
This is exactly the interpretation Jacobs et al. (1991) give to MoE in Part 1: each expert defines a Gaussian centred at its output, the gating network provides the mixing weights, and the combined model is a mixture of Gaussians. The term in the MoE error function is the Gaussian component (with unit variance, normalising constant dropped), and is the mixture.
Numerical check
Using our two forecasters with unit variance and mixing weights , :
This is the marginal likelihood — the total probability of observing under the mixture. It is dominated by forecaster A’s contribution ( out of ), reflecting A’s much better fit to the data.
Notice that we computed this same quantity (up to the normalising constant) in the Bayes’ theorem section as the denominator . The difference is the normalising constant : . Both are computing the total probability of the data under the mixture — one with the constant, one without.
10. The Log-Likelihood of a Mixture
The loss function in MoE is the negative log of the mixture probability:
This is the negative log-likelihood of a mixture model. It is equation (1.3) in Jacobs et al. (1991), and equation (7) in Jordan & Jacobs (1993).
For a dataset of observations, the total log-likelihood is the sum over all data points:
Here the logarithm product rule (Section 6, property 2) has already done its work: the log of the product of independent likelihoods becomes a sum of individual log-likelihoods.
The log-sum problem
Notice that the logarithm sits outside the sum over experts: . This is fundamentally different from , which would be much easier to work with. The log of a sum does not simplify nicely — we cannot separate the parameters of different experts.
Compare:
- Log of a product (easy): — the terms separate.
- Log of a sum (hard): — the terms stay coupled.
This is the central computational difficulty of mixture models. The EM algorithm, which we will derive in Part 1, exists specifically to solve this problem. It introduces hidden indicator variables that tell us which expert generated each data point, converting the log-of-a-sum into a sum-of-logs.
Numerical check
If we had used forecaster A alone: . If forecaster B alone: . The mixture loss (3.426) is slightly worse than using A alone (2.919) because the mixture includes B, which hurts the likelihood. But the mixture model has a crucial advantage: it can learn which forecaster to trust for each input, achieving better overall performance across many data points from different regions.
11. Conditional Probability and Nested Decisions
In the flat MoE, the gating network makes a single decision: which expert to use. In the Hierarchical Mixture of Experts (HME), decisions are nested: first pick a branch, then pick an expert within that branch. To handle this, we need conditional probability.
The conditional probability of event given that event has occurred is:
Rearranging gives the multiplication rule:
This says: the probability of both and happening is the probability that happens, times the probability that happens given that has already happened.
Nested decisions in our example
Suppose we organise our two forecasters into a hierarchy. We have a “weather service” that first picks a forecasting agency (branch), then picks a specific forecaster within that agency.
- Branch decision: Pick agency 1 with probability or agency 2 with probability .
- Within-branch decision: If agency 1 is picked, choose forecaster A with probability or forecaster B with probability .
The probability of reaching forecaster A through agency 1 is a nested decision — two choices in sequence:
This is the multiplication rule in action. The notation means “the probability of choosing forecaster A given that we are in agency 1.” The vertical bar always means “given that” — it separates what we are asking about from what we are conditioning on.
The law of total probability with nested decisions
If the observation can be generated through multiple paths (agency 1 → forecaster A, agency 1 → forecaster B, agency 2 → forecaster A, etc.), the total probability of sums over all paths:
This is the law of total probability applied to a two-level hierarchy. Each path has probability (by the multiplication rule), and generates with density . We sum over all possible paths.
In Part 1, this is exactly equation (4) from Jordan & Jacobs (1993) — the probability model for the Hierarchical Mixture of Experts.
Numerical check
Suppose agency 2 has two different forecasters, C and D, with , , predictions , . Using unit-variance Gaussians:
Path probabilities and likelihoods for :
- Agency 1, forecaster A:
- Agency 1, forecaster B:
- Agency 2, forecaster C:
- Agency 2, forecaster D:
Total: .
Forecaster A through agency 1 dominates, contributing out of — about 96% of the total probability. The hierarchy has identified the right path.
12. The Multinomial Distribution
When the gating network picks one of experts, it is making a multinomial (or categorical) decision. A multinomial distribution generalises a coin flip to outcomes. Instead of just Heads or Tails, we have outcomes with probabilities where .
For our two-forecaster example (), the multinomial reduces to a single coin flip: choose forecaster A with probability or forecaster B with probability . But in MoE, the gating network may choose among many experts — 4, 8, or even thousands. The multinomial distribution handles any number of outcomes.
The probability of outcome in a single draw is simply:
That is literally it. The term “multinomial” sounds imposing, but for a single draw it is just “pick option with probability .” The softmax function (Section 3) is the standard way to parameterise a multinomial distribution — it takes arbitrary scores and produces valid multinomial probabilities.
Where this appears in MoE
In the HME, there are multiple multinomial decisions:
- The top-level gating network makes a multinomial decision over branches: pick branch with probability .
- Each lower-level gating network makes a multinomial decision over experts within a branch: pick expert with probability .
Each of these is a separate multinomial distribution, parameterised by its own softmax. The nested structure — first a multinomial over branches, then a conditional multinomial over experts — is what creates the tree-structured probability model from Section 11.
Numerical check
With 4 forecasters and scores , , , :
Sum: .
Check: . This is a valid multinomial distribution over 4 outcomes. The gating network rolls this 4-sided weighted die to pick an expert.
13. The Gradient of the Mixture Loss
To train MoE, we need the derivative of the loss with respect to each expert’s prediction , where is the mixture likelihood. Let us derive this step by step.
Step 1. Since , by the chain rule:
Step 2. Since and only the -th term depends on :
Step 3. For the Gaussian (dropping the constant), we need the derivative. Let , so . By the chain rule, . Now:
The and the from the power rule cancel, and the comes from differentiating with respect to . So:
Step 4. Combining steps 1–3:
Step 5. Recognise as the posterior probability from Bayes’ theorem (Section 8):
This is equation (1.5) from Jacobs et al. (1991). The gradient for expert is proportional to two things: the posterior probability (how responsible this expert is for the data point) and the prediction error (how far off the expert’s prediction is). If is small — meaning the expert is not responsible — the gradient is small and the expert barely updates. If is large, the expert gets a strong push toward the target. This is the mathematical mechanism behind expert specialisation.
Numerical check
Using , , , , :
The negative gradient means the loss decreases when increases — forecaster A should increase its prediction from 20 toward the target 22. This is the correct direction.
For forecaster B:
The gradient for B is . The gradient for A is . Which gradient is “larger”? This brings us to a subtle but important distinction.
14. Gradient Magnitude vs. Signed Value
Gradients are signed quantities — they can be positive or negative. The sign tells you the direction: a negative gradient means “increase the parameter to decrease the loss,” and a positive gradient means “decrease the parameter to decrease the loss.” But when we ask which expert is learning faster, we care about the magnitude (absolute value) of the gradient, not its sign.
The magnitude of a number is — its distance from zero on the number line, ignoring the sign:
Note that as signed numbers, but — negative five is smaller than positive three, but it is farther from zero.
Why this matters in MoE
From our gradient computation:
As signed numbers, . But the magnitudes tell a different story:
The magnitude of A’s gradient () is 312 times larger than B’s (). Forecaster A is learning 312 times faster from this data point.
When you read in Part 1 that “expert 2 gets a larger gradient magnitude,” this means — the absolute value is bigger, meaning a stronger push. The sign tells the direction of the push; the magnitude tells how hard the push is. In MoE, the posterior controls the magnitude: an expert with high posterior gets a large and learns aggressively; an expert with near-zero posterior barely moves.
15. Indicator Variables and the Complete-Data Log-Likelihood
The EM algorithm, which we will derive in Part 1, relies on a clever trick: imagine that for each data point, we know which expert generated it. We represent this knowledge with indicator variables.
Define if forecaster generated the observation, and otherwise. Exactly one indicator is 1 for each data point — the data came from exactly one forecaster.
If we knew the ‘s, the complete-data log-likelihood would be:
The second step uses the logarithm product rule from Section 6: .
Compare this to the incomplete-data log-likelihood from Section 10:
The indicator variables have performed a crucial transformation: the logarithm has moved inside the summation. Instead of (log of a sum — hard), we have (sum of logs — easy). The parameters of different experts now appear in separate terms, making maximisation straightforward.
Numerical check
Suppose we knew forecaster A generated the data (, ):
If instead forecaster B generated it (, ):
The complete-data likelihood sharply favours the hypothesis that A generated the data ( vs. ).
Of course, we do not know the ‘s. The EM algorithm handles this by replacing with its expected value — the posterior probability from Bayes’ theorem (Section 8). For our numbers: and . This is the E step. The M step then maximises the expected complete-data log-likelihood with respect to the model parameters. The full derivation appears in Part 1.
16. The Matrix Inverse Update: Sherman-Morrison-Woodbury
The on-line version of MoE (Section 11 of Part 1) updates expert parameters after each data point, without re-processing the entire dataset. This requires maintaining a running estimate of an inverse matrix. The Sherman-Morrison-Woodbury formula tells us how to update a matrix inverse when a small change is made to the original matrix.
The problem
Suppose we have a matrix and we know its inverse . Now changes by a small amount — specifically, a rank-1 update: , where and are column vectors and is their outer product (a matrix where entry is ).
Computing from scratch is expensive — it takes operations for an matrix. But if we already know , the Sherman-Morrison formula gives us in only operations:
The formula says: start with the old inverse , then subtract a correction term. The correction term involves multiplying the old inverse by and — all or cheaper operations.
Verification
We can verify this by checking that . Multiply the original matrix by the proposed inverse:
Distribute:
Since , the first term is . Simplify in the second term. In the fourth term, is a scalar, call it . So:
Combine the last three terms (they all contain ):
The expression in parentheses is .
So the result is . The formula is correct.
Numerical check
Let , so .
Let and . Then .
So .
Applying the formula:
Verify: . Correct.
How this appears in MoE
In the on-line MoE algorithm, each expert maintains a matrix (the inverse of a weighted sum of outer products ). When a new data point arrives, must be updated. The Sherman-Morrison formula gives:
This has exactly the form of the Sherman-Morrison formula: the old inverse () minus a correction term involving the outer product . The parameter is a decay factor that down-weights old data, and is the posterior probability that controls how much this data point affects expert ‘s update. Without this formula, the on-line algorithm would need to re-invert a matrix at every time step — making it impractical.
Summary
We have built a complete mathematical toolkit for Mixture of Experts. Squared error measures prediction quality. Expected value averages over stochastic selections, giving us the competitive error function. The softmax function converts arbitrary scores into valid probabilities for the gating network. The Gaussian density connects squared error to probability, and the likelihood function flips the perspective from “how probable is the data” to “how plausible is the model.” The logarithm converts products into sums, and the negative log-likelihood gives us a loss function to minimise. Bayes’ theorem updates prior gating probabilities into posterior responsibilities — the mechanism that makes experts specialise. Mixture models describe the generative process, and the log-of-a-sum structure of the mixture log-likelihood creates the computational challenge that the EM algorithm solves. Conditional probability and the multiplication rule handle nested decisions in the hierarchical architecture, while the multinomial distribution formalises the gating network’s multi-way choices. Differentiating the mixture loss yields the posterior-weighted gradient, where gradient magnitude — not signed value — determines how fast each expert learns. Indicator variables with the complete-data log-likelihood show how the EM algorithm brings the logarithm inside the sum. And the Sherman-Morrison-Woodbury formula enables efficient on-line updates without re-inverting matrices.
With these tools in hand, we are ready for Part 1, where we put them all together to derive the Mixture of Experts framework from scratch — cooperative and competitive errors, the mixture-of-Gaussians interpretation, hierarchical mixtures, and the EM algorithm.
Previous: MoE Load Balancing from Scratch
Next: Mixture of Experts from Scratch — Part 1
Enjoyed this post?
Subscribe to get notified when I publish new posts. No spam, unsubscribe anytime.