# Differentiable Expectation-Maximisation and Applications to Gaussian Mixture Model Optimal Transport

arXiv: https://arxiv.org/abs/2509.02109 · HTML: https://mlsanity.com/paper/2509.02109/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 47 of its 111 references that are in the map.
- First posted on arXiv 2 Sep 2025.

## Builds on (most cited first, 30 of 47)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks](https://mlsanity.com/paper/1511.06434/) (2015)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [A Neural Algorithm of Artistic Style](https://mlsanity.com/paper/1508.06576/) (2015)
- [Joint Distribution Optimal Transportation for Domain Adaptation](https://mlsanity.com/paper/1705.08848/) (2017)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Truncated Back-propagation for Bilevel Optimization](https://mlsanity.com/paper/1810.10667/) (2018)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](https://mlsanity.com/paper/2103.03606/) (2021)
- [Generative Modeling using the Sliced Wasserstein Distance](https://mlsanity.com/paper/1803.11188/) (2018)
- [Kernel Operations on the GPU, with Autodiff, without Memory Overflows](https://mlsanity.com/paper/2004.11127/) (2020)
- [DeepGMR: Learning Latent Gaussian Mixture Models for Registration](https://mlsanity.com/paper/2008.09088/) (2020)
- [Learning with minibatch Wasserstein : asymptotic and gradient properties](https://mlsanity.com/paper/1910.04091/) (2019)
- [Variational inference via Wasserstein gradient flows](https://mlsanity.com/paper/2205.15902/) (2022)
- [Smooth and Sparse Optimal Transport](https://mlsanity.com/paper/1710.06276/) (2017)
- [A Wasserstein-type distance in the space of Gaussian Mixture Models](https://mlsanity.com/paper/1907.05254/) (2019)
- [Faster Wasserstein Distance Estimation with the Sinkhorn Divergence](https://mlsanity.com/paper/2006.08172/) (2020)
- [Differential Properties of Sinkhorn Approximation for Learning with Wasserstein Distance](https://mlsanity.com/paper/1805.11897/) (2018)
- [Nonsmooth Implicit Differentiation for Machine Learning and Optimization](https://mlsanity.com/paper/2106.04350/) (2021)
- [Global analysis of Expectation Maximization for mixtures of two Gaussians](https://mlsanity.com/paper/1608.07630/) (2016)
- [Deep Gaussian Mixture Models](https://mlsanity.com/paper/1711.06929/) (2017)
- [Wasserstein barycenters are NP-hard to compute](https://mlsanity.com/paper/2101.01100/) (2021)

## Abstract (arXiv)

The Expectation-Maximisation (EM) algorithm is a central tool in statistics and machine learning, widely used for latent-variable models such as Gaussian Mixture Models (GMMs). Despite its ubiquity, EM is typically treated as a non-differentiable black box, preventing its integration into modern learning pipelines where end-to-end gradient propagation is essential. In this work, we present and compare several differentiation strategies for EM, from full automatic differentiation to approximate methods, assessing their accuracy and computational efficiency. As a key application, we leverage this differentiable EM in the computation of the Mixture Wasserstein distance MW2 between GMMs, allowing MW2 to be used as a differentiable loss in imaging and machine learning tasks. To complement our practical use of MW2, we contribute a novel stability result which provides theoretical justification for the use of MW2 with EM, and also introduce a novel unbalanced variant of MW2. Numerical experiments on barycentre computation, colour and style transfer, image generation, and texture synthesis illustrate the versatility of the proposed approach in different settings.