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

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.