Learning Mixture Density via Natural Gradient Expectation Maximization

We present *natural gradient expectation maximization* (nGEM), a novel approach to learning mixture density networks for conditional density estimation. Mixture density networks are conventionally optimized via the negative log-likelihood (NLL) objective, which suffers from slow convergence and mode collapse. We instead cast mixture density networks as latent-variable models that can be optimized in an expectation maximization framework, and derive tractable factorized natural gradient updates drawing upon theoretical ties between expectation maximization and natural gradient descent. Empirically, we show that nGEM achieves up to 10$\times$ faster convergence with almost zero extra computational overhead, yet scales well to practical and challenging tasks where NLL otherwise fails for mode collapse.