Gaussian Mixture Model (GMM) (spec model) #
This file defines a basic GMM with nComponents multivariate Gaussians over nFeatures:
- mixing weights $\pi$,
- means $\mu$, and
- covariances $\Sigma$.
gmmForwardSpec computes per-component log scores for a single input:
$$ \log \pi_k + \log \mathcal{N}(x \mid \mu_k, \Sigma_k). $$
PyTorch analogies:
torch.distributions.MultivariateNormalfor $\mathcal{N}(x \mid \mu, \Sigma)$,torch.distributions.MixtureSameFamilyfor mixture distributions,torch.softmaxfor turning per-component log-probabilities into responsibilities.
Mixture weights must be positive and normalized within the scalar backend's validation tolerance.
Covariances must be symmetric positive definite. Parameters outside this domain are reported as
none. Scores use the stored weights directly: when the weights sum to one, their exponential sum
is a probability density. The default normalization tolerance also applies to exact scalar
backends; callers needing a stricter gate can check mixtureWeightsValidSpec with zero tolerance.
Determinants and inverses are defined via NN.Spec.Core.Tensor.Numerics; those
definitions are intended for small feature dimensions and proof/reference usage, not
high-performance clustering on large matrices.
References (background, not required to read the code):
- Dempster, Laird, Rubin (1977), "Maximum Likelihood from Incomplete Data via the EM Algorithm": https://www.jstor.org/stable/2984875
- Bishop (2006), "Pattern Recognition and Machine Learning", Chapter 9 (Mixture Models and EM): https://www.microsoft.com/en-us/research/people/cmbishop/prml-book/
Implementation status #
No API builder implements this model. It is exercised numerically in
NN/Tests/Runtime/Floats/TorchLeanOpsCheck.lean; no theorem is proved about it.
Parameters #
Parameters of a Gaussian mixture model (GMM).
- weights : TorchLean.Tensor α [nComponents]
Mixing weights $\pi_k$. Model operations require every weight to be strictly positive and their compensated total to be within
Context.defaultEpsilonof one. - means : TorchLean.Tensor α [nComponents, nFeatures]
Component means $\mu_k$.
- covariances : TorchLean.Tensor α [nComponents, nFeatures, nFeatures]
Component covariance matrices $\Sigma_k$ (typically symmetric positive definite).
Instances For
Whether a matrix is symmetric under the scalar backend's equality operation.
Instances For
Executable Sylvester-criterion check for a symmetric positive-definite covariance matrix.
The matrix must be symmetric and every nonempty leading principal minor must be positive. This is the domain on which the Gaussian density, inverse, and logarithmic determinant used below have their usual meaning.
Instances For
Check that every mixture weight is positive and their total is within tolerance of one.
Floating-point division usually cannot represent the exact uniform weight 1 / n. A check against
exactly one can therefore reject gmmInitSpec before its first forward pass. The compensated total
limits error from the reduction, and the tolerance accounts for rounding in the weights themselves.
The admissible tolerance is in [0, 1), and the comparison uses absolute error. The default
Context.defaultEpsilon is the backend's configured numerical safeguard. It also applies to exact
scalar backends. Zero requires the computed compensated total to equal one; callers can use this
stricter check before invoking model operations, which use the default tolerance.
Instances For
Check the dimensions, mixing weights, and covariance matrices before evaluating the model.
The model needs at least one component and one feature. Its weights must be positive and satisfy the configured normalization tolerance, and every covariance must be symmetric positive definite. Accepted weights are used as stored; this check does not rescale them to sum to exactly one.
Instances For
Per-component log scores for a single input.
Given $x \in \mathbb{R}^d$, each component contributes
$$ \log \pi_k-\frac12\left( (x-\mu_k)^\mathsf{T}\Sigma_k^{-1}(x-\mu_k) +\log\det\Sigma_k+d\log(2\pi) \right). $$
This is the natural "logit vector" for responsibilities. If you want posterior probabilities
$P(z=k\mid x)$, apply gmmExpectationSpec (a last-axis softmax).
PyTorch analogy: the returned vector is like per-component logProb values before the final
mixture logsumexp.
Instances For
E-step responsibilities for a single input.
Mathematically:
$$ \gamma_k = P(z=k\mid x) = \operatorname{softmax}_k\!\left( \log \pi_k+\log\mathcal{N}(x\mid\mu_k,\Sigma_k) \right). $$
PyTorch analogy: torch.softmax(component_log_probs, dim=-1) where the logits are the
per-component log-probabilities.
Instances For
Batched forward pass: apply gmmForwardSpec to each sample in a batch.
Invalid model parameters return none, including for an empty batch.
Instances For
Backward/VJP (for gmmForwardSpec) #
gmmForwardSpec is vector-valued: it returns one log score per component.
The gradients below are the VJP for that vector function. In particular, responsibilities $\gamma=\operatorname{softmax}(\text{component log scores})$ do not appear in these formulas by themselves.
Responsibilities show up when you differentiate a scalar objective that aggregates components,
like the mixture log-likelihood $\operatorname{logsumexp}(\text{component log scores})$.
In that case, you compute the derivative with respect to the component log scores first
(which will involve $\gamma$), then feed that vector into
gmmBackwardSpec.
Gradient/VJP with respect to weights $\pi$ for the output of gmmForwardSpec.
For $y_k=\log\pi_k+\cdots$, we have $\partial y_k/\partial\pi_k=1/\pi_k$.
Instances For
Gradient/VJP with respect to means $\mu$ for the output of gmmForwardSpec.
For a single component:
$$ \frac{\partial}{\partial\mu}\log\mathcal{N}(x\mid\mu,\Sigma) =\frac12\left(\Sigma^{-1}+\Sigma^{-\mathsf{T}}\right)(x-\mu). $$
For a valid symmetric covariance this reduces to the familiar $\Sigma^{-1}(x-\mu)$.
Instances For
Gradient/VJP with respect to the input $x$ for the output of gmmForwardSpec.
For one component:
$$ \frac{\partial}{\partial x}\log\mathcal{N}(x\mid\mu,\Sigma) =-\frac12\left(\Sigma^{-1}+\Sigma^{-\mathsf{T}}\right)(x-\mu). $$
We sum the contributions from all components, weighted by the upstream gradient $g_k$.
Instances For
Gradient/VJP with respect to covariances $\Sigma$ for the output of gmmForwardSpec.
For one component:
$$ \frac{\partial}{\partial\Sigma}\log\mathcal{N}(x\mid\mu,\Sigma) =\frac12\left( \Sigma^{-\mathsf{T}}(x-\mu)(x-\mu)^\mathsf{T}\Sigma^{-\mathsf{T}} -\Sigma^{-\mathsf{T}} \right). $$
Instances For
Gradients for a GMMSpec, one field per differentiable component.
- weightsGradient : TorchLean.Tensor α [nComponents]
Gradient with respect to the mixture weights.
- meansGradient : TorchLean.Tensor α [nComponents, nFeatures]
Gradient with respect to the component means.
- covariancesGradient : TorchLean.Tensor α [nComponents, nFeatures, nFeatures]
Gradient with respect to the component covariance matrices.
- inputGradient : TorchLean.Tensor α [nFeatures]
Gradient with respect to the observed point.
Instances For
Backward/VJP for gmmForwardSpec.
Every component gradient needs a covariance inverse, so the whole record is optional: a singular covariance makes the derivative undefined rather than zero.
Instances For
Default initialization for a GMM.
This is kept simple and deterministic:
- uniform weights,
- zero means,
- identity covariances.
Instances For
Numerically stable log-sum-exp reduction: $\log\!\left(\sum_i \exp(\mathtt{log\_probs}[i])\right)$.
This is the standard $m+\log\!\left(\sum_i\exp(x_i-m)\right)$ trick, where $m=\max_i x_i$.
Instances For
Mixture log-likelihood $\log p(x)$ computed via log-sum-exp over components.
Mathematically, $$ \log p(x)=\log\!\left( \sum_k \exp\!\left(\log p(x\mid z_k)+\log\pi_k\right) \right). $$
This is the log of a normalized mixture density when the stored weights sum to one.
The default domain check allows Context.defaultEpsilon error in that sum; accepted weights
are used directly without renormalization.
Instances For
Classical training: EM for Gaussian mixtures #
For a GMM, “training” is typically done with the Expectation–Maximization (EM) algorithm:
- E-step: compute responsibilities $r_{ik}=P(z=k\mid x_i)$ for each sample/component.
- M-step: update $\pi$, $\mu$, and $\Sigma$ from the weighted sufficient statistics.
This file already provides gmmExpectationSpec (responsibilities for one sample). The helpers
below lift that to a batched dataset and implement a deterministic EM update step.
Numerical notes:
- If a component gets (near) zero total responsibility ($N_k\approx 0$), we keep that component’s parameters unchanged (otherwise we’d divide by zero).
- We add a small diagonal “jitter,” $\mathtt{Context.defaultEpsilon}\,I$, to covariances to keep them well-behaved.
Batched responsibilities: apply gmmExpectationSpec to each sample.
Invalid model parameters return none, including for an empty dataset.
Instances For
One EM step for a batched dataset.
An empty dataset preserves a valid model; invalid parameters still return none.
Instances For
Sum of negative gmmLogLikelihoodSpec scores over the dataset.
An empty dataset has score zero for a valid model; invalid parameters still return none.
Instances For
Run epochs EM steps (deterministic).
Zero epochs preserve a valid model; invalid initial parameters still return none.