Wasserstein Barycenter Model Ensembling

1 INTRODUCTION . Model ensembling consists in combining many models into a stronger , more robust and more accurate model . Ensembling is ubiquitous in machine learning and yields improved accuracies across multiple prediction tasks such as multi-class or multi-label classification . For instance in deep learning , output layers of Deep Neural Networks ( DNNs ) , such as softmaxes or sigmoids , are usually combined using a simple arithmetic or geometric mean . The arithmetic mean rewards confidence of the models while the geometric means seeks the consensus across models . What is missing in the current approaches to models ensembling , is the ability to incorporate side information such as class relationships represented by a graph or via an embedding space . For example a semantic class can be represented with a finite dimensional vector in a pretrained word embedding space such as GloVe ( Pennington et al. , 2014 ) . The models ’ predictions can be seen as defining a distribution in this label space defined by word embeddings : if we denote pi to be the confidence of a model on a bin corresponding to a word having an embedding xi , the distribution on the label space is therefore p = ∑ i piδxi . In order to find the consensus between many models predictions , we propose to achieve this consensus within this representation in the label space . In contrast to arithmetic and geometric averaging , which are limited to the independent bins ’ confidence , this has the advantage of carrying the semantics to model averaging via the word embeddings . More generally this semantic information can be encoded via cost a matrix C , where Cij encodes the dissimilarity between semantic classes i and j , and C defines a ground metric on the label space . To achieve this goal , we propose to combine model predictions via Wasserstein ( W. ) barycenters ( Agueh & Carlier , 2011 ) , which enables us to balance the confidence of the models and the semantic side information in finding a consensus between the models . Wasserstein distances are a naturally good fit for such a task , since they are defined with respect to a ground metric in the label space of the models , which carry such semantic information . Moreover they enable the possiblity of ensembling predictions defined on different label sets , since the Wasserstein distance allows to align and compare those different predictions . Since their introduction in ( Agueh & Carlier , 2011 ) W. barycenter computations were facilitated by entropic regularization ( Cuturi , 2013 ) and iterative algorithms that rely on iterative Bregman projections ( Benamou et al. , 2015 ) . Many applications have used W. barycenters in Natural Language Processing ( NLP ) , clustering and graphics . We show in this paper that W. barycenters are effective in model ensembling and in finding a semantic consensus , and can be applied to a wide range of problems in machine learning ( Table 1 ) . The paper is organized as follows : In Section 2 we revisit geometric and arithmetic means from a geometric viewpoint , showing that they are ` 2 and Kullback Leibler divergence KL ( extended KL divergence ) barycenters respectively . We give a brief overview of optimal transport metric and W. barycenters in Section 3 . We highlight the advantages of W. barycenter ensembling in terms of semantic smoothness and diversity in Section 4 . Related work on W. barycenters in Machine learning are presented in Section 5 . Finally we show applications of Wasserstein ensembling on attribute based classification , multi-label learning and image captioning in Section 6 . 2 WASSERSTEIN BARYCENTERS FOR MODEL ENSEMBLING . Normalized and Unnormalized predictions Ensembling . In deep learning , predictions on a label space of fixed size M are usually in one of two forms : a ) normalized probabilities : in a multiclass setting , the neural network outputs a probability vector ( normalized through softmax ) , where each bin corresponds to a semantic class ; b ) unnormalized positive scores : in a multi-label setting , the outputs of M independent logistic units are unnormalized positive scores , where each unit corresponds to the presence or the absence of a semantic class . Model ensembling in those two scenarios has long history in deep learning and more generally in machine learning ( Breiman , 1996 ; Freund & Schapire , 1999 ; Wolpert , 1992 ) as they lead to more robust and accurate models . As discussed in the introduction , two methods have been prominent in model ensembling due to their simplicity : majority vote using the arithmetic mean of predictions , or consensus based using the geometric mean . Revisiting Arithmetic and Geometric Means from a geometric viewpoint . Givenm predictions µ ` , and weights λ ` ≥ 0 such that ∑m ` =1 λ ` = 1 , the weighted arithmetic mean is given by µ̄a =∑m ` =1 λ ` µ ` , and the weighted geometric mean by µ̄g = Π m ` =1 ( µ λ ` ` ) . It is instructive to reinterpret the arithmetic and geometric mean as weighted Frechet means ( Definition 1 ) ( Zemel & Panaretos , 2017 ) . Definition 1 . [ Weighted Frechet Mean ] Given a distance d and { ( λ ` , µ ` ) , λ ` > 0 , µ ` ∈ RM+ } ` =1 ... m , the Frechet mean is defined as follows : µ̄ = arg minρ ∑m ` =1 λ ` d ( ρ , µ ` ) . It is easy to prove ( Appendix F ) that the arithmetic mean corresponds to a Frechet mean for d = ‖.‖22 ( the ` 2 Euclidean distance ) . A less known fact is that the geometric mean corresponds to a Frechet Mean for d = K̃L , where K̃L is the extended KL divergence to unnormalized measures : K̃L ( p , q ) = ∑ i pi log ( pi qi ) − pi + qi . We give proofs and properties of arithmetic and geometric mean in Appendix F. Following this geometric viewpoint , in order to incorporate the semantics of the target space in model ensembling , we need to use a distance d that takes advantage of the underlying geometry of the label space via a cost matrix C when comparing positive measures . Optimal transport ( OT ) metrics such as Wasserstein-2 have this property since they are built on an explicit cost matrix defining pairwise distance between the semantic classes . In this paper we propose to use the Frechet means with Wasserstein distance ( d = W 22 ) for model ensembling , i.e . use Wasserstein barycenters ( Agueh & Carlier , 2011 ) for model ensembling : µ̄w = arg min ρ m∑ ` =1 λ ` W 2 2 ( ρ , µ ` ) . Intuitively , the barycenter looks for a distribution ρ ( a histogram ) that is close to all the base distributions µ ` in the Wasserstein sense . In our context transporting the consensus ρ to each individual model µ ` should have a minimal cost , where the cost is defined by the distance in the word embedding space . 3 WASSERSTEIN BARYCENTERS . Wasserstein distances were originally defined between normalized probability vectors ( Balanced OT ) ( Villani , 2008 ; Peyré & Cuturi , 2017 ) , but they have been extended to deal with unnormalized measures and this problem is referred to as unbalanced OT ( Chizat et al. , 2018 ; Frogner et al. , 2015 ) . Motivated by the multi-class and the multi-label ensembling applications , in the following we present a brief overview of W. barycenters in the balanced and unbalanced cases . 3.1 OPTIMAL TRANSPORT METRICS . Balanced OT . Given p ∈ ∆N , where ∆N = { p ∈ RN , pk ≥ 0 , ∑N k=1 pk = 1 } , p represents histograms on source label space ΩS = { xi ∈ Rd , i = 1 . . . N } , for e.g words embeddings . Consider similarly q ∈ ∆M representing histograms whose bins are defined on a target label space ΩT = { yj ∈ Rd , j = 1 . . .M } . Consider a cost function c ( x , y ) , ( for example c ( x , y ) = ‖x− y‖2 ) . Let C be the matrix in ∈ RN×M such that Cij = c ( xi , yj ) . 1N denotes a vector with all ones . Let γ ∈ RN×M be a coupling matrix whose marginals are p and q such that : γ ∈ Π ( p , q ) = { γ ∈ RN×M , γ1M = p , γ > 1N = q } . The optimal transport metric is defined as follows : W ( p , q ) = min γ∈Π ( p , q ) { 〈C , γ〉 = ∑ ij Cijγij . } ( 1 ) When c ( x , y ) = ‖x− y‖22 , this distance corresponds to the so called Wasserstein−2 distance W 22 . Unbalanced OT . When p and q are unnormalized and have different total masses , optimal transport metrics have been extended to deal with this unbalanced case . The main idea is in relaxing the set Π ( p , q ) using a divergence such as the extended KL divergence : K̃L . ( Chizat et al. , 2018 ) define for λ > 0 the following generalized Wasserstein distance between unnormalized measures : Wunb ( p , q ) = min γ 〈C , γ〉+ λK̃L ( γ1M , p ) + λK̃L ( γ > 1N , q ) . ( 2 ) 3.2 BALANCED AND UNBALANCED WASSERSTEIN IN MODELS ENSEMBLING . Throughout the paper we consider m discrete prediction vectors µ ` ∈ RN ` + , ` = 1 . . .m defined on a discrete space ( word embeddings for instance ) ΩS ` = { x ` i ∈ Rd , i = 1 . . . N ` } . We refer to ΩS ` as source spaces . Our goal is to find a consensus prediction µ̄w ∈ RM+ defined on a discrete target space ΩT = { yj ∈ Rd , j = 1 . . .M } . Let C ` ∈ RN ` ×M be the cost matrices , C ` , i , j = c ( x ` i , yj ) . Balanced W. Barycenters : Normalized predictions . The W. barycenter ( Agueh & Carlier , 2011 ) of normalized predictions is defined as follows : µ̄w = arg minρ ∑m ` =1 λ ` W ( ρ , µ ` ) , for the Wasserstein distance W defined in equation ( 1 ) . Hence one needs to solve the following problem , for m coupling matrices γ ` , ` = 1 . . .m : min ρ min γ ` ∈Π ( µ ` , ρ ) , ` =1 ... m m∑ ` =1 λ ` 〈C ` , γ ` 〉 . ( 3 ) Unbalanced W. Barycenters : Unnormalized predictions . Similarly the W. barycenter of unnormalized predictions is defined as follows : µ̄w = arg minρ ∑m ` =1 λ ` Wunb ( ρ , µ ` ) , for the Generalized Wasserstein distance Wunb defined in equation ( 2 ) . Hence the unbalanced W. barycenter problem ( Chizat et al. , 2018 ) amounts to solving , for m coupling matrices γ ` , ` = 1 . . .m : min ρ min γ ` , ` =1 ... m m∑ ` =1 λ ` ( 〈C ` , γ ` 〉+ λK̃L ( γ ` 1M , µ ` ) + λK̃L ( γ > ` 1N ` , ρ ) ) . ( 4 )

Summary: This paper has a simple message. When predicting families (weight vectors) of labels, it makes sense to use an ensemble of predictors and average them using a Wasserstein barycenter, where the ground metric is defined using some a priori knowledge on the labels, here usually distances between word embeddings or more elaborate metrics (or kernels K, as described in p.8). Such barycenters can be easily computed using an algorithm proposed by Benamou et al. 18. When these histograms are not normalized (e.g. their count vectors do not sum to the same quantity) then, as shown by Frogner, Zhang et al, an alternative penalized formulation of OT can be studied, solved numerically with a modified Sinkhorn algorithm, which also leads to a simple W barycenter algorithm as shown by Chizat et al.