Cramer Type Distances for Learning Gaussian Mixture Models by Gradient Descent

TL;DR

This paper introduces a new gradient-compatible distance function called Sliced Cramér 2-distance for learning Gaussian Mixture Models, enabling seamless integration with neural networks and offering theoretical guarantees.

Contribution

It derives a closed-form formula for univariate GMMs and proposes a new distance for multivariate GMMs that is easy to compute and compatible with gradient descent.

Findings

Closed-form formula for univariate GMMs.

Proposed Sliced Cramér 2-distance for multivariate GMMs.

Demonstrated effectiveness in distributional reinforcement learning.

Abstract

The learning of Gaussian Mixture Models (also referred to simply as GMMs) plays an important role in machine learning. Known for their expressiveness and interpretability, Gaussian mixture models have a wide range of applications, from statistics, computer vision to distributional reinforcement learning. However, as of today, few known algorithms can fit or learn these models, some of which include Expectation-Maximization algorithms and Sliced Wasserstein Distance. Even fewer algorithms are compatible with gradient descent, the common learning process for neural networks. In this paper, we derive a closed formula of two GMMs in the univariate, one-dimensional case, then propose a distance function called Sliced Cram\'er 2-distance for learning general multivariate GMMs. Our approach has several advantages over many previous methods. First, it has a closed-form expression for the univariate case and is easy to compute and implement using common machine learning libraries (e.g., PyTorch and TensorFlow). Second, it is compatible with gradient descent, which enables us to integrate GMMs with neural networks seamlessly. Third, it can fit a GMM not only to a set of data points, but also to another GMM directly, without sampling from the target model. And fourth, it has some theoretical guarantees like global gradient boundedness and unbiased sampling gradient. These features are especially useful for distributional reinforcement learning and Deep Q Networks, where the goal is to learn a distribution over future rewards. We will also construct a Gaussian Mixture Distributional Deep Q Network as a toy example to demonstrate its effectiveness. Compared with previous models, this model is parameter efficient in terms of representing a distribution and possesses better interpretability.