summaries: The paper proposes a new conformalized procedure for computing uncertainty sets in classification tasks. The key feature of the method is that the size of the uncertainty sets are regularized via a penalty on the size. The issue of large uncertainty sets produced by conformalized procedures is an interesting one, which the paper does well to highlight. The proposed solution of using an additive regularizer is reasonable, and appears to be effective for sensible choices of the hyper-parameters. However, the paper has some significant weaknesses.

summaries: In this paper, the authors propose a regularized conformal score for use in a conformal prediction framework. This regularizer is motivated by the instabilities of top-p variations on conformal scores (cf. Romano et. al., 2020) and the resulting high-variance in output conformal prediction set sizes. The proposed regularizer smooths top-p scores with top-k scores, which empirically results in more robust predictive sets. The authors also perform a large-scale evaluation on ImageNet with modern architectures, which serves as a helpful benchmark for conformal prediction algorithms.

summaries: Prediction sets are used to quantify the uncertainty of classification. The naive approach which include the labels until a pre-specified coverage probability is satisfied often leads to large prediction sets. Adaptive Prediction Sets (APS) can output prediction sets with desired coverage but set sizes are still not satisfyingly small and the results are unstable, especially when many probability estimations fall into the tail of the distribution.

abstractText: Convolutional image classifiers can achieve high predictive accuracy, but quantifying their uncertainty remains an unresolved challenge, hindering their deployment in consequential settings. Existing uncertainty quantification techniques, such as Platt scaling, attempt to calibrate the network’s probability estimates, but they do not have formal guarantees. We present an algorithm that modifies any classifier to output a predictive set containing the true label with a user-specified probability, such as 90%. The algorithm is simple and fast like Platt scaling, but provides a formal finite-sample coverage guarantee for every model and dataset. Our method modifies an existing conformal prediction algorithm to give more stable predictive sets by regularizing the small scores of unlikely classes after Platt scaling. In experiments on both Imagenet and Imagenet-V2 with ResNet-152 and other classifiers, our scheme outperforms existing approaches, achieving coverage with sets that are often factors of 5 to 10 smaller than a stand-alone Platt scaling baseline.

authors:
Maxime Cauchois
Suyash Gupta
John Duchi

title: Knowing what you know: valid and validated confidence sets in multiclass and multilabel prediction

authors:
Adam Fisch
Tal Schuster
Tommi S. Jaakkola
Regina Barzilay

title: Efficient conformal prediction via cascaded inference with expanded admission

title: Uncertainty in Deep Learning

venue: University of Cambridge,

authors:
Leying Guan
Rob Tibshirani

title: Prediction and outlier detection in classification problems

authors:
Chuan Guo
Geoff Pleiss
Yu Sun
Kilian Q Weinberger

title: On calibration of modern neural networks. 2017

authors:
Chirag Gupta
Arun K Kuchibhotla
Aaditya K Ramdas

title: Nested conformal prediction and quantile out-of-bag ensemble methods

venue: arXiv, pp

authors:
Heinrich Jiang
Been Kim
Melody Guan
Maya Gupta

title: To trust or not to trust a classifier

venue: In Advances in Neural Information Processing Systems,

authors:
Volodymyr Kuleshov
Nathan Fenner
Stefano Ermon

title: Accurate uncertainties for deep learning using calibrated regression. 2018

authors:
Balaji Lakshminarayanan
Alexander Pritzel
Charles Blundell

title: Simple and scalable predictive uncertainty estimation using deep ensembles

authors:
Jing Lei
Larry Wasserman

title: Distribution-free prediction bands for non-parametric regression

venue: Journal of the Royal Statistical Society: Series B (Statistical Methodology),

authors:
Jing Lei
Max G’Sell
Alessandro Rinaldo
Ryan J. Tibshirani
Larry Wasserman

title: Distributionfree predictive inference for regression

venue: Journal of the American Statistical Association,

authors:
Haoxiang Li
Zhe Lin
Xiaohui Shen
Jonathan Brandt
Gang Hua

title: A convolutional neural network cascade for face detection

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

authors:
Qing Li
Weidong Cai
Xiaogang Wang
Yun Zhou
David Dagan Feng
Mei Chen

title: Medical image classification with convolutional neural network

venue: In International Conference on Control Automation Robotics & Vision,

authors:
Alexander Selvikvåg Lundervold
Arvid Lundervold

title: An overview of deep learning in medical imaging focusing on MRI

venue: Zeitschrift für Medizinische Physik,

title: Bayesian methods for adaptive models

venue: PhD thesis, California Institute of Technology,

authors:
Soundouss Messoudi
Sylvain Rousseau
Sebastien Destercke

title: Deep conformal prediction for robust models

venue: In Information Processing and Management of Uncertainty in Knowledge-Based Systems,

title: Bayesian Learning for Neural Networks

venue: Springer Science & Business Media,

authors:
Jeremy Nixon
Michael W Dusenberry
Linchuan Zhang
Ghassen Jerfel
Dustin Tran

title: Measuring calibration in deep learning

venue: In CVPR Workshops,

authors:
Harris Papadopoulos
Kostas Proedrou
Vladimir Vovk
Alex Gammerman

title: Inductive confidence machines for regression

venue: In Machine Learning: European Conference on Machine Learning ECML

authors:
Sangdon Park
Osbert Bastani
Nikolai Matni
Insup Lee

title: PAC confidence sets for deep neural networks via calibrated prediction

authors:
Adam Paszke
Sam Gross
Francisco Massa
Adam Lerer
James Bradbury
Gregory Chanan
Trevor Killeen
Zeming Lin
Natalia Gimelshein
Luca Antiga

title: Pytorch: An imperative style, highperformance deep learning library

authors:
Tim Pearce
Mohamed Zaki
Alexandra Brintrup
Andy Neely

title: High-quality prediction intervals for deep learning: A distribution-free, ensembled approach

venue: In International Conference on Machine Learning,

title: Probabilistic outputs for support vector machines and comparisons to regularized likelihood methods

venue: Advances in Large Margin Classifiers,

authors:
Joaquin Quinonero-Candela
Carl Edward Rasmussen
Fabian Sinz
Olivier Bousquet
Bernhard Schölkopf

title: Evaluating predictive uncertainty challenge

venue: In Machine Learning Challenges Workshop,

authors:
Muhammad Imran Razzak
Saeeda Naz
Ahmad Zaib

title: Deep learning for medical image processing: Overview, challenges and the future

venue: In Classification in BioApps,

authors:
Carlos Riquelme
George Tucker
Jasper Snoek

title: Deep Bayesian bandits showdown: An empirical comparison of Bayesian deep networks for Thompson sampling. 2018

authors:
Yaniv Romano
Evan Patterson
Emmanuel Candès

title: Conformalized quantile regression

authors:
Yaniv Romano
Matteo Sesia
Emmanuel J. Candès

title: Classification with valid and adaptive coverage

authors:
Mauricio Sadinle
Jing Lei
Larry Wasserman

title: Least ambiguous set-valued classifiers with bounded error levels

authors:
Murat Sensoy
Lance Kaplan
Melih Kandemir

title: Evidential deep learning to quantify classification uncertainty

authors:
Ryan J Tibshirani
Rina Foygel Barber
Emmanuel Candes
Aaditya Ramdas

title: Conformal prediction under covariate shift

title: Conditional validity of inductive conformal predictors

venue: In Proceedings of the Asian Conference on Machine Learning,

authors:
Vladimir Vovk
Alexander Gammerman
Craig Saunders

title: Machine-learning applications of algorithmic randomness

authors:
Vladimir Vovk
David Lindsay
Ilia Nouretdinov
Alex Gammerman

title: Mondrian confidence machine

venue: Technical report,

authors:
Vladimir Vovk
Alex Gammerman
Glenn Shafer

title: Algorithmic Learning in a Random World

authors:
Chong Zhang
Wenbo Wang
Xingye Qiao

title: On reject and refine options in multicategory classification

heading: 1 INTRODUCTION

text:
Imagine you are a doctor making a high-stakes medical decision based on diagnostic information from a computer vision classifier. What would you want the classifier to output in order to make the best decision? This is not a casual hypothetical; such classifiers are already used in medical settings (e.g., Razzak et al., 2018; Lundervold & Lundervold, 2019; Li et al., 2014). A maximumlikelihood diagnosis with an accompanying probability may not be the most essential piece of information. To ensure the health of the patient, you must also rule in or rule out harmful diagnoses. In other words, even if the most likely diagnosis is a stomach ache, it is equally or more important to rule out stomach cancer. Therefore, you would want the classifier to give you—in addition to an estimate of the most likely outcome—actionable uncertainty quantification, such as a set of predictions that provably covers the true diagnosis with a high probability (e.g., 90%). This is called a prediction set (see Figure 1). Our paper describes a method for constructing prediction sets from any pre-trained image classifier that are formally guaranteed to contain the true class with the desired probability, relatively small, and practical to implement. Our method modifies a conformal predictor (Vovk et al., 2005) given in Romano et al. (2020) for the purpose of modern image classification in order to make it more stable in the presence of noisy small probability estimates. Just as importantly, we provide extensive evaluations and code for conformal prediction in computer vision.
Formally, for a discrete response Y ∈ Y = {1, . . . ,K} and a feature vector X ∈ Rd, we desire an uncertainty set function, C(X), mapping a feature vector to a subset of {1, . . . ,K} such that
P (Y ∈ C(X)) ≥ 1− α, (1)
for a pre-specified confidence level α such as 10%. Conformal predictors like our method can modify any black-box classifier to output predictive sets that are rigorously guaranteed to satisfy the desired coverage property shown in Eq. (1). For evaluations, we focus on Imagenet classification
∗Equal contribution. Blog: https://people.eecs.berkeley.edu/˜angelopoulos/blog/ posts/conformal-classification
using convolutional neural networks (CNNs) as the base classifiers, since this is a particularly challenging testbed. In this setting, X would be the image and Y would be the class label. Note that the guarantee in Eq. (1) is marginal over X and Y—it holds on average, not for a particular image X .
A first approach toward this goal might be to assemble the set by including classes from highest to lowest probability (e.g., after Platt scaling and a softmax function; see Platt et al., 1999; Guo et al., 2017) until their sum just exceeds the threshold 1 − α. We call this strategy naive and formulate it precisely in Algorithm 1. There are two problems with naive: first, the probabilities output by CNNs are known to be incorrect (Nixon et al., 2019), so the sets from naive do not achieve coverage. Second, image classification models’ tail probabilities are often badly miscalibrated, leading to large sets that do not faithfully articulate the uncertainty of the model; see Section 2.3. Moreover, smaller sets that achieve the same coverage level can be generated with other methods.
The coverage problem can be solved by picking a new threshold using holdout samples. For example, with α =10%, if choosing sets that contain 93% estimated probability achieves 90% coverage on the holdout set, we use the 93% cutoff instead. We refer to this algorithm, introduced in Romano et al. (2020), as Adaptive Prediction Sets (APS). The APS procedure provides coverage but still produces large sets. To fix this, we introduce a regularization technique that tempers the influence of these noisy estimates, leading to smaller, more stable sets. We describe our proposed algorithm, Regularized Adaptive Prediction Sets (RAPS), in Algorithms 2 and 3 (with APS as a special case). As we will see in Section 2, both APS and RAPS are always guaranteed to satisfy Eq. (1)—regardless of model and dataset. Furthermore, we show that RAPS is guaranteed to have better performance than choosing a fixed-size set. Both methods impose negligible computational requirements in both training and evaluation, and output useful estimates of the model’s uncertainty on a new image given, say, 1000 held-out examples.
In Section 3 we conduct the most extensive evaluation of conformal prediction in deep learning to date on Imagenet and Imagenet-V2. We find that RAPS sets always have smaller average size than naive and APSsets. For example, using a ResNeXt-101, naive does not achieve coverage, while APS and RAPS achieve it almost exactly. However, APS sets have an average size of 19, while RAPS sets have an average size of 2 at α = 10% (Figure 2 and Table 1). We will provide an accompanying codebase that implements our method as a wrapper for any PyTorch classifier, along with code to exactly reproduce all of our experiments.

heading: 1.1 RELATED WORK

text:
Reliably estimating predictive uncertainty for neural networks is an unsolved problem. Historically, the standard approach has been to train a Bayesian neural network to learn a distribution over network weights (Quinonero-Candela et al., 2005; MacKay, 1992; Neal, 2012; Kuleshov et al., 2018; Gal, 2016). This approach requires computational and algorithmic modifications; other approaches avoid these via ensembles (Lakshminarayanan et al., 2017; Jiang et al., 2018) or approximations of Bayesian inference (Riquelme et al., 2018; Sensoy et al., 2018). These methods also have major practical limitations; for example, ensembling requires training many copies of a neural network adversarially. Therefore, the most widely used strategy is ad-hoc traditional calibration of the softmax scores with Platt scaling (Platt et al., 1999; Guo et al., 2017; Nixon et al., 2019).
This work develops a method for uncertainty quantification based on conformal prediction. Originating in the online learning literature, conformal prediction is an approach for generating predictive sets that satisfy the coverage property in Eq. (1) (Vovk et al., 1999; 2005). We use a convenient data-splitting version known as split conformal prediction that enables conformal prediction meth-
ods to be deployed for essentially any predictor (Papadopoulos et al., 2002; Lei et al., 2018). While mechanically very different from traditional calibration as discussed above, we will refer to our approach as conformal calibration to highlight that the two methodologies have overlapping but different goals.
Conformal prediction is a general framework, not a specific algorithm—important design decisions must be made to achieve the best performance for each context. To this end, Romano et al. (2020) and Cauchois et al. (2020) introduce techniques aimed at achieving coverage that is similar across regions of feature space, whereas Vovk et al. (2003); Hechtlinger et al. (2018) and Guan & Tibshirani (2019) introduce techniques aimed at achieving equal coverage for each class. While these methods have conceptual appeal, thus far there has been limited empirical evaluation of this general approach for state-of-the-art CNNs. Concretely, the only works that we are aware of that include some evaluation of conformal methods on ImageNet—the gold standard for benchmarking computer vision methods—are Hechtlinger et al. (2018), Park et al. (2019), Cauchois et al. (2020), and Messoudi et al. (2020), although in all four cases further experiments are needed to more fully evaluate their operating characteristics for practical deployment. At the heart of conformal prediction is the conformal score - a measure of similarity between labeled examples which is used to compare a new point to among those in a hold out set. Our theoretical contribution can be summarized as a modification of the conformal score from Romano et al. (2020) to have smaller, more stable sets. Lastly, there are alternative approaches to returning prediction sets not based on conformal prediction (Pearce et al., 2018; Zhang et al., 2018). These methods can be used as input to a conformal procedure to potentially improve performance, but they do not have finite-sample coverage guarantees when used alone.

heading: 2 METHODS

text:
In developing uncertainty set methods to improve upon naive, we are guided by three desiderata. First and most importantly, the coverage desideratum says the sets must provide 1−α coverage, as discussed above. Secondly, the size desideratum says we want sets of small size, since these convey more detailed information and may be more useful in practice. Lastly, the adaptiveness desideratum says we want the sets to communicate instance-wise uncertainty: they should be smaller for easy test-time examples than for hard ones; see Figure 1 for an illustration. Coverage and size are obviously competing objectives, but size and adaptiveness are also often in tension. The size desideratum seeks small sets, while the adaptiveness desideratum seeks larger sets when the classi-
Algorithm 1 Naive Prediction Sets Input: α, sorted scores s, associated permutation of classes I , boolean rand
1: procedure NAIVE(α, s, I, rand) 2: L← 1 3: while ∑L i=1 si < 1− α do . Stop if 1− α probability exceeded 4: L← L+ 1 5: if rand then . Break ties randomly (explained in Appendix B) 6: U ← Unif(0, 1) 7: V ← ( ∑L i=1 si − (1− α))/sL 8: if U ≤ V then 9: L← L− 1
10: return { I1, ..., IL } Output: The 1− α prediction set, { I1, ..., IL }
fier is uncertain. For example, always predicting a set of size five could achieve coverage, but it is not adaptive. As noted above, both APSand RAPS achieve correct coverage, and we will show that RAPS improves upon APS according to the other two desiderata.
We now turn to the specifics of our proposed method. We begin in Subsection 2.1 by describing an abstract data-splitting procedure called conformal calibration that enables the near-automatic construction of valid predictive sets (that is, sets satisfying Eq. (1)). Subsequently, in Subsection 2.2, we provide a detailed presentation of our procedure, with commentary in Section 2.3. In Subsection 2.4 we discuss the optimality of our procedure, proving that it is at least as good as the procedure that returns sets of a fixed size, unlike alternative approaches.

heading: 2.1 CONFORMAL CALIBRATION

text:
We first review a general technique for producing valid prediction sets, following the articulation in Gupta et al. (2019). Consider a procedure that outputs a predictive set for each observation, and further suppose that this procedure has a tuning parameter τ that controls the size of the sets. (In RAPS, τ is the cumulative sum of the sorted, penalized classifier scores.) We take a small independent conformal calibration set of data, and then choose the tuning parameter τ such that the predictive sets are large enough to achieve 1 − α coverage on this set. See Figure 3 for an illustration. This calibration step yields a choice of τ , and the resulting set is formally guaranteed to have coverage 1− α on a future test point from the same distribution; see Theorem 1 below. Formally, let (Xi, Yi)i=1,...,n be an independent and identically distributed (i.i.d.) set of variables that was not used for model training. Further, let C(x, u, τ) : Rd × [0, 1] × R → 2Y be a setvalued function that takes a feature vector x to a subset of the possible labels. The second argument u is included to allow for randomized procedures; let U1, . . . , Un be i.i.d. uniform [0, 1] random variables that will serve as the second argument for each data point. Suppose that the sets are indexed by τ such that they are nested, meaning larger values of τ lead to larger sets:
C(x, u, τ1) ⊆ C(x, u, τ2) if τ1 ≤ τ2. (2) To find a function that will achieve 1 − α coverage on test data, we select the smallest τ that gives at least 1 − α coverage on the conformal calibration set, with a slight correction to account for the finite sample size:
τ̂ccal = inf
{ τ : |{i : Yi ∈ C(Xi, Ui, τ)}|
n ≥ d(n+ 1)(1− α)e n
} . (3)
The set function C(x, u, τ) with this data-driven choice of τ is guaranteed to have correct finitesample coverage on a fresh test observation, as stated formally next. Theorem 1 (Conformal calibration coverage guarantee). Suppose (Xi, Yi, Ui)i=1,...,n and (Xn+1, Yn+1, Un+1) are i.i.d. and let C(x, u, τ) be a set-valued function satisfying the nesting property in Eq. (2). Suppose further that the sets C(x, u, τ) grow to include all labels for large enough τ : for all x ∈ Rd, C(x, u, τ) = Y for some τ . Then for τ̂ccal defined as in Eq. (3), we have the following coverage guarantee:
P ( Yn+1 ∈ C(Xn+1, Un+1, τ̂ccal) ) ≥ 1− α.
This is the same coverage property as Eq. (1) in the introduction, written in a more explicit manner. The result is not new—a special case of this result leveraging sample-splitting first appears in the regression setting in Papadopoulos et al. (2002), and the core idea of conformal prediction was introduced even earlier; see (Vovk et al., 2005).
As a technical remark, the theorem also holds if the observations to satisfy the weaker condition of exchangeability; see Vovk et al. (2005). In addition, for most families of set-valued functions C(x, u, τ) there is a matching upper bound:
P ( Yn+1 ∈ C(Xn+1, Un+1, τ̂ccal) ) ≤ 1− α+ 1
n+ 1 .
Roughly speaking, this will hold whenever the sets grow smoothly in τ . See Lei et al. (2018) for a formal statement of the required conditions.

heading: 2.2 OUR METHOD

text:
Conformal calibration is a powerful general idea, allowing one to achieve the coverage desideratum for any choice of sets C(x, u, τ). Nonetheless, this is not yet a full solution, since the quality of the resulting prediction sets can vary dramatically depending on the design of C(x, u, τ). In particular, we recall the size and adaptiveness desiderata from Section 1—we want our uncertainty sets to be as small as possible while faithfully articulating the instance-wise uncertainty of each test point. In this section, we explicitly give our algorithm, which can be viewed as a special case of conformal calibration with the uncertainty sets C designed to extract information from CNNs. Our algorithm has three main ingredients. First, for a feature vector x, the base model computes class probabilities π̂x ∈ Rk, and we order the classes from most probable to least probable. Then, we add a regularization term to promote small predictive sets. Finally, we conformally calibrate the penalized prediction sets to guarantee coverage on future test points.
Formally, let ρx(y) = ∑K y′=1 π̂x(y
′)I{π̂x(y′)>π̂x(y)} be the total probability mass of the set of labels that are more likely than y. These are all the labels that will be included before y is included. In addition, let ox(y) = |{y′ ∈ Y : π̂x(y′) ≥ π̂x(y)}| be the ranking of y among the label based on the probabilities π̂. For example, if y is the third most likely label, then ox(y) = 3.1 We take
C∗(x, u, τ) := { y : ρx(y) + π̂x(y) · u+ λ · (ox(y)− kreg)+︸ ︷︷ ︸
regularization
≤ τ } , (4)
where (z)+ denotes the positive part of z and λ, kreg ≥ 0 are regularization hyperparameters that are introduced to encourage small set sizes. See Figure 3 for a visualization of a RAPS predictive set and Appendix E for a discussion of how to select kreg and λ.
Since this is the heart of our proposal, we carefully parse each term. First, the ρx(y) term increases as y ranges from the most probable to least probable label, so our sets will prefer to include the y that are predicted to be the most probable. The second term, π̂x(y) · u, is a randomized term to handle the fact that the value will jump discretely with the inclusion of each new y. The randomization term can never impact more than one value of y: there is at most one value of y such that y ∈ C(x, 0, τ) but y /∈ C(x, 1, τ). These first two terms can be viewed as the CDF transform after arranging the classes from most likely to least likely, randomized in the usual way to result in a continuous uniform random variable (cf. Romano et al., 2020). We discuss randomization further in Appendix B.
Lastly, the regularization promotes small set sizes: for values of y that occur farther down the ordered list of classes, the term λ · (ox(y)− kreg)+ makes that value of y require a higher value of τ before it is included in the predictive set. For example, if kreg = 5, then the sixth most likely value of y has an extra penalty of size λ, so it will never be included until τ exceeds ρx(y) + π̂x(y) · u + λ, whereas it enters when τ exceeds ρx(y) + π̂x(y) · u in the nonregularized version. Our method has the following coverage property: Proposition 1 (RAPS coverage guarantee). Suppose (Xi, Yi, Ui)i=1,...,n and (Xn+1, Yn+1, Un+1) are i.i.d. and let C∗(x, u, τ) be defined as in Eq. (4). Suppose further that π̂x(y) > 0 for all x and y. Then for τ̂ccal defined as in Eq. (3), we have the following coverage guarantee:
1− α ≤ P ( Yn+1 ∈ C∗(Xn+1, Un+1, τ̂ccal) ) ≤ 1− α+ 1
n+ 1 .
1For ease of notation, we assume distinct probabilities. Else, label-ordering ties should be broken randomly.
Algorithm 2 RAPS Conformal Calibration
Input: α; s ∈ [0, 1]n×K , I ∈ {1, ...,K}n×K , and one-hot y ∈ {0, 1}K corresponding respectively to the sorted scores, the associated permutation of indexes, and labels for each of n examples in the calibration set; kreg; λ; boolean rand
1: procedure RAPSC(α,s,I ,y,λ) 2: for i ∈ {1, · · · , n} do 3: Li ← { j : Ii,j = yi } 4: Ei ← ΣLij=0si,j + λ(Li − kreg + 1)+ 5: if rand then 6: U ∼ Unif(0, 1) 7: Ei ← Ei − si,Li + U ∗ si,Li 8: τ̂ccal ← the d(1− α)(1 + n)e largest value in {Ei}ni=1 9: return τ̂ccal
Output: The generalized quantile, τ̂ccal . The value in Eq. (3)
Algorithm 3 RAPS Prediction Sets Input: α, sorted scores s and the associated permutation of classes I for a test-time example, τ̂ccal
from Algorithm 2, kreg , λ, boolean rand 1: procedure RAPS(α, s, I, τ̂ccal, kreg, λ, rand) 2: L← | j ∈ Y : Σji=0si + λ(L− kreg)+ ≤ τ̂ccal |+ 1 3: V ← (τ̂ccal − ΣL−1i=0 si − λ(L− kreg)+ + sL−1)/sL−1 4: if rand & V ≤ U ∼ Unif(0, 1) then 5: L← L− 1 6: return C = { I1, ...IL } . The L most likely classes
Output: The 1− α confidence set, C . The set in Eq. (4)
Note that the first inequality is a corollary of Theorem 1, and the second inequality is a special case of the remark in Section 2.1. The restriction that π̂x(y) > 0 is not necessary for the first inequality.

heading: 2.3 WHY REGULARIZE?

text: In our experiments, the sets from APS are larger than necessary, because APS is sensitive to the noisy probability estimates far down the list of classes. This noise leads to a permutation problem of unlikely classes, where ordering of the classes with small probability estimates is determined mostly by random chance. If 5% of the true classes from the calibration set are deep in the tail due to the permutation problem, APS will choose large 95% predictive sets; see Figure 2. The inclusion of the RAPS regularization causes the algorithm to avoid using the unreliable probabilities in the tail; see Figure 4. We discuss how RAPS improves the adaptiveness of APS in Section 4 and Appendix E.

heading: 2.4 OPTIMALITY CONSIDERATIONS

text:
To complement these experimental results, we now formally prove that RAPS with the correct regularization parameters will always dominate the simple procedure that returns a fixed set size. (Section 3.5 shows the parameters are easy to select and RAPS is not sensitive to their values). For a feature vector x, let ŷ(j)(x) be the label with the jth highest predicted probability. We define the top-k predictive sets to be {ŷ(1)(x), . . . , ŷ(k)(x)}. Proposition 2 (RAPS dominates top-k sets). Suppose (Xi, Yi, Ui)i=1,...,n and (Xn+1, Yn+1, Un+1) are i.i.d. draws. Let k∗ be the smallest k such that the top-k predictive sets have coverage at least d(n+ 1)(1−α)e/n on the conformal calibration points (Xi, Yi)i=1,...,n. Take C∗(x, u, τ) as in Eq. (4) with any kreg ≤ k∗ and λ = 1. Then with τ̂ccal chosen as in Eq. (3), we have
C∗(Xn+1, Un+1, τ̂ccal) ⊆ {ŷ(1)(x), . . . , ŷ(k∗)(x)}.
In words, the RAPS procedure with heavy regularization will be at least as good as the top-k procedure in the sense that it has smaller or same average set size while maintaining the desired coverage level. This is not true of either the naive baseline or the APS procedure; Table 2 shows that these two procedures usually return predictive sets with size much larger than k∗.
3 EXPERIMENTS
In this section we report on experiments that study the performance of the predictive sets from naive, APS, and RAPS, evaluating each based on the three desiderata above. We begin with a brief preview of the experiments. In Experiment 1, we evaluate naive, APS, and RAPS on ImagenetVal. Both APS and RAPS provided almost exact coverage, while naive sets had coverage slightly below the specified level. APS has larger sets on average than naive and RAPS. RAPS has a much smaller average set size than APS and naive. In Experiment 2, we repeat Experiment 1 on Imagenet-V2, and the conclusions still hold. In Experiment 3, we produce histograms of set sizes for naive, APS, and RAPS for several different values of λ, illustrating a simple tradeoff between set size and adaptiveness. In Experiment 4, we compute histograms of RAPS sets stratified by image difficulty, showing that RAPS sets are smaller for easier images than for difficult ones. In Experiment 5, we report the performance of RAPS with many values of the tuning parameters.
In our experiments, we use nine standard, pretrained Imagenet classifiers from the torchvision repository (Paszke et al., 2019) with standard normalization, resize, and crop parameters. Before applying naive, APS, or RAPS, we calibrated the classifiers using the standard temperature scaling/Platt scaling procedure as in Guo et al. (2017) on the calibration set. Thereafter, naive, APS, and RAPS were applied, with RAPS using a data-driven choice of parameters described in Appendix E. We use the randomized versions of these algorithms—see Appendix B for a discussion.

heading: 3.1 EXPERIMENT 1: COVERAGE VS SET SIZE ON IMAGENET

text: In this experiment, we calculated the coverage and mean set size of each procedure for two different choices of α. Over 100 trials, we randomly sampled two subsets of Imagenet-Val: one conformal calibration subset of size 20K and one evaluation subset of size 20K. The median-of-means over trials for both coverage and set size are reported in Table 1. Figure 2 illustrates the performances of naive, APS, and RAPS; RAPS has much smaller sets than both naive and APS, while achieving coverage. We also report results from a conformalized fixed-k procedure, which finds the smallest fixed set size achieving coverage on the holdout set, k∗, then predicts sets of size k∗ − 1 or k∗ on new examples in order to achieve exact coverage; see Algorithm 4 in Appendix E.

heading: 3.2 EXPERIMENT 2: COVERAGE VS SET SIZE ON IMAGENET-V2

text:
The same procedure as Experiment 1 was repeated on Imagenet-V2, with exactly the same normalization, resize, and crop parameters. The size of the calibration and evaluation sets was 5K, since Imagenet-V2 is a smaller dataset. The result shows that our method can still provide coverage even for models trained on different distributions, as long as the conformal calibration set comes from the new distribution. The variance of the coverage is higher due to having less data.
3.3 EXPERIMENT 3: SET SIZES OF NAIVE , APS, AND RAPS ON IMAGENET
We investigate the effect of regularization in more detail. For three values of λ, we collected the set sizes produced by each of naive, APS, and RAPS and report their histograms in Figure 4.
3.4 EXPERIMENT 4: ADAPTIVENESS OF RAPS ON IMAGENET
We now show that RAPS sets are smaller for easy images than hard ones, addressing the adaptiveness desideratum. Table 4 reports the size-stratified coverages of RAPS at the 90% level with kreg = 5 and different choices of λ. When λ is small, RAPS allows sets to be large. But when λ = 1, RAPS
clips sets to be a maximum of size 5. Table 7 (in the Appendix) stratifies by image difficulty, showing that RAPS sets are small for easy examples and large for hard ones. Experiments 3 and 4 together illustrate the tradeoff between adaptiveness and size: as the average set size decreases, the RAPS procedure truncates sets larger than the smallest fixed set that provides coverage, taming the heavy tail of the APS procedure. Since RAPS with large λ undercovers hard examples, it must compensate by taking larger sets for easy examples to ensure the 1− α marginal coverage guarantee. However, the size only increases slightly since easy images are more common than hard ones, and the total probability mass can often exceed τ̂ccal by including only one more class. If this behavior is not desired, we can instead automatically pick λ to optimize the adaptiveness of RAPS; see Section 4.

heading: 3.5 EXPERIMENT 5: CHOICE OF TUNING PARAMETERS

text: While any value of the tuning parameters λ and kreg lead to coverage (Proposition 1), some values will lead to smaller sets. In Experiments 1 and 2, we chose kreg and λ adaptively from data (see Appendix E), achieving strong results for all models and choices of the coverage level. Table 3 gives the performance of RAPS with many choices of kreg and λ for ResNet-152.

heading: 4 ADAPTIVENESS AND CONDITIONAL COVERAGE

text:
In this section, we point to a definition of adaptiveness that is more natural for the image classification setting than the existing notion of conditional coverage. We show that APS does not satisfy conditional coverage, and that RAPS with small λ outperforms it in terms of adaptiveness.
We say that a set-valued predictor C : Rd → 2Y satisfies exact conditional coverage if P (Y ∈ C(X) | X = x) = 1 − α for each x. Distribution-free guarantees on conditional coverage are impossible (Vovk, 2012; Lei & Wasserman, 2014), but many algorithms try to satisfy it approximately (Romano et al., 2019; 2020; Cauchois et al., 2020). In a similar spirit, Tibshirani et al. (2019) suggest a notion of local conditional coverage, where one asks for coverage in a neighborhood of each point, weighted according to a chosen kernel. Cauchois et al. (2020) introduce the worst-case slab metric for measuring violations of the conditional coverage property. We present a different way of measuring violations of conditional coverage.
Proposition 3. Suppose P (Y ∈ C(X) | X = x) = 1− α for each x ∈ Rd. Then, P (Y ∈ C(X) | {|C(X)| ∈ A}) = 1− α for any A ⊂ {0, 1, 2, . . . }.
In words, if conditional coverage holds, then coverage holds after stratifying by set size. Based on this result, In Appendix E, we introduce the size-stratified coverage violation criterion, a simple and pragmatic way of quantifying adaptiveness. Then, we automatically tune λ on this metric so RAPS markedly outperforms the adaptiveness of APS (see Table 8).

heading: 1 11.2 10.2 7.0 3.6 2.9 2.3 2.1 2.3 2.2 2.2

heading: 2 11.2 10.2 7.1 3.7 3.0 2.4 2.1 2.3 2.2 2.2

heading: 5 11.2 10.2 7.2 3.9 3.4 2.9 2.6 2.5 2.5 2.5

heading: 10 11.2 10.2 7.4 4.5 4.0 3.6 3.4 3.4 3.4 3.4

heading: 50 11.2 10.6 8.7 7.2 7.0 6.9 6.9 6.9 6.9 6.9

text:
In Table 4, we report on the coverage of APS and RAPS, stratified by the size of the prediction set. Turning our attention to the λ = 0 column, we see that when APS outputs a set of size 101− 1000, APS has coverage 97%, substantially higher than 90% nominal rate. By Proposition 3, we conclude that APS is not achieving exact conditional coverage, because the scores are far from the oracle probabilities. The APS procedure still achieves marginal coverage by overcovering hard examples and undercovering easy ones, an undesirable behavior. Alternatively, RAPS can be used to regularize the set sizes—for λ = .001 to λ = .01 the coverage stratified by set size is more balanced. In summary, even purely based on the adaptiveness desideratum, RAPS with light regularization is preferable to APS. Note that as the size of the training data increases, as long as π̂ is consistent, naive and APS will become more stable, and so we expect less regularization will be needed.
Lastly, we argue that conditional coverage is a poor notion of adaptiveness when the best possible model (i.e., one fit on infinite data) has high accuracy. Given such a model, the oracle procedure from Romano et al. (2020) would return the correct label with probability 1 − α and the empty set with probability α. That is, having correct conditional coverage for high-signal problems where Y is perfectly determined by X requires a perfect classifier. In our experiments on ImageNet, APS does not approximate this behavior. Therefore, conditional coverage isn’t the right goal for prediction sets with realistic sample sizes. Proposition 3 suggests a relaxation. We could require that we have the right coverage, no matter the size of the prediction set: P (Y ∈ C(X) | {|C(x)| ∈ A}) ≥ 1− α for any A ⊂ {0, 1, 2, . . . }; Appendix E.2 develops this idea. We view this as a promising way to reason about adaptiveness in high-signal problems such as image classification.

heading: 5 DISCUSSION

text: For classification tasks with many possible labels, our method enables a researcher to take any base classifier and return predictive sets guaranteed to achieve a pre-specified error level, such as 90%, while retaining small average size. It is simple to deploy, so it is an attractive, automatic way to quantify the uncertainty of image classifiers—an essential task in such settings as medical diagnostics, self-driving vehicles, and flagging dangerous internet content. Predictive sets in computer vision (from RAPS and other conformal methods) have many further uses, since they systematically identify hard test-time examples. Finding such examples is useful in active learning where one only has resources to label a small number of points. In a different direction, one can improve efficiency of a classifier by using a cheap classifier outputting a prediction set first, and an expensive one only when the cheap classifier outputs a large set (a cascade; see, e.g., Li et al. 2015), and Fisch et al. (2021) for an implementation of conformal prediction in this setting. One can also use predictive sets during model development to identify failure cases and outliers and suggest strategies for improving its performance. Prediction sets are most useful for problems with many classes; returning to our initial medical motivation, we envision RAPS could be used by a doctor to automatically screen for a large number of diseases (e.g. via a blood sample) and refer the patient to relevant specialists.

heading: A PROOFS

text:
Theorem 1. Let s(x, u, y) = infτ{y ∈ C(x, u, τ)}, and let si = s(Xi, Ui, Yi) for i = 1, . . . , n. Then
{y : s(x, u, y) ≤ τ} = {y : y ∈ C(x, u, τ)}
because C(x, u, τ) is a finite set growing in τ by the assumption in Eq. (2). Thus,
{τ : |{i : si ≤ τ} | ≥ d(1−α)(n+ 1)e} = { τ : |{i : Yi ∈ C(Xi, Ui, τ)}|
n ≥ d(n+ 1)(1− α)e n
} .
Considering the left expression, the infimum over τ of the set on the left hand side is the d(1 − α)(n + 1)e smallest value of the si, so this is the value of τ̂ccal. Since s1, . . . , sn, s(Xn+1, Un+1, Yn+1) are exchangeable random variables, |{i : s(Xn+1, Un+1, Yn+1) > si}| is stochastically dominated by the discrete uniform distribution on {0, 1, . . . , n}. We thus have that
P (Yn+1 /∈ C(Xn+1, Un+1, τ̂ccal)) = P (s(Xn+1, Un+1, Yn+1) > τ̂ccal) = P (|{i : s(Xn+1, Un+1, Yn+1) > si}| ≥ d(n+ 1)(1− α)e)
= P
( |{i : s(Xn+1, Un+1, Yn+1) > si}|
n+ 1 ≥ d(n+ 1)(1− α)e n+ 1 ) ≤ α.
Proposition 1. The lower bound follows from Theorem 1. To prove the upper bound, using the result from Theorem 2.2 of Lei et al. (2018) it suffices to show that the variables s(Xi, Ui, Yi) = inf{τ : Yi ∈ C(Xi, Ui, τ)} are almost surely distinct. To this end, note that that
s(Xi, Ui, Yi) = ρXi(Yi) + π̂Xi(Yi) · Ui + λ(oXi(Yi)− kreg)+,
and due to the middle term of the sum, these values are distinct almost surely provided π̂Xi(Yi) > 0.
Proposition 2. We first show that τ̂ccal ≤ 1 + k∗ − kreg . Note that since at least d(1 − α)(n + 1)e of the conformal calibration points are covered by a set of size k∗, at least d(1 − α)(n + 1)e of the Ei in Algorithm 2 are less than or equal to 1 + k∗ − kreg. Thus, by the definition of τ̂ccal, we have that it is less than or equal to 1 + k∗ − kreg. Then, note that by the definition of C∗ in Eq. (4), we have that
|C∗(Xn+1, Un+1, τ̂ccal)| ≤ k∗.
as long as τ̂ccal ≤ 1+k∗−kreg , since for the k∗+1 most likely class, the sum in Eq. (4) will exceed λ · (1 + k∗ − kreg) = (1 + k∗ − kreg) ≥ τ̂ccal, and so the k∗ + 1 class will not be in the set.
Proposition 3. Suppose P (Y ∈ C(X) | X = x) = 1− α for each x ∈ Rd. Then, P (Y ∈ C(X) | |C(X)| ∈ A) = ∫ x P (Y ∈ C(x) | X = x})I{|C(x)|∈A}dP (x)
P (|C(X)| ∈ A)
=
∫ x (1− α)I{|C(x)|∈A}dP (x)
P (|C(X)| ∈ A) = 1− α.

heading: B RANDOMIZED PREDICTORS

text:
The reader may wonder why we choose to use a randomized procedure. The randomization is needed to achieve 1 − α coverage exactly, which we will explain via an example. Note that the randomization is of little practical importance, since the predictive set output by the randomized procedure will differ from the that of the non-randomized procedure by at most one element.
Turning to an example, assume for a particular input image we expect a set of size k to have 91% coverage, and a set of size k − 1 to have 89% coverage. In order to achieve our desired coverage of 90%, we randomly choose size k or k − 1 with equal probability. In general, the probabilities will not be equal, but rather chosen so the weighted average of the two coverages is exactly 90%. If a user of our method desires deterministic sets, it is easy to turn off this randomization with a single flag, resulting in slightly conservative sets.
C IMAGENET AND IMAGENETV2 RESULTS FOR α = 5%
We repeated Experiments 1 and 2 with α = 5%. See the results in Tables 5 and 6.

heading: D COVERAGE AND SIZE CONDITIONAL ON IMAGE DIFFICULTY

text:
In order to probe the adaptiveness properties of APS and RAPS we stratified coverage and size by image difficulty (the position of the true label in the list of most likely to least likely classes, based on the classifier predictions) in Table 7. With increasing λ, coverage decreases for more difficult images and increases for easier ones. In the most difficult regime, even though APS can output large sets, those sets still rarely contain the true class. This suggests regularization is a sensible way to stabilize the sets. As a final word on Table 7, notice that as λ increases, coverage improves for the more common medium-difficulty examples, although not for very rare and difficult ones.
E CHOOSING kreg AND λ TO OPTIMIZE SET SIZE AND ADAPTIVENESS
This section describes two procedures for picking kreg and λ that optimize for set size or adaptiveness, outperforming APS in both cases.
E.1 OPTIMIZING SET SIZE WITH RAPS
Algorithm 4 Adaptive Fixed-K
Input: α; I ∈ {1, ...,K}n×K , and one-hot y ∈ {0, 1}K corresponding respectively to the classes from highest to lowest estimated probability mass, and labels for each of n examples in the dataset
1: procedure GET-KSTAR(α,I ,y) 2: for i ∈ {1, · · · , n} do 3: Li ← { j : Ii,j = yi } 4: k̂∗ ← the d(1− α)(1 + n)e largest value in {Li}ni=1 5: return k̂∗
Output: The estimate of the smallest fixed size set that achieves coverage, k̂∗
To produce Tables 1, 5, 2, and 6, we chose kreg and λ adaptively. This required an extra data splitting step, where a small amount of tuning data { xi, yi }m i=1
were used to estimate k∗, and then kreg is set to k∗. Takingm ≈ 1000 was sufficient, since the algorithm is fairly insensitive to kreg (see Table 3). Then, k̂∗ was calculated with Algorithm 4. We produced the Imagenet V2 tables with m = 1000 and the Imagenet tables with m = 10000.
After choosing k̂∗, we chose λ to have small set size. We used the same tuning data to pick k̂∗ and λ for simplicity (this does not invalidate our coverage guarantee since conformal calibration still uses fresh data). A coarse grid search on λ sufficed, since small parameter variations have little impact on RAPS. For example, we chose the λ ∈ {0.001, 0.01, 0.1, 0.2, 0.5} that achieved the smallest size on the m holdout samples in order to produce Tables 1, 5, 2, and 6. We include a subroutine that automatically chooses k̂∗ and λ to optimize size in our GitHub codebase.
E.2 OPTIMIZING ADAPTIVENESS WITH RAPS
In this appendix, we show empirically that RAPS with an automatically chosen set of kreg and λ improves the adaptiveness of APS. Recall our discussion in Section 4 and Proposition 3, wherein we propose size-stratified coverage as a useful definition of adaptiveness in image classification. After picking kreg as in Appendix E, we can choose λ using the same tuning data to optimize this notion of adaptiveness.
We now describe a particular manifestation of our adaptiveness criterion that we will use to optimize λ. Consider disjoint set-size strata {Si}i=si=1, where ⋃j=s j=1 Si = {1, . . . , |Y|}. Then define the indexes
of examples stratified by the prediction set size of each example from algorithm C as Jj = { i :
|C(Xi, Yi, Ui)| ∈ Sj }
. Then we can define the size-stratified coverage violation of an algorithm C on strata {S}i=si=1 as
SSCV ( C, {S}j=sj=1 ) = sup
j ∣∣∣∣∣ |{i : Yi ∈ C(Xi,Yi,Ui), i ∈ Jj}||Jj| − (1− α) ∣∣∣∣∣. (5)
In words, Eq. (5) is the worst-case deviation of C from exact coverage when it outputs sets of a certain size. Computing the size-stratified coverage violation thus only requires post-stratifying the results of C on a set of labeled examples. If conditional coverage held, the worst stratum coverage violation would be 0 by Proposition 3.
To maximize adaptiveness, we’d like to choose λ to minimize the size-stratified coverage violation of RAPS. Write Cλ to mean the RAPS procedure for a fixed choice of kreg and λ. Then we would like to pick
λ = arg min λ′
SSCV(Cλ′ , {S}j=sj=1). (6)
In our experiments, we choose a relatively coarse partitioning of the possible set sizes: 0-1, 2-3, 4- 10, 11-100, and 101-1000. Then, we chose the λ ∈ {0.00001, 0.0001, 0.0008, 0.001, 0.0015, 0.002} which minimized the size-stratified coverage violation on the tuning set. The results in Table 8 show RAPS always outperforms the adaptiveness of APS on the test set, even with this coarse, automated choice of parameters. The table reports the median size-stratified coverage violation over 10 independent trials of APS and RAPS with automated parameter tuning.

heading: F COMPARISON WITH LEAST AMBIGUOUS SET-VALUED CLASSIFIERS

text:
In this section, we compare RAPS to the Least Ambiguous Set-valued Classifier (LAC) method introduced in Sadinle et al. (2019), an alternative conformal procedure that is designed to have small sets. The LAC method provable gives the smallest possible average set size in the case where the input probabilities are correct, with the idea that these sets should be small even when the estimated probabilities are only approximately correct. In the notation of this paper, the LAC method considers nested sets of the following form:
CLAC(x, τ) := {y : π̂x(y) ≥ 1− τ}, which can be calibrated using as before in using τ̂ccal from Eq. (3).
We first compare naive, APS, RAPS, and LAC in terms of power and coverage in Table 9. In this experiment, we tuned RAPS to have small set size as described in Appendix E.1. We see that LAC also achieves correct coverage, as expected since it is a conformal method and satisfies the guarantee from Theorem 1. We further see that it has systematically smaller sets that RAPS, although the difference is slight compared to the gap between APS and RAPS or APS and LAC.
We next compare RAPS to LAC in terms of adaptiveness, tuning RAPS as in Section E.2. First, in Table 10, we report on the coverage of LAC for images of different difficulties, and see that LAC has dramatically worse coverage for hard images than for easy ones. Comparing this to RAPS in Table 7, we see that RAPS also has worse coverage for more difficult images, although the gap is much smaller for RAPS. Next, in Table 11, we report on the SSCV metric for of adaptiveness (and conditional coverage) for APS, RAPS, and LAC. We find that APS and RAPS have much better adaptiveness than LAC, with RAPS being the overall winner. The results of all of these comparisons are expected: LAC is not targetting adpativeness and instead trying to achieve the smallest possible set size. It succeeds at its goal, sacrificing adaptiveness to do so.

heading: 1 15668 1.00 1.5


summaries: This work introduces a new Wasserstein-2 barycenter computation method. The authors first derive the dual formulation of the Wasserstein-2 barycenter problem, and then parametrize the convex potentials by ICNNs. The congruent and conjugacy conditions are enforced by regularization terms, respectively. They then show that the algorithm can find a good barycenter if the objective function is properly minimized.

summaries: The paper derives the barycenter mapping problem as an optimization over *congruent* convex functions---each convex potential corresponding to a component distribution.  Congruency is a property on the set of optimal potential functions that ties them together.  However, this optimization is quite challenging and so the paper derives an principled objective function that includes two regularization terms.  The first regularization term encourages congruency of the set of convex functions and can be seen as a variational bound on an ideal congruency regularization.  The second regularization term encourages the pairs of convex functions to be conjugate.  The paper proves that the optimal solution of this objective is the true potentials and thus no bias is introduced.  The proposed approach is demonstrated on the tasks of generative modeling (2-256 dimensions), posterior inference, and color pallete barycenters (3D)

summaries: The paper considers the Wasserstein Barycenter problems in the continuous setting. In particular, the authors propose an algorithm to compute the Wasserstein-2 barycenter when only samples from the marginals are accessible. Some theoretical analysis of this method is presented. Several numerical examples are carried out to compare this method with two other recently proposed methods.

abstractText: Wasserstein barycenters provide a geometric notion of the weighted average of probability measures based on optimal transport. In this paper, we present a scalable algorithm to compute Wasserstein-2 barycenters given sample access to the input measures, which are not restricted to being discrete. While past approaches rely on entropic or quadratic regularization, we employ input convex neural networks and cycle-consistency regularization to avoid introducing bias. As a result, our approach does not resort to minimax optimization. We provide theoretical analysis on error bounds as well as empirical evidence of the effectiveness of the proposed approach in low-dimensional qualitative scenarios and high-dimensional quantitative experiments.

authors:
Martial Agueh
Guillaume Carlier

title: Barycenters in the wasserstein space

venue: SIAM Journal on Mathematical Analysis,

authors:
Pedro C Álvarez-Esteban
E Del Barrio
JA Cuesta-Albertos

title: Matrán. A fixed-point approach to barycenters in wasserstein space

venue: Journal of Mathematical Analysis and Applications,

authors:
Brandon Amos
Lei Xu
J Zico Kolter

title: Input convex neural networks

venue: In Proceedings of the 34th International Conference on Machine Learning-Volume

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

venue: Communications on pure and applied mathematics,

authors:
Bob Carpenter
Andrew Gelman
Matthew D Hoffman
Daniel Lee
Ben Goodrich
Michael Betancourt
Marcus Brubaker
Jiqiang Guo
Peter Li
Allen Riddell

title: Stan: A probabilistic programming language

venue: Journal of statistical software,

authors:
Sinho Chewi
Tyler Maunu
Philippe Rigollet
Austin J Stromme

title: Gradient descent algorithms for bures-wasserstein barycenters

venue: arXiv preprint arXiv:2001.01700,

authors:
Samuel Cohen
Michael Arbel
Marc Peter Deisenroth

title: Estimating barycenters of measures in high dimensions

venue: arXiv preprint arXiv:2007.07105,

title: Sinkhorn distances: Lightspeed computation of optimal transport

venue: In Advances in neural information processing systems,

authors:
Marco Cuturi
Arnaud Doucet

title: Fast computation of wasserstein barycenters

authors:
DC Dowson
BV Landau

title: The fréchet distance between multivariate normal distributions

venue: Journal of multivariate analysis,

authors:
Pavel Dvurechenskii
Darina Dvinskikh
Alexander Gasnikov
Cesar Uribe
Angelia Nedich

title: Decentralize and randomize: Faster algorithm for wasserstein barycenters

venue: In Advances in Neural Information Processing Systems,

authors:
Jiaojiao Fan
Amirhossein Taghvaei
Yongxin Chen

title: Scalable computations of wasserstein barycenter via input convex neural networks

venue: arXiv preprint arXiv:2007.04462,

title: On conjugate convex functions

venue: Canadian Journal of Mathematics,

authors:
Aude Genevay
Marco Cuturi
Gabriel Peyré
Francis Bach

title: Stochastic optimization for largescale optimal transport

authors:
Sham Kakade
Shai Shalev-Shwartz
Ambuj Tewari

title: On the duality of strong convexity and strong smoothness: Learning applications and matrix regularization. Unpublished Manuscript, http://ttic

title: On the translocation of masses

venue: Management Science,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

authors:
Alexander Korotin
Vladimir V’yugin
Evgeny Burnaev

title: Integral mixability: a tool for efficient online aggregation of functional and probabilistic forecasts

venue: arXiv preprint arXiv:1912.07048,

authors:
Lingxiao Li
Aude Genevay
Mikhail Yurochkin
Justin Solomon

title: Continuous regularized wasserstein barycenters

venue: arXiv preprint arXiv:2008.12534,

authors:
Ashok Vardhan Makkuva
Amirhossein Taghvaei
Sewoong Oh
Jason D Lee

title: Optimal transport mapping via input convex neural networks

title: Existence and uniqueness of monotone measure-preserving maps

venue: Duke Mathematical Journal,

title: Wasserstein style transfer

venue: arXiv preprint arXiv:1905.12828,

title: Fast exact multiplication by the hessian

venue: Neural computation,

authors:
Gabriel Peyré
Marco Cuturi

title: Computational optimal transport

venue: Foundations and Trends® in Machine Learning,

authors:
Julien Rabin
Gabriel Peyré
Julie Delon
Marc Bernot

title: Wasserstein barycenter and its application to texture mixing

venue: In International Conference on Scale Space and Variational Methods in Computer Vision,

authors:
Julien Rabin
Sira Ferradans
Nicolas Papadakis

title: Adaptive color transfer with relaxed optimal transport

venue: In 2014 IEEE International Conference on Image Processing (ICIP),

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

title: Large-scale optimal transport and mapping estimation

venue: arXiv preprint arXiv:1711.02283,

authors:
Justin Solomon
Fernando De Goes
Gabriel Peyré
Marco Cuturi
Adrian Butscher
Andy Nguyen
Tao Du
Leonidas Guibas

title: Convolutional wasserstein distances: Efficient optimal transportation on geometric domains

venue: ACM Transactions on Graphics (TOG),

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

venue: The Journal of Machine Learning Research,

authors:
Matthew Staib
Sebastian Claici
Justin M Solomon
Stefanie Jegelka

title: Parallel streaming Wasserstein barycenters

title: Topics in optimal transportation

venue: Number 58. American Mathematical Soc.,

authors:
Jonathan Weed
Francis Bach

title: Sharp asymptotic and finite-sample rates of convergence of empirical measures in wasserstein

venue: distance. Bernoulli,

authors:
Yujia Xie
Minshuo Chen
Haoming Jiang
Tuo Zhao
Hongyuan Zha

title: On scalable and efficient computation of large scale optimal transport

venue: Proceedings of Machine Learning Research,

heading: 1 INTRODUCTION

text:
Wasserstein barycenters have become popular due to their ability to represent the average of probability measures in a geometrically meaningful way. Techniques for computing Wasserstein barycenters have been successfully applied to many computational problems. In image processing, Wasserstein barycenters are used for color and style transfer (Rabin et al., 2014; Mroueh, 2019), and texture synthesis (Rabin et al., 2011). In geometry processing, shape interpolation can be done by computing barycenters (Solomon et al., 2015). In online machine learning, barycenters are used for aggregating probabilistic predictions of experts (Korotin et al., 2019b). Within the context of Bayesian inference, the barycenter of subset posteriors converges to the full data posterior, thus enabling efficient computational methods based on finding the barycenters (Srivastava et al., 2015; 2018).
Fast and accurate barycenter algorithms exist for discrete distributions (see Peyré et al. (2019) for a survey), while for continuous distributions the situation is more difficult and remains unexplored until recently (Li et al., 2020; Fan et al., 2020; Cohen et al., 2020). The discrete methods scale poorly with the number of support points of the barycenter and thus cannot approximate continuous barycenters well, especially in high dimensions.
In this paper, we present a method to compute Wasserstein-2 barycenters of continuous distributions based on a novel regularized dual formulation where the convex potentials are parameterized by input convex neural networks (Amos et al., 2017). Our algorithm is straightforward without introducing bias (e.g. Li et al. (2020)) or requiring minimax optimization (e.g. Fan et al. (2020)). This is made possible by combining a new congruence regularizing term combined with cycle-consistency regularization (Korotin et al., 2019a). As we will show in the analysis, thanks to the properties of
Wasserstein-2 distances, the gradients of the resulting convex potentials “push” the input distributions close to the true barycenter, allowing good approximation of the barycenter.

heading: 2 PRELIMINARIES

text:
We denote the set of all Borel probability measures on RD with finite second moment by P2(RD). We use P2,ac(RD) ⊂ P2(RD) to denote the subset of all absolutely continuous measures (w.r.t. the Lebesgue measure).
Wasserstein-2 distance. For P,Q ∈ P2(RD), the Wasserstein-2 distance is defined by
W22(P,Q) def = min
π∈Π(P,Q) ∫ RD×RD ‖x− y‖2 2 dπ(x, y), (1)
where Π(P,Q) is the set of probability measures on RD×RD whose marginals are P,Q, respectively. This definition is known as Kantorovich’s primal form of transport distance (Kantorovitch, 1958).
The Wasserstein-2 distance W2 is well-studied in the theory of optimal transport (Brenier, 1991; McCann et al., 1995). In particular, it has a dual formulation (Villani, 2003):
W22(P,Q) = ∫ RD ‖x‖2 2 dP(x) + ∫ RD ‖y‖2 2 dQ(y)− min ψ∈Conv [ ∫ RD ψ(x)dP(x) + ∫ RD ψ(y)dQ(y) ] , (2)
where the minimum is taken over all the convex functions (potentials) ψ : RD → R ∪ {∞}, and ψ(y) = maxx∈RD ( 〈x, y〉−ψ(x) ) : RD → R∪{∞} is the convex conjugate of ψ (Fenchel, 1949), which is also a convex function. The optimal potential ψ∗ is defined up to an additive constant.
Brenier (1991) shows that if P does not give mass to sets of dimensions at most D − 1, then the optimal plan π is uniquely determined by π∗ = [idRD , T ∗]]P, where T ∗ : RD → RD is the unique solution to the Monge’s problem
T ∗ = arg min T]P=Q ∫ RD ‖x− T (x)‖2 2 dP(x). (3)
The connection between T ∗ and the dual formulation (2) is that T ∗ = ∇ψ∗, where ψ∗ is the optimal solution of (2). Additionally, if Q does not give mass to sets of dimensions at most D − 1, then T ∗ is invertible and
T ∗(x) = ∇ψ∗(x) = (∇ψ∗)−1(x), (T ∗)−1(y) = ∇ψ∗(y) = (∇ψ∗)−1(y). In particular, the above discussion applies to the case where P,Q ∈ P2,ac(RD).
Wasserstein-2 barycenter. Let P1, . . . ,PN ∈ P2,ac(RD). Then, their barycenter w.r.t. weights α1, . . . , αN (αn > 0 and ∑N n=1 αn = 1) is
P def= arg min P∈P2(RD) N∑ n=1 αnW22(Pn,P). (4)
Throughout this paper, we assume that at least one of P1, . . . ,PN ∈ P2,ac(RD) has bounded density. Under this assumption, P is unique and absolutely continuous, i.e., P ∈ P2,ac(RD), and it has bounded density (Agueh & Carlier, 2011, Definition 3.6 & Theorem 5.1).
For n ∈ {1, 2, . . . , N}, let (ψ∗n, ψ∗n) be the optimal pair of (mutually) conjugate potentials that transport Pn to P, i.e., ∇ψ∗n]Pn = P and∇ψ∗n]P = Pn. Then {ψ∗n} satisfy
N∑ n=1 αn∇ψ∗n(x) = x and N∑ n=1 αnψ∗n(x) = ‖x‖2 2 + c. (5)
for all x ∈ RD (Agueh & Carlier, 2011; Álvarez-Esteban et al., 2016). Since optimal potentials are defined up to a constant, for convenience, we set c = 0. The condition (5) serves as the basis for our algorithm for computing Wasserstein-2 barycenters. We say that potentialsψ1, . . . , ψN are congruent w.r.t. weights α1, . . . , αn if their conjugate potentials satisfy (5), i.e., ∑D n=1 αnψn(x) = ‖x‖2 2 for all x ∈ RD.

heading: 3 RELATED WORK

text:
Most algorithms in the field of computational optimal transport are designed for the discrete setting where the input distributions have finite support; see the recent survey by Peyré et al. (2019) for discussion. A particular popular line of algorithms are based on entropic regularization that gives rise to the famous Sinkhorn iteration (Cuturi, 2013; Cuturi & Doucet, 2014). These methods are typically limited to a support of 105−106 points before the problem becomes computationally infeasible. Similarly, discrete barycenter methods (Cuturi & Doucet, 2014), particularly the ones that rely on a fixed support for the barycenter (Dvurechenskii et al., 2018; Staib et al., 2017), cannot provide precise approximation of continuous barycenters in high dimensions, since a large number of samples is needed; see experiments in Fan et al. (2020, §4.3) for an example. Thus we focus on the existing literature in the continuous setting.
Computation of Wasserstein-2 distances and maps. Genevay et al. (2016) demonstrate the possibility of computing Wasserstein distances given only sample access to the distributions by parameterizing the dual potentials as functions in the reproducing kernel Hilbert spaces. Based on this realization, Seguy et al. (2017) propose a similar method but use neural networks to parameterize the potentials, using entropic or L2 regularization w.r.t. P×Q to keep the potentials approximately conjugate. The transport map is recovered from optimized potentials via barycentric projection.
As we note in §2, W2 enjoys many useful theoretical properties. For example, the optimal potential ψ∗ is convex, and the corresponding optimal transport map is given by ∇ψ∗. By exploiting these properties, Makkuva et al. (2019) propose a minimax optimization algorithm for recovering transport maps, using input convex neural networks (ICNNs) (Amos et al., 2017) to approximate the potentials.
An alternative to entropic regularization is the cycle-consistency regularization proposed by Korotin et al. (2019a). It uses the property that the gradients of optimal dual potentials are inverses of each other. The imposed regularizer requires integration only over the marginal measures P and Q, instead of over P × Q as required by entropy-based alternatives. Their method converges faster than the minimax method since it does not have an inner optimization cycle.
Xie et al. (2019) propose using two generative models with a shared latent space to implicitly compute the optimal transport correspondence between P and Q. Based on the obtained correspondence, the authors are able to compute the optimal transport distance between the distributions.
Computation of Wasserstein-2 barycenters. A few recent techniques tackle the barycenter problem (4) using continuous rather than discrete approximations of the barycenter:
• MEASURE-BASED (GENERATIVE) OPTIMIZATION: Problem (4) optimizes over probability measures. This can be done using the generic algorithm by Cohen et al. (2020) who employ generative networks to compute barycenters w.r.t. arbitrary discrepancies. They test their method with the maximum mean discrepancy (MMD) and Sinkhorn divergence. This approach suffers from the usual limitations of generative models such as mode collapse. Applying it to W2 barycenters requires estimation of W22(Pn,P). Fan et al. (2020) test this approach using the minimax method by Makkuva et al. (2019), but they end up with a challenging min-max-min problem. • POTENTIAL-BASED OPTIMIZATION: Li et al. (2020) recover the optimal potentials {ψ∗n} via a non-minimax regularized dual formulation. No generative model is needed: the barycenter is recovered by pushing forward measures using gradients of potentials or by barycentric projection.

heading: 4 METHODS

text: Inspired by Li et al. (2020) we use a potential-based approach and recover the barycenter by using gradients of the potentials as pushforward maps. The main differences are: (1) we restrict the potentials to be convex, (2) we enforce congruence via a regularizing term, and (3) our formulation does not introduce bias, meaning the optimal solution of our formulation gives the true barycenter.

heading: 4.1 DERIVING THE DUAL PROBLEM

text:
Let P be the true barycenter. Our goal is to recover the optimal potentials {ψ∗n, ψ∗n} mapping the input measures Pn into P.
To start, we express the barycenter objective (4) after substituting the dual formulation (2):
N∑ n=1 αnW22(Pn,P) = [ N∑ n=1 αn ∫ RD ‖x‖2 2 dPn(x) ] + ∫ RD ‖y‖2 2 dP(y)−
min {ψn}∈Conv [ N∑ n=1 αn ∫ RD ψn(x)dPn(x) + N∑ n=1 αn ∫ RD ψn(y)dP(y) ]
(6)
The minimum is attained not just among convex potentials {ψn}, but among congruent potentials (see discussion under (5)); thus, we can add the constraint that {ψn} are congruent to (6). Hence, N∑ n=1 αnW22(Pn,P) = [ N∑ n=1 αn ∫ RD ‖x‖2 2 dPn(x) ] − min {ψn} congruent [ N∑ n=1 αn ∫ RD ψn(y)dPn(y)︸ ︷︷ ︸
MultiCorr({αn,Pn}|{ψn})
] . (7)
To transition from (6) to (7), we used the fact that for congruent {ψn} we have∑N n=1 αnψn(x) = ‖x‖2 2 , so ∑N n=1 ∫ RD αnψn(y)dP(y) = ∫ RD ‖y‖2 2 dP(y).
We call the value inside the minimum in (7) the multiple correlation of {Pn} with weights {αn} w.r.t. potentials {ψn}. Notice that the true barycenter P appears nowhere on the right side of (7). Thus the optimal potentials {ψ∗n} can be recovered by solving the following
min {ψn} congruent MultiCorr({αn,Pn}|{ψn}) = min {ψn} congruent [ N∑ n=1 αn ∫ RD ψn(y)dPn(y) ] . (8)

heading: 4.2 IMPOSING THE CONGRUENCE CONDITION

text:
It is challenging to impose the congruence condition on convex potentials. What if we relax the congruence condition? The following theorem bounds how close a set of convex potentials {ψn} is to {ψ∗n} in terms of the difference of multiple correlation. Theorem 4.1. Let P ∈ P2,ac(RD) be the barycenter of P1, . . . ,PN ∈ P2,ac(RD) w.r.t. weights α1, . . . , αN . Let {ψ∗n} be the optimal congruent potentials of the barycenter problem. Suppose we have B-smooth1 convex potentials {ψn} for some B ∈ [0,+∞], and denote ∆ = MultiCorr({αn,Pn} | {ψn})−MultiCorr({αn,Pn} | {ψ∗n}). Then,
∆ + ∫ RD N∑ n=1 [ αnψn(y)− ‖y‖2 2 ] dP(y)︸ ︷︷ ︸
Congruence mismatch
≥ 1 2B N∑ n=1 αn‖∇ψ∗n(x)−∇ψn(x)‖2Pn . (9)
Here ‖ · ‖µ denotes the norm induced by inner product in Hilbert space L2(RD → RD, µ). We call the second term on the left of (9) the congruence mismatch.
We prove this in Appendix B. Note that if the congruence mismatch is non-positive, then
∆ ≥ 1 2B N∑ n=1 αn‖∇ψ∗n(x)−∇ψn(x)‖2Pn ≥ 1 B N∑ n=1 αnW22(∇ψn]Pn,P), (10)
where the last inequality of (10) follows from (Korotin et al., 2019a, Lemma A.2). From (10), we conclude that for all n ∈ {1, . . . , N}, we have W22(∇ψn]Pn,P) ≤ B∆αn . This shows that if the congruence mismatch is non-positive, then ∆, the difference in multiple correlation, provides
1We say that a diffirentiable function f : RD → R is B-smooth if its gradient ∇f is B-Lipschitz.
an upper bound for the Wasserstein-2 distance between the true barycenter and each pushforward ∇ψn]Pn. This justifies the use of ∇ψn]Pn to recover the barycenter. Notice for optimal potentials, the congruence mismatch is zero.
Thus to penalize positive congruence mismatch, we introduce a regularizing term
RP1({αn}, {ψn}) def = ∫ RD [ N∑ n=1 αnψn(y)− ‖y‖2 2 ] + dP(y). (11)
Because we take the positive part of the integrand of (9) to get (11) and that the right side of (9) is non-negative, we have[
MultiCorr({αn,Pn} | {ψn}) + 1 · RP1({αn}, {ψn}) ] −MultiCorr({αn,Pn} | {ψ∗n}) ≥ 0
for all convex potentials {ψn}. On the other hand, for optimal potentials {ψn} = {ψ∗n}, the inequality turns into equality, implying that adding the regularizing term 1 · RP1({αn}, {ψn}) to (8) will not introduce bias – the optimal solution still yields {ψ∗n}.
However, evaluating (11) exactly requires knowing the true barycenter P a priori. To remedy this issue, one may replace P with another absolutely continuous measure τ · P̂ (τ ≥ 1 and P̂ is a probability measure) whose density bounds that of P from above almost everywhere. In this case,
τ · RP̂1({αn}, {ψn}) = τ · ∫ RD [ N∑ n=1 αnψn(y)− ‖y‖2 2 ] + dP̂ ≥ RP1({αn}, {ψn}). (12)
Hence we obtain the following regularized version of (8) where {ψ∗n} is the optimal solution:
min {ψn}∈Conv
[ MultiCorr({αn,Pn} | {ψn}) + τ · RP̂1({αn}, {ψn}) ] . (13)
Selecting a measure τ · P̂ is not obvious. Consider the case when {Pn} are supported on compact sets X1, . . . ,XN ⊂ RD and P1 has density upper bounded by h < ∞. In this scenario, the barycenter density is upper bounded by h · α−D1 (Álvarez-Esteban et al., 2016, Remark 3.2). Thus, the measure τ · P̂ supported on ConvexHull(X1, . . . ,XN ) with this density is an upper bound for P. We will address the question of how to choose τ, P̂ properly in practice in §4.4.

heading: 4.3 ENFORCING CONJUGACY OF POTENTIALS PAIRS

text:
Throughout this subsection, we assume the upper bound finite measure τ · P̂ of the P is known. The optimization problem (13) involves not only the potentials {ψn}, but also their conjugates {ψn}. This brings practical difficulty since evaluating conjugate potentials is hard (Korotin et al., 2019a).
Instead we parameterize potentials ψn and ψn separately using input convex neural networks (ICNN) as ψ†n and ψ ‡ n respectively. We add an additional cycle-consistency regularizer to enfore the conjugacy of ψ†n and ψ ‡ n as in Korotin et al. (2019a). This regularizer is defined as
RPn2 (ψ†n, ψ ‡ n) def = ∫ RD ‖∇ψ‡n ◦ ∇ψ†n(x)− x‖22 dPn(x) = ‖∇ψ ‡ n ◦ ∇ψ†n − idRD‖2Pn .
Note thatRPn2 (ψ†n, ψ ‡ n) = 0 this condition is necessary for ψ†n and ψ ‡ n to be conjugate with each other. Also, it is a sufficient condition for convex functions to be conjugates up to an additive constant.
We use one-sided regularization. In our case, computing the regularizer of the other direction ‖∇ψ†n ◦∇ψ ‡ n− idRD‖2P is infeasible, since P is unknown. If fact, Korotin et al. (2019a) demonstrates that such one-sided condition is sufficient.
In this way we use 2N input convex neural networks for {ψ†n, ψ ‡ n}. By adding the new cycle consistency regularizer into (13), we obtain our final objective:
min {ψ†n,ψ‡n} Approximate multiple correlation︷ ︸︸ ︷ N∑ n=1 [ αn ∫ RD [〈x,∇ψ†n(x)〉−ψ ‡ n(∇ψ†n(x))︸ ︷︷ ︸
≈ψ‡n(x)
]dPn(x) ] +τ ·RP̂1({ψ
‡ n})︸ ︷︷ ︸
Congruence reg.
+λ N∑ n=1 αnRPn2 (ψ†n, ψ ‡ n)︸ ︷︷ ︸
Cycle regularizer
. (14)
Note that we express the aproximate multiple correlation by using both potentials {ψ†n} and {ψ ‡ n}. This is done to eliminate the freedom of an additive constant on {ψ†n} that is not addressed by cycle regularization. We denote the entire objective as MultiCorr ( {Pn} | {ψ†}, {ψ‡}; τ, P̂, λ ) . Analogous to Theorem 4.1, we have following result showing that this new objective enjoys the same properties as the unregularized version from (8).
Theorem 4.2. Let P ∈ P2,ac(RD) be the barycenter of P1, . . . ,PN ∈ P2,ac(RD) w.r.t. weights α1, . . . , αN . Let {ψ∗n} be the optimal congruent potentials of the barycenter problem. Suppose we have τ, P̂ such that τ ≥ 1 and τ · P̂ ≥ P. Suppose we have convex potentials {ψ†n} and β‡-strongly convex and B‡-smooth convex potentials {ψ‡n} with 0 < β‡ ≤ B‡ <∞ and λ > B †
2(β‡)2 . Then
MultiCorr ( {αn,Pn} | {ψ†n}, {ψ ‡ n}; τ, P̂, λ ) ≥ MultiCorr ( {αn,Pn} | {ψ∗n} ) . (15)
Denote ∆ = MultiCorr ( {αn,Pn} | {ψ†n}, {ψ ‡ n}; τ, P̂, λ ) −MultiCorr ( {αn,Pn} | {ψ∗n} ) . Then for all n ∈ {1, . . . , N}, we have
W22 ( ∇ψ†n]Pn,P ) ≤ 2∆ αn ·
(√ 1
β‡ +
√ 1
λ(β‡)2 − B†2
)2 = O(∆). (16)
Informally, Theorem 4.2 states that the better we solve the regularized dual problem, (14) the closer we expect each ∇ψ†n]Pn to be to the true barycenter P in W2. It follows from (15) that our final objective (14) is unbiased: the optimal solution is obtained by {ψ∗n, ψ∗n}.

heading: 4.4 PRACTICAL ASPECTS AND OPTIMIZATION PROCEDURE

text:
In practice, even if the choice of τ, P̂ does not satisfy τ · P̂ ≥ P, we observe the pushforward measures∇ψ†n]Pn often converge to P. To partially bridge the gap between theory and practice, we dynamically update the measure P̂ so that after each optimization step we set (for γ ∈ [0, 1])
P̂′ := γ · P̂ + (1− γ) · N∑ n=1 αn · [ ∇ψ†]Pn ] ,
i.e., the probability measure P̂′ is a mixture of the given initial measure P̂ and the current barycenter estimates {∇ψ†]Pn}. For the initial P̂ one may use the barycenter of {N (µPn ,ΣPn)}. It can be efficiently computed via an iterative fixed point algorithm (Álvarez-Esteban et al., 2016; Chewi et al., 2020). During the optimization, these estimates become closer to the true barycenter and can thus improve the congruence regularizer (12).
We use mini-batch stochastic gradient descent to solve (14) where the integration is done by MonteCarlo sampling from input measures {Pn} and regularization measure P̂, similar to Li et al. (2020). We provide the detailed optimization procedure (Algorithm 1) and discuss its computational complexity in Appendix A. In Appendix C.3, we demonstrate that the impact of the considered regularization on our model: we show that cycle consistency and the congruence condition of the potentials are well satisfied.

heading: 5 EXPERIMENTS

text:
The code is written on PyTorch framework and is publicly available at
https://github.com/iamalexkorotin/Wasserstein2Barycenters.
We compare our method [CW2B] with the potential-based method [CRWB] by Li et al. (2020) (with Wasserstein-2 distance and L2-regularization) and with the measure-based generative method [SCW2B] by Fan et al. (2020). All considered methods recover 2N potentials {ψ†n, ψ ‡ n} ≈ {ψ∗n, ψ∗n} and approximate the barycenter as pushforward measures {∇ψ†n]Pn}. Regularization in [CRWB] allows access to the joint density of the transport plan, a feature of their method that we do not consider here. The method [SCW2B] additionally outputs a generated barycenter g]S ≈ P where g is the generative network and S is the input noise distribution.
To assess the quality of the computed barycenter, we consider the unexplained variance percentage defined as UVP(P̃) = 100 W 2 2(P̃,P)
1/2Var(P)%. When UVP ≈ 0%, P̃ is a good approximation of P. For values ≥ 100%, the distribution P̃ is undesirable: a trivial baseline P0 = δEP[y] achieves UVP(P
0) = 100%. Evaluating UVP in high dimensions is infeasible: empirical estimates of W22 are unreliable due to high sample complexity (Weed et al., 2019). To overcome this issue, for barycenters given by ∇ψ†n]Pn we use L2-UVP defined by
L2-UVP(∇ψ†n,Pn) def = 100 ‖∇ψ†n −∇ψ∗n‖2Pn Var(P) %
[ ≥ UVP(∇ψ†n]Pn) ] , (17)
where the inequality in brackets follows from (Korotin et al., 2019a, Lemma A.2). We report the weighted average of L2-UVP of all pushforward measures w.r.t. the weights αn. For barycenters given in an implicit form g]S, we compute the Bures-Wasserstein UVP defined by
BW22-UVP(g]S) def = 100 BW22(g]S,P) 1 2 Var(P) %
[ ≤ UVP(g]S) ] , (18)
where BW22(P,Q) = W22 ( N (µP,ΣP),N (µQ,ΣQ) ) is the Bures-Wasserstein metric and we use µP,ΣP to denote the mean and the covariance of a distribution P (Chewi et al., 2020). It is known that BW22 lower-bounds W22 (Dowson & Landau, 1982), so the inequality in the brackets of (18) follows. A detailed discussion of the adopted metrics is given in Appendix C.2.

heading: 5.1 HIGH-DIMENSIONAL LOCATION-SCATTER EXPERIMENTS

text:
In this section, we consider N = 4 with (α1, . . . , α4) = (0.1, 0.2, 0.3, 0.4) as weights. We consider the location-scatter family of distributions (Álvarez-Esteban et al., 2016, §4) whose true barycenter can be computed. Let P0 ∈ P2,ac and define the following location-scatter family of distributions F(P0) = {fS,u]P0 | S ∈ M+D×D, u ∈ RD}, where fS,u : RD → RD is a linear map fS,u(x) = Sx+ u with positive definite matrix S ∈M+D×D. When {Pn} ⊂ F(P0), their barycenter P is also an element of F(P0) and can be computed via fixed-point iterations (Álvarez-Esteban et al., 2016). Figure 1a shows a 2-dimensional location-scatter family generated by using the Swiss roll distribution as P0. The true barycenter is shown in Figure 1b. The generated barycenter g]S of [SCW2B] is given in Figure 1c. The pushforward measures ∇ψ†n]Pn of each method are provided in Figures 1d, 1e, 1f, respectively. In this example, the pushforward measures ∇ψn]Pn all reasonably approximate P, whereas the generated barycenter g]S of [SCW2B] (Figure 1c) visibly underfits.
For quantitative comparison, we consider two choices for P0: the D-dimensional standard Gaussian distribution and the uniform distribution on [− √ 3,+ √
3]D. Each Pn is constructed as fSTn ΛSn,0]P0 ∈ F(P0), where Sn is a random rotation matrix and Λ is diagonal with entries [ 12b 0, 12b 1, . . . , 2] where
b = D−1 √
4. We consider only centered distributions (i.e. zero mean) because the barycenter of non-centered {Pn} ∈ P2,ac(RD) is the barycenter of {P′n} shifted by ∑N n=1 αnµPn , where {P′n} are centered copies of {Pn} (Álvarez-Esteban et al., 2016). Results are shown in Table 1 and 2. In these experiments, our method outperforms [CRWB] and [SCW2B]. For [CRWB], dimension ∼16 is the breakpoint: the method does not scale well to higher dimensions. [SCW2B] scales with the increasing dimension better, but its errors L2-UVP and BW22-UVP are twice as high as ours. This is likely due to the generative approximation and the difficult min-max-min optimization in [SCW2B]. For completeness, we also compare our algorithm to the proposed in Cuturi & Doucet (2014) which approximates the barycenter by a discrete distribution on a fixed number of free-support points. In our experiment, similar to Li et al. (2020), we set 5000 as the support size. As expected, the BW22-UVP error of the method increases drastically as the dimension grows and the method is outperformed by our approach.
To show the scalability of our method with the number of input distributions N , we conduct an analogous experiment with a high-dimensional location-scatter family for N = 20. We set αn = 2nN(N+1) for n = 1, 2, ..., 20 and choose the uniform distribution on [− √ 3,+ √
3]D as P0 and construct distributions Pn ∈ F(P0) as before. The results for dimensions 32, 64 and 128 are provided in Table 3. Similar to the results from Tables 1 and 2, we see that our method outperforms the alternatives.
Metric Method D=2 4 8 16 32 64 128 256
BW22-UVP, % [FCWB], Cuturi & Doucet (2014) 0.64 0.77 1.22 3.75 8.92 14.3 18.46 21.64
[SCW2B], (Fan et al., 2020) 0.12 0.10 0.19 0.29 0.46 0.6 1.38 2.9
L2-UVP, % (potentials)
0.17 0.12 0.2 0.31 0.47 0.62 1.21 1.52 [CRWB], (Li et al., 2020) 0.58 1.83 8.09 21.23 55.17 > 100
[CW2B], ours 0.17 0.08 0.06 0.1 0.2 0.25 0.42 0.82
Table 2: Comparison of UVP for the case {Pn} ⊂ F(P0), P0 = Uniform
( [− √ 3,+ √ 3]D ) , N = 4.
Metric Method D=32 64 128
BW22-UVP, % [FCWB], Cuturi & Doucet (2014) 14.09 26.21 38.43
[SCW2B], (Fan et al., 2020) 0.62 0.93 1.83
L2-UVP, % (potentials)
0.60 0.86 1.52 [CW2B], ours 0.31 0.58 1.45
Table 3: Comparison of UVP for the case {Pn} ⊂ F(P0), P0 = Uniform
( [− √ 3,+ √ 3]D ) , N = 20.

heading: 5.2 SUBSET POSTERIOR AGGREGATION

text:
We apply our method to aggregate subset posterior distributions. The barycenter of subset posteriors converges to the true posterior (Srivastava et al., 2018). Thus, computing the barycenter of subset
posteriors is an efficient alternative to obtaining a full posterior in the big data setting (Srivastava et al., 2015; Staib et al., 2017; Li et al., 2020).
Analogous to (Li et al., 2020), we consider Poisson and negative binomial regressions for predicting the hourly number of bike rentals using features such as the day of the week and weather conditions.2 We consider the posterior on the 8-dimensional regression coefficients for both Poisson and negative binomial regressions. We randomly split the data into N = 5 equally-sized subsets and obtain 105 samples from each subset posterior using the Stan library (Carpenter et al., 2017). This gives the discrete uniform distributions {Pn} supported on the samples. As the ground truth barycenter P, we consider the full dataset posterior also consisting of 105 points.
We use BW22-UVP(P̃,P) to compare the estimated barycenter P̃ (pushforward measure ∇ψ†n]Pn or generated measure g]S) with the true barycenter. The results are in Table 4. All considered methods perform well (UVP< 2%), but our method outperforms the alternatives.

heading: 5.3 COLOR PALETTE AVERAGING

text:
For qualitative study, we apply our method to aggregating color palettes of images. For an RGB image I , its color palette is defined by the discrete uniform distribution P(I) of all its pixels ∈ [0, 1]3. For 3 images {In} we compute the barycenter P of each color palette Pn = P(In) w.r.t. uniform weights αn = 1 3 . We apply each computed potential∇ψ † n pixel-wise to In to obtain the “pushforward” image ∇ψ†n]In. These “pushforward” images should be close to the barycenter P of {Pn}.
The results are provided in Figure 2. Note that the image∇ψ†1]I1 inherits certain attributes of images I2 and I3: the sky becomes bluer and the trees becomes greener. On the other hand, the sunlight in images∇ψ†2]I2,∇ψ † 3]I3 has acquired an orange tint, thanks to the dominance of orange in I1.

heading: ACKNOWLEDGMENTS

text:
The Skoltech Advanced Data Analytics in Science and Engineering Group acknowledges the support of Russian Foundation for Basic Research grant 20-01-00203, Skoltech-MIT NGP initiative and thanks the Skoltech CDISE HPC Zhores cluster staff for computing cluster provision.
The MIT Geometric Data Processing group acknowledges the generous support of Army Research Office grant W911NF2010168, of Air Force Office of Scientific Research award FA9550-19-1-031, of National Science Foundation grant IIS-1838071, from the CSAIL Systems that Learn program, from the MIT–IBM Watson AI Laboratory, from the Toyota–CSAIL Joint Research Center, from a gift from Adobe Systems, from an MIT.nano Immersion Lab/NCSOFT Gaming Program seed grant, and from the Skoltech–MIT Next Generation Program.
2http://archive.ics.uci.edu/ml/datasets/Bike+Sharing+Dataset

heading: A THE ALGORITHM

text:
The numerical procedure for solving our final objective (14) is given below.
Algorithm 1: Numerical Procedure for Optimizing Multiple Correlations (14) Input :Distributions P1, . . . ,PN with sample access;
Weights α1, . . . , αN ≥ 0 with ∑N n=1 αn = 1; Regularization distribution P̂′ given by a sampler; Congruence regularizer coefficient τ ≥ 1; Balancing coefficient γ ∈ [0, 1]; Cycle-consistency regularizer coefficient λ > 0; 2N ICNNs {ψθn , ψωn}; Batch size K > 0;
for t = 1, 2, . . . do 1. Sample batches Xn ∼ Pn for all n = 1, . . . , N ; 2. Compute the pushforwards Yn = ∇ψθn]Xn for all n = 1, . . . , N ; 3. Sample batch Y0 ∼ P̂; 4. Compute the Monte-Carlo estimate of the congruence regularizer:
LCongruence := 1 K · N∑ n=1 γn ∑ y∈Yn [ N∑ n′=1 αn′ψωn′ (y)− ‖y‖2 2 ] + ,
where γ0 = γ and γn = αn · (1− γ) for n = 1, 2, . . . , N ; 5. Compute the Monte-Carlo estimate of the cycle-consistency regularizer:
LCycle := 1
K N∑ n=1 αn [ ∑ x∈Xn ‖∇ψωn ◦ ∇ψθn(x)− x‖22 ] ;
6. Compute the Monte-Carlo estimate of multiple correlations:
LMultiCorr := N∑ n=1 [ αn · 1 K ∑ x∈Xn [ 〈x,∇ψθn(x)〉 − ψωn(∇ψθn(x))] ] ;
7. Compute the total loss:
LTotal := LMultiCorr + λ · LCycle + τ · LCongruence;
8. Perform a gradient step over {θn, ωn} by using ∂LTotal∂{θn,ωn} ; end
Parametrization of the potentials. To parametrize potentials {ψθn , ψωn}, we use DenseICNN (dense input convex neural network) with quadratic skip connections; see (Korotin et al., 2019a, Appendix B.2). As an initialization step, we pre-train the potentials to satisfy
ψθn(x) ≈ ‖x‖2
2 and ψωn(y) ≈
‖y‖2
2 .
Such pre-training provides a good start for the networks: each ψθn is approximately conjugate to the corresponding ψωn . On the other hand, the initial networks {ψθn} are approximate congruent according to (5).
Computational Complexity. For a single training iteration, the time complexity of both forward (evaluation) and backward (computing the gradient with respect to the parameters) passes through the objective function (14) is O(NT ). Here N is the number of input distributions and T is the time taken by evaluating each individual potential (parameterized as a neural network) on a batch of points sampled from either Pn or P̂. This claim follows from the well-known fact that gradient evaluation ∇θhθ(x) of hθ : RD → R, when parameterized as a neural network, requires time proportional
to the size of the computational graph. Hence, gradient computation requires computational time proportional to the time for evaluating the function hθ(x) itself. The same holds when computing the derivative with respect to x. Then, for instance, computing the term ∇ψ‡n ◦ ∇ψ†n(x) in (14) takes O(T ) time. The gradient of this term with respect to θ also takes O(T ) time: Hessian-vector products that appear can be calculated in O(T ) time using the famous Hessian trick, see Pearlmutter (1994).
In practice, we compute all the gradients using automatic differentiation. We empirically measured that for our DenseICNN potentials, the computation of their gradient w.r.t. input x, i.e., ∇ψ†(x), requires roughly 3-4x more time than the computation of ψ†(x).

heading: B PROOFS

text:
In this section, we prove our main Theorems 4.1 and 4.2.
We use L2(RD → RD, µ) to denote the Hilbert space of functions f : RD → RD with integrable square w.r.t. a probability measure µ. The corresponding inner product for f1, f2 ∈ L2(RD → RD, µ) is denoted by
〈f1, f2〉µ def = ∫ RD 〈f1(x), f2(x)〉dµ(x),
where 〈f1(x), f2(x)〉 is the Euclidean dot product. We use ‖ · ‖µ = √ 〈·, ·〉µ to denote the norm induced by the inner product in L2(RD → RD, µ). We also recall a useful property of lower semi-continuous convex function ψ : RD → R:
∇ψ(x) = arg max y∈RD
[ 〈y, x〉 − ψ(y) ] , (19)
which follows from the fact that
ŷ = arg max y∈RD
[ 〈y, x〉 − ψ(y) ] ⇐⇒ x−∇ψ(ŷ) = 0.
We begin with the proof of Theorem 4.1.
Proof. We consider the difference between the estimated correlations and true ones:
∆ = N∑ n=1 αn ∫ RD ψn(x)dPn(x)− N∑ n=1 αn ∫ RD ψ∗n(x)dPn(x) =
N∑ n=1 αn ∫ RD [ 〈∇ψn(x), x〉 − ψn ( ∇ψn(x)) ] dPn(x)−
N∑ n=1 αn ∫ RD [ 〈∇ψ∗n(x), x〉 − ψ∗n ( ∇ψ∗n(x)) ] dPn(x), (20)
where we twice use (19) for f = ψn and f = ψ∗n. We note that N∑ n=1 αn ∫ RD 〈∇ψ∗n(x), x〉dPn(x) = N∑ n=1 αn ∫ RD 〈y,∇ψ∗n(y)〉dP(y) =
∫ RD 〈y, N∑ n=1 αn∇ψ∗n(y)〉dP(y) = ∫ RD 〈y, y〉dP(y) = ‖idRD‖2P, (21)
where we use of change-of-variable formula for∇ψ∗n]Pn = P and (5). Analogously, N∑ n=1 αn ∫ RD ψ∗n ( ∇ψ∗n(x))dPn(x) = N∑ n=1 αn ∫ RD ψ∗n ( y)dP(y) =
∫ RD N∑ n=1 αnψ∗n ( y)dP(y) = ∫ RD ‖y‖2 2 dP(y) = 1 2 ‖idRD‖2P. (22)
Since each ψn is B-smooth, we conclude that ψn is 1B -strongly convex, see (Kakade et al., 2009). Thus, we have
ψn ( ∇ψ∗n(x))) ≥
ψn ( ∇ψn(x))) + 〈∇ψ†n ◦ ∇ψ†n(x)︸ ︷︷ ︸
=x
,∇ψ∗n(x)−∇ψn(x)〉+ 1
2B ‖∇ψ∗n(x)−∇ψn(x)‖2 =
ψn ( ∇ψn(x))) + 〈x,∇ψ∗n(x)−∇ψn(x)〉+ 1
2B ‖∇ψ∗n(x)−∇ψn(x)‖2, (23)
or equivalently
−ψn ( ∇ψn(x))) ≥ −ψn ( ∇ψ∗n(x)))+ 〈x,∇ψ∗n(x)−∇ψn(x)〉+ 1
2B ‖∇ψ∗n(x)−∇ψn(x)‖2. (24)
We integrate (24) w.r.t. Pn and sum over n = 1, 2, . . . , N with weights αn:
− N∑ n=1 αn ∫ RD ψn ( ∇ψn(x))dPn(x) ≥
− N∑ n=1 αn ∫ RD ψn ( ∇ψ∗n(x))dPn(x) + N∑ n=1 αn〈x,∇ψ∗n(x)〉Pn − N∑ n=1 αn〈x,∇ψn(x)〉Pn +
N∑ n=1 αn 1 2B ‖∇ψ∗n(x)−∇ψn(x)‖2Pn =
− ∫ RD N∑ n=1 αnψn ( y)dP(y) + N∑ n=1 αn〈x,∇ψ∗n(x)〉Pn −
N∑ n=1 αn〈x,∇ψn(x)〉Pn + N∑ n=1 αn 1 2B ‖∇ψ∗n(x)−∇ψn(x)‖2Pn . (25)
We note that − ∫ RD N∑ n=1 αnψn ( y)dP(y) = ∫ RD [‖y‖2 2 − N∑ n=1 αnψn ( y) ] dP(y)− ∫ RD ‖y‖2 2 dP(y)
∫ RD [‖y‖2 2 − N∑ n=1 αnψn ( y) ] dP(y)− 1 2 ‖idRD‖2P. (26)
Now we substitute (25), (26), (21) and (22) into (20) to obtain (9).
Next, we prove Theorem 4.2.
Proof. Since ψ‡n is β‡ strongly convex, its conjugate ψ‡n is 1 β‡ -smooth, i.e. has 1 β‡
-Lipschitz gradient ∇ψ‡n (Kakade et al., 2009). Thus, for all x, x′ ∈ RD:
‖∇ψ‡n(x)−∇ψ‡n(x′)‖2 ≤ ( 1
β‡ )2 · ‖x− x′‖2.
We substitute x′ = ∇ψ‡n ◦ ∇ψ†n(y) = ( ∇ψ‡n )−1 ◦ ∇ψ†n(y) and obtain: ‖∇ψ†n(x)−∇ψ‡n(x)‖2 ≤ ( 1
β‡ )2‖x−∇ψ‡n ◦ ∇ψ†n(x)‖2. (27)
Since the function ψ‡n is B‡-smooth, we have for all x ∈ RD:
ψ‡n(∇ψ†n(x)) ≤ ψ ‡ n(∇ψ‡n(x)) + 〈∇ψ ‡ n ◦ ∇ψ‡n(x)︸ ︷︷ ︸
=x
,∇ψ†n(x)−∇ψ‡n(x)〉+ B‡
2 ‖∇ψ†n(x)−∇ψ‡n(x)‖2,
that is equivalent to:
〈x,∇ψ†n(x)〉 − ψ ‡ n(∇ψ†n(x)) ≥ 〈x,∇ψ‡n(x)〉 − ψ ‡ n(∇ψ‡n(x))︸ ︷︷ ︸
ψ‡n(x)
−B ‡
2 ‖∇ψ†n(x)−∇ψ‡n(x)‖2. (28)
We combine (28) with (27) to obtain
〈x,∇ψ†n(x)〉 − ψ ‡ n(∇ψ†n(x)) ≥ ψ‡n(x)−
B‡
2(β‡)2 · ‖idRD −∇ψ‡n ◦ ∇ψ†n‖2. (29)
For every n = 1, 2, . . . , N we integrate (29) w.r.t. Pn and sum up the corresponding cycle-consistency regularization term:∫
RD
[ 〈x,∇ψ†n(x)〉 − ψ ‡ n(∇ψ†n(x))]dPn(x) + λ · ‖∇ψ
‡ n ◦ ∇ψ†n − idRD‖2Pn ≥∫
RD ψ‡(x)dPn(x) +
( λ− B † 2(β‡)2 ) · ‖∇ψ‡n ◦ ∇ψ†n − idRD‖2Pn︸ ︷︷ ︸
RPn2 (ψ † n,ψ ‡ n)
. (30)
We sum (30) for n = 1, 2, . . . , N w.r.t. weights αn to obtain:
N∑ n=1 αn ∫ RD [ 〈x,∇ψ†n(x)〉 − ψ ‡ n(∇ψ†n(x))]dPn(x) + λ N∑ n=1 αnRPn2 (ψ†n, ψ ‡ n) ≥
N∑ n=1 αn ∫ RD
ψ‡(x)dPn(x)︸ ︷︷ ︸ MultiCorr({αn,Pn}|{ψ‡n}) +
N∑ n=1 αn ( λ− B † 2(β‡)2 ) · RPn2 (ψ†n, ψ ‡ n).
We add τ · RP̂1({ψ ‡ n}) to both sides of (31) to get
MultiCorr ( {αn,Pn} | {ψ†n}, {ψ ‡ n}; τ, P̂, λ ) ≥ MultiCorr({αn,Pn} | {ψ‡n}) +
τ · RP̂1({ψ ‡ n}) + N∑ n=1 αn ( λ− B † 2(β‡)2 ) · RPn2 (ψ†n, ψ ‡ n). (31)
We substract MultiCorr({αn,Pn} | {ψ‡n}) from both sides and use Theorem 4.1 to obtain
∆ ≥ − ∫ RD N∑ n=1 [ αnψ ‡ n(y)− ‖y‖2 2 ] dP(y) + β‡ 2 N∑ n=1 αn‖∇ψ∗n(x)−∇ψ‡n(x)‖2Pn + (32)
τ · RP̂1({ψ ‡ n}) + N∑ n=1 αn ( λ− B † 2(β‡)2 ) · RPn2 (ψ†n, ψ ‡ n) ≥ (33)
N∑ n=1 αn ( λ− B † 2(β‡)2 ) · RPn2 (ψ†n, ψ ‡ n) + β‡ 2 N∑ n=1 αn‖∇ψ∗n(x)−∇ψ‡n(x)‖2Pn . (34)
In transition from (33) to (34), we explot the fact that the sum of the first term of (32) with the regularizer τ · RP̂1({ψ ‡ n}). Since λ > B † 2(β‡)2 , from (34) we immediately conclude ∆ ≥ 0; i.e., the multiple correlations upper bound (15) holds true. On the other hand, for every n = 1, 2, . . . , N we have
‖∇ψ∗n(x)−∇ψ‡n(x)‖2Pn ≤ 2∆ αnβ‡ and ‖∇ψ‡n ◦ ∇ψ†n − idRD‖2Pn ≤
2∆
αn · (λ− B † 2(β‡)2 ) . (35)
We combine the second part of (35) with (27) integrated w.r.t. Pn:
‖∇ψ‡n −∇ψ†n‖2Pn ≤ 2∆
αn · (λ(β‡)2 − B † 2 ) . (36)
Finally, we use the triangle inequality for ‖ · ‖Pn and conclude ‖∇ψ∗n −∇ψ†n‖Pn ≤ ‖∇ψ‡n −∇ψ†n‖Pn + ‖∇ψ‡n −∇ψ∗n‖Pn ≤√
2∆ αn · (√ 1 β‡ + √ 1
λ(β‡)2 − B†2
) , (37)
i.e.,
W22(∇ψ†n]Pn,P) ≤ ‖∇ψ∗n −∇ψ†n‖2Pn ≤ 2∆ αn · (√ 1 β‡ + √ 1
λ(β‡)2 − B†2
)2 = O(∆),
where the first inequality follows from (Korotin et al., 2019a, Lemma A.2).

heading: C EXPERIMENTAL DETAILS AND EXTRA RESULTS

text:
In this section, we provide experimental details and additional results. In Subsection C.1, we demonstrate qualitative results of computed barycenters in the 2-dimensional space. In Subsection C.2, we discuss used metrics in more detail. In Subsection C.4, we list the used hyperparameters of our method (CW2B) and methods [SCW2B], [CRWB].
C.1 ADDITIONAL TOY EXPERIMENTS IN 2D
We provide additional qualitative examples of computed barycenters of probability measures on R2. In Figure 3, we consider the location-scatter family F(P0) with P0 = Uniform[− √ 3, √
3]D. In principle, all the methods capture the true barycenter. However, the generated distribution g]S of [SCW2B] (Figure 3c) provides samples that lies outside of the actual barycenter’s support (Figure 3b). Also, in [CRWB] method, one of the potentials’ pushforward measure (top-right in Figure 3e) has visual artifacts.
In Figure 4, we consider the Gaussian Mixture example by (Fan et al., 2020). The barycenter computed by [SCW2B] method (Figure 4b) suffers from the behavior similar to mode collapse.
C.2 METRICS
The unexplained variance percentage (UVP) (introduced in Section 5) is a natural and straightforward metric to assess the quality of the computed barycenter. However, it is difficult to compute in high dimensions: it requires computation of the Wasserstein-2 distance. Thus, we use different but highly related metrics L2-UVP and BW22-UVP. To access the quality of the recovered potentials {ψ†n} we use L2-UVP defined in (17). L2-UVP compares not just pushforward distribution ∇ψ†n]Pn with the barycenter P, but also the resulting transport map with the optimal transport map ∇ψ∗n. It bounds UVP(∇ψ†n]Pn) from above, thanks to (Korotin et al., 2019a, Lemma A.2). Besides, L2-UVP naturally admits unbiased Monte Carlo estimates using random samples from Pn.
For measure-based optimization method, we also evaluate the quality of the generated measure g]S using Bures-Wasserstein UVP defined in (18). For measures P,Q whose covariance matrices are not degenerate, BW22 is given by
BW22(P,Q) = 1
2 ‖µP − µQ‖2 + [1 2 Tr ΣP + 1 2 Tr ΣQ − Tr(Σ 1 2 P ΣQΣ 1 2 P ) 1 2 ] .
Bures-Wasserstein metric compares P,Q by considering only their first and second moments. It is known that BW22(P,Q) is a lower bound for W22(P,Q), see (Dowson & Landau, 1982). Thus, we have BW22-UVP(g]S) ≤ UVP(g]S). In practice, to compute BW22-UVP(g]S), we estimate means and covariance matrices of distributions by using 105 random samples.
C.3 CYCLE CONSISTENCY AND CONGRUENCE IN PRACTICE
To assess the effect of the regularization of cycle consistency and the congruence condition in practice, we run the following sanity checks.
For cycle consistency, for each input distribution Pn we estimate (by drawing samples from Pn) the value ‖∇ψ‡n ◦ ∇ψ†n(x) − x‖2Pn/Var(Pn). This metric can be viewed as an analog of the L
2-UVP that we used for assessing the resulting transport maps. In all the experiments, this value does not exceed 2%, which means that cycle consistency and hence conjugacy are satisfied well.
For the congruence condition, we need to check that ∑N n=1 αnψ † n(x) = ‖x‖2/2. However, we do not know any straightforward metric to check this exact condition that is scaled properly by the variance of the distributions. Thus, we propose to use an alternative metric to check a slightly weaker condition on gradients, e.g., that ∑N n=1 αn∇ψ†n(x) = x. This is weaker due to the ambiguity of the additive
constants. For this we can compute ‖ ∑N n=1 αn∇ψ†n(x) − x‖2P/Var(P), where the denominator is
the variance of the true barycenter. We computed this metric and found that it is also less than 2% in all the cases, which means that congruence condition is mostly satisfied.
C.4 TRAINING HYPERPARAMETERS
The code is written using the PyTorch framework. The networks are trained on a single GTX 1080Ti.
C.4.1 WASSERSTEIN-2 CONTINUOUS BARYCENTERS (CW2B, OUR METHOD)
Regularization. We use τ = 5 and P̂ = N (0, ID) in our congruence regularizer τ · RP̂1 . We use λ = 10 for the cycle regularization λ · RPn2 for all n = 1, 2, . . . , N .
Neural Networks (Potentials). To approximate potentials {ψ†n, ψ ‡ n} in dimension D, we use
DenseICNN[2; max(64, 2D),max(64, 2D),max(32, D)] with CELU activation function. DenseICNN is an input-convex dense architecture with additional convex quadratic skip connections. Here 2 is the rank of each input-quadratic skip-connection’s Hessian matrix. Each following number max(·, ·) represents the size of a hidden dense layer in the sequantial part of the network. For detailed discussion of the architecture see (Korotin et al., 2019a, Section B.2).
Training process. We perform training according to Algorithm 1 of Appendix A. We set batch size K = 1024 and balancing coefficient γ = 0.2. We use Adam optimizer by (Kingma & Ba, 2014) with a fixed learning rate 10−3. The total number of iterations is set to 50000.
C.4.2 SCALABLE COMPUTATION OF WASSERSTEIN BARYCENTERS (SCW2B)
Generator Neural Network. For the input noise distribution of the generative model we use S = N (0, ID). For the generative network g : RD → RD we use a fully-connected sequential ReLU network with hidden layer sizes
[max(100, 2D),max(100, 2D),max(100, 2D)].
Before the main optimization, we pre-train the network to satisfy g(z) ≈ z for all z ∈ RD. This has been empirically verified as a better option than random initialization of network’s weights.
Neural Networks (Potentials). We used exactly the same networks as in Subsection C.4.1.
Training process. We perform training according to the min-max-min procedure described by (Fan et al., 2020, Algorithm 1). The batch size is set to 1024. We use Adam optimizer by (Kingma & Ba, 2014) with fixed learning rate 10−3 for potentials and 10−4 for generative network g. The number of iterations of the outer cycle (min-max-min) number of iterations is set to 15000. Following (Fan et al., 2020), we use 10 iterations per the middle cycle (min-max-min) and 6 iterations per the inner cycle (min-max-min).
C.4.3 CONTINUOUS REGULARIZED WASSERSTEIN BARYCENTERS (CRWB)
Regularization. [CRWB] method uses regularization to keep the potentials conjugate. The authors impose entropy or L2 regularization w.r.t. some proposal measure P̂; see (Li et al., 2020, Section 3) for more details. Following the source code provided by the authors, we use L2 regularization (empirically shown as a more stable option than entropic regularization). The regularization measure P̂ is set to be the uniform measure on a box containing the support of all the source distributions, estimated by sampling. The regularization parameter is set to 10−4.
Neural Networks (Potentials). To approximate potentials {ψ†n, ψ ‡ n} in dimension D, we use fullyconnected sequential ReLU neural networks with layer sizes given by [max(128, 4D),max(128, 4D),max(128, 4D)].
We have also tried using DenseICNN architecture, but did not experience any performance gain.
Training process. We perform training according to (Li et al., 2020, Algorithm 1). We set batch size to 1024. We use Adam optimizer by (Kingma & Ba, 2014) with fixed learning rate 10−3. The total number of iterations is set to 50000.

title: CONTINUOUS WASSERSTEIN-2 BARYCENTER ESTIMATION WITHOUT MINIMAX OPTIMIZATION