Generalized Jensen-Shannon Divergence Loss for Learning with Noisy Labels

1 Introduction . Labeled datasets , even the systematically annotated ones , contain noisy labels [ 1 ] . Therefore , designing noise-robust learning algorithms are crucial for the real-world tasks . An important avenue to tackle noisy labels is to devise noise-robust loss functions [ 2 , 3 , 4 , 5 ] . Similarly , in this work , we propose two new noise-robust loss functions based on two central observations as follows . Observation I : Provably-robust loss functions can underfit the training data [ 2 , 3 , 4 , 5 ] . Observation II : Standard networks show low consistency around noisy data points 1 , see Figure 1 . We first propose to use Jensen-Shannon divergence ( JS ) as a loss function , which we crucially show interpolates between the noise-robust mean absolute error ( MAE ) and the cross entropy ( CE ) that better fits the data through faster convergence . Figure 2 illustrates the CE-MAE interpolation . Regarding Observation II , we adopt the generalized version of Jensen-Shannon divergence ( GJS ) to encourage predictions on perturbed inputs to be consistent , see Figure 3 . Notably , Jensen-Shannon divergence has previously shown promise for test-time robustness to domain shift [ 6 ] , here we further argue for its training-time robustness to label noise . The key contributions of this work2 are : • We make a novel observation that a network predictions ’ consistency is reduced for noisylabeled data when overfitting to noise , which motivates the use of consistency regularization . • We propose using Jensen-Shannon divergence ( JS ) and its multi-distribution generalization ( GJS ) as loss functions for learning with noisy labels . We relate JS to loss functions that are based on the noise-robustness theory of Ghosh et al . [ 2 ] . In particular , we prove that JS generalizes CE and MAE . Furthermore , we prove that GJS generalizes JS by incorporating consistency regularization in a single principled loss function . • We provide an extensive set of empirical evidences on several datasets , noise types and rates . They show state-of-the-art results and give in-depth studies of the proposed losses . 1we call a network consistent around a sample ( x ) if it predicts the same class for x and its perturbations ( x̃ ) . 2implementation available at https : //github.com/ErikEnglesson/GJS 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 2 Generalized Jensen-Shannon Divergence . We propose two loss functions , the Jensen-Shannon divergence ( JS ) and its multi-distribution generalization ( GJS ) . In this section , we first provide background and two observations that motivate our proposed loss functions . This is followed by definition of the losses , and then we show that JS generalizes CE and MAE similarly to other robust loss functions . Finally , we show how GJS generalizes JS to incorporate consistency regularization into a single principled loss function . We provide proofs of all theorems , propositions , and remarks in this section in Appendix C . 2.1 Background & Motivation . Supervised Classification . Assume a general function class3 F where each f ∈ F maps an input x ∈ X to the probability simplex ∆K−1 , i.e . to a categorical distribution over K classes y ∈ Y = { 1 , 2 , . . . , K } . We seek f∗ ∈ F that minimizes a risk RL ( f ) = ED [ L ( e ( y ) , f ( x ) ) ] , for some loss function L and joint distribution D over X×Y , where e ( y ) is a K-vector with one at index y and zero elsewhere . In practice , D is unknown and , instead , we use S = { ( xi , yi ) } Ni=1 which are independently sampled from D to minimize an empirical risk 1N ∑N i=1 L ( e ( yi ) , f ( xi ) ) . Learning with Noisy Labels . In this work , the goal is to learn from a noisy training distribution Dη where the labels are changed , with probability η , from their true distribution D. The noise is called instance-dependent if it depends on the input , asymmetric if it dependents on the true label , and symmetric if it is independent of both x and y . Let f∗η be the optimizer of the noisy distribution risk RηL ( f ) . A loss function L is then called robust if f∗η also minimizes RL . The MAE loss ( LMAE ( e ( y ) , f ( x ) ) : = ‖e ( y ) − f ( x ) ‖1 ) is robust but not CE [ 2 ] . Issue of Underfitting . Several works propose such robust loss functions and demonstrate their efficacy in preventing noise fitting [ 2 , 3 , 4 , 5 ] . However , all those works have observed slow convergence of such robust loss functions leading to underfitting . This can be contrasted with CE that has fast convergence but overfits to noise . Ghosh et al . [ 2 ] mentions slow convergence of MAE and GCE [ 3 ] extensively analyzes the undefitting thereof . SCE [ 4 ] reports similar problems for the reverse cross entropy and proposes a linear combination with CE . Finally , Ma et al . [ 5 ] observe the same problem and consider a combination of “ active ” and “ passive ” loss functions . Consistency Regularization . This encourages a network to have consistent predictions for different perturbations of the same image , which has mainly been used for semi-supervised learning [ 7 ] . Motivation . In Figure 1 , we show the validation accuracy and a measure of consistency during training with the CE loss for varying amounts of noise . First , we note that training with CE loss eventually overfits to noisy labels . Figure 1a , indicates that the higher the noise rate , the more accuracy drop when it starts to overfit to noise . Figure 1 ( b-c ) shows the consistency of predictions for correct and noisy labeled examples of the training set , with the consistency measured as the ratio of examples that have the same class prediction for two perturbations of the same image , see 3e.g . softmax neural network classifiers in this work Appendix B.6 for more details . A clear correlation is observed between the accuracy and consistency of the noisy examples . This suggests that maximizing consistency of predictions may improve the robustness to noise . Next , we define simple loss functions that ( i ) encourage consistency around data points and ( ii ) alleviate the “ issue of underfitting ” by interpolating between CE and MAE . 2.2 Definitions . DJS . Let p ( 1 ) , p ( 2 ) ∈ ∆K−1 have corresponding weights π = [ π1 , π2 ] T ∈ ∆ . Then , the JensenShannon divergence between p ( 1 ) and p ( 2 ) is DJSπ ( p ( 1 ) , p ( 2 ) ) : = H ( m ) − π1H ( p ( 1 ) ) − π2H ( p ( 2 ) ) = π1DKL ( p ( 1 ) ‖m ) + π2DKL ( p ( 2 ) ‖m ) ( 1 ) with H the Shannon entropy , and m = π1p ( 1 ) + π2p ( 2 ) . Unlike Kullback–Leibler divergence ( DKL ( p ( 1 ) ‖p ( 2 ) ) ) or cross entropy ( CE ) , JS is symmetric , bounded , does not require absolute continuity , and has a crucial weighting mechanism ( π ) , as we will see later . DGJS . Similar to DKL , DJS satisfies DJSπ ( p ( 1 ) , p ( 2 ) ) ≥ 0 , with equality iff p ( 1 ) = p ( 2 ) . For DJS , this is derived from Jensen ’ s inequality for the concave Shannon entropy . This property holds for finite number of distributions and motivates a generalization of DJS to multiple distributions [ 8 ] : DGJSπ ( p ( 1 ) , . . . , p ( M ) ) : = H ( M∑ i=1 πip ( i ) ) − M∑ i=1 πiH ( p ( i ) ) = M∑ i=1 πiDKL ( p ( i ) ∥∥∥ M∑ j=1 πjp ( j ) ) ( 2 ) where M is the number of distributions , and π = [ π1 , . . . , πM ] T ∈ ∆M−1 . Loss functions . We aim to use DJS and DGJS divergences , to measure deviation of the predictive distribution ( s ) , f ( x ) , from the target distribution , e ( y ) . Without loss of generality , hereafter , we dedicate p ( 1 ) to denote the target distribution . JS loss , therefore , can take the form of DJSπ ( e ( y ) , f ( x ) ) . Generalized JS loss is a less straight-forward construction since DGJS can accommodate more predictive distributions . While various choices can be made for these distributions , in this work , we consider predictions associated with different random perturbations of a sample , denoted by A ( x ) . This choice , as shown later , implies an interesting analogy to consistency regularization . The choice , also entails no distinction between the M − 1 predictive distributions . Therefore , we consider π2 = · · · = πM = 1−π1M−1 in all our experiments . Finally , we scale the loss functions by a constant factor Z = − ( 1−π1 ) log ( 1−π1 ) . As we will see later , the role of this scaling is merely to strengthen the already existing and desirable behaviors of these losses as π1 approaches zero and one . Formally , we have JS and GJS losses : LJS ( y , f , x ) : = DJSπ ( e ( y ) , f ( x̃ ) ) Z , LGJS ( y , f , x ) : = DGJSπ ( e ( y ) , f ( x̃ ( 2 ) ) , . . . , f ( x̃ ( M ) ) ) Z ( 3 ) with x̃ ( i ) ∼ A ( x ) . Next , we study the connection between JS and losses which are based on the robustness theory of Ghosh et al . [ 2 ] . 2.3 JS ’ s Connection to Robust Losses . Cross Entropy ( CE ) is the prevalent loss function for deep classifiers with remarkable successes . However , CE is prone to fitting noise [ 9 ] . On the other hand , Mean Absolute Error ( MAE ) is theoretically noise-robust [ 2 ] . Evidently , standard optimization algorithms struggle to minimize MAE , especially for more challenging datasets e.g . CIFAR-100 [ 3 , 5 ] . Therefore , there have been several proposals that combine CE and MAE , such as Generalized CE ( GCE ) [ 3 ] , Symmetric CE ( SCE ) [ 4 ] , and Normalized CE ( NCE+MAE ) [ 5 ] . The rationale is for CE to help with the learning dynamics of MAE . Next , we show JS has CE and MAE as its asymptotes w.r.t . π1 . Proposition 1 . Let p ∈ ∆K−1 , then lim π1→0 LJS ( e ( y ) , p ) = H ( e ( y ) , p ) , lim π1→1 LJS ( e ( y ) , p ) = 1 2 ‖e ( y ) − p‖1 where H ( e ( y ) , p ) is the cross entropy of e ( y ) relative to p. Figure 2 depicts how JS interpolates between CE and MAE for π1 ∈ ( 0 , 1 ) . The proposition reveals an interesting connection to state-of-the-art robust loss functions , however , there are important differences . SCE is not bounded ( so it can not be used in Theorem 1 ) , and GCE is not symmetric , while JS and MAE are both symmetric and bounded . In Appendix B.3 , we perform a dissection to better understand how these properties affect learning with noisy labels . GCE is most similar to JS and is compared further in Appendix B.4 . A crucial difference to these other losses is that JS naturally extends to multiple predictive distributions ( GJS ) . Next , we show how GJS generalizes JS by incorporating consistency regularization .

Summary: The paper proposed to use JS divergence as a loss function (JS) to address the underfit problem of provably-robust loss functions. Further, the authors generalize the JS loss function (GJS) to include multiple distributions to encourage consistency during learning with noisy labels. They demonstrate that JS interpolates MAS and CE loss, and they proved the noisy robustness property of the JS and GJS. Extensive experiments are provided to show the effectiveness of the proposed loss functions.