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: The paper proposes a framework based on Wasserstein barycenter to ensemble learning models for a multiclass or a multilabel learning problem. The paper has theoretically shown that the model ensembling using Wasserstein barycenters preserves accuracy, and has a higher entropy than the individual models. Experimental results in the context of attribute-based classification, multilabel learning, and image captioning generation have shown the effectiveness of Wasserstein-based ensembling in comparison to geometric or arithmetic mean ensembling.