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: This paper proposes a very interesting approach for training DNNs on binary (and multiclass) classification tasks with critical under-represented classes, which is very typical in domains like healthcare. In this approach, the task is posed as a constrained optimization problem and a novel constraint is designed that emphasizes on reducing the false positive rate while maintaining a high true positive rate, which is then added to the loss function and solved using an Augmented Lagrangian Method. The main contribution of the paper is proposing a novel constraint optimization approach that can be added to any previous work, and that improves upon the state of the art results.