PaperlayerView exact revision on arXiv

Source-linked AI summary

Maximum Mean Discrepancy Gradient Flow

Michael Arbel, Anna Korba, Adil Salim, Arthur Gretton

arXiv:1906.04370v2stat.MLcs.LG

TL;DR

The paper studies how to transport an initial probability distribution toward a target using a Wasserstein gradient flow of the MMD, without requiring an analytical target density. It establishes convergence conditions, proposes noise injection to regularize the flow, and develops a practical sample-based discretization. The resulting framework provides theoretical convergence guarantees and a simple implementation connected to particle transport, sampling, and neural-network optimization.

  • Problem

    The paper addresses distribution transport when the target may be available only through samples, while global convergence is difficult because the MMD functional is non-convex.

  • Method

    The paper constructs continuous and discrete Wasserstein MMD flows, analyzes their convergence, and regularizes them by injecting noise into the gradient.

  • Results

    The paper establishes conditions for convergence to the global optimum and reports that noise injection converges in an illustrative case where the original MMD flow fails.

  • Takeaways & Limitations

    MMD flow offers a sample-accessible transport and sampling procedure with a practical discretized implementation and connections to particle-based neural-network optimization.

  • Takeaways & Limitations

    The discrete flow is generally intractable because its updates require exact knowledge of the evolving distribution and kernel gradient.

Abstract

from arXiv · show

We construct a Wasserstein gradient flow of the maximum mean discrepancy (MMD) and study its convergence properties. The MMD is an integral probability metric defined for a reproducing kernel Hilbert space (RKHS), and serves as a metric on probability measures for a sufficiently rich RKHS. We obtain conditions for convergence of the gradient flow towards a global optimum, that can be related to particle transport when optimizing neural networks. We also propose a way to regularize this MMD flow, based on an injection of noise in the gradient. This algorithmic fix comes with theoretical and empirical evidence. The practical implementation of the flow is straightforward, since both the MMD and its gradient have simple closed-form expressions, which can be easily estimated with samples.

1 Introduction

The paper defines an MMD Wasserstein gradient flow and establishes conditions for its global convergence. It also introduces noise regularization and a sample-based implementation relevant to particle transport and sampling.

  • The flow connects theoretically to particle transport interpretations of GAN training and neural-network optimization.
  • The paper establishes conditions under which MMD gradient flow converges to its global optimum.
  • MMD flow can approximate target distributions from samples, unlike Langevin Monte Carlo and SVGD, which require an analytical target density.
  • The convergence analysis addresses non-convexity and provides a barrier, a sufficient condition for global optimality, and convergence rates under that condition.
  • Noise injection regularizes the MMD flow and yields convergence in an illustrative case where the original flow fails.
  • The method has an O(MN + N^2) implementation using N samples from ν and M samples from µ, requiring kernel-gradient evaluations on those samples.

2 Gradient flow of the MMD in W2

The paper constructs continuous and discrete Wasserstein gradient flows for the MMD, transporting an initial distribution toward a target while decreasing the discrepancy. Under kernel regularity assumptions, the flow is well posed, its energy decreases, and the Euler scheme approximates the continuous dynamics for small step sizes.

  • MMD formulation: For a characteristic kernel with Lipschitz gradient, the MMD defines a distance on probability measures with finite second moment.
  • MMD formulation: The MMD differential is the witness function, whose induced potentials represent target-generated confinement and particle interactions.
  • Continuous-time flow: The continuous flow transports ν0 toward µ through the continuity equation with velocity Vt(x) = −∇fµ,νt(x).
  • Continuous-time flow: The nonlinear velocity field depends on the evolving distribution, yielding a McKean–Vlasov particle process with a unique solution under Assumption (A).
  • Continuous-time flow: The MMD functional decreases along the continuous flow and acts as a Lyapunov functional for its dynamics.
  • Euler scheme: The forward-Euler discretization converges to the continuous flow on bounded time intervals as the step size tends to zero.
  • Euler scheme: For γ ≤ 2/3L, the discrete sequence F(νn) is decreasing, although exact updates generally require knowledge of the full evolving distribution.

3 Convergence properties of the MMD flow

The MMD flow decreases the objective but need not reach a global optimum because the functional lacks displacement convexity. Global convergence is obtained under additional boundedness conditions, while a noise-regularized flow offers more tractable convergence conditions.

  • Optimization in a (W2) non-convex setting: The MMD objective decreases along the flow, but its limit can be a local minimum rather than a global optimum.The equilibrium condition alone does not guarantee global optimality for characteristic kernels.
  • Optimization in a (W2) non-convex setting: F is merely Λ-convex because a term in the Λ lower bound can become negative, so displacement convexity does not generally hold.The opposing-sign kernel terms prevent convexity arguments from compensating for one another.
  • Optimization in a (W2) non-convex setting: Theorem 6 gives an upper bound on the asymptotic objective for discrete gradient descent, but the bound can be loose when the average K̄ is negative.When K̄ is non-negative, the usual convergence rate is recovered.
  • A condition for global convergence: A Łojasiewicz-type analysis proves convergence to the global optimum when the weighted negative Sobolev distance between µ and ν_t remains bounded.This distance is linked to the squared weighted Sobolev seminorm governing the objective's time derivative.
  • A condition for global convergence: The weighted negative Sobolev distance linearizes the Wasserstein distance and acts as the boundedness quantity required by the convergence proof.The paper notes that guaranteeing this condition in general remains difficult and may require regularization.

4 A practical algorithm to descend the MMD flow

The paper regularizes noisy MMD updates to help escape local optima and implements the flow with interacting particles and sampled target data. The resulting scheme has finite-sample convergence guarantees and performs well empirically.

  • 4.1 A noisy update as a regularization: Noise injection evaluates the witness gradient at perturbed points, helping escape local minima where gradients vanish on the particle support.The update differs from diffusion because the MMD global optimum remains a fixed point under injection.
  • 4.2 The sample-based approximate scheme: The particle algorithm approximates the flow using samples from the initial and target distributions, with Gaussian perturbations and empirical measures.When βn = 0, the update becomes sample-based gradient descent over particles.
  • 4.2 The sample-based approximate scheme: The algorithm costs O((M + N)N) time and O(M + N) memory per iteration, falling to O(M + N) time with random-feature kernels.This contrasts with the cubic cost of KSD flow, which requires solving a linear system each iteration.
  • 4.2 The sample-based approximate scheme: The particle approximation converges toward the population flow as the numbers of particles and target samples increase, including the noiseless case.Theorem 9 controls propagation of chaos at each iteration, and the dependence on M supports sample-based use when only target samples are available.
  • Experiments: Training error decreases below 10^-5 with noise injection, producing better validation error than noiseless MMD, diffusion, and KSD in the student-teacher experiment.Noise injection also performs best across a wider noise range, whereas diffusion works only in a narrow low-noise region.

5 Conclusion

The paper introduces MMD flow as a Wasserstein gradient flow over probability distributions, together with a discretized implementation and regularization scheme. It establishes theoretical properties and convergence guarantees while identifying future applications in sampling and neural-network optimization.

  • 5 Conclusion: The practical implementation uses a space-time discretization of the flow rather than requiring exact access to the evolving distribution.The paper develops a finite-sample algorithm that converges toward the time-discretized flow as sample size increases.
  • 5 Conclusion: Future work targets deeper regularization theory and applications to sampling and optimization for large neural networks.These directions follow the paper’s stated scope rather than an established empirical result.
  • 5 Conclusion: The analysis relies on Wasserstein geometry, continuity equations, RKHS structure, and displacement-convexity concepts.The paper distinguishes displacement geodesics used for the flow from mixture geodesics naturally associated with MMD distances.
  • 5 Conclusion: MMD flow transports probability distributions using the negative Wasserstein gradient of an MMD-based functional.The MMD compares distributions through differences in RKHS mean embeddings and can be viewed as an integral probability metric.
  • 5 Conclusion: The continuous flow decreases its objective over time, making the functional a Lyapunov function for the distributional dynamics.Its time-discretized counterpart provides a gradient-descent interpretation on probability space.

A.5 Comparison with the Kullback Leilber divergence flow

The appendix contrasts MMD flow with KL-based Wasserstein flows and sample-based sampling methods. It highlights MMD’s non-convexity and states regularity conditions supporting well-posedness and energy decay.

  • Continuity equation and McKean Vlasov process: Gaussian-kernel interaction energy is attractive at long distances and repulsive at short distances, helping prevent particle collapse.This role is compared with the entropy term in KL-based Langevin dynamics.
  • Convergence to a global minimum: KL gradient flow converges toward its target under mild assumptions because KL is displacement convex along Wasserstein geodesics.The MMD lacks displacement convexity in general, making global convergence harder to establish.
  • Sampling algorithms derived from gradient flows: MMD flow can operate with samples from the target, unlike density-based Langevin methods that require analytical access to ∇log µ.This places MMD flow in the sample-based setting when only target samples are available.
  • Kernel assumptions: Under the kernel assumptions, Lipschitz regularity of ∇k guarantees existence and uniqueness of the MMD flow solution.The gradient field is Lipschitz in both position and distribution under the stated assumptions.
  • Convergence to a global minimum: The MMD functional decreases along the flow through an energy identity, providing decay of the objective over time.The proof uses the first variation represented by the witness-gradient field and λ-displacement convexity.

C.2 Time-discretized flow

The time-discretized MMD flow is analyzed as an approximation to the continuous process. Under Lipschitz assumptions and sufficiently small step size, its Wasserstein error remains controlled over finite time horizons and the objective decreases.

  • C.2 Time-discretized flow: The discrete scheme iteratively applies the MMD gradient field at the current distribution with step size γ.The analysis couples continuous and discrete trajectories through one-step approximation errors.
  • C.2 Time-discretized flow: For γ ≤ 1/(4L), the Wasserstein distance between the continuous and discretized flows is bounded over finite time.The bound follows from Lipschitz control, Taylor expansion, and recursive error estimates.
  • C.2 Time-discretized flow: The coupling analysis yields recursive Wasserstein error inequalities whose growth is controlled over iterations up to time T.The discrete-continuous discrepancy includes accumulated one-step errors of order γ^2 before applying the recursion.
  • C.2 Time-discretized flow: The finite-time approximation bounds depend on the kernel’s Lipschitz regularity and constants determined by the time horizon.These assumptions support the estimates used for the continuous and discrete processes.
  • C.2 Time-discretized flow: The discrete update decreases the MMD functional under the fixed-step analysis.The proof examines the push-forward path generated by the current gradient field and establishes absolute continuity along that path.

D.1 Equilibrium condition

The equilibrium condition for MMD flow is characterized by the witness function lying in the null space of a Hilbert-Schmidt operator. This generalizes a finite-dimensional result and clarifies when stationarity need not imply global optimality.

  • Equilibrium characterization: Condition (11) is equivalent to fµ,ν∗ belonging to null(Dν∗).The operator Dν∗ is symmetric, non-negative, and Hilbert-Schmidt.
  • Equilibrium characterization: Proposition 13 generalizes the equivalence between the equilibrium condition and null-space membership from finite to infinite dimensions.The finite-dimensional equivalence was previously established for KSD flow.
  • Convexity analysis: The proof of Λ-displacement convexity evaluates the second derivative of F along constant-speed Wasserstein geodesics.Under Assumption (A), the second derivative is bounded below by a quadratic form Λ(ρt,Vt).

D.3 Descent up to a barrier

The analysis establishes descent and convergence guarantees for discrete MMD flow despite nonconvexity, using Λ-displacement convexity and a barrier term along Wasserstein geodesics. The resulting bounds depend on step size, smoothness, and the accumulated barrier.

  • Discrete descent: The discrete flow satisfies a descent inequality combining the objective gap, a gradient-related term, Wasserstein distance decrease, and the barrier K(ρn).For sufficiently small γ, smoothness allows the additional right-hand terms to be controlled.
  • Discrete descent: The convergence analysis averages K(ρj) over iterations to obtain the final barrier-dependent result.The average barrier is denoted by K̄.
  • Sobolev distance: The negative weighted Sobolev distance is defined as a dual norm and can take infinite values.Its definition is used to formulate convergence estimates for the flow.
  • Gaussian example: In Gaussian examples with a quadratic-plus-linear kernel, boundedness of the negative Sobolev distance requires covariance singular values to remain bounded away from zero.The same condition extends to higher dimensions through the singular-value representation of the covariance matrix.

D.6 Lojasiewicz-type inequalities for F under different metrics

The paper relates MMD flows under alternative metrics to local linearizations of familiar divergences and transport distances. These relationships motivate convergence conditions based on bounded linearized distances, while leaving some metric connections unresolved.

  • Alternative metrics: The hybrid metric interpolates between squared Wasserstein distance and squared Fisher-Rao distance.The Wasserstein endpoint has β = 1, α = 0, while the Fisher-Rao endpoint has β = 0, α = 1.
  • KL and Fisher-Rao linearization: The square root of the χ2 divergence linearizes the square root of KL divergence locally.A second-order expansion gives KL(ν∥ν+ϵ(ν′−ν)) = 1/2χ2(ν′∥ν)ϵ2 + o(ϵ2).
  • Wasserstein linearization: The negative weighted Sobolev distance linearizes the Wasserstein-2 distance under suitable regularity and density conditions.For bounded densities, the local Wasserstein displacement equals the negative Sobolev distance.
  • Noisy flow: For the noisy discrete update, each sample is moved by the gradient evaluated at a Gaussian-perturbed location.The update uses xn+1 = xn −γV(xn + βnun), with un standard Gaussian noise.

E.2 Sample-based approximate scheme

The sample-based approximation tracks interacting particles through empirical distributions and compares them with independent samples from the target and the current distribution. Its error analysis controls sampling, target, and gradient-related terms recursively.

  • Empirical approximation: The approximate scheme represents the interacting particle system by the empirical distribution of N particles.An independent empirical distribution from i.i.d. samples is used for comparison.
  • Error decomposition: The expected approximation error is recursively decomposed into errors from empirical-current discrepancies and target sampling.The terms Ei and Gi respectively measure these two sources.
  • Variance control: The analysis controls particle variance using the kernel gradient’s Lipschitz property and a variance recursion.A bounded noise sequence and finite initial variance are assumed.
  • Variance control: Under the stated assumptions, the variance bound grows at most exponentially with time and the Lipschitz constant.For n ≤ 2T/γ, the bound is proportional to (B + var(ν0)^1/2)e^(LT).
  • Error control: The gradient field is jointly Lipschitz in particle position and distribution, enabling recursive control of the approximation error.This property is invoked to bound the interaction error between empirical and population dynamics.

F Connection with Neural Networks

The paper connects MMD gradient flow to optimization over probability distributions of neural-network parameters, especially in the infinite-width limit. Under well-specified conditions, optimizing the MMD functional is equivalent to optimizing the network loss.

  • F Connection with Neural Networks: In the infinite-neuron limit, averaged neural networks become expectations over a probability distribution ν of network parameters.The empirical version represents ν with N particles and optimizes them by gradient descent.
  • F Connection with Neural Networks: The neural-network objective can be lifted to an optimization problem over probability distributions in P2(X), linking particle optimization to MMD flow.The kernel is constructed from expected products of network outputs.
  • F Connection with Neural Networks: When the model is well specified, the centered network loss matches the MMD functional F for a data-induced kernel.More generally, Proposition 20 relates the two losses when a global minimizer exists.
  • F Connection with Neural Networks: In the well-specified case, optimizing F is equivalent to optimizing the network loss, so MMD gradient flow can solve the lifted neural-network problem.The student-teacher setting is one concrete example where this well-specified condition holds.
  • F Connection with Neural Networks: The particle discretization uses a rescaled gradient and converges to the continuous MMD flow as N →∞ and γ →0.The discretization corresponds to the zero-noise update and can estimate the kernel from minibatch samples.

G Numerical Experiments

The experiments apply MMD flow to student-teacher networks and to matching Gaussian distributions. Noise injection helps escape local optima and reach the global target in the Gaussian example, although it slows convergence.

  • G Numerical Experiments: The student network uses an empirical distribution of N particles over neural-network parameters, while the teacher distribution is represented by fixed particles.The reported experiment uses M = 1 teacher particle and N = 1000 student particles.
  • G Numerical Experiments: The network experiment uses ReLU subnetworks and a fixed Gaussian function G with bandwidth σ = 2, avoiding the homogeneity condition that would guarantee global optimality.With identity G and no bias, the setting from prior work is partially 1-homogeneous; the experiment instead studies the non-guaranteed case.
  • G Numerical Experiments: The noisy student-teacher algorithm can be implemented with automatic differentiation by replacing the MMD loss with an auxiliary loss computed from two forward passes.One pass computes predictions without parameter gradients, while a noisy pass supplies the update gradient.
  • G Numerical Experiments: In Gaussian matching, noise-free updates fail to reach the global optimum: some samples move away from the target, others collapse near the origin, and error remains above 10^-3.The experiment starts from N(10, 0.5), targets N(0, 1), and uses 1000 samples from each distribution.
  • G Numerical Experiments: With noise, the training error falls below 10^-4 and oscillates between 10^-8 and 10^-4, while the noise-free flow converges faster.Noise is used through iteration 5000 and then set to zero; the oscillation is attributed to the fixed step size.

H Auxiliary results

The auxiliary results establish regularity properties of the MMD witness and functional, including differentiability, Lipschitz control, path derivatives, and mixture convexity. They also provide assumptions and discrete estimates used in the flow analysis.

  • H Auxiliary results: Under Assumption (A), the unnormalised witness function is differentiable, its gradient is Lipschitz in points and measures, and each gradient component belongs to the RKHS.The bound scales with point distances and Wasserstein distances between the corresponding measures.
  • H Auxiliary results: The discrete algorithms initialize particles from ν0, sample target points and Gaussian perturbations, update particles, and adapt the noise level through a rule h.Algorithm 2 is the noisy MMD flow, while Algorithm 1 gives the general noisy-flow template.
  • H Auxiliary results: Along pushforward paths ρt = (ψ + tφ)#q, the MMD functional is differentiable and its derivative is expressed through the witness gradient and velocity φ.Under additional assumptions, the functional is twice differentiable along such paths.
  • H Auxiliary results: The analysis uses smoothness and Hessian assumptions to control derivatives of the MMD functional along paths and derive bounds involving kernel Hessians and the witness function.These controls support the subsequent flow estimates under the stated assumptions.
  • H Auxiliary results: The MMD functional is mixture convex, providing a convexity property for mixtures of probability distributions even when displacement convexity is unavailable.The result follows by expanding the RKHS norm representation of F.
  • H Auxiliary results: The discrete Gronwall lemma bounds recurrences of the form a_n+1 ≤ (1 + γA)a_n + b by b(e^nγA − 1)/(γA).This estimate is used as a standard tool for controlling accumulated discretization errors.
Loading 1906.04370v2…