Constrained Optimization to Train Neural Networks on Critical and Under-Represented Classes

1 Introduction . Deep Neural Networks ( DNNs ) perform extremely well in many classification tasks when sufficiently large and representative datasets are available for training . However , in many real world applications , it is not uncommon to encounter highly-skewed class distributions , i.e. , majority of the data belong to only a few classes while some classes are represented with scarce instances . Training DNNs on such imbalanced datasets leads to models that are biased toward majority classes with poor prediction accuracy for samples of minority class . While this is problematic for all such applications , it poses an even greater issue for “ critical ” applications where misclassifying samples belonging to the minority class can have severe consequences . One domain where such applications are common and machine learning is having an important impact is medical imaging . In medical imaging , applications with data imbalance are ubiquitous [ 14 ] and costs of making some types of mistakes are more severe than others . For instance , in a diagnosis application , discarding a 1Code is available at : https : //github.com/salusanga/alm-dnn . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . cancer case as healthy ( False Negative ) is more costly than classifying a healthy subject as having cancer ( False Positive ) . While the latter creates burden for the subject as well as health-care system through additional tests that may be invasive and expensive , the former , i.e . failure to identify a cancerous case , would delay the diagnosis and jeopardise the treatment success . In such applications , binary classifiers are operated at high True Positive Rates ( TPRs ) even when this means having higher False Positive Rates ( FPRs ) . To make matters more complicated , there are usually significantly fewer samples to represent critical classes , where mistakes are more severe . For instance , in [ 6 ] authors found out that in prostate cancer screening only 30 % of even the most suspicious cases identified with initial testing actually have cancer . Such class imbalance increases the FPR even higher in “ critical ” applications , because the models tend to misclassify minority classes more often . Useful algorithms need to achieve low FPR at high TPR operating points , even under class imbalance . While various methods for learning with imbalanced datasets exist , to the best of our knowledge , these methods do not take into account the fact that “ critical ” applications need to be operated at high accuracy for the critical classes . We believe that for such applications ensuring low misclassification rate for the non-critical samples and high accuracy for the critical classes should be the main goal , to make binary classifiers useful in practice . This motivates us to design new strategies for training DNNs for classification . Contribution : In this paper , we pose the training of a DNN for binary classification under class imbalance as a constrained optimization problem and propose a novel constraint that can be used with existing loss functions . We define the constraint using Mann-Whitney statistics [ 16 ] in order to maximize the AUC , but in an asymmetric way to favor reduction of false positives at high true positive ( or low false negative ) rates . Then , we transfer the constrained problem to its dual unconstrained optimization problem using an Augmented Lagrangian method ( ALM ) [ 2 ] . We optimize the resulting loss function using stochastic gradient descent . Unlike the existing methods that directly optimize AUC , we incorporate AUC optimization in a principled way into a constrained optimization framework . We finally present two possible extensions of the proposed constraint for multi-class classification problems . We present an extensive evaluation of the proposed method for image-based binary and multi-class classification problems on three datasets : an in-house medical dataset for prostate cancer , CIFAR10 , and CIFAR100 [ 11 ] . In all datasets , we perform experiments by simulating different class imbalance ratios . In our experiments , we apply the proposed constraint to 9 different baseline loss functions , most of which were proposed to handle class imbalance . We compare the results with the baselines without any constraint . The results demonstrate that the proposed method improves the baselines in majority of the cases . 2 Related work . Various methods have already been proposed to learn better models with class-imbalanced datasets . We group the existing methods into three categories : cost sensitive training-based methods , samplingbased and classifier-based methods . Here , we focus on the first one and present related work for the other groups in the supplementary material for space reasons . Cost sensitive training-based methods : This family of methods aims at handling class imbalance by designing an appropriate loss function to be used during training . In particular , they design loss functions to give more emphasis to the samples belonging to minority class , or the class with higher associated risk , than the majority ones during training [ 22 ] . [ 25 ] proposes a loss function , which we refer to as Weighted BCE ( W-BCE ) , where samples from minority class are multiplied by a constant to introduce more cost to misclassification of those samples . [ 10 ] proposes a function that aims to learn more discriminative latent representations by enforcing DNNs to maintain inter-cluster and inter-class margins , where clusters are formed using k-means clustering . They demonstrate that the tighter constraint inherently reduces class imbalance . In a more recent work , [ 5 ] proposes a loss function called class-balanced binary cross-entropy ( CB-BCE ) to weight BCE inversely proportionally to the class frequencies to amplify the loss for samples from the minority class . In a similar vein , [ 13 ] modifies BCE and propose symmetric focal loss ( S-FL ) by multiplying it with the inverse of the prediction probability to introduce more cost to the samples that DNNs are not very confident . [ 15 ] introduces symmetric margin loss ( S-ML ) by introducing a margin to the BCE loss . [ 12 ] investigates different loss functions such as S-FL and S-ML , and propose their asymmetric versions , A-FL and A-ML , by introducing a margin for samples from the minority class to handle class imbalance . In a different line of work , [ 23 ] proposes a method called mean squared false error by performing simple yet effective modification to the mean squared error ( MSE ) loss . Unlike MSE , which computes an average error from all samples without considering their classes , this loss computes a mean error for each class and averages them . [ 3 ] proposes a label-distribution-aware margin ( LDAM ) loss motivated by minimizing a margin-based generalization bound , optionally coupled with a training schedule that defers re-weighting until after the initial stage . [ 19 ] introduces balanced meta-softmax for long-tailed recognition , which accommodates the label distribution shift between training and testing , as well as a meta sampler that learns to re-sample training set by meta-learning . [ 21 ] proposes a loss that ignores the gradient from samples of large classes for the rare ones , making the training more fair . A particular group within the cost sensitive training-based methods focuses on optimizing AUC and our method falls into this group . AUC optimization is an ideal choice for class imbalance since AUC is not sensitive to class distributions [ 4 ] . [ 18 ] proposes a support vector machine ( SVM ) based loss function that maximizes AUC and demonstrates its effectiveness for the class imbalance . [ 26 ] approaches the class imbalance problem from online learning perspective and proposes an AUC optimization-based loss function . [ 7 ] proposes a one-pass method for AUC optimization that does not require storing data unlike the previous online methods . Another online AUC optimization method proposed by [ 24 ] formulates AUC optimization as a convex-concave saddle point problem . Despite their usefulness , all aforementioned AUC optimization-based methods were applied to linear predictive models , as this allows to simplify the Mann-Whitney statistics [ 16 ] for the definition of AUC , and their performance on DNNs is unknown . [ 20 ] applies online AUC optimization on a small dataset for breast cancer detection where they also mention that extension to larger datasets may not be feasible . In a very recent work called mini-batch AUC ( MBAUC ) [ 8 ] , authors extend AUC optimization to non-linear models with DNNs by optimizing AUC with mini-batches and demonstrate its effectiveness on various datasets . The proposed constrained optimization method differs from the existing works in that it enforces maximal AUC as a constraint in a way that favors reducing FPR at high TPR and can be used with existing loss functions . 3 Background - Augmented Lagrangian method ( ALM ) . A generic optimization problem for an objective function F ( θ ) subject to the constraints C ( θ ) = { c1 ( θ ) , ... , cm ( θ ) } can be expressed as [ 2 , 17 ] : arg min θ∈Θ F ( θ ) ; subject to C ( θ ) ( 1 ) Augmented Lagrangian method ( ALM ) [ 1 ] , also known as methods of multipliers , converts the constrained optimization problem in Eq . ( 1 ) to an unconstrained optimization problem . ALM is proposed to overcome the limitations of two earlier methods called quadratic penalty method and method of Lagrangian multipliers which suffer from training instability and non-convergence due to the difficulty of convexifying loss functions2 . In ALM , the penalty concept is merged with the primal-dual philosophy of classic Lagrangian function . In such methods , the penalty term is added not to the objective function F ( θ ) but rather to its Lagrangian function , thus forming the Augmented Lagrangian Function : Lµ ( θ , λ ) = F ( θ ) + µ m∑ i=1 ∥∥ci ( θ ) ∥∥2 + m∑ i=1 λici ( θ ) ( 2 ) In practice , this method consists in iteratively solving a sequence of problems as : max λk min θ Lµk ( θ , λk ) , θ ∈ Θ ( 3 ) Where { λk } is a bounded sequence in Rm , updated as λk+1i = λki + µci ( θ ) . { µk } is a positive penalty parameter sequence , with 0 < µk ≤ µk+1 , µk → ∞ , which may be either pre-selected or generated during the computation according to a defined scheme . In ALM , increasing µk indefinitely 2Please consult supplementary material for more details of quadratic penalty method and method of Lagrangian multipliers . is not necessary as in the quadratic penalty method . Thus , it does not suffer from training instabilities due to the constraint prevailing F ( θ ) . Furthermore , it does not require convexity assumption as in the method of Lagrange multipliers to ensure convergence [ 1 ] . 4 Proposed methods . 4.1 Proposed constraint for binary classification . Let F ( θ ) be a generic loss function that is used to train classification DNNs , fθ ( . ) , for binary problems . Let us also define p , { xp1 , · · · , x p |p| } and n , { x n 1 , · · · , xn|n| } as the sets of samples of the positive ( critical ) and negative classes , respectively . Note that we choose p as the minority class in our description , i.e. , |p| < |n| , and assume that this is a critical class associated with higher risk of making a mistake . We define our constrained optimization problem as follows : arg min θ F ( θ ) subject to |n|∑ k=1 max ( 0 , − ( fθ ( x p j ) − fθ ( x n k ) ) + δ ) = 0 , j ∈ { 1 , ... , |p| } , ( 4 ) where fθ ( x ) indicates output probability of the DNN on input x . Note that the constraint states that the output of the DNN for each sample of critical class should be larger than the outputs of all of the negative samples by a margin δ . Satisfying the constraint would directly ensure maximal AUC [ 16 ] . We define the equivalent unconstrained version of Eq . ( 4 ) by writing it in the form given in Eq . ( 2 ) Lµ ( θ , λ ) = F ( θ ) + µ ∑|p| j=1 q 2 j 2 · |p| · |n| + ∑|p| j=1 λj · qj |p| · |n| ( 5 ) where , qj = ∑|n| k=1 max ( 0 , − ( fθ ( x p j ) − fθ ( xnk ) ) + δ ) , µ is the penalty coefficient corresponding to the quadratic penalty term , λj is the estimate of Lagrange multiplier corresponding to each positive training sample j , and δ is the margin that we determine using a validation dataset . The crucial aspect of the formulation is the asymmetry between positive and negative classes . A constraint is defined for each sample from the positive class , thus each positive sample gets a separate Lagrange multiplier . This form prioritizes the reduction of FPR at high TPR values as illustrated next . We use Algorithm 1 to estimate the parameters of a DNN , θ , using the proposed loss function in Eq . 5 . The parameters θ are updated with every batch using gradient descent with learning rate α. Concurrently , µ is increased using a multiplicative coefficient ρ only when a chosen metric on validation is not improved , by a margin to avoid training instabilities . The validation metric ( ValMETRIC ) is : 1 ) Validation AUC for the binary setup ; 2 ) Validation Accuracy for the multi-class version . We update λj in each iteration for each positive sample x p j . Algorithm 1 ALM for Training DNNs Input : θ ( 0 ) , µ ( 0 ) , λ ( 0 ) j , ρ ; for t = 1 , . . . , T do for each mini-batch of xB with size B do yB = f ( XB ) ; Calculate q ( t ) j ; . ∀j ∈ [ 1 , B ] and yj = y+ θ ( t+1 ) ← θ ( t ) − α · ∇θLµ ( θ ( t ) , λ ) ; λ ( t+1 ) j ← λ ( t ) j + µ ( t ) · q ( t ) j ; . ∀j ∈ [ 1 , B ] and yj = y+ end for if V alMETRIC ( t ) < V alMETRIC ( t−1 ) then µ ( t+1 ) ← µ ( t ) · ρ ; else µ ( t+1 ) ← µ ( t ) ; end if end for Return θ ( t+1 )

Summary: Summary This paper proposes a new optimization approach to classification problems with class imbalance, where correctly classifying the minority class is more important than correctly classifying the majority class. One specific motivation is in medical imaging diagnostics, where misclassifying the “critical” minority class (e.g., diseased) is a more severe mistake than misclassifying the majority class (e.g., healthy). The authors therefore use an augmented Lagrangian method (ALM) approach to minimize the FPR for the majority while maintaining a high TPR for the critical class in a binary classification problem, as well as suggesting an optimization formulation for multi-class classification tasks. The ALM approach tries to enforce a margin between the critical and non-critical samples in the output space, and is structured to penalize a critical misclassification more than a non-critical misclassification.  The problem is certainly important for medical diagnostic tools, as well as other problems where errors have different costs for different classes. The proposed method is a useful approach for these problems, and the theoretical analysis of the asymmetry and required imbalance of the problem is clear and interesting. The experiments show some benefit of the ALM-based approach and its ability to be coupled with many different loss functions. The writing is clear overall, but with various grammatical mistakes throughout (and particularly in the supplement). My major concerns are the significance of the result and the setting of the CIFAR experiments.