Technical Summary: A Gradient Flow Perspective on Minimum MMD Estimation

Problem Statement

Minimum Maximum Mean Discrepancy (MMD) estimation has emerged as a robust, likelihood-free alternative to Maximum Likelihood Estimation (MLE) for parameter inference. It minimizes the discrepancy between a parametric model family $\{P_\theta : \theta \in \Theta\}$ and an unknown data-generating distribution $Q$. While the statistical properties of minimum MMD estimators (robustness to misspecification, likelihood-free computation) are well-established, the associated optimization problem, $\arg \min_{\theta \in \Theta} \text{MMD}^2(P_\theta, Q)$, remains poorly understood.

Existing optimization methods typically rely on standard (stochastic) gradient descent (GD). However, the MMD objective is generally non-convex, meaning standard GD is only guaranteed to converge to stationary points, not global minimizers. The only existing theoretical convergence guarantees for minimum MMD optimization (e.g., Chérief-Abdellatif and Alquier, 2022) rely on convexity assumptions that rarely hold in practice, as illustrated by simple Gaussian mixture models where the objective landscape is highly non-convex. The central open problem addressed is whether a gradient-based optimization scheme can be designed for minimum MMD estimation that offers convergence guarantees beyond convex objectives.

Methodology

The authors propose a Preconditioned Gradient Descent (PGD) scheme, inspired by recent advances in MMD gradient flows on the space of probability measures. The methodology bridges the gap between nonparametric gradient flows and parametric optimization through three key components:

1. Adaptive Nonparametric MMD Descent

The authors first analyze a nonparametric descent scheme where the distribution $P_t$ evolves directly in the space of probability measures $\mathcal{P}_2(\mathbb{R}^d)$ via the continuity equation:
$$P_{t+1} = (\text{Id} - \gamma_t \nabla f_{\ell_t, P_t, Q}) \# P_t$$
where $f_{\ell, P, Q}$ is the witness function and $\nabla f$ is the gradient field. Crucially, they introduce an adaptive kernel lengthscale schedule $(\ell_t)_{t \in \mathbb{N}}$, where the lengthscale starts large and monotonically decreases to a target $\ell > 0$. This "annealing" of the kernel is shown to improve the optimization landscape, smoothing the objective to avoid local traps. Under a trajectory-wise gradient-dominance condition, they prove that this adaptive nonparametric scheme converges globally to the target distribution $Q$ (Theorem 3.1).

2. Projection onto Parametric Subspace

To apply this to parametric models $P_\theta = (G_\theta)\#\rho$, the authors design a descent direction that approximates the nonparametric flow. Instead of computing the standard gradient of the MMD with respect to $\theta$, they project the nonparametric gradient field $\nabla f_{\ell_t, P_{\theta_t}, Q} \circ G_{\theta_t}$ onto the linear subspace of directions realizable by infinitesimal updates of the parameter $\theta$.
The update direction $\Delta \theta_t$ is obtained by solving a least-squares problem:
$$\Delta \theta_t = \arg \min_{u \in \mathbb{R}^p} \mathbb{E}_{z \sim \rho} \left[ \| J_\theta G_{\theta_t}(z) u - \nabla f_{\ell_t, P_{\theta_t}, Q}(G_{\theta_t}(z)) \|^2 \right]$$
This yields a closed-form solution involving a preconditioning matrix (the empirical Fisher-like matrix of the generator Jacobian) and a ridge regularization term $\lambda_t$:
$$\Delta \theta_t^{PGD} = \left( \mathbb{E}[J_\theta G_{\theta_t}^\top J_\theta G_{\theta_t}] + \lambda_t I \right)^{-1} \mathbb{E}[J_\theta G_{\theta_t}^\top \nabla f_{\ell_t, P_{\theta_t}, Q}(G_{\theta_t})]$$
This differs from Natural Gradient Descent (NGD), which uses a preconditioner derived from the geometry of the MMD itself; the PGD preconditioner is derived from minimizing the discrepancy between the parametric and nonparametric update directions.

3. Theoretical Guarantees

The authors establish global convergence for the PGD scheme (Theorem 4.6) under two main conditions:

1. Gradient Dominance: A condition relating the norm of the gradient field to the MMD value (or the gap to the optimal MMD in misspecified cases).
2. Projection Residual: A condition (Assumption 3) ensuring that the nonparametric gradient field is not orthogonal to the tangent space of the parametric model, and that the projection error is bounded.

The theorem proves that under these conditions, the PGD scheme drives the MMD objective to the global minimum (or the best achievable approximation in the misspecified case) even in non-convex settings, provided the kernel lengthscale is adaptively scheduled.

Key Results

The paper provides both theoretical and empirical evidence supporting the PGD scheme:

• Theoretical Convergence: Theorem 4.6 establishes asymptotic global convergence for PGD without assuming convexity of the MMD objective. This is a significant departure from prior work which required convexity.
• Nonparametric Convergence: Theorem 3.1 proves that adaptive lengthscale nonparametric MMD descent converges to the target distribution, validating the "annealing" intuition.
• Empirical Performance:
  • Gaussian Mixtures: In a two-mode Gaussian mixture example, standard GD gets trapped in local minima, while PGD successfully converges to the global minimizer.
  • G-and-k Distribution: On this likelihood-free benchmark, PGD and NGD outperform standard GD. PGD reaches the optimum in fewer iterations and handles challenging parameters better than GD.
  • Lotka-Volterra Model: Under data corruption (outliers), PGD consistently recovers the global minimizer across various corruption levels, whereas GD and NGD fail under poor initialization or high corruption.
  • Composite Goodness-of-Fit Test: In a toggle-switch gene expression model, PGD achieves well-calibrated Type-I error and higher statistical power than GD, even when GD uses a multi-start scheme (400 initializations) to mitigate non-convexity. PGD achieves this with a single initialization.
• Ablation Studies: Experiments indicate that both the adaptive lengthscale and the preconditioner contribute to performance, with their relative importance varying by problem (e.g., lengthscale is critical for Lotka-Volterra, while the preconditioner is more vital for G-and-k).

Significance and Claims

The paper claims to bridge the gap between the statistical theory of minimum MMD estimation and its practical computation. Its primary significance lies in providing the first global convergence guarantees for minimum MMD estimators without assuming convexity of the objective function.

By leveraging the geometric insights of MMD gradient flows and introducing an adaptive lengthscale schedule, the authors demonstrate that it is possible to design optimization schemes that escape local minima in non-convex landscapes. The proposed PGD scheme not only offers theoretical convergence guarantees under explicit (though non-trivial) conditions but also empirically outperforms standard gradient descent and natural gradient descent in robustness and efficiency across diverse statistical inference tasks, including those involving model misspecification and data corruption.

The authors remain modest regarding limitations, noting that their analysis relies on a priori conditions on the adaptive lengthscale and the projection residual, which are difficult to verify in practice. They also acknowledge that their convergence results are asymptotic and do not yet provide non-asymptotic rates with respect to both sample size and iteration count. Future work is suggested to extend these techniques to other minimum-distance estimation methods.