Stretch: EM Algorithm
The problem, from first principles
04-gaussian-mixture built two individual moves: given parameters, compute responsibilities (E-step); given responsibilities, re-estimate parameters (M-step). Neither move alone fits a model — you have to actually alternate them, starting from some initial guess, and keep going until the parameters settle. That's the same "build the primitive, then assemble the loop" split 03-best-split-minimal-tree used for a single tree split versus full recursion, and 04-full-boosting-loop used for one boosting round versus the whole training loop.
The other missing piece is knowing whether the loop is actually helping. With gradient descent you'd watch a loss go down; here the natural quantity to watch is the mixture's total log-likelihood of the data — how well the current parameters explain what was actually observed. EM comes with an unusually strong guarantee about that quantity, one most iterative algorithms in this curriculum can't make: each E-step/M-step pair can only increase it, never decrease it.
From theory to code
Implement gmm_log_likelihood(input, weights, means, variances), which scores how well a given set of mixture parameters explains input, and fit_gmm(input, n_components, max_iter, seed=None), which initializes parameters and repeatedly calls 04-gaussian-mixture's gmm_e_step/gmm_m_step, tracking gmm_log_likelihood after each iteration. Theory below derives the log-likelihood formula as a log-sum-exp over the same per-component scores gmm_e_step already computes; the loop itself is just calling the two existing functions in sequence.
Constraints
- Reuse
04-gaussian-mixture'sgmm_e_step/gmm_m_step— don't reimplement either. gmm_log_likelihoodreturns a single float: the total log-likelihood ofinputunder the given(weights, means, variances).- Must be numerically stable — a point far from every component's mean must not produce
-infornan. fit_gmminitializes means fromn_componentsdistinct rows ofinput(sampled without replacement, usingseedfor reproducibility), uniform starting weights, and every component starting withinput's overall per-feature variance.fit_gmmreturns a dict with"weights","means","variances"(final, post-loop),"responsibilities"(from the last E-step), and"log_likelihood_history"(array, one entry per iteration, recorded after that iteration's M-step).- The loop runs exactly
max_iteriterations, always E-step then M-step.
Hints
Open one at a time. Each gives away a little more than the last.
Hint 1
gmm_log_likelihood needs the same log(weights[j]) + gaussian_log_likelihood(input[i], means[j], variances[j]) matrix gmm_e_step builds internally. The difference is what you do with each row afterward: gmm_e_step normalizes a row into a probability distribution; here you need the row's un-normalized total, in log space.
Hint 2
"Sum a row of log-values, in log space, stably" is log-sum-exp: subtract the row's max, exponentiate, sum, take the log, then add the max back. Skipping the max-subtraction is exactly what breaks on a point whose true probability under every component underflows to exactly 0.0.
Hint 3
For fit_gmm's initialization, rng.choice(n_samples, n_components, replace=False) picks distinct row indices — index input with them to get starting means. The loop body itself is just two calls: feed the current parameters into gmm_e_step to get responsibilities, then feed those into gmm_m_step to get the next parameters, and score the result with gmm_log_likelihood before moving to the next iteration.
Keywords
No results yet
Run your code to check it against the examples, or submit to run every test.