Mixture Density Networks (MDN)¶

← Evidential Regression | Normalizing Flows →

Mixture Density Networks model the conditional distribution \(p(y \mid x)\) as a mixture of Gaussian components, capturing multimodality, heteroscedastic noise, and complex distributional shapes.

Mathematical Background¶

The NLL is:

computed via log-sum-exp for numerical stability.

See the MixtureDensityLoss API for constructor parameters and the MixtureDensityLoss alias.

Usage¶

from torchregress.losses import MixtureDensityLoss

loss_fn = MixtureDensityLoss(
    n_components=3,    # K = number of Gaussian components
    n_features=2,      # d = dimensionality of targets
    covariance_type="diagonal",  # or "full"
    min_std=1e-3,
)

Model Output Size¶

# K=3, d=2 → output_dim = 3(1 + 4) = 15
model = nn.Sequential(
    nn.Linear(in_dim, 128), nn.ReLU(),
    nn.Linear(128, 15),
)

Full covariance: precision or covariance factor¶

With covariance_type="full" the network outputs, per component, a lower-triangular \(T_k\). full_parameterization sets what it means:

The precision form is linear in \(T_k\) and needs no triangular solve, which makes it much easier to optimise when there are many targets. It is not better everywhere. Harness run #6 (density_parity, test NLL, mean of 3 seeds, lower is better) against run #5 (covariance form), with sbi's MDN on the same splits:

Training¶

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(100):
    params = model(x_train)           # (batch, output_dim)
    loss = loss_fn(params, y_train)   # MDN NLL
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Inference Methods¶

predict_mean_std¶

Gaussian-approximation summary (mean and variance of the mixture):

mean, std = loss_fn.predict_mean_std(model(x_test))

predict_interval¶

MC-based prediction intervals that respect multimodality:

lower, upper = loss_fn.predict_interval(
    model(x_test), confidence=0.95, n_samples=10_000,
)

sample¶

Draw samples from the learned mixture:

samples = loss_fn.sample(model(x_test), n_samples=1000)
# samples shape: [n_samples, batch, n_features]

Practical Tips¶

log_weights, means, scales = loss_fn._extract_distribution_parameters(params)
active = (log_weights.exp() > 0.05).float().mean()  # should stay > 0.5

Limitations & Optimization Challenges¶

While MDNs are highly expressive, they present several training and optimization challenges:

1. Numerical Instability & Log-Sum-Exp: Computing the NLL directly via \(p(y \mid x) = \sum_k \pi_k \mathcal{N}(y \mid \mu_k, \sigma_k^2)\) involves exponentials that can easily underflow or overflow. torchregress internally utilizes the log-sum-exp trick: $\(\log p(y \mid x) = \text{log-sum-exp}\left(\log \pi_k + \log \mathcal{N}(y \mid \mu_k, \sigma_k^2)\right)\)$ Always use a non-zero min_std (e.g., \(10^{-3}\)) to prevent variance collapse (\(\sigma_k^2 \rightarrow 0\)), which causes division by zero.

2. Computational Scaling: The forward pass cost is \(\mathcal{O}(K \cdot d)\) for diagonal covariance and \(\mathcal{O}(K \cdot d^2)\) for full covariance per sample. NLL evaluation adds determinant and solve operations that scale as \(\mathcal{O}(d^3)\) for full covariance. For large \(K\) or high-dimensional targets, prefer diagonal covariance or consider flow-based alternatives.

3. Mode Collapse: The network may collapse into using only a single mixture component, effectively ignoring the others (\(\pi_k \approx 0\) for \(k > 1\)). To monitor this, check active components during validation: $\(\text{active\_fraction} = \frac{1}{B} \sum_{i=1}^B \sum_{k=1}^K \mathbf{1}_{\pi_k^{(i)} > 0.05}\)$ Mode collapse is especially common with large \(K\) (> 10) or strong regularization.

4. Initialization Sensitivity: If initialized poorly, components can "claim" the same regions of target space. Initialize the final output layer weight matrices with small random values to diversify components early in training.

5. Dtype Sensitivity: The log-sum-exp trick is numerically sensitive in float16 or bfloat16. Use float32 for the NLL computation and only cast down after the loss is computed. Mixed-precision training with MDNs requires careful gradient scaling.

6. Covariance Type Constraints: With "full" covariance the head has \(K \cdot d(d+1)/2\) factor outputs and memory grows as \(\mathcal{O}(K \cdot d^2)\) per sample. The default precision parameterisation trains on 16 targets (scm20d above); for very large \(d\) prefer diagonal components or a flow.

Next steps¶

• Normalizing Flows — more flexible distributions for complex targets
• EIV + MDN — marginalize over input noise with MDN heads
• MDN Ensembles — mixture-of-mixtures for decomposed uncertainty
• Multimodal comparison — benchmark MDN vs flows vs SLS

References¶