🖊️ ing down my thoughts here

a tiny introduction to MOEs - part 1

when we first learn about regression, we're taught to model the world with a single, beautiful assumption: for a given input x, the output y comes from a unimodal distribution. maybe it's a gaussian. the mean and variance of this distribution are functions of our input, and we can write this as:


p(y|x)=𝒩(y|fμ(x),diag(σ+(fσ(x))))


this looks so clean. but it's a lie. a fucking beautiful lie. it works great when the relationship is one-to-one, but what if it's not?


the one-to-many problem: a messy reality

real life is often messy. what if for a certain input x, there are two equally valid outputs, say, y1 and y2? a single unimodal model would get pulled in both directions, and its mean would land right in the middle, between the two true modes. it's the dreaded regression to the mean. a blurry, useless prediction.

this is the one-to-many problem. a single model is just not enough.


the solution: a team of fucking experts

instead of using one expert to do everything, why don't we hire a goddamn team? this is the idea behind mixtures of experts (moe). we assume the output distribution isn't unimodal, but a mixture of K simpler, unimodal distributions.


the full probability distribution looks like this:


p(y|x)=k=1Kp(y|x,z=k)p(z=k|x)


let's break down this powerful equation:


1. the experts (p(y|x,z=k)):

each term p(y|x,z=k) is an expert model. it's a specialist. each expert is its own unimodal distribution, with its own set of parameters. for a gaussian moe, the k-th expert looks like this:


p(y|x,z=k)=𝒩(y|fμ,k(x),diag(fσ,k(x)))


each expert has its own mean function fμ,k and variance function fσ,k. this is crucial. it allows each expert to specialize.


2. the gating function (p(z=k|x)):

this is the master strategist. the gating function decides which expert to use for a given input x. it assigns a responsibility score to each expert. it's a classifier using a softmax function to produce probabilities.


p(z=k|x)=softmax(fz(x))k


the vector fz(x) contains scores for each expert, and the softmax turns them into probabilities that sum to 1.


training and conditional computation

we can train an moe using stochastic gradient descent (sgd). the loss function is the negative log-likelihood of the data. the chain rule takes care of the backpropagation through both the gating network and the experts.


(θ)=i=1Nlog(k=1Kp(yi|xi,z=k,θk)p(z=k|xi,θg))


what's really fucking cool is this is an example of conditional computation. for a given input, we don't have to run all the experts. we only activate the most promising ones. this makes moes incredibly efficient, especially at a large scale, like in google's gemini.


types of moes

the experts and gating function can be anything.

in the end, the key takeaway is this: for a one-to-many problem, don't average. specialize. a team of experts will always be more effective than a single one trying to be everything to everyone. it's the goddamn difference between a useless middle-of-the-road prediction and a powerful, multi-modal one that actually works.

we will meet again to discuss more about moes very soon...

References