Optimal Transport for Long-Tailed Recognition with Learnable Cost Matrix

1 INTRODUCTION . Classification problems in the real world are generally challenged by the long-tailed label distribution , i.e. , having a small number of samples for a majority of labels , and a dominant number of samples for a minority of labels ( Van Horn & Perona , 2017 ; Buda et al. , 2018 ; Liu et al. , 2019 ) on the training data set . It is also known as imbalanced recognition , which has been widely studied in the past decades ( Cardie & Howe , 1997 ; Chawla et al. , 2002 ; Qiao & Liu , 2009 ; Cui et al. , 2019 ) . These distribution biases pose a significant challenge to predictive modeling ; conceivably , models often suffer from poor generalisation and undesirable estimation bias ( Cao et al. , 2019 ; Kang et al. , 2019 ; Zhou et al. , 2020 ) . Recently , a renewed interest in the problem of long-tail recognition has emerged following the context of neural networks , as numerous publications in the literature endeavour to resolve the problem albeit in different ways including decouple ( Kang et al. , 2019 ) , meta-learning ( Ren et al. , 2020 ; Wang et al. , 2020a ; Li et al. , 2021 ) , post-hoc correction ( Tang et al. , 2020 ; Hong et al. , 2020 ) , etc ( Liu et al. , 2019 ; Cao et al. , 2019 ; Tang et al. , 2020 ) . One of the representative methods of post-hoc correction , Logit Adjustment Menon et al . ( 2020 ) , provides a statistical correction to the prediction , receiving widespread attention for its simplicity and validity , but the downside is that its optimisation is conducted on individual samples , without any indication that rectified marginal distribution satisfies the desired distribution we are striving for . Figuring out how to solve the above problem , our explicit modeling of the problem mathematically turns into an equational constraint , meanwhile to minimise the difference between refined distribution and the original one , this minimisation is motivated upon the inner-product similarity . A little further , the resulting problem can be linked to OT . Drawing on this linkage , we develop it further by proposing a linear mapping to automatically learn cost matrix , thereby circumventing the requirement for expert knowledge to configure this matrix . In short , we contribute as follows : • We propose an alternative direction based on convex optimisation to do post-hoc correction , which goes beyond previous direction from the statistical view . • Imposing marginal distributions to align ideal marginal distributions , we derive an optimisation problem tied to the OT that is solved using Sinkhorn , a highly efficient algorithm . More further , for better learning of the cost matrix , we present a linear mapping enabling very elegant learning with one-layer neural network . • The experimental evidence shows the high efficiency of the algorithm , and the best performance achieved on three benchmarks , also supports our claim that addressing the post-hoc problem via OT is helpful and effective . 2 PRELIMINARIES . In this section , we begin by giving a notational definition , followed by an introduction to the longtailed recognition problem , and finally , we briefly review the OT , and give a quick recap of Logit Adjustment Menon et al . ( 2020 ) . Notations : In what follows , for two matrices X , Y ∈ RN×K , we denote 〈X , Y 〉 =∑N n=1 ∑K k=1XnkYnk as the Frobenius dot-product . δ ( · ) stands for the Dirac function , p ( · ) represents the probability distribution . U ( r , c ) = { P ∈ RN×K+ |P1K = r , P ᵀ1N = c } , where 1N and 1K are N -dimension and K-dimension vector whose elements are all 1. r and c refer to the vectors of size N and K respectively , U ( r , c ) include all matrices with row and column sums r and c respectively 2.1 PROBLEM FORMULATION . Having a collection of training samples { ( xsn , ysn ) } Ns n=1 , validation samples { ( xvn , yvn ) } Nv n=1 and test samples { ( xtn , yn ) t } Nt n=1 for classification with K labels and input x ∈ Rd , long-tailed recognition assumes that the class-prior distribution for training data p ( ys ) is different from that for validation data p ( yv ) and test data p ( yt ) . Specifically , long-tailed recognition means the distribution p ( ys ) is highly skewed , some classes have the dominant number of samples , while tailed labels own a very small number of samples . We can use imbalance ratio to measure the skewness in training data set , which can be defined as R = N s max Nsmin , where Nsmax and N s min denote the largest and smallest number of samples in the training data set , respectively . In this paper , we assume that the marginal distribution of the test set is known , we consider it as an implicit prior knowledge to be applied . Stepping back , even if we do not know the marginal distribution of the test dataset in advance . There are still ways to estimate the marginal distribution of the test dataset relatively precisely , such as methods in Hendrycks et al . ( 2018 ) ; Azizzadenesheli et al . ( 2019 ) . Obviously , most models trained on imbalanced training data set would suffer from extremely limited generalisation ability . Hence the ultimate goal is to learn a model that minimises the empirical risk , which can be formulated as : J ( Φ ( xsn ) , ysn ) = 1 Ns Ns∑ n=1 L ( Φ ( xsn ) , ysn ) , ( 1 ) where Φ ( xsn ) ∈ RK denotes logits with associated sample , Φ ( · ) : Rd → RK represents the mapping via neural networks , L stands for the loss function , typically cross entropy loss function for classification problem . 2.2 REMINDERS ON OPTIMAL TRANSPORT . OT is used to calculate the cost of transporting one probability measure to another . We next present a brief introduction to OT with the aim of helping us better view the long-tailed problem from an OT perspective . Consider probability measures r and c of the two random variables X and Y respectively , when the cost function C ( X , Y ) : X × Y → R+ is defined to denote the cost of transporting X to Y , the OT distance can be defined as : d ( r , c ) = min π∈Π ( r , c ) ∫ X×Y C ( x , y ) π ( x , y ) dxdy , ( 2 ) where Π ( r , c ) = { ∫ Y π ( x , y ) dy = r , ∫ X π ( x , y ) dx = c } is the joint probability measure with r and c. When we extend the above continuous case to the discrete situation , we consider following discrete distributions : r = N∑ i=1 pi ( xi ) δ ( xi ) c = K∑ j=1 pi ( yj ) δ ( yj ) ( 3 ) where pi ( xi ) and pi ( yj ) represent the probability mass to the sample xi and yj respectively . In this context , OT distance can be expressed as : dM ( r , c ) = min P∈U ( r , c ) 〈P , M〉 . ( 4 ) where M stands for the cost matrix constructed by Mij = C ( xi , yj ) . The goal of OT is to find a transportation matrix P that minimizes the distance dM ( r , c ) As we can see , OT is a distance measure between two probability distributions under some cost matrix ( Villani , 2008 ) . However , when we use network simplex or interior point methods to solve the above optimization problem , it often comes at the cost of heavy computational demands . To tackle this issue , OT with entropy constraint is proposed to allow the optimisation at small computational cost in sufficient smoothness ( Cuturi , 2013 ) . By adding a Lagrangian multiplier to the entropy constraint , the new formulation can be defined as follows : dλM ( r , c ) = 〈P λ , M〉 where P λ = arg min P∈U ( r , c ) 〈P , M〉 − λh ( P ) , ( 5 ) where λ ∈ [ 0 , +∞ ] , h ( P ) = − ∑N n=1 ∑K k=1Pnk logPnk , d λ M ( r , c ) is also known as dual-Sinkhorn divergence , besides , it can be calculated with matrix scaling algorithms for cheaper computational demand . The following lemma guarantees the convergence and uniqueness of the solution of the algorithm . Lemma 1 For λ > 0 , the solution P λ is unique and has the form P λ = diag ( u ) Kdiag ( v ) , where u and v are two non-negative vectors uniquely defined up to a multiplicative factor andK = e−M/λ is the element-wise exponential of −M/λ . The above lemma states the uniqueness of P λ ( Sinkhorn , 1974 ) , and P λ can be efficiently computed via Sinkhorn ’ s fixed point iteration u , v ← r./Kv , c./Kᵀu . 2.3 A QUICK RECAP OF LOGIT ADJUSTMENT . We give a brief introduction to Logit Adjustment ( Menon et al. , 2020 ; Hong et al. , 2020 ) . For the model Φ ( · ) , it is trained by the standard cross-entropy loss function on imbalanced training data set , and evaluated on test data . In this algorithm , the test logit is adjusted as follows : Φ ( xtn ) = Φ ( x t n ) − log p ( ys ) ( 6 ) This simple procedure is derived from the Bayes optimal rule . It is apparent that Logit Adjustment involves a post hoc correction on an individual sample , which does not necessarily guarantee that the marginal distribution of the whole dataset matches the desired distribution . 3 METHODOLOGY . The first part of this section explores post-hoc correction from an OT perspective , proceeds to the automatic learning of the cost matrix via linear mapping . Lastly , we demonstrate how it can be achieved simply with one-layer neural network . 3.1 POST-HOC CORRECTION FORMALISED FROM AN OT PERSPECTIVE . Since Logit Adjustment applies adjustment at the individual sample level . It doesn ’ t assure that the marginal distribution of the overall data set fulfils our desired distribution . In this respect , we clearly put the constraint into an equation : Y ᵀ1N = µ , ( 7 ) where Y ∈ RN×K indicates the refined prediction value in matrix form , µ represents the expected distribution on the test set . Alternatively , it is desirable to preserve another characteristic of Y , namely , remaining almost as similar to the original prediction as possible . We consider inner-product based similarity to measure the similarity , which is a straightforward yet useful similarity measure . In mathematical terms , this can be expressed as follows : maximize Y 〈C ( Ẑ ) , Y 〉 , ( 8 ) where Ẑ represents the original prediction in matrix form , C ( · ) denotes to perform some transformation to Ẑ , it can be some simple function , like Logarithmic function log ( z ) , exponential function zα . Here we select − log ( · ) as the cost function . This choice was driven by the requirement that the cost matrix must be positive definite , whereas the transformation of the original prediction by − log ( · ) satisfies this condition , in addition , as log likelihood represents the local probability density of the associated samples , it can also be used to substitute Ẑ for the similarity approximation . In brief , the resulting numerical form can be put in formal terms as follows : minimize Y 〈− log ( Ẑ ) , Y 〉 ( 9 ) subject to Y ᵀ1N = µ , Y 1K = 1N . ( 10 ) Extra constraint on Y are imposed simply cos the tuned estimation has to fulfil the basic probabilistic requirement that its sum is one . Comparing Eq . ( 9-10 ) with Eq . ( 4 ) , we can see that if we substitute P with Y and substitute r and c with 1N and µ respectively , we find that the above optimization problem is actually a special case of OT . In preliminaries , the entropy regularised OT is introduced ( EOT ) , by adding entropy regularisation to OT , the given equation can be solved efficiently by Sinkhorn algorithm . Specifically , the optimisation problem is : minimize Y 〈− log ( Ẑ ) , Y 〉+ λY log ( Y ) ( 11 ) subject to Y ᵀ1N = µ , Y 1K = 1N . ( 12 ) The associated algorithmic flow is outlined in detail in Algorithm 1 . Algorithm 1 : Solve OT-related algorithm efficiently in the post-hoc correction via Sinkhorn Algorithm . Input : Cost matrixM , trade-off parameter λ , max number of iterations NT , iteration number t , error threshold , current error δ , two probability vectors r and c , |·| denotes the vector norm . Result : Refined predictions Y 1 InitialiseK = e−M/λ , uold = 1N , v = 1K , t = 0 ; 2 while t ≤ NT and δ ≤ do 3 u = r./Kv ; 4 v = c./Kᵀu ; 5 σ = |uold − u| ; 6 uold = u ; 7 t = t+ 1 ; 8 end 9 Output Y = diag ( u ) Kdiag ( v ) Assigning λ with 1 , it is observed that we equate our objective function to the KL divergence , thus illustrating the extensive nature of our approach . DARP ( Kim et al. , 2020a ) has previously applied it to long-tailed semi-supervised classification . Remark We would like to illustrate the non-applicable scenarios of our method . Firstly , our method requires a large number of samples for evaluation . This is because if the batch size is small , we can not guarantee that the desired marginal distribution can be satisfied within the batch . In some online scenarios , the sample-wise correction method is more suitable . In addition , our method assumes that the marginal distribution is already known , and in the paper we assume in particular that it is consistent with a uniform distribution .

Summary: This paper contributes an extension of post hoc correction of long-tailed recognition with Optimal Transport (OT). Unlike the previous work (e.g. logit adjustment) which focuses on sample-wise correction, this work, on the other hand, considers the marginal distribution of the overall data for correction. The method is further extended by learning a cost matrix. Their experiments show the effectiveness of optimal transport in the long-tail recognition problem via comprehensive comparison with previous works in terms of performance and efficiency.