Mixture Density Network Loss

The problem, from first principles

Mixture Density Network (MDN) loss is used when the target distribution is multimodal. Instead of predicting a single value, the model predicts a mixture of Gaussians.

For K components, the model outputs:

• pi_k: mixture weights (sum to 1)
• mu_k: means
• sigma_k: standard deviations

The loss is the negative log-likelihood:

Where N is the Gaussian probability density.

From theory to code

Implement:

mdn_loss(y, pi, mu, sigma)

Where:

• y: targets (shape N,)
• pi: mixture weights (shape N, K)
• mu: means (shape N, K)
• sigma: standard deviations (shape N, K)

Constraints

• All inputs must have compatible shapes.
• sigma > 0.
• sum(pi) = 1 for each sample.
• Return scalar loss.