summaries: This article proposes a study on the rates of estimation of Brenier maps using the plug in "barycentric mapping" estimator. As a consequence of this study authors obtain rates of convergence for $W_2^{2}$ in various context (depending on the regularity of the measures and on the setting: discrete-discrete of semi-discrete). They apply their results for Wasserstein barycenter estimation and Nonparametric independence testing.

summaries: This paper analyze the statistical behavior of the general plug-in estimators of OT defined by barycentric projections. The authors provide a thorough analysis of the rate of convergence for the transport cost and the transport map. They also consider kernel smoothed plug-in estimators and relate its rate of convergence to the smoothness of the densities, which alleviates the curse of dimensionality suffered by the plug-in estimator.

summaries: Summary =======  Context: --------  This paper tackles the problem of estimating the rate of convergence of empirical Monge maps (through its Kantorovich barycentric map approximator) to the true Monge map for  absolutely continuous Lebesgue measures.    Contribution ------------  - The authors provide new rates of convergence of the barycentric map using standard  empirical plug-in estimators (O(n^{-2/d})) - If the measures are regular enough (Holder smooth), the dependence on the dimension completely vanishes (which was established for Gaussian distributions https://arxiv.org/pdf/1905.10155.pdf) - The authors propose to take advantage of this smoothness result to sample from smoothened distributions (using a Kernel density estimator) and provide the sufficient conditions to obtain a desired rate. Higher smoothness however requires more samples: a tradeoff between statistical accuracy and computational cost cannot be avoided.

summaries: In standard OT problem, marginal distributions of the data are unknown. The current work considers plug-in estimation using Barycentric projection and derived rate of estimation. The idea of Barycentric projection is to plug the estimated marginal distributions of the data in the OT problem. Conditional mean of given us estimated and rate of estimation is derived. Various plug-in estimators of density of are considered such as empirical CDF and kernel density estimator.

summaries: This paper studies the problem of estimating an optimal transport map between two probability distributions on Euclidean space, given IID samples from each of the distributions. The main result (Theorem 2.1) gives a general upper bound, in terms of the dual representation of the 2-Wasserstein distance. This is used to give upper bounds, in probability, on the rate of decay of the estimation risk both for the plugging in empirical distributions in the absence of smoothness assumptions and for kernel estimates under smoothness assumptions. Finally, implications for two applications (Wasserstein barycenter estimation and nonparametric independence testing) are discussed.

abstractText: Optimal transport maps between two probability distributions μ and ν on R have found extensive applications in both machine learning and statistics. In practice, these maps need to be estimated from data sampled according to μ and ν. Plugin estimators are perhaps most popular in estimating transport maps in the field of computational optimal transport. In this paper, we provide a comprehensive analysis of the rates of convergences for general plug-in estimators defined via barycentric projections. Our main contribution is a new stability estimate for barycentric projections which proceeds under minimal smoothness assumptions and can be used to analyze general plug-in estimators. We illustrate the usefulness of this stability estimate by first providing rates of convergence for the natural discretediscrete and semi-discrete estimators of optimal transport maps. We then use the same stability estimate to show that, under additional smoothness assumptions of Sobolev type or Besov type, kernel smoothed or wavelet based plug-in estimators respectively speed up the rates of convergence and significantly mitigate the curse of dimensionality suffered by the natural discrete-discrete/semi-discrete estimators. As a by-product of our analysis, we also obtain faster rates of convergence for plug-in estimators of W2(μ, ν), the Wasserstein distance between μ and ν, under the aforementioned smoothness assumptions, thereby complementing recent results in Chizat et al. (2020). Finally, we illustrate the applicability of our results in obtaining rates of convergence for Wasserstein barycenter between two probability distributions and obtaining asymptotic detection thresholds for some recent optimaltransport based tests of independence.

authors:
Pankaj K Agarwal
R Sharathkumar

title: Approximation algorithms for bipartite matching with metric and geometric costs

venue: In Proceedings of the forty-sixth annual ACM symposium on Theory of computing,

authors:
Martial Agueh
Guillaume Carlier

title: Barycenters in the wasserstein space

venue: SIAM Journal on Mathematical Analysis,

authors:
Luigi Ambrosio
Nicola Gigli
Giuseppe Savaré

title: Gradient flows in metric spaces and in the space of probability measures

venue: Lectures in Mathematics ETH Zürich. Birkhäuser Verlag, Basel, second edition,

authors:
Ery Arias-Castro
David Mason
Bruno Pelletier

title: On the estimation of the gradient lines of a density and the consistency of the mean-shift algorithm

venue: The Journal of Machine Learning Research,

authors:
Arnab Auddy
Nabarun Deb
Sagnik Nandy

title: Exact detection thresholds for chatterjee’s correlation

venue: arXiv preprint arXiv:2104.15140,

authors:
Franck Barthe
Charles Bordenave

title: Combinatorial optimization over two random point sets. In Séminaire de Probabilités XLV, volume 2078 of Lecture Notes in Math., pages 483–535

authors:
Jean-David Benamou
Yann Brenier

title: A computational fluid mechanics solution to the monge-kantorovich mass transfer problem

venue: Numerische Mathematik,

authors:
Jean-David Benamou
Brittany D Froese
Adam M Oberman

title: Numerical solution of the optimal transportation problem using the monge–ampère equation

venue: Journal of Computational Physics,

title: Convergence rates for discretized monge–ampère equations and quantitative stability of optimal transport

venue: Foundations of Computational Mathematics,

authors:
Espen Bernton
Pierre E Jacob
Mathieu Gerber
Christian P Robert

title: Inference in generative models using the Wasserstein distance

venue: arXiv preprint arXiv:1701.05146,

authors:
T.B. Berrett
R.J. Samworth

title: Nonparametric independence testing via mutual information

authors:
Jérémie Bigot
Raúl Gouet
Thierry Klein
Alfredo López

title: Geodesic PCA in the Wasserstein space by convex PCA

venue: In Annales de l’Institut Henri Poincaré, Probabilités et Statistiques,

authors:
Jérémie Bigot
Raúl Gouet
Thierry Klein
Alfredo Lopez

title: Upper and lower risk bounds for estimating the wasserstein barycenter of random measures on the real line

venue: Electronic journal of statistics,

authors:
Jérémie Bigot
Thierry Klein

title: Characterization of barycenters in the wasserstein space by averaging optimal transport

venue: maps. ESAIM: Probability and Statistics,

authors:
Adrien Blanchet
Guillaume Carlier

title: Optimal transport and cournot-nash equilibria

venue: Mathematics of Operations Research,

authors:
Melf Boeckel
Vladimir Spokoiny
Alexandra Suvorikova

title: Multivariate brenier cumulative distribution functions and their application to non-parametric testing

venue: arXiv preprint arXiv:1809.04090,

authors:
Emmanuel Boissard
Thibaut Le Gouic
Jean-Michel Loubes

title: Distribution’s template estimate with wasserstein

venue: metrics. Bernoulli,

authors:
Nicolas Bonneel
Michiel Van De Panne
Sylvain Paris
Wolfgang Heidrich

title: Displacement interpolation using lagrangian mass transport

venue: In Proceedings of the 2011 SIGGRAPH Asia Conference,

title: Polar factorization and monotone rearrangement of vector-valued functions

venue: Comm. Pure Appl. Math.,

title: Boundary regularity of maps with convex potentials

title: The regularity of mappings with a convex potential

venue: J. Amer. Math. Soc.,

venue: II. Ann. of Math. (2),

authors:
Guillaume Carlier
Ivar Ekeland

title: Matching for teams

venue: Economic theory,

authors:
Guillaume Carlier
Adam Oberman
Edouard Oudet

title: Numerical methods for matching for teams and wasserstein barycenters

venue: ESAIM: Mathematical Modelling and Numerical Analysis,

authors:
Rick Chartrand
Brendt Wohlberg
Kevin Vixie
Erik Bollt

title: A gradient descent solution to the monge-kantorovich problem

venue: Applied Mathematical Sciences,

authors:
Victor Chernozhukov
Alfred Galichon
Marc Hallin
Marc Henry

title: Monge-Kantorovich depth, quantiles, ranks and signs

venue: Ann. Statist.,

authors:
Sinho Chewi
Tyler Maunu
Philippe Rigollet
Austin J Stromme

title: Gradient descent algorithms for bures-wasserstein barycenters

venue: In Conference on Learning Theory,

authors:
Pierre-André Chiappori
Robert J McCann
Lars P Nesheim

title: Hedonic price equilibria, stable matching, and optimal transport: equivalence, topology, and uniqueness

venue: Economic Theory,

authors:
Lenaic Chizat
Gabriel Peyré
Bernhard Schmitzer
François-Xavier Vialard

title: Scaling algorithms for unbalanced optimal transport problems

venue: Mathematics of Computation,

authors:
Lenaic Chizat
Pierre Roussillon
Flavien Léger
François-Xavier Vialard
Gabriel Peyré

title: Faster Wasserstein distance estimation with the sinkhorn divergence

authors:
Evgenii Chzhen
Christophe Denis
Mohamed Hebiri
Luca Oneto
Massimiliano Pontil

title: Leveraging Labeled and Unlabeled Data for Consistent Fair Binary Classification

venue: NeurIPS

authors:
Sebastian Claici
Edward Chien
Justin Solomon

title: Stochastic wasserstein barycenters

venue: In International Conference on Machine Learning,

title: Numerical analysis of wavelet methods, volume 32 of Studies in Mathematics and its Applications

authors:
Nicolas Courty
Rémi Flamary
Amaury Habrard
Alain Rakotomamonjy

title: Joint distribution optimal transportation for domain adaptation

venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,

authors:
Nicolas Courty
Rémi Flamary
Devis Tuia
Alain Rakotomamonjy

title: Optimal transport for domain adaptation

venue: IEEE transactions on pattern analysis and machine intelligence,

title: Sinkhorn distances: Lightspeed computation of optimal transport

venue: Advances in neural information processing systems,

authors:
Marco Cuturi
Arnaud Doucet

title: Fast computation of wasserstein barycenters

venue: In International conference on machine learning,

authors:
Bharath Bhushan Damodaran
Benjamin Kellenberger
Rémi Flamary
Devis Tuia
Nicolas Courty

title: Deepjdot: Deep joint distribution optimal transport for unsupervised domain adaptation

venue: In Proceedings of the European Conference on Computer Vision (ECCV),

authors:
Nabarun Deb
Bhaswar B Bhattacharya
Bodhisattva Sen

title: Efficiency lower bounds for distribution-free hotelling-type two-sample tests based on optimal transport

venue: arXiv preprint arXiv:2104.01986,

authors:
Nabarun Deb
Promit Ghosal
Bodhisattva Sen

title: Measuring association on topological spaces using kernels and geometric graphs

venue: arXiv preprint arXiv:2010.01768,

authors:
Nabarun Deb
Bodhisattva Sen

title: Multivariate rank-based distribution-free nonparametric testing using measure transportation

venue: Journal of the American Statistical Association,

authors:
Steffen Dereich
Michael Scheutzow
Reik Schottstedt

title: Constructive quantization: approximation by empirical measures

venue: Ann. Inst. Henri Poincaré Probab. Stat.,

authors:
David L. Donoho
Iain M. Johnstone
Gérard Kerkyacharian
Dominique Picard

title: Density estimation by wavelet thresholding

authors:
Dominique Drouet Mari
Samuel Kotz

title: Correlation and dependence. Imperial College Press, London; distributed by World Scientific Publishing Co., Inc., River

authors:
Ivar Ekeland
Alfred Galichon
Marc Henry

title: Optimal transportation and the falsifiability of incompletely specified economic models

authors:
Tarek A El Moselhy
Youssef M Marzouk

title: Bayesian inference with optimal maps

authors:
Sira Ferradans
Nicolas Papadakis
Gabriel Peyré
Jean-François Aujol

title: Regularized discrete optimal transport

venue: SIAM Journal on Imaging Sciences,

authors:
Aden Forrow
Jan-Christian Hütter
Mor Nitzan
Philippe Rigollet
Geoffrey Schiebinger
Jonathan Weed

title: Statistical optimal transport via factored couplings

authors:
Nicolas Fournier
Arnaud Guillin

title: On the rate of convergence in Wasserstein distance of the empirical measure

venue: Probab. Theory Related Fields,

authors:
Terry L Friesz
J Enrique Fernandez

title: A model of optimal transport maintenance with demand responsiveness

venue: Transportation Research Part B: Methodological,

authors:
Charlie Frogner
Chiyuan Zhang
Hossein Mobahi
Mauricio Araya-Polo
Tomaso Poggio

title: Learning with a Wasserstein loss

venue: In Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 2,

authors:
Kenji Fukumizu
Bharath K Sriperumbudur
Arthur Gretton
Bernhard Schölkopf

title: Characteristic kernels on groups and semigroups

venue: In NIPS,

authors:
Harold N Gabow
Robert E Tarjan

title: Faster scaling algorithms for network problems

venue: SIAM Journal on Computing,

title: Optimal transport methods in economics

authors:
Promit Ghosal
Bodhisattva Sen

title: Multivariate ranks and quantiles using optimal transportation and applications to goodness-of-fit testing

venue: arXiv preprint arXiv:1905.05340,

title: On Hölder continuity-in-time of the optimal transport map towards measures along a curve

venue: Proc. Edinb. Math. Soc. (2),

authors:
Evarist Giné
Richard Nickl

title: Uniform central limit theorems for kernel density estimators

title: Mathematical foundations of infinite-dimensional statistical models

venue: Cambridge Series in Statistical and Probabilistic Mathematics,

authors:
Joan Glaunes
Alain Trouvé
Laurent Younes

title: Diffeomorphic matching of distributions: A new approach for unlabelled point-sets and sub-manifolds matching

venue: In Proceedings of the 2004 IEEE Computer Society Conference on Computer Vision and Pattern Recognition,

authors:
Ian J Goodfellow
Jean Pouget-Abadie
Mehdi Mirza
Bing Xu
David Warde-Farley
Sherjil Ozair
Aaron Courville
Yoshua Bengio

title: Generative adversarial networks

venue: Advances in Neural Information Processing Systems

authors:
Raghuraman Gopalan
Ruonan Li
Rama Chellappa

title: Domain adaptation for object recognition: An unsupervised approach

venue: In 2011 International Conference on Computer Vision,

authors:
Paula Gordaliza
Eustasio Del Barrio
Gamboa Fabrice
Jean-Michel Loubes

title: Obtaining fairness using optimal transport theory

authors:
A. Gretton
R. Herbrich
A. Smola
O. Bousquet
B. Schölkopf

title: Kernel methods for measuring independence

venue: J. Mach. Learn. Res.,

authors:
Arthur Gretton
Olivier Bousquet
Alex Smola
Bernhard Schölkopf

title: Measuring statistical dependence with hilbert-schmidt norms

venue: In International conference on algorithmic learning theory,

authors:
Arthur Gretton
Kenji Fukumizu
Choon Hui Teo
Le Song
Bernhard Schölkopf
Alexander J Smola

title: A kernel statistical test of independence

venue: In Nips,

authors:
Florian Gunsilius
Susanne M Schennach

title: Independent nonlinear component analysis

venue: Technical report, cemmap working paper,

title: On the convergence rate of potentials of brenier maps

venue: Econometric Theory,

authors:
Marc Hallin
Eustasio Del Barrio
Juan Cuesta-Albertos
Carlos Matrán

title: Distribution and quantile functions, ranks and signs in dimension d: A measure transportation approach

venue: The Annals of Statistics,

title: Uniform convergence rates for kernel estimation with dependent data

authors:
Wolfgang Härdle
Gerard Kerkyacharian
Dominique Picard
Alexander Tsybakov

title: Wavelets, approximation, and statistical applications, volume 129 of Lecture Notes in Statistics

authors:
Ruth Heller
Yair Heller
Malka Gorfine

title: A consistent multivariate test of association based on ranks of distances

authors:
Jan-Christian Hütter
Philippe Rigollet

title: Minimax estimation of smooth optimal transport maps

authors:
Roy Jonker
Anton Volgenant

title: A shortest augmenting path algorithm for dense and sparse linear assignment problems

authors:
Julie Josse
Susan Holmes

title: Measuring multivariate association and beyond

venue: Stat. Surv.,

authors:
Kirthevasan Kandasamy
Willie Neiswanger
Jeff Schneider
Barnabás Póczos
Eric P. Xing

title: Neural architecture search with bayesian optimisation and optimal transport

venue: In Proceedings of the 32nd International Conference on Neural Information Processing Systems,

title: On a problem of Monge

venue: Zap. Nauchn. Sem. S.-Peterburg. Otdel. Mat. Inst. Steklov. (POMI), 312(Teor. Predst. Din. Sist. Komb. i Algoritm. Metody

title: On the translocation of masses

venue: C. R. (Doklady) Acad. Sci. URSS (N.S.),

authors:
G. Kerkyacharian
D. Picard

title: Density estimation in Besov spaces

venue: Statist. Probab. Lett.,

authors:
Ilmun Kim
Sivaraman Balakrishnan
Larry Wasserman

title: Minimax optimality of permutation tests

venue: arXiv preprint arXiv:2003.13208,

authors:
Sanggyun Kim
Rui Ma
Diego Mesa
Todd P Coleman

title: Efficient bayesian inference methods via convex optimization and optimal transport

venue: IEEE International Symposium on Information Theory,

authors:
Diederik P. Kingma
Max Welling

title: Auto-Encoding Variational Bayes

venue: In 2nd International Conference on Learning Representations,

authors:
J. Komlós
P. Major
G. Tusnády

title: An approximation of partial sums of independent RV’s, and the sample DF

venue: II. Z. Wahrscheinlichkeitstheorie und Verw. Gebiete,

authors:
Gil Kur
Fuchang Gao
Adityanand Guntuboyina
Bodhisattva Sen

title: Convex regression in multidimensions: Suboptimality of least squares estimators

venue: arXiv preprint arXiv:2006.02044,

authors:
Thibaut Le Gouic
Jean-Michel Loubes

title: Existence and consistency of wasserstein barycenters

venue: Probability Theory and Related Fields,

authors:
Tong Li
Ming Yuan

title: On the optimality of gaussian kernel based nonparametric tests against smooth alternatives

venue: arXiv preprint arXiv:1909.03302,

authors:
Wenbo Li
Ricardo H Nochetto

title: Quantitative stability and error estimates for optimal transport plans

venue: IMA Journal of Numerical Analysis,

authors:
Yujia Li
Kevin Swersky
Rich Zemel

title: Generative moment matching networks

authors:
Giulia Luise
Alessandro Rudi
Massimiliano Pontil
Carlo Ciliberto

title: Differential properties of sinkhorn approximation for learning with Wasserstein distance

authors:
Valentina Masarotto
Victor M Panaretos
Yoav Zemel

title: Procrustes metrics on covariance operators and optimal transportation of gaussian processes

venue: Sankhyā A,

title: Existence and uniqueness of monotone measure-preserving maps

venue: Duke Math. J.,

authors:
Gonzalo Mena
Jonathan Niles-Weed

title: Statistical bounds for entropic optimal transport: Sample complexity and the central limit theorem

title: A multiscale approach to optimal transport

venue: In Computer Graphics Forum,

authors:
Quentin Merigot
Boris Thibert

title: Optimal transport: discretization and algorithms. Handbook of Numerical Analysis 22 – Geometric PDES (arXiv

title: Ondelettes et opérateurs

venue: I. Actualités Mathématiques. [Current Mathematical Topics]. Hermann, Paris,

authors:
Shakir Mohamed
Balaji Lakshminarayanan

title: Learning in implicit generative models

venue: In Proceedings of the International Conference in Learning Representations,

title: Integral probability metrics and their generating classes of functions

venue: Adv. in Appl. Probab.,

title: On non-parametric estimates of density functions and regression

venue: Teor. Verojatnost. i Primenen.,

authors:
Jonathan Niles-Weed
Philippe Rigollet

title: Estimation of Wasserstein distances in the spiked transport model

venue: arXiv preprint arXiv:1909.07513,

authors:
Nadia Oudjane
Christian Musso

title: L2-density estimation with negative kernels

venue: ISPA

authors:
Nicolas Papadakis
Gabriel Peyré
Edouard Oudet

title: Optimal transport with proximal splitting

title: On estimation of a probability density function and mode

venue: Ann. Math. Statist.,

authors:
Gabriel Peyré
Marco Cuturi

title: Computational optimal transport: With applications to data science

venue: Foundations and Trends® in Machine Learning,

title: The monge–kantorovich mass transference problem and its stochastic applications

venue: Theory of Probability & Its Applications,

authors:
Aaditya Ramdas
Nicolás García Trillos
Marco Cuturi

title: On Wasserstein two-sample testing and related families of nonparametric tests

venue: Entropy,

title: A dynamical systems framework for intermittent data assimilation

venue: BIT Numerical Mathematics,

title: A nonparametric ensemble transform method for bayesian inference

venue: SIAM Journal on Scientific Computing,

authors:
Thomas Rippl
Axel Munk
Anja Sturm

title: Limit laws of the empirical Wasserstein distance: Gaussian distributions

venue: J. Multivariate Anal.,

authors:
Tim Salimans
Han Zhang
Alec Radford
Dimitris Metaxas

title: Improving GANs using optimal transport

venue: arXiv preprint arXiv:1803.05573,

title: Optimal transport for applied mathematicians

authors:
Vivien Seguy
Bharath Bhushan Damodaran
Rémi Flamary
Nicolas Courty
Antoine Rolet
Mathieu Blondel

title: Large-scale optimal transport and mapping estimation

title: Distribution-free consistent independence tests via center-outward ranks and signs

authors:
Hongjian Shi
Marc Hallin
Mathias Drton
Fang Han

title: Rate-optimality of consistent distribution-free tests of independence based on center-outward ranks and signs

title: Weak and strong uniform consistency of the kernel estimate of a density and its derivatives

title: Inference for empirical Wasserstein distances on finite spaces

venue: J. R. Stat. Soc. Ser. B. Stat. Methodol.,

authors:
Bharath K. Sriperumbudur
Kenji Fukumizu
Arthur Gretton
Bernhard Schölkopf
Gert R.G. Lanckriet

title: On the empirical estimation of integral probability metrics

authors:
Bharath K. Sriperumbudur
Arthur Gretton
Kenji Fukumizu
Bernhard Schölkopf
Gert R.G. Lanckriet

title: Hilbert space embeddings and metrics on probability measures

authors:
Sanvesh Srivastava
Volkan Cevher
Quoc Dinh
David Dunson

title: Wasp: Scalable bayes via barycenters of subset posteriors

venue: In Artificial Intelligence and Statistics,

authors:
Sanvesh Srivastava
Cheng Li
David B Dunson

title: Scalable bayes via barycenter in wasserstein space

authors:
Zhengyu Su
Yalin Wang
Rui Shi
Wei Zeng
Jian Sun
Feng Luo
Xianfeng Gu

title: Optimal mass transport for shape matching and comparison

title: Measuring and testing dependence by correlation of distances

title: The transportation cost from the uniform measure to the empirical measure in dimension ≥ 3

title: Theory of function spaces. III, volume 100 of Monographs in Mathematics

title: Introduction to nonparametric estimation

title: Asymptotic statistics, volume 3 of Cambridge Series in Statistical and Probabilistic Mathematics

title: Weak convergence and empirical processes

title: Topics in optimal transportation, volume 58 of Graduate Studies in Mathematics

title: Optimal transport, volume 338 of Grundlehren der Mathematischen Wissenschaften [Fundamental Principles of Mathematical Sciences

title: Approximation of the delta function by wavelets

title: Estimation of smooth densities in Wasserstein distance

heading: 1 Introduction

text:
Given two random variables X ∼ µ and Y ∼ ν, where µ, ν are probability measures on Rd, d ≥ 1, the problem of finding a “nice" map T0(·) such that T0(X) ∼ ν has numerous applications in machine learning such as domain adaptation and data integration [34, 35, 38, 48, 61, 112], dimension reduction [12, 66, 90], generative models [60, 81, 88, 110], to name a few. Of particular interest is the case when T0(·) is obtained by minimizing a cost function, a line of work initiated by Gaspard Monge [97] in 1781 (see (1.1) below), in which case T0(·) is termed an optimal transport (OT) map and has applications in shape matching/transfer problems [29, 47, 107, 121], Bayesian statistics [46, 75, 80, 108], econometrics [15, 28, 45, 50, 54], nonparametric statistical inference [39– 41, 113, 114]; also see [111, 128, 129] for book-length treatments on the subject. In this paper, we will focus on the OT map obtained using the standard squared Euclidean cost function, i.e.,
T0 := argmin T :T#µ=ν
E‖X − T (X)‖2, (1.1)
where T#µ = ν means T (X) ∼ ν for X ∼ µ. The estimation of T0 has attracted a lot of interest in recent years due to its myriad applications (as stated above) and interesting geometrical properties (see [19, 56, 91] and Definition 1.1 below). In practice, the main hurdle in constructing estimators for T0 is that the explicit forms of the measures µ, ν are unknown; instead only random samples
X1, . . . , Xm ∼ µ and Y1, . . . , Yn ∼ ν
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
are available. A natural strategy in this scenario is to estimate T0 using T̃m,n, where T̃m,n is computed as in (1.1) with µ and ν replaced by µ̃m and ν̃n which are empirical approximations of µ and ν based on X1, . . . , Xm and Y1, . . . , Yn respectively (see Definition 1.2). Such estimators are often called plug-in estimators and have been used extensively; see [7, 30, 67, 93, 94, 102, 116].
The main goal of this paper is to study the rates of convergence of general plug-in estimators of T0 under a unified framework. We show that when µ̃m and ν̃n are chosen as µ̂m and ν̂n respectively, where µ̂m and ν̂n are the standard empirical distributions supported on m and n atoms, i.e.,
µ̂m := 1
m m∑ i=1 δXi and ν̂n := 1 n n∑ j=1 δYj , (1.2)
T̃m,n (appropriately defined using Definition 1.2) converges at a rate of m−2/d + n−2/d for d ≥ 4 in the sense of (1.8). This rate happens to be minimax optimal under minimal smoothness assumptions (see [72, Theorem 6]) but suffers from the curse of dimensionality. We next show that, if µ and ν are known to admit sufficiently smooth densities, it is possible to apply kernel or wavelet based smoothing techniques on µ̂m and ν̂n to obtain plug-in estimators that mitigate the aforementioned curse of dimensionality.
Our next contribution pertains to the estimation of W 22 (µ, ν) (the squared Wasserstein distance), see (1.3) below, a quantity of independent interest in statistics and machine learning with applications in structured prediction [51, 89], image analysis [18, 59], nonparametric testing [16, 106], generative modeling [10, 96], etc. In this paper, we also obtain rates of convergence for plug-in estimators W 22 (µ̃m, ν̃n) of W 2 2 (µ, ν). We show that kernel smoothing µ̂m and ν̂n can be used to obtain plug-in estimators of W 22 (µ, ν) that mitigate the curse of dimensionality as opposed to a direct plug-in approach using µ̂m and ν̂n (as used in [30, Theorem 2]). This provides an answer to the open question of estimating W 22 (µ, ν) when µ, ν admit smooth densities laid out in [30].

heading: 1.1 Background on optimal transport

text:
In this section, we present some basic concepts and results associated with the OT problem that will play a crucial role in the sequel. Let Pac(Rd) denote the set of all Lebesgue absolutely continuous probability measures on Rd andP2(Rd) be the set of probability measures with finite second moments. Then the 2-Wasserstein distance (squared) between µ, ν ∈ P2(Rd) is defined as:
W 22 (µ, ν) := min π∈Π(µ,ν)
∫ ‖x− y‖2 dπ(x, y), (1.3)
where Π(µ, ν) is the set of probability measures on Rd×Rd with marginals µ and ν. The optimization problem in (1.3) is often called the Kantorovich relaxation (see [76, 77]) of the optimization problem in (1.1). The existence of a minimizer in (1.3) follows from [129, Theorem 4.1].
Proposition 1.1 (Brenier-McCann polar factorization theorem, see [91, 128]). Suppose µ ∈ Pac(Rd). Then there exists a µ-a.e. (almost everywhere) unique function T0(·) : Rd → Rd, which is the gradient of a real-valued d-variate convex function, say ϕ0(·) : Rd → R, such that T0#µ = ν. Further, the distribution defined as π(A×B) = µ(A ∩ (T0)−1(B)) for all Borel sets A,B ⊆ Rd is the unique minimizer in (1.3) provided µ, ν ∈ P2(Rd).
Definition 1.1 (OT map and potential function). The function T0 : Rd → Rd in Proposition 1.1 which satisfies T0#µ = ν will be called the OT map from µ to ν. A convex function ϕ0(·) in Proposition 1.1 satisfying∇ϕ0 = T0 will be termed an OT potential.
The next and final important ingredient is the alternate dual representation of (1.3) which gives: 1
2 W 22 (µ, ν) =
1
2
∫ ‖x‖2 dµ(x) + 1
2
∫ ‖y‖2 dν(y)−min
f∈F Sµ,ν(f), where (1.4) Sµ,ν(f) = ∫ f dµ+ ∫ f∗ dν. (1.5)
Here F denotes the space of convex functions on Rd which are also elements of L1(µ) and f∗(·) is the standard Legendre-Fenchel dual defined as:
f∗(x) := sup y∈Rd [y>x− f(y)], for x ∈ dom(f). (1.6)

heading: 1.2 Estimating OT map via barycentric projection

text:
Recall the setting from the Introduction. Let µ̃m, ν̃n ∈ P2(Rd). Here µ̃m, ν̃n need not be absolutely continuous and can be very general. Intuitively, µ̃m and ν̃n can be viewed as some empirical approximation of µ and ν respectively.
Example 1.2 (Simple choices of µ̃m and ν̃n). Let X1, . . . , Xm i.i.d.∼ µ and Y1, . . . , Yn i.i.d.∼ ν; in which case a natural choice would be to set µ̃m = µ̂m and ν̃n = ν̂n where µ̂m and ν̂n are the empirical distributions of X1, . . . , Xm and Y1, . . . , Yn respectively, as defined in (1.2). This is the standard choice adopted in the discrete-discrete Kantorovich relaxation; see [104, Section 2.3]. Another popular choice is µ̃m = µ̂m, ν̃n = ν or µ̃m = µ, ν̃n = ν̂n. This is the semi-discrete Kantorovich problem and is popular when one of the measures is fully specified; see [26, 55].
A natural way to estimate T0(·), as defined in (1.1), would be to approximate it using the OT map from µ̃m to ν̃n. However as µ̃m and ν̃n may not be elements of Pac(Rd), Proposition 1.1 does not apply and an OT map may not exist from µ̃m to ν̃n. Such is the case in Example 1.2 in the discrete-discrete case when m 6= n. To circumvent this issue, we leverage the notion of barycentric projections (see [3, Definition 5.4.2]) defined below:
Definition 1.2 (Barycentric projection). Define the set
Γ̃min := argmin π∈Π(µ̃m,ν̃n)
∫ ‖x− y‖2 dπ(x, y).
The optimization problem above is the plug-in analog of the optimization problem on the right hand side of (1.3). Given any γ ∈ Γ̃min, define the barycentric projection of γ as the conditional mean of y given x under γ, i.e.,
T̃m,n(x) ≡ T̃ γm,n(x) := ∫ y y dγ(x, y)∫ y dγ(x, y) , for x ∈ supp (µ̃m) . (1.7)
In general, Γ̃min need not be a singleton which is why we index the barycentric projection T̃ γm,n(·) by γ ∈ Γ̃min. Note that T̃ γm,n(·) need not be a transport map; however, if an OT map exists then it must be equal to T̃ γm,n(·) (µ̃m-a.e.). Our goal is to obtain stochastic upper bounds for
sup γ∈Γ̃min ∫ ∥∥T̃ γm,n(x)− T0(x)∥∥2 dµ̃m(x). (1.8) In addition, our proof techniques also yield rates of convergence for∣∣W 22 (µ̃m, ν̃n)−W 22 (µ, ν)∣∣. (1.9) In this paper, we will focus on d ≥ 2. Due to the canonical ordering of R, the case d = 1 can be handled easily using the classical Hungarian embedding theorem [82].

heading: 1.3 Contributions

text:
1. We provide a new and flexible stability estimate Theorem 2.1 which yields a unified approach
to obtaining rates of convergence for general plug-in estimators of the OT map T0(·). Unlike existing stability estimates, Theorem 2.1 holds for the barycentric projection (which is the same as the OT map when it exists) and does not require any smoothness assumptions on µ̃m, ν̃n or T̃ γm,n(·); also see Remark 2.1 for a comparison with the existing literature.
2. In Sections 2.1 and 2.2, we use Theorem 2.1 to bound (1.8) and (1.9):
• In Section 2.1, we show that in both the discrete-discrete and semi-discrete Kantorovich relaxation problems (see Example 1.2), the rate of convergence of (1.8) is m−2/d + n−2/d for d ≥ 4 when T0 is assumed to be Lipschitz (see Theorem 2.2), which is the minimax rate (see [72, Theorem 6]). To the best of our knowledge, rates of convergence for these natural estimators weren’t previously established in the literature.
• In Section 2.2 and Appendix A, we show that the curse of dimensionality in the above rates can be mitigated provided µ and ν admit (uniform) Sobolev smooth densities (see Section 2.2) or Besov smooth densities (see Appendix A). In Section 2.2, our plug-in estimator is obtained by choosing µ̃m (and ν̃n) as the convolution of µ̂m (and ν̂n) and a smooth kernel with an appropriate bandwidth. Under this choice, the rate of convergence in (1.8) is m−( s+2 d ∧ 1 2 ) +
n−( s+2 d ∧ 1 2 ), where s denotes the degree of Sobolev smoothness (see Theorem 2.5). Clearly, if 2(s + 2) ≥ d, the rate of convergence becomes dimension-free and mitigates the curse of dimensionality. We also show the same rates of convergence mentioned above hold for (1.9) (see e.g., Proposition 2.6) which makes a strong case in favor of incorporating smoothness in the construction of plug-in estimators as was conjectured in [30]. In Appendix A, our plug-in estimator is obtained using natural wavelet based density estimators. The rate of convergence in (1.8) turns out to be n− 1+s d+2s where s denotes the degree of Besov smoothness (see Theorem A.1). Note that by choosing s large enough, the exponent in the rate can be made arbitrarily close to 1/2, thereby reducing the curse of dimensionality.
3. In Section 2.3, we use a discretization technique from [131] to construct discrete approximations to the smoothed µ̃m and ν̃n from the previous paragraph that in turn yield computable plug-in estimators for T0 (provided one can sample from µ̃m and ν̃n) that also achieve the same statistical guarantees as the smoothed plug-in estimator from Section 2.2 (see Theorem 2.7). However the number of atoms required in the discretizations and correspondingly the computational complexity increases with the degree of smoothness; this highlights a statistical and computational trade-off.
4. We provide implications of our results in popular applications of OT such as estimating the barycenter of two multivariate probability distributions (see Theorem B.1 in Appendix B.1) and in nonparametric independence testing (see Theorem B.3 in Appendix B.2).

heading: 1.4 Related work

text:
Many recent works have focused on obtaining consistent estimators of T0 using the plug-in principle, see [26, 55] (in the semi-discrete problem) and [41, 68, 132] (in the discrete-discrete problem). In [55], the authors studied the rate of convergence of the semi-discrete optimal transport map from ν (absolutely continuous) to µ̂m. This paper complements the aforementioned papers by studying the rates of convergence for general plug-in estimators in a unified fashion. In two other papers [9, Theorem 1.1] and [87, Section 4], the authors use a “Voronoi tessellation" approach to estimate T0, however the rates obtained in this paper, even in the absence of smoothness, are strictly better than those in [9, 87]. Perhaps the most closely related paper to ours would be [67]. In [67], the author uses variational techniques to arrive at stability estimates while we exploit the Lipschitz nature of the OT map (see Definition 1.1). Further the rates in this paper have exponents s+2d ∧ 1 2 which are strictly better than the exponents s+22(s+2)+d obtained in [67, Proposition 1] under the same smoothness assumptions (Sobolev type of order s, see Definition 2.4). In another line of work [72], the authors use theoretical wavelet based estimators (not of the plug-in type) of T0 to obtain nearly minimax optimal rates of convergence. However these estimators, by themselves, are not transport maps between two probability measures, which makes them harder to interpret. In contrast, our focus is on obtaining rates of convergence for plug-in estimators, which are transport maps between natural approximations of µ and ν. Such plug-in type strategies are a lot more popular in computational OT [7, 30, 67, 93, 94, 102, 116].
In terms of obtaining rates of convergence for (1.9), some attempts include [109, 116] where parametric rates are obtained when µ, ν are known to be finitely supported or are both Gaussian. In a related problem, bounds for W 22 (µ̂m, µ) were obtained in [6, 42, 49, 100, 123, 131]. Using these bounds, for m = n, it is easy to get a n−1/d rate of convergence for (1.9). This rate was recently improved to n−2/d in [30] under no smoothness assumptions. Our rates coincide with the n−2/d rate from [30] under no smoothness assumptions. But further, we show in this paper that the curse of dimensionality in the above rate can be mitigated by incorporating smoothness into the plug-in procedure.

heading: 2 Main results

text:
Recall ϕ0(·) from Definition 1.1. The following is our main result.
Theorem 2.1 (Stability estimate). Suppose that µ, ν ∈ Pac(Rd) ∩ P2(Rd) and µ̃m, ν̃n ∈ P2(Rd). Assume that T0(·) (as defined in (1.1)) is L-Lipschitz (L > 0). Then,
sup γ∈Γ̃min
∫ ‖T̃ γm,n(x)− T0(x)‖2 dµ̃m(x) ≤ Lmax {∣∣∣∣ ∫ Ψ∗µ̃m,ν̃n d(ν̃n − ν†m)∣∣∣∣, ∣∣∣∣ ∫ Ψ∗µ̃m,ν†m d(ν̃n − ν†m) ∣∣∣∣}
+ 2L ∫ ϕ∗0(y) d(ν̃n − ν†m)(y), (2.1)
where ν†m := T0#µ̃m, ϕ ∗ 0(·) is defined as in (1.6), and with S·,·(·) defined as in (1.5), Ψµ̃m,ν̃n(·) := argminf∈F Sµ̃m,ν̃n(f), Ψµ̃m,ν†m(·) := argminf∈F Sµ̃m,ν†m(f), and D denotes the space of realvalued convex functions on Rd.
The proof of Theorem 2.1 (see Appendix C.1) starts along the same lines as the proof of the curvature estimate in [56, Proposition 3.3]. This is followed by some careful manipulations of W 22 (·, ·) (as in (1.3)) and an application of the conditional version of Jensen’s inequality, see (C.3). The final step of the proof uses the dual representation in (1.4) with techniques similar to some intermediate steps in the proof of [92, Proposition 2] and [30, Lemma 3].
Remark 2.1 (Comparison with other stability estimates). Theorem 2.1 provides some important advantages to existing stability estimates in the literature. One of the earliest results in this direction can be found in [56, Proposition 3.3] but their bound involves a push-forward constraint which makes it hard to use for rate of convergence analysis. A bound similar to Theorem 2.1 is presented in [55, Lemma 5.1] but there the authors assume the existence of an OT map from µ̃m to ν̃n. Therefore, it does not apply to the discrete-discrete problem where µ̃m = µ̂m and ν̃n = ν̂n with m 6= n. Overcoming all these limitations is an important contribution of Theorem 2.1 and allows us to deal with popular plug-in estimators all in one go. The stability estimate in [72, Proposition 10] on the other hand requires µ̃m, ν̃n to be sufficiently smooth and hence it does not hold for discrete-discrete or semi-discrete plug-in estimators (see Example 1.2). Further their result requires all the measures involved to be compactly supported unlike the much milder requirements of Theorem 2.1. However, a shortcoming of Theorem 2.1 is that it is hard to obtain rates faster than n−1/2 using it directly, whereas [72] can obtain rates arbitrarily close to n−1. This is a price we pay for analyzing natural and popular plug-in estimators as opposed to the (more intractable) wavelet based estimators in [72].
Remark 2.2 (How to use Theorem 2.1 to obtain rates of convergence?). Note that the second term on the right hand side of (2.1), under appropriate moment assumptions, is Op(m−1/2 +n−1/2) (free of dimension) by a direct application of Markov’s inequality. We therefore focus on the first term. By (1.5), Ψ∗µ̃m,ν̃n(·), Ψ ∗ µ̃m,ν † m
(·) ∈ F . Further, by Caffarelli’s regularity theory [20–22], depending on the “smoothness" of µ̃m, ν̃n, it can be shown that there exists a further class of functions Fs (see Remarks 2.3 and 2.6) such that Ψ∗µ̃m,ν̃n(·), Ψ ∗ µ̃m,ν † m
(·) ∈ F ∩ Fs. Thus, we can bound the first term on the right hand side of (2.1) as:
max {∣∣∣∣ ∫ Ψ∗µ̃m,ν̃n d(ν̃n − ν†m)∣∣∣∣, ∣∣∣∣ ∫ Ψ∗µ̃m,ν†m d(ν̃n − ν†m) ∣∣∣∣} ≤ sup
f∈F∩Fs ∣∣∣∣ ∫ f d(ν̃n − ν†m)∣∣∣∣. (2.2) The right hand side of (2.2) can now be bounded using the corresponding Dudley’s entropy integral bounds using empirical process techniques; see [126, Lemmas 19.35-19.37].
To conclude, the two main steps in our strategy are identifying the family of functions Fs and computing Dudley’s entropy integral. Further, the more the smoothness of µ̃m, ν̃n, the smaller is the class of functions Fs and smaller the supremum on the right hand side of (2.2). This shows why better rates can be expected under smoothness assumptions.

heading: 2.1 Natural non-smooth plug-in estimators

text:
In this case, we discuss the rates of convergence for the discrete-discrete problem and the semi-discrete problem, where no smoothness is available on µ̃m and ν̃n.
Theorem 2.2. Suppose that T0(·) is L-Lipschitz, ν is compactly supported and E exp(t‖X1‖α) <∞ for some t > 0, α > 0.
(Discrete-discrete): Set µ̃m = µ̂m and ν̃n = ν̂n. Then the following holds:
sup γ∈Γ̃min
∫ ‖T̃ γm,n(x)− T0(x)‖2 dµ̃m(x) = Op ( r (m,n) d × (log (1 + max{m,n})) td,α ) , (2.3)
where r(m,n)d :=  m−1/2 + n−1/2 for d = 2, 3, m−1/2 log (1 +m) + n−1/2 log (1 + n) for d = 4, m−2/d + n−2/d for d ≥ 5,
(2.4)
and
td,α :=  (4α)−1(4 + ((2α+ 2dα− d) ∨ 0)) for d < 4, (α−1 ∨ 7/2)− 1 for d = 4, 2(1 + d−1) for d > 4.
The same bound holds for |W 22 (µ̃m, ν̃n)−W 22 (µ, ν)| without assuming T0(·) is Lipschitz. (Semi-discrete): Set µ̃m = µ, ν̃n = ν̂n or µ̃m = µ̂m, ν̃n = ν. Then the left hand side of (2.3) is Op(r (n,n) d × (log (1 + n))td,α) or Op(r (m,m) d × (log (1 +m))td,α) respectively.
A stronger result can be proved if both µ and ν are compactly supported.
Corollary 2.3. Consider the setting from Theorem 2.2 and assume further that µ is compactly supported. Then, with r(m,n)d defined as in (2.4), we have:
E [ sup
γ∈Γ̃min
∫ ‖T̃ γm,n(x)− T0(x)‖2 dµ̃m(x) ] ≤ Cr(m,n)d ,
for some constant C > 0, in both the discrete-discrete and semi-discrete settings from Theorem 2.2.
A brief description of the proof technique of Theorem 2.2 using Theorem 2.1 is provided in Remark 2.3 below, and the actual proof is presented in Appendix C.1.
Remark 2.3 (Proof technique). The proof of Theorem 2.2 proceeds via the strategy outlined in Remark 2.2. We first show that Fs (see Remark 2.2) can be chosen as a certain sub-class of convex functions which are in L2(ν). We then use Dudley’s entropy integral type bounds which in turn requires the bracketing entropy [126, Page 270] of Fs, recently proved in [83, Equation 26]. This strategy is slightly different from that used in the proof of [30, Theorem 2], where the authors assume that µ is compactly supported whereas we only assume the finiteness of E exp(t‖X1‖α) for some t > 0, α > 0. The compactness assumption on µ allows one to further restrict Fs to the class of Lipschitz functions. This additional restriction does not seem to be immediate without the compactness assumption.
As discussed in Section 1.3, the exponents obtained in Theorem 2.2 are minimax optimal, up to multiplicative logarithmic factors, under bare minimal smoothness assumptions (see [72, Theorem 6]). To the best of our knowledge, rates for the discrete-discrete case for m 6= n and those for the semi-discrete case were not known previously in the literature. Our rates are also strictly better than those (for different estimators, based on space tessellations) obtained in [9, 87] and require less stringent assumptions than those in [30]. In the next section, we show how smoothness assumptions can be leveraged to mitigate the curse of dimensionality in Theorem 2.2.

heading: 2.2 Smooth kernel based plug-in estimator: mitigating the curse of dimensionality

text:
In this section, we focus on kernel based density estimators for the probability densities associated with µ and ν (see [57, 58, 99, 103, 115]). We will show, using Theorem 2.1, that the corresponding estimators of T0(·) achieve (near) dimension-free rates under sufficient smoothness assumptions. We first introduce the Sobolev class of functions which we will exploit in this subsection to construct estimators that achieve rates of convergence which mitigate the curse of dimensionality under sufficient smoothness.
Definition 2.4 (Uniform Sobolev class of functions). Let Ω ⊆ Rd and f(·) be uniformly continuous on Ω and admits uniformly continuous derivatives up to order s on Ω for some s ∈ N. For any m := (m1, . . . ,md) ∈ Nd, let
∂mf := ∂
∂m1x1 . . .
∂
∂mdxd f, |m| := d∑ i=1 mi.
For any k ≤ s, we further define, ‖f‖Ck(Ω) := ∑ |m|≤k ‖∂mf‖L∞(Ω).
The space Cs(Ω) is defined as the set of functions f(·) for which ‖f‖Ck(Ω) <∞ for all k ≤ s.
For this subsection, assume that µ and ν admit Sobolev smooth densities fµ(·) and fν(·) in the uniform norm (see Definition 2.4 above). Given Ω ⊆ Rd and s ∈ N, let Cs(Ω) denote the set of Sobolev smooth functions on Ω of order s.
Assumption (A1) (Regularity of the densities). Suppose that
1. fµ and fν are supported on compact and convex subsets of Rd, say X and Y respectively. 2. There exists s,M > 0 such that fµ(·) ∈ Cs(X ;M) and fν(·) ∈ Cs(Y;M) where Cs(X ;M) is
the space of real valued functions supported on X such that for all f(·) ∈ Cs(X ;M), we have M−1 ≤ f(x) ≤ M for all x ∈ X and ‖f‖Cs(X ) ≤ M . Here ‖·‖Cs(X ) is the standard uniform Sobolev norm as defined in Definition 2.4. The space Cs(Y;M) is defined analogously.
We now define our estimators for fµ(·) and fν(·) using the standard kernel density estimation technique (see [125, Section 1.2]). Set
f̂µ(x) := 1
mhdm m∑ i=1 Kd ( Xi − x hm ) , (2.5)
for some bandwidth parameter hm > 0 and d-variate kernel Kd(·). We assume that Kd(·) is the dfold product of univariate kernels, i.e., there exists a kernelK(·) such that for u = (u1, . . . , ud) ∈ Rd, Kd(u) = ∏d i=1K(ui). We define f̂ν(·) similarly with the same univariate kernel and bandwidth.
Assumption (A2) (Regularity of the kernel). Assume that K(·) is a symmetric, bounded, s+ 1 times differentiable kernel on Rd with all s+ 1 derivatives bounded and integrable. Further, suppose that K(·) is of order 2s+ 2, i.e.,∫ ujK(u) du = 1(j = 0), for j = {0, 1, 2, . . . , 2s+1}, and ∫ |u|2s+2|K(u)| du <∞.
The above assumptions on K(·) are standard for estimating smooth densities and their derivatives of different orders in the kernel density estimation literature; see e.g. [4, 57, 58, 69, 125]. There are several natural ways to construct kernels satisfying Assumption (A2), see [125, Section 1.2.2]; an example is also provided in Example 2.4 below.
Example 2.4 (Example of a kernel satisfying Assumption (A2)). Let ψm(·) be the m-th Hermite polynomial on R (see [84]). Then the kernel function defined as
K(u) := 2s+2∑ m=0 ψm(0)ψm(u) exp(−u2/2)
satisfies Assumption (A2).
It is evident from Assumption (A2) that K(·) may take some negative values, in which case, f̂µ(·) (respectively f̂ν(·)) may not be a probability density. Consequently the barycentric projection (see Definition 1.2) between f̂µ(·) and f̂ν(·) is not well-defined. We get around this by projecting
f̂µ(·) and f̂ν(·) on an appropriate space of “smooth" probability densities (see (2.6)), via an integral probability metric (see Definition 2.5 below; also see [98, 105, 117] for examples, computational procedures and applications of such metrics).
Definition 2.5 (Integral probability metric). Given a classH of bounded functions on Rd and two probability densities g1(·) and g2(·) on Rd, the integral probability metric/distance between g1(·) and g2(·) with respect toH is defined as
dIP(g1, g2;H) := sup ψ(·)∈H ∣∣∣∣ ∫ ψ(x)(g1(x)− g2(x)) dx∣∣∣∣. Sufficient conditions onH for dIP(·, ·;H) to be a metric on the space of probability measures (not on the space of probability densities as they can be altered on set of Lebesgue measure 0 without altering the underlying probability measures) on Rd have been discussed in [98]. Observe that the measure dIP(g1, g2;H) is well defined even when g1(·) and g2(·) are not probability densities.
In Theorem 2.5 below, we useH = Cs+2(X ,M ′). Note that any function in Cs+2(X ,M ′) can be extended to a function in Cs+2(Rd;M ′) (see [72, Theorem 23] and [124, Theorem 1.105]). The fact that this choice of F results in a metric follows from the argument in [98, Page 8].
We are now in a position to describe the projection estimators for fµ(·) and fν(·), and the rates achieved by the corresponding plug-in estimator.
Theorem 2.5. Assume that T0(·) is L-Lipschitz and fµ, fν are Lebesgue densities satisfying Assumption (A1). Also suppose that K(·) satisfies Assumption (A2). Define hm := m− 1 d+2s logm , hn := n − 1d+2s log n and T := ∫ |Kd(u)| du+ 1. Fix any M ′ > 0. Consider any probability density f̃M ′
µ (·) ∈ Cs(X ;TM) (where M is defined as in Assumption (A1)) which satisfies
dIP
( f̃M ′
µ , f̂µ;C s+2(X ;M ′) ) ≤ inf f(·)∈Cs(X ;TM) f≥0, ∫ f=1 dIP ( f̂µ, f ;C s+2(X ;M ′) ) + r (m,n) d,s (2.6)
where r(m,n)d,s is defined as in (2.7) and dIP(·, ·;Cs+2(X ;M ′)) is the integral probability metric defined in Definition 2.5. We define f̃M ′
ν (·) analogously as in (2.6) with X , f̂µ(·) replaced by Y , f̂ν(·). Then the following conclusions hold.
1. SetM ′ := 8(1+TM). If µ̃m and ν̃n are the probability measures corresponding to the probability densities f̃M ′ µ (·) and f̃M ′ ν (·), then the following holds for some constant C > 0:
E [ sup
γ∈Γ̃min
∫ ‖T̃ γm,n(x)− T0(x)‖2 dµ̃m(x) ] ≤ Cr(m,n)d,s ,
where r(m,n)d,s :=  m−1/2 + n−1/2 for d < 2(s+ 2), m−1/2 (log (1 +m)) d + n−1/2 (log (1 + n))
d for d = 2(s+ 2), m− s+2 d + n− s+2 d for d ≥ 2(s+ 2). (2.7)
The same bound also holds for E|W 22 (µ̃m, ν̃n)−W 22 (µ, ν)|.
2. f̂µ(·) satisfies
lim n→∞ max
{ P ( ‖f̂µ‖Cs(X̃ ) ≥ TM ) ,P ( sup x∈X̃ |f̂µ(x)− fµ(x)| ≥ ε )} = 0 (2.8)
for any ε > 0, where X̃ is any compact subset of X o. The same conclusion holds for f̂ν(·) with X replaced by Y .
In Theorem 2.5, we have shown that the plug-in estimator for T0(·) using f̃M ′ µ (·) and f̃M ′
ν (·) (with M ′ = 8(1 + TM)) achieves rates that mitigate the curse of dimensionality under sufficient smoothness. In fact, f̃M ′
µ (·) can be viewed as an approximate minimizer of dIP(f̂µ, ·;Cs+2(X ,M ′)) over an appropriate class of Sobolev smooth probability densities. This is carried out because f̂µ(·) by itself may not be a probability density.
Further note that µ̃m, ν̃n as specified in Theorem 2.5 are both smooth, and consequently Γ̃min is a singleton and the supremum in Theorem 2.5 can be dropped. A brief description of the proof technique for Theorem 2.5 is presented in Remark 2.6 below and the actual proof is given in Appendix C.1.
Remark 2.6 (Proof technique). The proof of Theorem 2.5 proceeds along the same lines as Remark 2.3. We first show that Fs (see Remark 2.2) can be chosen as a certain subset of Cs+2(Y◦). We then use Dudley’s entropy integral type bounds which in turn requires the bracketing entropy [126, Page 270] of the class of compactly supported Sobolev smooth functions which can be found in [127, Corollary 2.7.2].
We now explain the implications of both the parts of Theorem 2.5 in the following two remarks.
Remark 2.7 (Mitigating the curse of dimensionality). Theorem 2.5 shows that, under enough smoothness, i.e., when 2(s+ 2) > d, both the upper bounds for (1.8) and (1.9) are Op(n−1/2). This shows that, for large dimensions, provided µ and ν admit smooth enough densities, it is possible to construct plug-in estimators that mitigate the curse of dimensionality. Note that a similar estimator was analyzed in [67, Proposition 1] when m = n. However, the rates obtained in Theorem 2.5 are strictly better than those in [67, Proposition 1]. For m = n, when d < 2(s + 2), [67] obtained a rate of n− s+2 2(s+2)+d which is worse than n−1/2 obtained in Theorem 2.5. For the other regimes, [67] obtains rates (up to log factors) of n−1/4 and n− 1
(s+2)(d+2(s+2)) which are both worse than the respective rates of n−1/2 and n− s+2 d in Theorem 2.5.
Remark 2.8 (Computational aspects of Theorem 2.5). Note that f̃M ′ µ (·) (with M ′ = 8(1 + TM)) is hard to compute whereas f̂µ(·) is computable easily in linear time. Note that if f̂µ(·) itself were a probability density in Cs(X ;TM), then we would have f̂µ = f̃M ′
µ . While Theorem 2.5 does not establish that, it does come close in part 2, from which we can easily derive the following:
lim n→∞
P(f̂µ(·) /∈ Cs(X̃ ;TM)) = 0.
The above shows that f̂µ(·) is indeed bounded below by (TM)−1 on X̃ (any compact subset of the interior of X ), and additionally belongs to Cs(X̃ ;TM) with probability converging to 1. This leads us to conjecture that the natural density version of f̂µ(·), i.e.,
max{f̂µ(·), 0}∫ max{f̂µ(x), 0} dx
should serve as a good proxy for f̃M ′
µ (·) and lead to rates of convergence that mitigate the curse of dimensionality. From a computational perspective, the density specified above is easy to simulate from using an accept-reject algorithm without computing the integral in the denominator (see [101, Algorithm 4.3]). However, our current proof technique does not provide rates of convergence for the above density estimator based on f̂µ(·).
Another important implication of Theorem 2.5 is the bound obtained on |W2(µ̃m, ν̃n)−W2(µ, ν)| when µ 6= ν. We first present the result and then describe the implication.
Proposition 2.6. Consider the setting in Theorem 2.5. Then, provided µ 6= ν, the following holds:
|W2(µ̃m, ν̃n)−W2(µ, ν)| = Op(r(m,n)d,s ).
Proposition 2.6 (see Appendix C.1 for a proof) shows an interesting distinction between the µ 6= ν case and the µ = ν case. For µ = ν, the best possible exponent is n− 1+s 2s+d for d ≥ 3 (see [131, Theorem 3] where the result was established under more general Besov smoothness assumptions). On the contrary, when µ 6= ν, Proposition 2.6 establishes a rate of n− s+2d for the Wasserstein distance which is strictly better than the minimax achievable rate mentioned above when µ = ν. This observation complements [30, Corollary 1] where the authors make a similar remark for the special case of s = 0.

heading: 2.3 Discretized plug-in estimator under smoothness assumptions

text:
In Section 2.1, we discussed how smoothness can be incorporated into the plug-in procedure to get faster rates of convergence. Such plug-in estimators are popular in the computational OT literature (see [7, 8, 25, 36]). However, even after f̃µ(·) ≡ f̃M ′ µ (·), f̃ν(·) ≡ f̃M ′
ν (·) are calculated, T̃ γm,n as in Theorem 2.5 cannot be computed explicitly from data if f̃µ(·) and f̃ν(·) are continuous densities. This is in contrast to T̃ γm,n from Theorem 2.2 in the discrete-discrete case which is explicitly computable using a standard linear program, but achieves worse rates of convergence. This is not unexpected. Thanks to the no free lunch principle, better statistical accuracy is naturally accompanied by heavier computational challenges. Therefore, our goal here is to construct estimators, under smoothness assumptions as in Section 2.2, which are computable in polynomial time (with complexity increasing with smoothness) provided f̃µ(·) and f̃ν(·) can be sampled from, and also attain rates that mitigate the curse of dimensionality.
Construction: We will illustrate the discretized estimator using the kernel based estimator from Section 2.2. Similar results also hold for the wavelet based estimator from Appendix A. Recall the kernel density estimators f̃µ(·) and f̃ν(·) (see (2.6)). Sample M ≥ 1 random points from both f̃µ(·) and f̃ν(·). Let µ̂m,M and ν̂n,M denote the standard empirical measures on the M points sampled from f̃µ(·) and f̃ν(·) respectively. Finally construct T̃m,n ≡ T̃ γm,n as in Definition 1.2 with µ̃m = µ̂m,M and ν̃n = ν̂n,M . It should be pointed out that a similar construction was also used in [131, Section 6] for estimating probability densities under the Wasserstein loss. Based on this construction, the main result of this section is as follows:
Theorem 2.7. Consider the setting in Theorem 2.5 and the same construction of T̃ γm,n as above. For simplicity, let’s also assume m = n. Accordingly set M = n s+2 2 . Then Γ̃min is a singleton and consequently the following conclusion holds for some constant C > 0:
E [∫ ‖T̃m,n(x)− T0(x)‖2 dµ̃m(x) ] ≤ Cr(n,n)d,s .
The same rates also hold for E|W 22 (µ̃m, ν̃n)−W 22 (µ, ν)|.
The proof of Theorem 2.7 is given in Appendix C.1. Once the empirical measures µ̂m,M and ν̂n,M have been obtained, an explicit computation of T̃m,n as described above requires O(M3) = O(n 3(s+2) 2 ) steps using the Hungarian algorithm, see [73]. This highlights the statistical versus computational trade-off, i.e., in order to mitigate the curse of dimensionality in convergence rates by exploiting smoothness, the computational complexity gets progressively worse by polynomial factors in n. It should be mentioned that (approximate) algorithms faster than the Hungarian algorithm stated above, can be found in [1, 36, 53] to name a few. Due to space constraints, we avoid a detailed discussion on this.
In the above construction, sampling from the smoothed kernel densities f̃µ(·) and f̃ν(·) is crucial. If we would simply draw M bootstrap samples from the empirical distributions µ̂m and ν̂n, the rates of convergence wouldn’t improve from those observed in Theorem 2.2 no matter how large M is.
Acknowledgments and Disclosure of Funding We would like to thank the reviewers for their constructive suggestions that greatly helped improve the quality of the paper. The third author is supported by NSF Grant DMS-2015376.

title: Rates of Estimation of Optimal Transport Maps using Plug-in Estimators via Barycentric Projections


summaries: The paper proposes to use infinitely wide CNNs to perform dataset distillation – obtaining a tiny dataset that enables high accuracy. The approach is an extension of a prior work (Kernel Inducing Points (KIP), etc.). Results are shown on MNIST, Fashion-MNIST, SVHN, CIFAR-10/100 datasets.

summaries: In this work the authors extend the Kernel Inducing Points (KIP) method for dataset distillation to a new set of kernel functions and achieve new state of the art results on dataset compression. Since the KIP algorithm is a key algorithmic mechanism in this work, I will summarize it: The KIP method optimizes what can be seen as a test residual of a target dataset fitted by a kernel regression model that is created using a support dataset. This means that given a target dataset distribution, one can differentiate through the kernel computations to evolve the support data points in the support dataset freely in the input space and optimize them. After optimization the support dataset (Which is usually taken to be much smaller than the target distribution) can be used as a distilled version of the dataset. This compressed dataset enables classification and training without training with extensively large datasets. In this work the authors use the KIP algorithm using a family of kernels corresponding to infinitely wide neural networks that were discovered in the realm of NTK. Such infinite NTK convolution kernels are known to provide state of the art results (when compared with other kernels) on classic classification tasks such as MNIST, CIFAR, etc. The biggest issue of such infinite-network kernels is that they require immense amounts of computation to construct the kernel. In this work, the authors construct such kernels at each training step (!) and then back-propagate through the kernels (!!) to optimize the datapoint. This impressive feat is made possible by considering 3 layer convolutional neural networks and working with a relatively small number of data-points (5K as opposed to the whole dataset). This is still highly computationally challenging and the authors create a careful implementation including an orchestration system working with 100s of GPUs for optimizing the support dataset. With this new kernel family the authors achieve state of the art results on dataset compression and analyze the compressed dataset using high dimensional statistics tools.

summaries: Paper mainly extends the dataset distillation algorithms KIP and LS (Nguyen et al. 2021) to infinitely wide neural networks. To that end, authors present a distributed framework that draws upon huge hardware resources and show improvement on the distillation performance. Paper also analyzes the synthesized data samples (images, labels) via multiple studies.

summaries: The paper performs thorough analysis and exploration on the implication of the proposed KIP and LS methods by Nyugen 2021 for dataset distillation. KIP and LS are based on the inf-width network limit as kernels, and devised based on kernel Ridge regression, and has been shown in prior work to have good dataset distillation performance. This paper greatly extends the analysis in Nyugen 2021, by investigating how KIP and LS can be used to achieve SOTA dataset distillation results in a kernel setting, and also sometimes in a neural network (NN) transfer setting. Furthermore, the authors provide impressive and insightful analyses on (1) how the dataset distillation performance behaves as knobs of the algorithms are tuned (2) what are the properties (spectral, intrinsic dim., visual) of the distilled data. Not every paper needs to propose a new method or new theory. Here, thorough analysis and understanding of an existing good method should still be appreciated.

summaries: This paper performs dataset distillation to improve the deep learning training efficiency. By using the algorithm made by [1] and some expanded methods, the authors achieved the state-of-the-arts accuracy on CIFAR-10, MNIST, Fashion-MNIST, CIFAR10, CIFAR100, and SVHN datasets.  [1] DATASET META-LEARNING FROM KERNEL RIDGE REGRESSION, ICLR 2021

abstractText: The effectiveness of machine learning algorithms arises from being able to extract useful features from large amounts of data. As model and dataset sizes increase, dataset distillation methods that compress large datasets into significantly smaller yet highly performant ones will become valuable in terms of training efficiency and useful feature extraction. To that end, we apply a novel distributed kernel-based meta-learning framework to achieve state-of-the-art results for dataset distillation using infinitely wide convolutional neural networks. For instance, using only 10 datapoints (0.02% of original dataset), we obtain over 65% test accuracy on CIFAR10 image classification task, a dramatic improvement over the previous best test accuracy of 40%. Our state-of-the-art results extend across many other settings for MNIST, Fashion-MNIST, CIFAR-10, CIFAR-100, and SVHN. Furthermore, we perform some preliminary analyses of our distilled datasets to shed light on how they differ from naturally occurring data.

authors:
Alex Krizhevsky
Ilya Sutskever
Geoffrey E Hinton

title: Imagenet classification with deep convolutional neural networks

venue: Advances in neural information processing systems,

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova. Bert

title: Pre-training of deep bidirectional transformers for language understanding

venue: arXiv preprint arXiv:1810.04805,

authors:
Tom B Brown
Benjamin Mann
Nick Ryder
Melanie Subbiah
Jared Kaplan
Prafulla Dhariwal
Arvind Neelakantan
Pranav Shyam
Girish Sastry
Amanda Askell

title: Language models are few-shot learners

authors:
Alexey Dosovitskiy
Lucas Beyer
Alexander Kolesnikov
Dirk Weissenborn
Xiaohua Zhai
Thomas Unterthiner
Mostafa Dehghani
Matthias Minderer
Georg Heigold
Sylvain Gelly

title: An image is worth 16x16 words: Transformers for image recognition at scale

venue: arXiv preprint arXiv:2010.11929,

authors:
Tongzhou Wang
Jun-Yan Zhu
Antonio Torralba
Alexei

title: A Efros. Dataset distillation

venue: arXiv preprint arXiv:1811.10959,

authors:
Ondrej Bohdal
Yongxin Yang
Timothy Hospedales

title: Flexible dataset distillation: Learn labels instead of images

venue: arXiv preprint arXiv:2006.08572,

authors:
Timothy Nguyen
Zhourong Chen
Jaehoon Lee

title: Dataset meta-learning from kernel ridge-regression

venue: In International Conference on Learning Representations,

authors:
Bo Zhao
Hakan Bilen

title: Dataset condensation with differentiable siamese augmentation

venue: In International Conference on Machine Learning,

authors:
Andrew Ilyas
Shibani Santurkar
Dimitris Tsipras
Logan Engstrom
Brandon Tran
Aleksander Madry

title: Adversarial examples are not bugs, they are features, 2019

authors:
Minyoung Huh
Pulkit Agrawal
Alexei A Efros

title: What makes imagenet good for transfer learning

venue: arXiv preprint arXiv:1608.08614,

authors:
Katherine L Hermann
Andrew K Lampinen

title: What shapes feature representations? exploring datasets, architectures, and training

venue: arXiv preprint arXiv:2006.12433,

authors:
Zalán Borsos
Mojmír Mutnỳ
Andreas Krause

title: Coresets via bilevel optimization for continual learning and streaming

venue: arXiv preprint arXiv:2006.03875,

authors:
Bo Zhao
Konda Reddy Mopuri
Hakan Bilen

title: Dataset condensation with gradient matching

authors:
Roman Novak
Lechao Xiao
Jaehoon Lee
Yasaman Bahri
Greg Yang
Jiri Hron
Daniel A. Abolafia
Jeffrey Pennington
Jascha Sohl-Dickstein

title: Bayesian deep convolutional networks with many channels are gaussian processes

authors:
Adrià Garriga-Alonso
Laurence Aitchison
Carl Edward Rasmussen

title: Deep convolutional networks as shallow gaussian processes

authors:
Sanjeev Arora
Simon S Du
Wei Hu
Zhiyuan Li
Russ R Salakhutdinov
Ruosong Wang

title: On exact computation with an infinitely wide neural net

venue: In Advances in Neural Information Processing Systems,

authors:
Zhiyuan Li
Ruosong Wang
Dingli Yu
Simon S Du
Wei Hu
Ruslan Salakhutdinov
Sanjeev Arora

title: Enhanced convolutional neural tangent kernels

authors:
Vaishaal Shankar
Alex Chengyu Fang
Wenshuo Guo
Sara Fridovich-Keil
Ludwig Schmidt
Jonathan RaganKelley
Benjamin Recht

title: Neural kernels without tangents

title: Approximation and learning with deep convolutional models: a kernel perspective

venue: arXiv preprint arXiv:2102.10032,

authors:
Sanjeev Arora
Simon S. Du
Zhiyuan Li
Ruslan Salakhutdinov
Ruosong Wang
Dingli Yu

title: Harnessing the power of infinitely wide deep nets on small-data tasks

authors:
Jaehoon Lee
Samuel S Schoenholz
Jeffrey Pennington
Ben Adlam
Lechao Xiao
Roman Novak
Jascha Sohl-Dickstein

title: Finite versus infinite neural networks: an empirical study

authors:
Roman Novak
Lechao Xiao
Jiri Hron
Jaehoon Lee
Alexander A. Alemi
Jascha Sohl-Dickstein
Samuel S. Schoenholz

title: Neural tangents: Fast and easy infinite neural networks in python

authors:
James Bradbury
Roy Frostig
Peter Hawkins
Matthew James Johnson
Chris Leary
Dougal Maclaurin
Skye Wanderman-Milne

title: JAX: composable transformations of Python+NumPy programs, 2018

venue: URL http://github.com/google/jax

authors:
Yann LeCun
Corinna Cortes
CJ Burges

title: Mnist handwritten digit database

venue: ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist,

authors:
Han Xiao
Kashif Rasul
Roland Vollgraf

title: Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017

authors:
Yuval Netzer
Tao Wang
Adam Coates
Alessandro Bissacco
Bo Wu
Andrew Y Ng

title: Reading digits in natural images with unsupervised feature learning

title: Learning multiple layers of features from tiny images

venue: Technical report,

authors:
Arthur Jacot
Franck Gabriel
Clement Hongler

title: Neural tangent kernel: Convergence and generalization in neural networks

authors:
Jaehoon Lee
Lechao Xiao
Samuel S. Schoenholz
Yasaman Bahri
Roman Novak
Jascha Sohl-Dickstein
Jeffrey Pennington

title: Wide neural networks of any depth evolve as linear models under gradient descent

authors:
Jiaoyang Huang
Horng-Tzer Yau

title: Dynamics of deep neural networks and neural tangent hierarchy

authors:
Ethan Dyer
Guy Gur-Ari

title: Asymptotics of wide networks from feynman diagrams

authors:
Anders Andreassen
Ethan Dyer

title: Asymptotics of wide convolutional neural networks

venue: arxiv preprint arXiv:2008.08675,

title: Non-Gaussian processes and neural networks at finite widths

venue: In Mathematical and Scientific Machine Learning Conference,

title: The intrinsic dimensionality of signal collections

venue: IEEE Transactions on Information Theory,

authors:
Francesco Camastra
Antonino Staiano

title: Intrinsic dimension estimation: Advances and open problems

venue: Information Sciences,

authors:
Elena Facco
Maria d’Errico
Alex Rodriguez
Alessandro Laio

title: Estimating the intrinsic dimension of datasets by a minimal neighborhood information

venue: Scientific reports,

authors:
A. Ansuini
A. Laio
J. Macke
D. Zoccolan

title: Intrinsic dimension of data representations in deep neural networks

venue: In NeurIPS,

authors:
Ilia Sucholutsky
Matthias Schonlau

title: Soft-label dataset distillation and text dataset distillation

venue: arXiv preprint arXiv:1910.02551,

title: Priors for infinite networks (tech

venue: rep. no. crg-tr-94-1). University of Toronto,

authors:
Jaehoon Lee
Yasaman Bahri
Roman Novak
Sam Schoenholz
Jeffrey Pennington
Jascha Sohl-dickstein

title: Deep neural networks as gaussian processes

authors:
Alexander G. de G. Matthews
Jiri Hron
Mark Rowland
Richard E. Turner
Zoubin Ghahramani

title: Gaussian process behaviour in wide deep neural networks

title: Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation

venue: arXiv preprint arXiv:1902.04760,

title: Wide feedforward or recurrent neural networks of any architecture are gaussian processes

authors:
Jiri Hron
Yasaman Bahri
Jascha Sohl-Dickstein
Roman Novak

title: Infinite attention: NNGP and NTK for deep attention networks

authors:
Lechao Xiao
Jeffrey Pennington
Samuel S Schoenholz

title: Disentangling trainability and generalization in deep learning

authors:
Ben Adlam
Jeffrey Pennington

title: The neural tangent kernel in high dimensions: Triple descent and a multi-scale theory of generalization

authors:
Aitor Lewkowycz
Yasaman Bahri
Ethan Dyer
Jascha Sohl-Dickstein
Guy Gur-Ari

title: The large learning rate phase of deep learning: the catapult mechanism

venue: arXiv preprint arXiv:2003.02218,

authors:
Aitor Lewkowycz
Guy Gur-Ari

title: On the training dynamics of deep networks with l_2 regularization

authors:
Ben Adlam
Jaehoon Lee
Lechao Xiao
Jeffrey Pennington
Jasper Snoek

title: Exploring the uncertainty properties of neural networks’ implicit priors in the infinite-width limit

authors:
Daniel S Park
Jaehoon Lee
Daiyi Peng
Yuan Cao
Jascha Sohl-Dickstein

title: Towards nngp-guided neural architecture search

venue: arXiv preprint arXiv:2011.06006,

authors:
Wuyang Chen
Xinyu Gong
Zhangyang Wang

title: Neural architecture search on imagenet in four {gpu} hours: A theoretically inspired perspective

authors:
Amir Zandieh
Insu Han
Haim Avron
Neta Shoham
Chaewon Kim
Jinwoo Shin

title: Scaling neural tangent kernels via sketching and random features

authors:
Joel Hestness
Sharan Narang
Newsha Ardalani
Gregory F. Diamos
Heewoo Jun
Hassan Kianinejad
Md. Mostofa Ali Patwary
Yang Yang
Yanqi Zhou

title: Deep learning scaling is predictable

venue: empirically. CoRR,

authors:
Jonathan S. Rosenfeld
Amir Rosenfeld
Yonatan Belinkov
Nir Shavit

title: A constructive prediction of the generalization error across scales

authors:
Jared Kaplan
Sam McCandlish
Tom Henighan
Tom B. Brown
Benjamin Chess
Rewon Child
Scott Gray
Alec Radford
Jeffrey Wu
Dario Amodei

title: Scaling laws for neural language models, 2020

authors:
Yasaman Bahri
Ethan Dyer
Jared Kaplan
Jaehoon Lee
Utkarsh Sharma

title: Explaining neural scaling laws

venue: arXiv preprint arXiv:2102.06701,

authors:
Ekin D. Cubuk
Barret Zoph
Dandelion Mane
Vijay Vasudevan
Quoc V. Le

title: Autoaugment: Learning augmentation strategies from data

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
Jascha Sohl-Dickstein
Roman Novak
Samuel S Schoenholz
Jaehoon Lee

title: On the infinite width limit of neural networks with a standard parameterization

venue: arXiv preprint arXiv:2001.07301,

heading: 1 Introduction

text:
Deep learning has become extraordinarily successful in a wide variety of settings through the availability of large datasets [Krizhevsky et al., 2012, Devlin et al., 2018, Brown et al., 2020, Dosovitskiy et al., 2020]. Such large datasets enable a neural network to learn useful representations of the data that are adapted to solving tasks of interest. Unfortunately, it can be prohibitively costly to acquire such large datasets and train a neural network for the requisite amount of time.
One way to mitigate this problem is by constructing smaller datasets that are nevertheless informative. Some direct approaches to this include choosing a representative subset of the dataset (i.e. a coreset) or else performing a low-dimensional projection that reduces the number of features. However, such methods typically introduce a tradeoff between performance and dataset size, since what they produce is a coarse approximation of the full dataset. By contrast, the approach of dataset distillation is to synthesize datasets that are more informative than their natural counterparts when equalizing for dataset size [Wang et al., 2018, Bohdal et al., 2020, Nguyen et al., 2021, Zhao and Bilen, 2021]. Such resulting datasets will not arise from the distribution of natural images but will nevertheless capture features useful to a neural network, a capability which remains mysterious and is far from being well-understood [Ilyas et al., 2019, Huh et al., 2016, Hermann and Lampinen, 2020].
The applications of such smaller, distilled datasets are diverse. For nonparametric methods that scale poorly with the training dataset (e.g. nearest-neighbors or kernel-ridge regression), having a reduced dataset decreases the associated memory and inference costs. For the training of neural networks, such distilled datasets have found several applications in the literature, including increasing the
∗Work done while at Google Research.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
effectiveness of replay methods in continual learning [Borsos et al., 2020] and helping to accelerate neural architecture search [Zhao et al., 2021, Zhao and Bilen, 2021].
In this paper, we perform a large-scale extension of the methods of Nguyen et al. [2021] to obtain new state-of-the-art (SOTA) dataset distillation results. Specifically, we apply the algorithms KIP (Kernel Inducing Points) and LS (Label Solve), first developed in Nguyen et al. [2021], to infinitely wide convolutional networks by implementing a novel, distributed meta-learning framework that draws upon hundreds of accelerators per training. The need for such resources is necessitated by the computational costs of using infinitely wide neural networks built out of components occurring in modern image classification models: convolutional and pooling layers (see §B for details). The consequence is that we obtain distilled datasets that are effective for both kernel ridge-regression and neural network training.
Additionally, we initiate a preliminary study of the images and labels which KIP learns. We provide a visual and quantitative analysis of the data learned and find some surprising results concerning their interpretability and their dimensional and spectral properties. Given the efficacy of KIP and LS learned data, we believe a better understanding of them would aid in the understanding of feature learning in neural networks.
To summarize, our contributions are as follows:
1. We achieve SOTA dataset distillation results on a wide variety of datasets (MNIST, FashionMNIST, SVHN, CIFAR-10, CIFAR-100) for both kernel ridge-regression and neural network training. In several instances, our results achieve an impressively wide margin over prior art, including over 25% and 37% absolute gain in accuracy on CIFAR-10 and SVHN image classification, respectively, when using only 10 images (Tables 1, 2, A11).
2. We develop a novel, distributed meta-learning framework specifically tailored to the computational burdens of sophisticated neural kernels (§2.1).
3. We highlight and analyze some of the peculiar features of the distilled datasets we obtain, illustrating how they differ from natural data (§4).
4. We open source the distilled datasets, which used thousands of GPU hours, for the research community to further investigate at https://github.com/google-research/ google-research/tree/master/kip.

heading: 2 Setup

text:
Background on infinitely wide convolutional networks. Recent literature has established that Bayesian and gradient-descent trained neural networks converge to Gaussian Processes (GP) as the number of hidden units in intermediary layers approaches infinity (see §5). These results hold for many different architectures, including convolutional networks, which converge to a particular GP in the limit of infinite channels [Novak et al., 2019, Garriga-Alonso et al., 2019, Arora et al., 2019]. Bayesian networks are described by the Neural Network Gaussian Process (NNGP) kernel, while gradient descent networks are described by the Neural Tangent Kernel (NTK). Since we are interested in synthesizing datasets that can be used with both kernel methods and common gradient-descent trained neural networks, we focus on NTK in this work.
Infinitely wide networks have been shown to achieve SOTA (among non-parametric kernels) results on image classification tasks [Novak et al., 2019, Arora et al., 2019, Li et al., 2019, Shankar et al., 2020, Bietti, 2021] and even rival finite-width networks in certain settings [Arora et al., 2020, Lee et al., 2020]. This makes such kernels especially suitable for our task. As convolutional models, they encode useful inductive biases of locality and translation invariance [Novak et al., 2019], which enable good generalization. Moreover, flexible and efficient computation of these kernels are possible due to the Neural Tangent library [Novak et al., 2020].
Specific models considered. The central neural network (and corresponding infinite-width model) we consider is a simple 4-layer convolutional network with average pooling layers that we refer to as ConvNet throughout the text. This architecture is a slightly modified version of the default model used by Zhao and Bilen [2021], Zhao et al. [2021] and was chosen for ease of baselining (see §A for details). In several other settings we also consider convolutional networks without pooling
layers ConvVec,2 and networks with no convolutions and only fully-connected layers FC. Depth of architecture (as measured by number of hidden layers) is indicated by an integer suffix.
Background on algorithms. We review the Kernel Inducing Points (KIP) and Label Solve (LS) algorithms introduced by Nguyen et al. [2021]. Given a kernel K, the kernel ridge-regression (KRR) loss function trained on a support dataset (Xs, ys) and evaluated on a target dataset (Xt, yt) is
L(Xs, ys) = 1 2 ÂÂÂÂÂyt −KXtXs(KXsXs + λI) −1 ys ÂÂÂÂÂ 2 2 , (1)
where if U and V are sets, KUV is the matrix of kernel elements (K(u, v))u∈U,v∈V . Here λ > 0 is a fixed regularization parameter. The KIP algorithm consists of minimizing (1) with respect to the support set (either just the Xs or along with the labels ys). Here, we sample (Xt, yt) from a target dataset D at every (meta)step, and update the support set using gradient-based methods. Additional variations include augmenting theXt or sampling a different kernelK (from a fixed family of kernels) at each step.
The Label Solve algorithm consists of solving for the least-norm minimizer of (1) with respect to ys. This yields the labels
y ∗ s = (KXtXs(KXsXs + λI) −1) + yt, (2)
where A+ denotes the pseudo-inverse of the matrix A. Note that here (Xt, yt) = D, i.e. the labels are solved using the whole target set.
In our applications of KIP and Label Solve, the target dataset D is always significantly larger than the support set (Xs, ys). Hence, the learned support set or solved labels can be regarded as distilled versions of their respective targets. We also initialize our support images to be a subset of natural images, though they could also be initialized randomly.
Based on the infinite-width correspondence outlined above and in §5, dataset distillation using KIP or LS that is optimized for KRR should extend to the corresponding finite-width neural network training. Our experimental results in §3 validate this expectation across many settings.

heading: 2.1 Client-Server Distributed Workflow

text:
We invoke a client-server model of distributed computation3, in which a server distributes independent workloads to a large pool of client workers that share a queue for receiving and sending work. Our distributed implementation of the KIP algorithm has two distinct stages:
Forward pass: In this step, we compute the support-support and target-support matrices K(Xs, Xs) and K(Xt, Xs). To do so, we partition Xs ×Xs and Xt ×Xs into pairs of images (x, x′), each with batch size B. We send such pairs to workers compute the respective matrix block K(x, x′). The server aggregates all these blocks to obtain the K(Xs, Xs) and K(Xt, Xs) matrices. Backward pass: In this step, we need to compute the gradient of the loss L (1) with respect to the support set Xs. We need only consider ∂L/∂Xs since ∂L/∂ys is cheap to compute. By the chain rule, we can write
∂L
∂Xs =
∂L ∂(K(Xs, Xs)) ∂K(Xs, Xs) ∂Xs +
∂L ∂(K(Xt, Xs)) ∂K(Xt, Xs) ∂Xs .
The derivatives of L with respect to the kernel matrices are inexpensive, since L depends in a simple way on them (matrix multiplication and inversion). What is expensive to compute is the derivative of the kernel matrices with respect to the inputs. Each kernel element is an independent function of the inputs and a naive computation of the derivative of a block would require forwardmode differentiation, infeasible due to the size of the input images and the cost to compute the individual kernel elements. Thus our main novelty is to divide up the gradient computation into backward differentiation sub-computations, specifically by using the built-in function jax.vjp in JAX [Bradbury et al., 2018]. Denoting K = K(Xs, Xs) or K(Xt, Xs) for short-hand, we divide the
2The abbreviation “Vec” stands for vectorization, to indicate that activations of the network are vectorized (flattened) before the top fully-connected layer instead of being pooled.
3Implemented using Courier available at https://github.com/deepmind/launchpad.

heading: 1 12.8±0.3 13.9±0.3 - 23.8 34.9±0.1 33.3±0.3CIFAR-100 10 25.2±0.3 32.3±0.3 - 39.2 47.9±0.2 49.5±0.3

text: matrix ∂L/∂K, computed on the server, into B ×B blocks corresponding to ∂L/∂K(x, x′), where x and x′ each have batch size B. We send each such block, along with the corresponding block of image data (x, x′), to a worker. The worker then treats the ∂L/∂K(x, x′) it receives as the cotangent vector argument of jax.vjp that, via contraction, converts the derivative of K(x, x′) with respect to x into a scalar. The server aggregates all these partial gradient computations performed by the workers, over all possible B ×B blocks, to compute the total gradient ∂L/∂Xs used to update Xs.

heading: 3 Experimental Results

heading: 3.1 Kernel Distillation Results

text:
We apply the KIP and LS algorithms using the ConvNet architecture on the datasets MNIST [LeCun et al., 2010], Fashion MNIST [Xiao et al., 2017], SVHN [Netzer et al., 2011], CIFAR-10 [Krizhevsky, 2009], and CIFAR-100. Here, the goal is to condense the train dataset down to a learned dataset of size 1, 10, or 50 images per class. We consider a variety of hyperparameter settings (image preprocessing method, whether to augment target data, and whether to train the support labels for KIP), the full details of which are described in §A. For space reasons, we show a subset of our results in Table 1, with results corresponding to the remaining set of hyperparameters left to Tables A3-A10. We highlight here that a crucial ingredient for our strong results in the RGB dataset setting is the use of regularized ZCA preprocessing. Note the variable effect that our augmentations have on performance (see the last two columns of Table 1): they typically only provides a benefit for a sufficiently large support set. This result is consistent with Zhao et al. [2021], Zhao and Bilen [2021], in which gains from augmentations are also generally obtained from larger support sets. We tried varying the fraction (0.25 and 0.5) of each target batch that is augmented at each training step and found that while 10 images still did best without augmentations, 100 and 500 images typically did slightly better with some partial augmentations (versus none or full). For instance, for 500 images on CIFAR-10, we obtained 81.1% test accuracy using augmentation rate 0.5. Thus, our observation is that given that larger support sets can distill larger target datasets, as the former increases in size, the latter can be augmented more aggressively for obtaining optimal generalization performance.
Remarkably, our results in Table 1 outperform all prior baselines across all dataset settings. Our results are especially strong in the small support size regime, with our 1 image per class results for KRR outperforming over 100 times as many natural images (see Table A1). We also obtain a significant margin over prior art across all datasets, with our largest margin being a 37% absolute gain in test accuracy for the SVHN dataset.

heading: 3.2 Kernel Transfer

text:
In §3.1, we focused on obtaining state-of-the-art dataset distillation results for image classification using a specific kernel (ConvNet). Here, we consider the variation of KIP in which we sample from a family of kernels (which we call sampling KIP). We validate that sampling KIP adds robustness to the learned images in that they perform well for the family of kernels sampled during training.
In Figure 1, we plot test performance of sampling KIP when using the kernels ConvNet, Conv-Vec3, and Conv-Vec8 (denoted by “all”) alongside KIP trained with just the individual kernels. Sampling KIP performs well at test time when using any of the three kernels, whereas datasets trained using a single kernel have a significant performance drop when using a different kernel.

heading: 3.3 Neural Network Transfer

text:
In this section, we study how our distilled datasets optimized using KIP and LS transfer to the setting of finite-width neural networks. The main results are shown in Table 2. The third column shows the best neural network performance obtained by training on a KIP dataset of the corresponding size with respect to some choice of KIP and neural network training hyperparameters (see §A for details). Since the datasets are optimized for kernel ridge-regression and not for neural network training itself, we expect some performance loss when transferring to finite-width networks, which we record in the fourth column. Remarkably, the drop due to this transfer is quite moderate or small and sometimes the transfer can even lead to gain in performance (see LS for SVHN dataset with 10 images per class).
Overall, our transfer to finite-width networks outperforms prior art based on DC/DSA [Zhao et al., 2021, Zhao and Bilen, 2021] in the 1 image per class setting for all the RGB datasets (SVHN, CIFAR-10, CIFAR-100). Moreover, for CIFAR-10, we outperform DC/DSA in all settings.
Figures 2 and 3 provide a closer look at KIP transfer changes under various settings. The first of these tracks how transfer performance changes when adding additional layers as function of the number of KIP training steps used. The normalization layers appear to harm performance for MSE loss, which can be anticipated from their absence in the KIP and LS optimization procedures. However they appear to provide some benefit for cross entropy loss. For Figure 3, we observe that as KIP training progresses, the downstream finite-width network’s performance also improves in general. A notable exception is observed when learning the labels in KIP, where longer training steps lead to deterioration of information useful to training finite-width neural networks. We also observe that as predicted by infinite-width theory [Jacot et al., 2018, Lee et al., 2019], the overall gap between KIP or LS performance and finite-width neural network decreases as the width increases. While our best performing transfer is obtained with width 1024, Figure 3 (middle) suggest that even with modest width of 64, our transfer can outperform prior art of 60.6% by Zhao and Bilen [2021].
Finally, in Figure 4, we investigate the performance of KIP images over the course of training of a single run, as compared to natural images, over a range of hyperparameters. We find the outperformance of KIP images above natural images consistent across hyperparameters and checkpoints. This suggests that our KIP images may also be effective for accelerated hyperparameter search, an application of dataset distillation explored in Zhao et al. [2021], Zhao and Bilen [2021].
Altogether, we find our neural network training results encouraging. First, it validates the applicability of infinite-width methods to the setting of finite width [Huang and Yau, 2020, Dyer and Gur-Ari, 2020, Andreassen and Dyer, 2020, Yaida, 2020, Lee et al., 2020]. Second, we find some of the transfer results quite surprising, including efficacy of label solve and the use of cross-entropy for certain settings (see §A for the full details).

heading: 4 Understanding KIP Images and Labels

text:
A natural question to consider is what causes KIP to improve generalization performance. Does it simplify support images, removing noise and minor sources of variation while keeping only the core features shared between many target images of a given class? Or does it make them more complex by producing outputs that combine characteristics of many samples in a single resulting collage? While
properly answering this question is subject to precise definitions of simplicity and complexity, we find that KIP tends to increase the complexity of pictures based on the following experiments:
• Visual analysis: qualitatively, KIP images tend to be richer in textures and contours. • Dimensional analysis: KIP produces images of higher intrinsic dimensionality (i.e. an
estimate of the dimensionality of the manifold on which the images lie) than natural images.
• Spectral analysis: unlike natural images, for which the bulk of generalization performance is explained by a small number of top few eigendirections, KIP images leverage the whole spectrum much more evenly.
Combined, these results let us conclude that KIP usually increases complexity, integrating features from many target images into much fewer support images.
Visual Analysis. A visual inspection of our learned data leads to intriguing observations in terms of interpretability. Figure 5 shows examples of KIP learned images from CIFAR-100. The resulting images are heterogeneous in terms of how they can be interpreted as distilling the data. For instance, the distilled apple image seems to consist of many apples nested within a possibly larger apple, whereas the distilled bottle image starts off as two bottles and before transforming into one, while other classes (like the beaver) are altogether visually indistinct. Investigating and quantifying aspects that make these images generalize so well is a promising avenue for future work. We show examples from other datasets in Figure A2.
In Figure 6, we compare MNIST KIP data learned with and without label learning. For the latter case with images and labels optimized jointly, while labels become more informative, encoding richer inter-class information, the images become less interpretable. This behavior consistently leads to superior KRR results, but appears to not be leveraged as efficiently in the neural network transfer setting (Table 2). Experimental details can be found in §A.
Dimensional Analysis. We study the intrinsic dimension of KIP images and find that they tend to grow. Intrinsic dimension (ID) was first defined by Bennett [1969] as “the number of free parameters required in a hypothetical signal generator capable of producing a close approximation to each signal in the collection”. In our context, it is the dimensionality of the manifold embedded into the image space which contains all the support images. Intuitively, simple datasets have a low ID as they can be described by a small number of coordinates on the low-dimensional data manifold.
ID can be defined and estimated differently based on assumptions on the manifold structure and the probability density function of the images on this manifold (see [Camastra and Staiano, 2016] for review). We use the “Two-NN” method developed by [Facco et al., 2017], which makes relatively few assumptions and allows to estimate ID only from two nearest-neighbor distances for each datapoint.
Figure 7 shows that the ID is increasing for the learned KIP images as a function of the training step across a variety of configurations (training with or without augmentations/label learning) and datasets. One might expect that a distillation procedure should decrease dimensionality. On the other hand, Ansuini et al. [2019] showed that ID increases in the earlier layers of a trained neural network. It remains to be understood if this latter observation has any relationship with our increased ID. Note that ZCA preprocessing, which played an important role for getting the best performance for our RGB datasets, increases dimensionality of the underlying data (see Figure A4).
Spectral Analysis. Another distinguishing property of KIP images is how their spectral components contribute to performance. In Figure 8, we spot how different spectral bands of KIP learned images affect test performance as compared to their initial natural images. Here, we use the FC2, Conv-Vec8, and ConvNet architectures. We note that for natural images (light bars), most of their performance is captured by the top 20% of eigenvalues. For KIP images, the performance is either more evenly distributed across the bands (FC and Conv-Vec8) or else is skewed towards the tail (ConvNet).

heading: 5 Related Work

text:
Dataset distillation was first studied in Wang et al. [2018]. The work of Sucholutsky and Schonlau [2019], Bohdal et al. [2020] build upon it by distilling labels. Zhao et al. [2021] proposes condensing a training set by harnessing a gradient matching condition. Zhao and Bilen [2021] takes this idea further by applying a suitable augmentation strategy. Note that while Zhao and Bilen [2021] is limited in augmentation expressiveness (they have to apply a single augmentation per training iteration), we can sample augmentations independently per image in our target set per train step. Our work together with Nguyen et al. [2021] are, to the best of our knowledge, the only works using kernel-based methods for dataset distillation on image classification datasets.
Our use of kernels stems from the correspondence between infinitely-wide neural networks and kernel methods [Neal, 1994, Lee et al., 2018, Matthews et al., 2018, Jacot et al., 2018, Novak et al., 2019, Garriga-Alonso et al., 2019, Arora et al., 2019, Yang, 2019a,b, Hron et al., 2020] and the extended correspondence with the finite width corrections [Dyer and Gur-Ari, 2020, Huang and Yau, 2020, Yaida, 2020]. These correspondences underlie the transferability of our KRR results to neural networks , and have been utilized in understanding trainability [Xiao et al., 2020], generalizations [Adlam and Pennington, 2020], training dynamics [Lewkowycz et al., 2020, Lewkowycz and Gur-Ari, 2020], uncertainty [Adlam et al., 2021], and demonstrated their effectiveness for smaller datasets [Arora et al., 2020] and neural architecture search [Park et al., 2020, Chen et al., 2021].

heading: 6 Conclusion

text:
We performed an extensive study of dataset distillation using the KIP and LS algorithms applied to convolutional architectures, obtaining SOTA results on a variety of image classification datasets. In some cases, our learned datasets were more effective than a natural dataset two orders of magnitude larger in size. There are many interesting followup directions and questions from our work:
First, integrating efficient kernel-approximation methods into our algorithms, such as those of Zandieh et al. [2021], will reduce computational burden and enable scaling up to larger datasets. In this direction, the understanding of how various resources (e.g. data, parameter count, compute) scale when optimizing for neural network performance has received significant attention as machine learning models continue to stretch computational limits [Hestness et al., 2017, Rosenfeld et al., 2020, Kaplan et al., 2020, Bahri et al., 2021]. Developing our understanding of how to harness smaller, yet more useful representations data would aid in such endeavors. In particular, it would be especially interesting to explore how well datasets can be compressed as they scale up in size.
Second, LS and KIP with label learning shows that optimizing labels is a very powerful tool for dataset distillation. The labels we obtain are quite far away from standard, interpretable labels and we feel their effectiveness suggests that understanding of how to optimally label data warrants further study.
Finally, the novel features obtained by our learned datasets, and those of dataset distillation methods in general, may reveal insights into interpretability and the nature of sample-efficient representations. For instance, observe the bee and bicycle images in Figure 5: the bee class distills into what appears to be spurious visual features (e.g. pollen), while the bicycle class distills to the essential contours of a typical bicycle. Additional analyses and explorations of this type could offer insights into the perennial question of how neural networks learn and generalize.
Acknowledgments We would like to acknowledge special thanks to Samuel S. Schoenholz, who proposed and helped develop the overall strategy for our distributed KIP learning methodology. We are also grateful to Ekin Dogus Cubuk and Manuel Kroiss for helpful discussions.

title: Dataset Distillation with Infinitely Wide Convolutional Networks