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: $\textbf{Summary and main contributions:}$  (1) The authors propose to use Jensen-Shannon divergence (JS) to learn with noisy labels, which is shown to interpolate between Cross-Entropy and MAE loss through a hyper-parameter.  (2) Based on the new observation that learning with noisy labels reduces the classifier's prediction consistency for noisy-labeled sets. To overcome this issue, the authors introduce multi-distribution generalization (GJS) which is shown to generalize JS with a consistency regularizer.  (3) Experiment on synthetic noisy labeled datasets and a real-world noisy dataset (WebVision) demonstrate the effectiveness of proposed methods.