Learning Mixture Density via Natural Gradient Expectation Maximization

TL;DR

This paper introduces nGEM, a natural gradient-based optimization method for mixture density networks, significantly improving convergence speed and scalability by leveraging information geometry and an EM framework.

Contribution

It develops the nGEM algorithm, integrating natural gradient descent with EM for mixture density networks, enhancing training efficiency and scalability.

Findings

nGEM achieves up to 10x faster convergence.

nGEM scales well to high-dimensional data.

It adds minimal computational overhead.

Abstract

Mixture density networks are neural networks that produce Gaussian mixtures to represent continuous multimodal conditional densities. Standard training procedures involve maximum likelihood estimation using the negative log-likelihood (NLL) objective, which suffers from slow convergence and mode collapse. In this work, we improve the optimization of mixture density networks by integrating their information geometry. Specifically, we interpret mixture density networks as deep latent-variable models and analyze them through an expectation maximization framework, which reveals surprising theoretical connections to natural gradient descent. We then exploit such connections to derive the natural gradient expectation maximization (nGEM) objective. We show that empirically nGEM achieves up to 10$\times$ faster convergence while adding almost zerocomputational overhead, and scales well to high-dimensional data where NLL otherwise fails.