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 paper studies learning in settings with imbalanced data or under-represented subgroups in the data, where the target objective is fairness-based. E.g. in the setting of imbalanced data, the target objective is to maximize the balanced test accuracy, and in the setting with minority subgroups, target objectives studied are balanced (w.r.t. subgroups) test accuracy, and Difference of Equal Opportunity (DEO) risk.   The paper is motivated by the fact that many existing methods for imbalanced datasets rely on training theoretically-motivated modifications to the cross-entropy loss, where there are multiplicative and additive adjustments applied to the predicted logits for each class. This paper proposes a bi-level optimization procedure to directly learn these modifications, by directly learning the parameters governing the multiplicative and additive adjustments to the class logits. The bi-level optimization optimizes a classifier on a loss defined with specific instantiations of the adjustment parameters in an inner loop, and the outer loop optimizes some desired test objective (e.g., class balanced loss) on a separate holdout validation set. The paper also proposes learning class-specific data augmentation, with the motivation that minority classes are regularized more strongly than majority classes.   The empirical results evaluate the algorithm on imbalanced class and imbalanced subgroup settings.n They show that their algorithms can outperform existing methods in several evaluation metrics.