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

Samuel Boïté; Eloi Tanguy; Julie Delon; Agnès Desolneux; Rémi Flamary

cs.LGSep 2, 2025

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

Authors: Samuel Boïté, Eloi Tanguy, Julie Delon, Agnès Desolneux, Rémi Flamary

Organizations: Université Paris Cité, CNRS, MAP5, F-75006 Paris, France · Centre Borelli, CNRS and ENS Paris-Saclay, F-91190 Gif-sur-Yvette, France · CMAP, CNRS, Ecole Polytechnique, Institut Polytechnique de Paris

Abstract

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 $\mathrm{MW}_2$ between GMMs, allowing $\mathrm{MW}_2$ to be used as a differentiable loss in imaging and machine learning tasks. To complement our practical use of $\mathrm{MW}_2$ , we contribute a novel stability result which provides theoretical justification for the use of $\mathrm{MW}_2$ with EM, and also introduce a novel unbalanced variant of $\mathrm{MW}_2$ . Numerical experiments on barycentre computation, colour and style transfer, image generation, and texture synthesis illustrate the versatility of the proposed approach in different settings.

Figures & tables

Panel of Figure 1 — Figure 1 : Left: sensitivity of the EM output to perturbation. Right: particle flow with repulsive energy.

Method	Time	Memory
Full Automatic Differentiation (AD)	$\mathcal{O}\bigl(T(nKd^{2}+Kd^{3})\bigr)$	$\mathcal{O}(TKd^{2}+nd)$
Approximate Implicit Gradient (AI)	$\mathcal{O}\bigl(TnKd^{2}+K^{3}d^{6}+nK^{2}d^{5}\bigr)$	$\mathcal{O}(nK^{2}d^{4})$
One-Step gradient (OS)	$\mathcal{O}\bigl(T(nKd^{2}+Kd^{3})\bigr)$	$\mathcal{O}(Kd^{2}+nd)$

Table 1 : Complexities of the backward passes (gradient computations) for

T

EM iterations on

n

points in

\mathbb{R}^{d}

with

K

components.

Panel of Figure 2 — Figure 2 : Sample complexity of $\mathrm{MW}_{2}^{2}$ for three specific GMMs of “low” to “high” mode separation. Curves show the median across $40$ repetitions, with shaded interquartile ranges.

Panel of Figure 3 — Figure 3 : Comparison of experimental setup and flows of $\mathcal{E}_{T}$ using different methods. The dark shades of purple correspond to earlier iterations, and the yellow shades to later iterations. (Warm-Start and AI are almost identical to AD)

Panel of Figure 4 — Figure 4 : Flow of $\mathcal{E}_{T}$ with the AD method and standard EM Algorithm 1 . We compare two settings: one with non-uniform GMM weights (NU) and one with uniform weights (U).

Figure 5 : Stochastic Flow of $\mathcal{E}_{T}$ for Algorithm 2 with the full automatic differentiation method. We vary the sub-sampling ratio $r\in(0,1]$ , which corresponds to performing EM on only $[r\times n]$ random points from the current point cloud at each step.

Figure 6 : Varying the number of samples

n

and the number of iterations

T

, we study the convergence of EM, the local contractivity of

F

, and the MSEs of the OS and AI gradients against the AD gradient.

Figure 7 : $\mathrm{EM}-\mathrm{MW}_{2}^{2}$ flow displacing particles in order to make their EM output approach a barycentre of three target GMMs.

Panel of Figure 8 — Figure 8 : Colour transfer experiments. Top row: $\mathrm{EM}-\mathrm{MW}_{2}^{2}$ from source to target. Bottom row: unbalanced colour transfer with regularisations $\lambda_{1}=10$ (source) and $\lambda_{2}=0.1$ (target).

Panel of Figure 9 — Figure 9 : Style transfer method inspired by [ GEB15 ] : setup and example result.

Panel of Figure 10 — Figure 10 : Choosing the number of components $K$ for texture synthesis.

Panel of Figure 11 — Figure 11 : Multi-scale texture synthesis with $K=4$ components for $8\times 8$ patches

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 12 — Figure 12 : Local behaviour of the EM map in a symmetric two-component example. The grey dashed line is the identity map.

Figure 13 : Two fixed GMMs

\mu

(top) and

\nu

(bottom) used in Section 3.2 . Columns correspond to low (

\sigma=10

), medium (

\sigma=0.5

), and high (

\sigma=0.1

) separation.

Figure 14 : 3D representation of the three GMMs used in Section 4.5 to compare EM gradient methods.

Figure 15 : Relative MSE of the full automatic differentiation gradient (AD) against the finite differences approximation (FD) for three different GMMs and varying the FD step size

\varepsilon_{\mathrm{FD}}

Sample	$T$	$\mathrm{MW}_{2}^{2}$ (warm)	$\mathrm{MW}_{2}^{2}$ (cold)	Difference
#1	1	676.5595	662.2502	14.3093
#1	2	674.8875	670.1658	4.7217
#1	5	674.6132	674.6144	0.0012
#1	10	674.6132	674.6132	0
#2	1	680.5627	622.3420	58.2207
#2	2	680.4759	629.5430	50.9329

Appendix

Table 2 : Comparison between warm-start and cold-start runs.

Figure 16 : Setup for the local minimum of

\mathrm{W}_{2}^{2}(\mu_{\alpha,\eta},\nu)

. The support of

\mu

is represented with blue squares, and its weights with vertical blue lines. For the target

\nu

, its support is red squares and its weights red lines. We consider a specific region where the points

x_{1}=\eta_{1}

and

x_{2}=\eta_{2}

stay within

(-\frac{1}{2},\frac{1}{2})

and the weights

a_{1}=\frac{1}{6}+\alpha_{1}

and

a_{2}=\frac{1}{6}+\alpha_{2}

stay within

[0,\frac{1}{3}]

, as represented by the orange rectangle. Likewise, the point

x_{3}=1+\eta_{3}

must stay within

(\frac{1}{2},\frac{3}{2})

and its weight

a_{3}=\frac{2}{3}-(\alpha_{1}+\alpha_{2})

must stay in

[\frac{1}{3},1]

, as shown with the purple rectangle.

Figure 17 : With the data

X_{\varepsilon}:=(x_{1},\cdots,x_{6})

, one iteration of the EM algorithm initialised at

\theta^{\star}

(corresponding to the GMM

\mu^{\star}

) yields approximately the same parameters

\theta^{\star}

. We shall see that the gradient of the energy

\mathcal{E}_{\mathrm{EM-}\mathrm{MW}_{2}^{2}}

X_{\varepsilon}

is approximately zero, illustrating the vanishing gradient phenomenon.

Panel of Figure 18 — Figure 18 : Local minimum in Example A.2 with fixed non-uniform weights.

Panel of Figure 19 — Figure 19 : Impact of the treatment of source weights in colour transfer.

Figure 20 : Comparison of balanced and unbalanced particle flows for non-uniform weights.

Figure 21 : Varying the number of components

K

, we study the convergence of EM, the local contractivity of

F

, and the MSEs of the OS and AI gradients.

Panel of Figure 22 — Figure 22 : Left: $\mathrm{EM}-\mathrm{MW}_{2}^{2}$ -Barycentre flow; right: input point clouds.

Figure 23 : 3D barycentre (left) and its projections (right).

Panel of Figure 24 — Figure 24 : Final GMMs in RGB space with (a) variable weights and (b) fixed weights. Target is in blue and optimised mixture is in red. In (a) we are stuck in a local minimum, in (b) we converged.

Panel of Figure 25 — Figure 25 : Colour transfer with (a) $K=1$ components and (b) $K=10$ components.

Panel of Figure 26 — Figure 26 : Colour transfer: barycentric transfer with entropic OT, and our method.

Panel of Figure 27 — Figure 27 : Four images and their corresponding log-likelihood vs. $K$ plots, for VGG layers $\ell\in\{1,\cdots,3\}$ .

Panel of Figure 28 — Figure 28 : Taking $K=1$ (top) gives results comparable to $K=3$ (bottom).

Panel of Figure 29 — Figure 29 : Evolution of the style transfer across iterations.

Panel of Figure 30 — Figure 30 : (a) Texture synthesis with $12\times 12$ patches, $K=4$ , mono-scale, (b) reference.

Panel of Figure 31 — Figure 31 : Image generation with the $\mathrm{MW}_{2}^{2}$ -GAN from Eq. 82 .

Explore similar work

Averaged Mirror Descent and Dual Gradient Methods: Convergent Algorithms for Entropic Gromov-Wasserstein Problems
Sep 25, 2026Joanna Marks, Gabriel Rioux, Riccardo PasseggeriGromov--Wasserstein Mirror Descent
⋆Denotes equal contribution to this work.
On the Wasserstein Gradient Flow Interpretation of Drifting Models
May 6, 2026Arthur Gretton, Li Kevin Wenliang, Alexandre Galashov +3Wasserstein Gradient Flows Generative Models
Google DeepMind
Distance-Matrix Wasserstein Statistics for Scalable Gromov--Wasserstein Learning
May 14, 2026Ao Xu, Tieru WuGromov--Wasserstein Wasserstein Distance
School of Artificial Intelligence Jilin University No. 2699, Qianjin Street, Chaoyang District Changchun 130012, China · Zhongguancun Academy Daniufang 2nd Ring Road, Haidian District2026 Beijing 100094, China · School of Artificial IntelligenceMay Jilin University No. 2699, Qianjin Street, Chaoyang District14 Changchun 130012, China

Sep 25, 2026, cs.LGJ/K move · Enter open · S save

Averaged Mirror Descent and Dual Gradient Methods: Convergent Algorithms for Entropic Gromov-Wasserstein Problems

Joanna Marks, Gabriel Rioux, Riccardo Passeggeri

⋆Denotes equal contribution to this work.

Gromov--Wasserstein Mirror Descent Convergence

Open paper AskPDF arXiv

The Gromov-Wasserstein (GW) distance measures the discrepancy between metric measure (mm) spaces and identifies optimal alignments between them based solely on their intrinsic structure. Since it identifies isomorphic mm spaces, it provides a natural notion of distance for heterogeneous datasets which may admit isomorphic representations. In order to accelerate computation of GW distances, many practitioners employ entropic regularization to obtain an Entropic GW (EGW) problem. The most popular EGW solver is the Mirror Descent (MD) algorithm, which reduces EGW computations to an iterative process where an entropic optimal transport (EOT) problem is solved at each iteration. Despite its widespread use, the convergence of MD for this problem has only been established for restricted classes of costs. On the other hand, a recently proposed dual gradient method is available for general costs, but requires a choice of step size which depends on the regularization parameter. To address these two issues, we introduce Averaged Mirror Descent (AMD), which averages consecutive MD steps, and prove its convergence for arbitrary costs. Then, we establish that the dual gradient method with a fixed step size also converges for arbitrary costs at the cost of a more complicated iteration. In both cases, we also account for inexact iterations which are inescapable in practice. We compare the empirical performance of these methods across various settings and, in particular, show that AMD and the dual gradient method both converge on an example where classical MD fails.

Abstract

Explore similar work

Averaged Mirror Descent and Dual Gradient Methods: Convergent Algorithms for Entropic Gromov-Wasserstein Problems

On the Wasserstein Gradient Flow Interpretation of Drifting Models

Distance-Matrix Wasserstein Statistics for Scalable Gromov--Wasserstein Learning