Fitting mixture weights¶

This section describes how to fit mixture weights for a convex combination of fixed candidate Gaussian mixtures. Given a reference mixture \(p\) and candidate mixtures \(q_1,\dots,q_N\), the goal is to choose non-negative weights \(w_i\) that sum to one so that the combined mixture \(q_{\mathbf{w}}=\sum_i w_i q_i\) approximates \(p\) as closely as possible. The primary formulation minimizes the forward KL divergence \(D_{\mathrm{KL}}(p\|q_{\mathbf{w}})\), with alternative reverse and bidirectional KL objectives discussed afterward.

Forward KL divergence formulation¶

Similar to the previous sections, let \(p\) be a fixed Gaussian mixture and let \(\{q_i\}_{i=1}^N\) denote a collection of fixed Gaussian mixtures. Specifically,

and, for each \(i \in \{1,\dots,N\}\),

Here, \(K_p\) is the number of components in \(p\), while \(K_i\) is the number of components in the mixture \(q_i\), which may vary with \(i\). All mixtures are fixed, meaning that the component weights, means, and covariances,

are treated as fixed quantities. The objective is to find a set of non-negative weights \(\{w_i\}_{i=1}^N\) that sum to one, such that the KL divergence from \(p\) to the mixture

is minimized. Formally, this can be expressed as the following optimization problem:

where

Importantly, the resulting mixture \(q_{\mathbf{w}}\) is itself a Gaussian mixture, with the number of components equal to the sum of the number of components across all \(q_i\):

Practical objective¶

Using the defenition of KL (see Kl estimation) divergence, the optimization problem in \(\eqref{eq:mixture-weight-optimization}\) can be rewritten as

Since \(p\) does not depend on \(\mathbf{w}\), the optimization problem is equivalent to

The objective function \(J(\mathbf{w})\) is the negative expected log-likelihood of the mixture \(q_{\mathbf{w}}\) under the distribution \(p\) and can generally not be expressed in closed form. However, it can be estimated using Monte Carlo sampling. Specifically, given \(M\) independent and identically distributed (iid) samples \(x^{(1)},\dots,x^{(M)}\) drawn from \(p\), we can construct the following estimator for \(J(\mathbf{w})\):

resulting in the final practical optimization problem

Alternative objective: reverse KL divergence¶

Alternatively, one could consider the reverse KL divergence from \(q_{\mathbf{w}}\) to \(p\):

In this case, the optimization problem would be

This optimization problem can also be estimated using Monte Carlo sampling, but it requires sampling from the mixture \(q_{\mathbf{w}}\), which itself depends on the optimization variable \(\mathbf{w}\). This can make the optimization more challenging, as the sampling distribution change as \(\mathbf{w}\) is updated.

However, the underlying components are fixed and it is possible to reuse samples from the candidate mixtures \(q_i\) to construct an estimator for the reverse KL divergence. For example, given \(M\) iid samples \(x_i^{(1)},\dots,x_i^{(M)}\) drawn from each candidate mixture \(q_i\), the following estimator for the reverse KL divergence can be constructed:

where

Alternative objective: bidirectional KL divergence¶

Another alternative is to consider a bidirectional KL divergence, which combines the forward and reverse KL divergences:

The optimization problem would essentially be a weighted combination of the two previous optimization problems, and can also be estimated using Monte Carlo sampling. In the case of \(\alpha=0.5\), the bidirectional KL divergence is symmetric and gives equal weight to both the forward and reverse KL divergences.

Alternative objective: Jensen-Shannon divergence¶

The Jensen-Shannon objective compares p and q_w through their midpoint mixture,

and minimizes

This gives a symmetric alternative to the directional KL objectives:

import gmm_divergence as gd

p = gd.GaussianMixture.from_components(
    [
        gd.Gaussian.univariate(mean=0.0, variance=0.5),
        gd.Gaussian.univariate(mean=2.0, variance=0.5),
    ],
    weights=[0.6, 0.4],
)
q1 = gd.Gaussian.univariate(mean=0.0, variance=0.5)
q2 = gd.Gaussian.univariate(mean=2.0, variance=0.5)

fit = gd.fit_mixture_weights(
    p,
    [q1, q2],
    objective=gd.fitting.JensenShannon(
        p_sampling=gd.sampling.Draw(10_000, rng=102), q_sampling=gd.sampling.Draw(10_000, rng=102)
    ),
)

For fitting objectives, p_sampling controls samples from the reference distribution and q_sampling controls one fixed batch per candidate distribution. Use gd.sampling.Samples(...) for precomputed reference samples and gd.sampling.SampleBatches(...) for precomputed candidate batches. gd.sampling.Stratified(...) can be used for either side; a single Gaussian is treated as a one-component mixture.

Example¶

The fit_mixture_weights function fits the weights of a mixture of candidate distributions \(q_i\) to a fixed reference mixture \(p\). For example:

import gmm_divergence as gd

p = gd.GaussianMixture.from_components(
    components=[
        gd.Gaussian.univariate(mean=0.0, variance=0.5),
        gd.Gaussian.univariate(mean=2.0, variance=0.5),
    ],
    weights=[0.6, 0.4],
)

q1 = gd.Gaussian.univariate(mean=0.0, variance=0.5)
q2 = gd.Gaussian.univariate(mean=2.0, variance=0.5)
result = gd.fit_mixture_weights(
    p,
    [q1, q2],
    method="simplex_slsqp",
    objective=gd.fitting.ForwardKL(sampling=gd.sampling.Draw(10_000, rng=102)),
)

assert result.converged
assert abs(result.weights[0] - 0.6) < 1e-2
assert abs(result.weights[1] - 0.4) < 1e-2

Here, the optimizer recovers the mixture weights of the reference distribution by combining the two candidate mixtures q1 and q2. The result keeps the scalar optimizer objective separate from the forward and reverse KL diagnostics, since the optimized objective depends on the selected fit direction.