Wasserstein Distributional Normalization : Nonparametric Stochastic Modeling for Handling Noisy Labels

We propose a novel Wasserstein distributional normalization ( WDN ) algorithm to handle noisy labels for accurate classification . In this paper , we split our data into uncertain and certain samples based on small loss criteria . We investigate the geometric relationship between these two different types of samples and enhance this relation to exploit useful information , even from uncertain samples . To this end , we impose geometric constraints on the uncertain samples by normalizing them into the Wasserstein ball centered on certain samples . Experimental results demonstrate that our WDN outperforms other state-of-the-art methods on the Clothing1M and CIFAR-10/100 datasets , which have diverse noisy labels . The proposed WDN is highly compatible with existing classification methods , meaning it can be easily plugged into various methods to improve their accuracy significantly . 1 INTRODUCTION . The successful results of deep neural networks ( DNNs ) on supervised classification tasks heavily rely on accurate and high-quality label information . However , annotating large-scale datasets is extremely expensive and a time-consuming task . Because obtaining high-quality datasets is very difficult , in most conventional works , training data have been obtained alternatively using crowd-sourcing platforms Yu et al . ( 2018 ) to obtain large-scaled datasets , which leads inevitable noisy labels in the annotated samples . While there are numerous methods that can deal with noisy labeled data , recent methods actively adopt the small loss criterion , which enables to construct classification models that are not susceptible to noise corruption . In this learning scheme , a neural network is trained using easy samples first in the early stages of training . Harder samples are then gradually selected to train mature models as training proceeds . Jiang et al . ( 2018 ) suggested collaborative learning models , in which a mentor network delivers the data-driven curriculum loss to a student network . Han et al . ( 2018 ) ; Yu et al . ( 2019 ) proposed dual networks to generate gradient information jointly using easy samples and employed this information to allow the networks to teach each other . Wei et al . ( 2020 ) adopted a disagreement strategy , which determines the gradient information to update based on disagreement values between dual networks . Han et al . ( 2020 ) implemented accumulated gradients to escape optimization processes from over-parameterization and to obtain more generalized results . In this paper , we tackle to solve major issues raised from the aforementioned methods based on the small-loss criterion , as follows . In comprehensive experiments , the aforementioned methods gain empirical insight regarding network behavior under noisy labels . However , theoretical and quantitative explanation have not been closely investigated . In contrast , we give strong theoretical/empirical explanations to understand the network under noisy labels . In particular , we present an in-depth analysis of small loss criteria in a probabilistic sense . We exploit the stochastic properties of noisy labeled data and develop probabilistic descriptions of data under the small loss criteria , as follows . Let P be a probability measure for the pre-softmax logits of the training samples , l be an objective function for classification , and 1 { · } be an indicator function . Then , our central object to deal with is a truncated measure defined as X ∼ µ|ζ = 1 { X ; l ( X ) > ζ } P P [ l ( X ) > ζ ] , Y ∼ ξ|ζ = 1 { X ; l ( Y ) ≤ζ } P P [ l ( Y ) ≤ ζ ] , ( 1 ) where X and Y , which are sampled from µ|ζ and ξ|ζ , denote uncertain and certain samples defined in the pre-softmax feature space1 ( i.e. , Rd ) , respectively . In equation 1 , µ and ξ denote the probability measures of uncertain and certain samples , respectively , and ζ is a constant . Most previous works have focused on the usage of Y and the sampling strategy of ζ , but poor generalization capabilities based on the abundance of uncertain samples X has not been thoroughly investigated , even though these samples potentially contain important information . To understand the effect of noisy labels on the generalized bounds , we provide the concentration inequality of uncertain measure µ , which renders the probabilistic relation between µ and ξ and learnability of the network under noisy labels . While most conventional methods Han et al . ( 2018 ) ; Wei et al . ( 2020 ) ; Li et al . ( 2019a ) ; Yu et al . ( 2019 ) require additional dual networks to guide misinformed noisy samples , the scalability is not guaranteed due to the existence of dual architectures , which have the same number of parameters as the base network . To alleviate this problem , we build a statistical machinery , which should be fully non-parametric , simple to implement , and computationally efficient to reduce the computational complexity of conventional approaches , while maintaining the concept of small-loss criterion . Based on the empirical observation of ill-behaved certain/uncertain samples , we propose the gradient flow in the Wasserstein space , which can be induced by simulating non-parametric stochastic differential equation ( SDE ) with respect to the Ornstein-Ulenbeck type to control the ill-behaved dynamics . The reason for selecting these dynamics will be thoroughly discussed in the following sections . Thus , key contributions of our work are as follows . • We theoretically verified that there exists a strong correlation between model confidence and statistical distance between X and Y . We empirically investigate that the classification accuracy worsens when the upper-bound of 2-Wasserstein distance W2 ( µ , ξ ) ≤ ε ( i.e. , distributional distance between certain and uncertain samples ) drastically increase . Due to the empirical nature of upper-bound ε , it can be used as an estimator to determine if a network suffers from over-parameterization . • Based on empirical observations , we develop a simple , non-parametric , and computationally efficient stochastic model to control the observed ill-behaved sample dynamics . As a primal object , we propose the stochastic dynamics of gradient flow ( i.e. , , Ornstein-Ulenbeck process ) to simulate simple/non-parametric stochastic differential equation . Thus , our method do not require any additional learning parameters . • We provide important theoretical results . First , the controllable upper-bound ε with the inverse exponential ratio is induced , which indicates that our method can efficiently control the diverging effect of Wasserstein distance . Second , the concentration inequality of transported uncertain measure is presented , which clearly renders the probabilistic relation between µ and ξ . 2 RELATED WORK . Curriculum Learning & Small-loss Criterion . To handle noisy labels , Han et al . ( 2018 ) ; Yu et al . ( 2019 ) ; Jiang et al . ( 2018 ) ; Wei et al . ( 2020 ) ; Lyu & Tsang ( 2020a ) ; Han et al . ( 2020 ) adopted curriculum learning or sample selection frameworks . However , these methods only consider a small number of selected samples , where large portion of samples are excluded at the end of the training . This inevitably leads to poor generalization capabilities . However , this conflicts with sample selection methods because a large portion of training samples are gradually eliminated . By contrast , our method can extract useful information from unselected samples X ∼ µ ( i.e. , uncertain samples ) and enhance these samples ( e.g. , X ′ ∼ Fµ ) for more accurate classification . Chen et al . ( 2019 ) iteratively apply cross-validation to randomly partitioned noisy labeled data to identify most samples that have correct labels . To generate such partitions , they adopt small-loss criterion for selecting samples . Loss Correction & Label Correction . Patrini et al . ( 2017a ) ; Hendrycks et al . ( 2018 ) ; Ren et al . ( 2018 ) either explicitly or implicitly transformed noisy labels into clean labels by correcting classification losses . Unlike these methods , our method transforms the holistic information from uncertain samples into certain samples , which implicitly reduces the effects of potentially noisy labels . While correction of label noisy by modifying the loss-dynamics do not perform well under extreme noise environments , Arazo et al . ( 2019 ) adopt label augmentation method called MixUp Zhang et al . ( 2018 ) . 1Due to the technical difficulties , we define our central objects on pre-softmax space rather than label space , i.e. , the space of σ ( X ) , σ ( Y ) , where σ indicates softmax function . Please refer to Appendix for more details . Distillation . Li et al . ( 2019b ) updated mean teacher parameters by calculating the exponential moving average of student parameters to mitigate the impact of gradients induced by noisy labels . Lukasik et al . ( 2020 ) deeply investigated the effects of label smearing for noisy labels and linked label smoothing to loss correction in a distillation framework . Similar to these methods , our method leverages the useful properties of distillation models . We set ν as a pivot measure , which guides our normalization functional Fµ for uncertain measures . This is similar to self-distillation because uncertain training samples are forced to be normalized to those of past states . Other methods . Lee et al . ( 2019 ) induced a robust generative classifier based on pre-trained deep models . Similar to our method , Damodaran et al . ( 2019 ) designed a constraint on the Wasserstein space and adopted an adversarial framework for classification models of noisy labeled data by implementing semantic Wasserstein distance . Pleiss et al . ( 2020 ) identify noisy labeled samples by considering AUM statistics which exploits differences in training dynamics of clean and mislabeled samples . In most recent work , Li et al . ( 2019a ) adopts semi-supervised learning ( SSL ) methods to deal with noisy labels where the student network utilizes both labeled/unlabeled samples to perform semi-supervised learning guided by the other teacher network . 3 DISTRIBUTIONAL NORMALIZATION . Because our main target object is a probability measure ( distribution ) , we first define an objective function in a distributional sense . Let l be cross entropy and r̂ be a corrupted label random vector for an unknown label transition matrix from a clean label r which is independent of X , with label transition matrix Q . Then , a conventional objective function for classification with noisy labels can be defined as follows : min µ J [ µ ] = min µ EX∼µ , r̂|Q [ l ( X ; r̂ ) ] . ( 2 ) However , due to the significant changes in label information , the conventional objective function defined in equation 2 can not be used for accurate classification . Instead of directly using uncertain samples X ∼ µ as in previous works , we normalize µ in the form of a metric ball and present a holistic constraint . For a clear mathematical description , we first introduce the following definition . Definition 1 . ( Wasserstein ambiguity set ) Let P2 ( Rd ) = { µ : Eµd2E ( x0 , x ) < ∞ , ∀x0 ∈ Rd } be a 2-Wasserstein space , where d denotes the number of classes , dE is Euclidean distance defined on Rd . Then , we define a Wasserstein ambiguity set ( i.e. , metric ball ) in this space as follows : BW2 ( ν , ε ) = { µ ∈ P2 ( Rd ) : W2 ( µ , ν ) ≤ ε } , ( 3 ) whereW2 denotes the 2-Wasserstein distance and ν is the pivot measure . Then , we propose a new objective function by imposing geometric constraints on µ as follows : min Fµ∈BW2 ( ν , ε ) , ξ J [ Fµ ] + J [ ξ ] = min θ EX∼Fµθ , r̂ [ l ( X ; r̂ ) ] + EX∼ξθ , r̂ [ l ( Y ; r̂ ) ] , ( 4 ) where F : P2 ( Rd ) → P2 ( Rd ) is a functional for probability measures , which assures the constraint on Fµ ( i.e. , Fµ ∈ BW2 ( ν , ε ) ) and our main objective . The right-hand side of equation equation 4 is equivalent vectorial form of distributional form in left-hand side . While our main objects are defined on pre-softmax , both probability measures µθ and ξθ is parameterized by neural network with parameters θ . This newly proposed objective function uses the geometrically enhanced version of an uncertain measure Fµ with a certain measure ξ . In equation 4 , probability measure ν is defined as follows : ν = arg minJ [ ξk ? ] , where ξk denotes a certain measure at the current k-th iteration and k ? ∈ Ik−1 = { 1 , · · · , k − 1 } . In other words , our method finds the best probability measure that represents all certain samples so far at training time , where the uncertain measures are transported to be lying in the Wasserstein ball centered on ν . In equation 4 , the Wasserstein constraint on Fµ enforces uncertain measures statistically resemble ν from a geometric perspective ( i.e. , W2 ( ν , Fµ ) ≤ ε ) . Now , an important question naturally stems from the aforementioned analysis : how can we select the optimal radius ε ? Clearly , finding an F that induces a small ε ≈ 0 is suboptimal because Fµ ≈ ν and using objective function J [ Fµ ≈ ν ] can lead to the following critical problem . As the optimization process proceeds , enhanced uncertain samples X ′ ∼ Fµ contribute less and less , because it is statistically identical to ν , meaning our objective in equation 4 would receive little benefits from these transported uncertain samples . By contrast , if we adopt a large radius for ε , enhanced uncertain samples will be statistically and geometrically unrelated to ν , which causes the normalized measure Fµ to yield large losses and violates our objective . To overcome two problems above and select the radius , we make a detour , i.e. , a Gaussian measure , for cutting the path between ν and Fµ ( i.e. , ν → N ( mν , Σν ) → Fµ ) rather than directly calculating the geodesic between ν and Fµ ( i.e. , ν → Fµ ) . Specifically , we decompose the original constraint in equation 4 into two terms using the triangle inequality of the Wasserstein distance : W2 ( ν , Fµ ) ≤ ε =W2 ( ν , N ( mν , Σν ) ) ︸ ︷︷ ︸ d1 : Intrinsic statistics +W2 ( N ( mν , Σν ) , Fµ ) ︸ ︷︷ ︸ d2 : Wasserstein Normalization . ( 5 ) The first intrinsic statistics term sets a detour point as a Gaussian measure , for which the mean and covariance are the same as those for ν ( i.e. , mν = EY∼ν [ Y ] and Σν = CovY∼ν [ Y ] ) . The Wasserstein upper bound of this term is only dependent on the statistical structure of ν because ( mν , Σν ) is dependent on ν . Thus , this term induces a data-dependent , non-zero constant upper bound whenever ν 6= N and can prevent the upper-bound from collapsing to ε → 0 , regardless of F . This gives huge advantage when dealing with ε because the first term can be considered a fixed constant during the training . The second normalization term represents our central objective . F facilitates geometric manipulation in the Wasserstein space and prevent uncertain measure µ from diverging , where µ is normalized onto the Wasserstein ambiguity BW2 ( ν , ε ) in Fig1 . The theoretical/numerical advantages of setting detour measure as Gaussian is well-explained following section .

Summary: The paper introduces a novel objective function by imposing geometric constraints on the logits of uncertain samples. The authors' approach is to map the distribution logits of uncertain samples onto the 2-Wasserstein ball centered on the measure of certain samples. To overcome the dilemma of selecting the ball radius, the authors propose a surrogate objective, namely Wasserstein Normalization. An SDE grad flow is proposed for solving the Wasserstein normalization. The paper also keeps the Gaussian parameters as moving average during training in light of batch normalization. The paper both theoretically and empirically validate their method.