AutoBalance: Optimized Loss Functions for Imbalanced Data

1 Introduction . Recently , deep learning , large datasets , and the evolution of computing power have led to unprecedented success in computer vision , and natural language processing [ 15 , 39 , 62 ] . This success is partially driven by the availability of high-quality datasets , built by carefully collecting a sufficient number of samples for each class . In practice , real-world datasets are frequently imbalanced and exhibit long-tailed behavior , necessitating a careful treatment of the minorities [ 20 , 53 , 23 ] . Indeed , modern classification tasks can involve thousands of classes , so it is perhaps intuitive that some classes should be over/under-represented compared to others . Besides class imbalance , minorities can also appear at the feature-level ; for instance , the specific values of the features of an example can vary depending on that example ’ s membership in certain sensitive or protected groups , e.g . race , gender , disabilities ( see also Figure 1a ) . In scenarios where imbalances are induced by heterogeneous client datasets ( e.g. , in the context of federated learning ) , addressing these imbalances can help ensure that a machine learning model works well for all clients , rather than just those that generate the majority of the training data . This rich set of applications motivate the careful treatment of imbalanced datasets . In the imbalanced classification literature , the recurring theme is maximizing a fairness-seeking objective , such as balanced accuracy . Unlike standard accuracy , which can be dominated by the majorities , a fairness-seeking objective seeks to promote examples from minorities , and downweigh 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . examples from majorities . Here , note that there is a distinction between the test and training objectives . While the overall goal is typically to maximize a non-differentiable objective such as balanced accuracy on the test set , during training , we use a differentiable proxy for this , such as weighted cross-entropy . Thus , the fundamental question of interest is : How to design a training loss to maximize a fairness-seeking objective on the test set ? A classical answer to this question is to use a Bayes-consistent loss functions . For instance , weighted cross-entropy ( e.g. , each class gets a different weight , see Sec . 2 ) is traditionally a good choice for optimizing weighted accuracy objectives . Unfortunately , this intuition starts to break down when the training problem is overparameterized , which is a common practice in deep learning : in essence , for large capacity deep nets , the training process can perfectly fit to the data , and training loss is no longer indicative of test error . In fact , recent works [ 8 , 38 ] show that weighted cross-entropy has minimal benefit to balanced accuracy , and instead alternative methods based on margin adjustment can be effective ( namely , by ensuring that minority classes are further away from decision boundary ) . These ideas led to the development of a parametric cross-entropy function ` ( y , f ( x ) ) = wy log ( 1 +∑ k 6=y e lk−ly ·e∆kfk ( x ) −∆yfy ( x ) ) , which allows for a personalized treatment of the individual classes via the design parameters ( wk , lk , ∆k ) Kk=1 [ 8 , 38 , 53 , 35 , 67 ] . Here , wk is the classical weighting term whereas lk and ∆k are additive and multiplicative logit adjustments . However , despite these developments , it is unclear how such parametric cross-entropy functions can be tuned for use for different fairness objectives , for example to tackle class or group imbalances . The works by [ 8 , 53 ] provide theoretically-motivated choices for ( wk , lk ) , while [ 38 ] argues that ( wk , lk ) is not as effective as ∆k in the interpolating regime of zero training error and proposes the simultaneous use of all three different parameter types . However , these works do not provide an optimized loss function that can be systematically tailored for different fairness objectives , such as balanced accuracy common in class imbalanced scenarios , or equal opportunity [ 23 , 17 ] which is relevant in group-sensitive settings . In this work , we address these shortcomings by designing the loss function within the optimization in a principled fashion , to handle different fairness-seeking objectives . Our main idea is to use bi-level optimization , where the model weights are optimized over the training data , and the loss function is automatically tuned by monitoring the validation loss . Our core intuition is that unlike training data , the validation data is difficult to fit and will provide a consistent estimator of the test objective . Contributions . Based on this high-level idea , this paper takes a step towards a systematic treatment of imbalanced learning problems with contributions along several fronts : state-of-the-art performance , data augmentation , applications to different imbalance types , and theoretical intuitions . Specifically : •We introduce AutoBalance —a bilevel optimization framework— that designs a fairness-seeking loss function by jointly training the model and the loss function hyperparameters in a systematic way ( Figure 1b , Section 2 ) . We introduce novel strategies that narrow down the search space to improve convergence and avoid overfitting . To further improve the performance , our design also incorporates data augmentation policies personalized to subpopulations ( classes or groups ) . We demonstrate the benefits of AutoBalance when optimizing various fairness-seeking objectives over the state-of-the-art , such as logit-adjustment ( LA ) [ 53 ] and label-distribution-aware margin ( LDAM ) [ 8 ] losses . The code is available online [ 42 ] . Algorithm 1 : AutoBalance via Bilevel Optimization Input : Model fθ with weights θ , dataset S = ST ∪ SV , step sizes ηα & ηθ , # iterations t2 > t1 1 Initialize α with ` fair ( · ) = ` train ( · ; α ) // Consistent initialization 2 Train θ for t1 iterations ( α is fixed ) // Search Phase : Starts with warm-up 3 for i← t1 to t2 do 4 Sample training batch BT from ST ; 5 BT ← A ( BT ) // Apply class-personalized augmentation 6 θ ← θ − ηθ∇θLBTtrain ( fθ ; α ) 7 Sample validation batch BV from SV ; 8 Compute hyper-gradient ∇αLBVfair ( fθ ) // via Approx . Implicit Differentiation 9 α← α− ηα∇LBVfair ( fθ ) // Update loss function hyper-parameters 10 end 11 Set α ? ← α , ST ← S , reset weights θ //Retraining Phase : Use all data and α ? 12 Train θ for t2 iterations using α ? Result : The final model θ ? ← θ and hyper-parameters α ? • Extensive experiments provide several takeaways ( Section 3 ) . First , AutoBalance discovers loss functions from scratch that are consistent with theory and intuition : hyperparameters of the minority classes evolve to upweight the training loss of minority to promote them . Second , the impact of individual design parameters in the loss function is revealed , with the additive adjustment lk and multiplicative adjustment ∆k synergistically improving the fairness objective . Third , personalized data augmentation can further improve the performance over a single generic augmentation policy . • Beyond class imbalance , we consider applications of loss function design to the group-sensitive setting ( Section 4 ) . Our experiments show that AutoBalance consistently outperforms various baselines , leading to a more efficient Pareto-frontier of accuracy-fairness tradeoffs . 1.1 Problem Setup for Class Imbalance . We first focus on the label-imbalance problem . The extension to the group-imbalanced setting ( approach , algorithms , evaluations ) is deferred to Section 4 . Let [ K ] denote the set { 1 , . . . , K } . Suppose we have a dataset S = ( xi , yi ) ni=1 sampled i.i.d . from a distribution D with input space X and K classes . For a training example ( x , y ) , x ∈ X is the input feature and y ∈ [ K ] is the output label . Let f : X → RK be a model that outputs a distribution over classes and let ŷf ( x ) = arg maxi∈ [ K ] f ( x ) . The standard classification error is denoted by E ( f ) = PD [ y 6= ŷf ( x ) ] . For a loss function ` ( y , ŷ ) ( e.g . cross-entropy ) , we similarly denote Population risk : L ( f ) = ED [ ` ( y , ŷf ( x ) ) ] and Empirical risk : LS ( f ) = 1 n n∑ i=1 ` ( yi , ŷf ( xi ) ) . X Setting : Imbalanced classes . Define the frequency of the k ’ th class via πk = P ( x , y ) ∼D ( y = k ) . Label/class-imbalance occurs when the class frequencies differ substantially , i.e. , maxi∈ [ K ] πi mini∈ [ K ] πi . Let us introduce Balanced risk : Lbal ( f ) = 1 K K∑ k=1 Lk ( f ) and Class-conditional risk : Lk ( f ) = EDk [ ` ( y , ŷf ( x ) ) ] . Similarly , let Ek ( f ) be the class-conditional classification error and Ebal ( f ) : = ( 1/K ) ∑K k=1 Ek ( f ) be the balanced error . In this setting , rather than the standard test error E ( f ) , our goal is to the minimize balanced error . At a high-level , we propose to do this by designing an imbalance-aware training loss that maximizes balanced validation accuracy . 2 Methods : Loss Functions , Search Space Design , and Bilevel Optimization . Our main goal in this paper is automatically designing loss functions to optimize target objectives for imbalanced learning ( e.g. , Settings A and B ) . We will employ a parametrizable family of loss functions that can be tailored to the needs of different classes or groups . Cross-entropy variations have been proposed by [ 44 , 35 , 16 ] to optimize balanced objectives . Our design space will utilize recent works which introduce Label-distribution-aware margin ( LDAM ) [ 8 ] , Logit-adjustment ( LA ) [ 53 ] , Class-dependent temperatures ( CDT ) [ 67 ] , Vector scaling ( VS ) [ 38 ] losses . Specifically , we build on the following parametric loss function controlled by three vectors w , l , ∆ ∈ RK : ` ( y , f ( x ) ) = wy log ( 1 + ∑ k 6=y elk−ly · e∆kfk ( x ) −∆yfy ( x ) ) . ( 2.1 ) Here , wy enables conventional weighted CE and ly , and ∆y are additive and multiplicative adjustments to the logits . This choice is same as the VS-loss introduced in [ 38 ] , which borrows the ∆ term from [ 67 ] and l term from [ 8 , 53 ] . [ 53 ] makes the observation that we can use l rather than w while ensuring Fisher consistency in balanced error . We make the following complementary observation . Lemma 1 Parametric loss function ( 2.1 ) is not consistent for standard or balanced errors if there are distinct multiplicative adjustments i.e . ∆i 6= ∆j for some i , j ∈ [ K ] . While consistency is a desirable property , it is intuitively more critical during the earlier phase of the training where the training risk is more indicative of the test risk . In the interpolating regime of zero-training error , [ 38 ] shows that w , l can be ineffective and multiplicative ∆-adjustment can be more favorable . Our algorithm will be initialized with a consistent weighted-CE ; however , we will allow the algorithm to automatically adapt to the interpolating regime by tuning l and ∆ . Proposed training loss function . For our algorithm , we will augment ( 2.1 ) with data augmentation that can be personalized to distinct classes . Let us denote the data augmentation policies by A = ( Ay ) Ky=1 where each Ay stochastically augments an input example with label y. Additionally , we clamp ∆i with the sigmoid function σ to limit its range to ( 0,1 ) to ensure non-negativity . To this end , our loss function for the lower-level optimization ( over training data ) is as follows : ` train ( y , x , f ; α ) = −EA [ wy log ( eσ ( ∆y ) fy ( Ay ( x ) ) +ly∑ i∈ [ K ] e σ ( ∆i ) fi ( Ay ( x ) ) +li ) ] . ( 2.2 ) Here , α is the set of hyperparameters of the loss function that we wish to optimize , specifically α = [ w , l , ∆ , param ( A ) ] . param ( A ) is the parameterization of the augmentation policies ( Ay ) y∈ [ K ] . Personalized data augmentation ( PDA ) . Remarkable benefits of data augmentation techniques provide a natural motivation to investigate whether one can benefit from learning class-personalized augmentation policies . The PDA idea relates to SMOTE [ 9 ] , where the minority class is over-sampled by creating synthetic examples . To formalize the benefits of PDA , consider a spherical augmentation strategy where Ay ( x ) samples a vector uniformly from an ` 2-ball of radius εy around x . As visualized in Figure 2 for a linear classifier , if the augmentation strengths of both classes are equal , the max-margin classifier is not affected by the application of the data augmentation and remains identical . Thus , augmentation has no benefit . However by applying a stronger augmentation on minority , the decision boundary is shifted to protect minority which can provably benefit the balanced accuracy [ 38 ] . The following intuitive observation links the PDA to parametric loss ( 2.1 ) . Lemma 2 Consider a binary classification task with labels 0 and 1 and a linearly separable training dataset . For any parametric loss ( 2.1 ) choices of ( li , ∆i , wi ) 1i=0 , there exists spherical augmentation strengths for minority/majority classes so that , without regularization , optimizing the logistic loss with personalized augmentations returns the same classifier as optimizing ( 2.1 ) . This lemma is similar in flavor to [ 31 ] , which considers a larger uncertainty set around the minority class . But , as discussed in the appendix , Lemma 2 is relevant in the overparameterized regime whereas the approach of [ 31 ] is ineffective for separable data [ 51 ] . Algorithmically , the augmentations that we consider are much more flexible than the ` p-balls of [ 31 ] and our experiments showcase the value of our approach in state-of-the-art multiclass settings . Besides , note that the ( theoretical ) benefits of PDA can go well-beyond Lemma 2 by leveraging the invariances [ 10 , 14 ] ( via rotation , translation ) .

Summary: This work introduces a novel algorithm (AutoBalance) to address the imbalance and group sensitive classification problems and maximise the fairness-seeking objectives. The algorithm uses a bi-level optimization framework to optimize the model weights over the training data while automatically fine-tuning the loss function by monitoring the validation data. In addition, AutoBalance uses data augmentation policies personalised to subpopulations in the data.