AutoBalance: Optimized Loss Functions for Imbalanced Data
Year: 2022

Abstract:
Imbalanced datasets are commonplace in modern machine learning problems. The presence of under-represented classes or groups with sensitive attributes results in concerns about generalization and fairness. Such concerns are further exacerbated by the fact that large capacity deep nets can perfectly fit the training data and appear to achieve perfect accuracy and fairness during training, but perform poorly during test. To address these challenges, we propose AutoBalance, a bi-level optimization framework that automatically designs a training loss function to optimize a blend of accuracy and fairness-seeking objectives. Specifically, a lower-level problem trains the model weights, and an upper-level problem tunes the loss function by monitoring and optimizing the desired objective over the validation data. Our loss design enables personalized treatment for classes/groups by employing a parametric cross-entropy loss and individualized data augmentation schemes. We evaluate the benefits and performance of our approach for the application scenarios of imbalanced and group-sensitive classification. Extensive empirical evaluations demonstrate the benefits of AutoBalance over state-of-the-art approaches. Our experimental findings are complemented with theoretical insights on loss function design and the benefits of train-validation split. All code is available open-source.

heading: 1 Introduction
text: 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.

heading: 1.1 Problem Setup for Class Imbalance
text: 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.

heading: 2 Methods: Loss Functions, Search Space Design, and Bilevel Optimization
text: 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).

heading: 2.1 Proposed Bilevel Optimization Method
text: We formulate the loss function design as a bilevel optimization over hyperparameters α and a hypothesis set F . We split the dataset S into training ST and validation SV sets with nT and nV
0 20 40 60 80 100 0
100
200
300
400
500 Validation Train
N um
be ro
fs am
pl es
Sorted class index (a) Train-validation sizes for CIFAR100-LT clusters
0 50 100 150 200 250 300
2
1
0
1
2
3 4 Class [0,19] Class [20,39] Class [40,59] Class [60,79] Class [80,99]
H yp
er -p
ar am
et er
va lu
e
Epoch (b) Evolution of ∆ when only training ∆
0 50 100 150 200 250 300
2
1
0
1
2
3 4 Class [0,19] Class [20,39] Class [40,59] Class [60,79] Class [80,99]
H yp
er -p
ar am
et er
va lu
e
Epoch (c) Evolution of l when only training l
0 50 100 150 200 250 300
2
1
0
1
2
3 4 Class [0,19] Class [20,39] Class [40,59] Class [60,79] Class [80,99]
H yp
er -p
ar am
et er
va lu
e
Epoch (d) Evolution of ∆ and l when optimizing jointly
Here, Lfair/LSVfair are the test/validation risks associated with `fair, e.g. Lfair = ED[`fair] as in Section 1.1. Algorithm 1 summarizes our approach and highlights the key components. The training loss `train(·;α) is also initialized to be consistent with Efair (e.g., same as `fair). In line with the literature on bilevel optimization, we will refer to the two minimizations of the validation and training losses in (2.3) as upper and lower level optimizations, respectively.
Implicit Differentiation and Warm-up training. For a loss function parameterα, the hyper-gradient can be written via the chain-rule ∂Lfair(θ
?) ∂α = ∂Lfair ∂α + ∂Lfair ∂θ? ∂θ? ∂α [46]. Here, θ ? is the solution of
the lower-level problem. We note that ∂Lfair/∂α = 0 since α does not appear within the upperlevel loss. Also observe that ∂Lfair(θ?)/∂θ? can be directly computed by taking the gradient. To compute ∂θ?/∂α, we follow the recent work [46] and employ the Implicit Function Theorem (IFT). If there exists a fixed point (θ?,α?) that satisfies ∂Ltrain(θ?,α?)/∂θ = 0 and regularity conditions are satisfied, then around α?, there exists a function θ(α) such that θ(α?) = θ? and we also have ∂θ∂α = ( ∂2Ltrain ∂θ2 ) −1 ∂2Ltrain ∂θ∂α . However, directly computing inverse Hessian ( ∂2Ltrain ∂θ2 )
−1 is usually time consuming or even impossible for modern neural networks which have millions of parameters. To compute the hyper-gradient while avoiding extensive computation, we approximate the inverse Hessian via the Neumann series, which is widely used for inverse Hessian estimation [43, 46]. Finally, the warm-up phase of our method (Line 2 of Algo. 1) is essential to guarantee that the IFT assumption ∂Ltrain(θ?,α?)
∂θ = 0 is approximately satisfied.
Why Bilevel Optimization? We choose differentiable optimization over alternative hyperparameter tuning methods because our hyperparameter space is continuous and potentially large (e.g. in the order of K = 8, 142 for iNaturalist). In our experiments, the runtime of our method was typically 4∼5 times that of standard training (with known hyperparameters). Intuitively, the runtime is at least twice due to our use of separate search and retraining phases. In the appendix, we also compare against alternative approaches (specifically SMAC of [26]) and found that our approach is faster and more accurate on CIFAR10-LT.

heading: 2.2 Reducing the Hyperparameter Search Space and the Benefits of Validation Set
text: Suppose we wish to optimize the hyperparameter α = (wy, ly,∆y)Ky=1 of the parametric loss (2.1). An important challenge is the dimensionality of α, which is proportional to the number of classes K, as we need a triplet (wy, ly,∆y) for each class. For instance, ImageNet has K = 1, 000
whereas iNaturalist has K = 8, 142 classes resulting in high-dimensional hyperparameters. In our experiments, we found that directly optimizing over such large spaces leads to convergence issues likely because of the difficulty of hypergradient estimation. Additionally, with large number of hyperparameters there is increased concern for validation overfitting. This is especially so for the tail classes (e.g. the smallest class in CIFAR100-LT has only 1 validation example with an 80-20% split). On the other hand, it is well-known in AutoML literature (e.g. neural architecture search [45], AutoAugment [12]) that designing a good search space is critical for attaining faster convergence and good validation accuracy. To this end, we propose subspace-based search spaces for hyperparameters (wy, ly,∆y). To explain the idea, consider the logit-adjustment parameters l = [l1 . . . lK ] and ∆ = [∆1 . . . ∆K ]. We propose representing these K dimensional vectors via K ′ < K dimensional embeddings l′,∆′ as follows
∆ = Dπ∆ ′ and l = Dπl′ where l′,∆′ ∈ RK ′ .
Here, Dπ ∈ RK×K ′
is a frequency-aware dictionary matrix that we design, and the range space ofDπ becomes the hyperparameter search space. In our algorithm, we cluster the classes in terms of their frequency and assign the same hyperparameter to classes with similar frequencies. To be concrete, if each cluster has size C, then K ′ = dK/Ce. For CIFAR10-LT, CIFAR100-LT, ImageNetLT and iNaturalist, we use C = {1, 10, 20, 40} respectively. In this scheme, each column d of the matrix Dπ is the indicator function of one of the K ′ clusters, i.e. di = 1 if ith class is within the cluster and 0 otherwise. Clusters are pictorially illustrated in Fig. 3a. Finally, we remark that the specific hyperparameter choices of CDT [67], LA [53], and VS [38] losses in the corresponding papers, can be viewed as specific instances of the above search space design. For instance, LA-loss chooses a scalar τ and sets li = τ log(πi). This corresponds to a dictionary containing a single column d ∈ RK with entries di = log(πi).
Why is train-validation split critical? In Figure 4 we plot the balanced errors of training/validation/test datasets at each epoch of the search phase. The training data is fixed whereas we evaluate different validation set sizes. The first finding is that training loss always overfits (dotted) until zero error whereas validation loss mildly overfits (dashed vs solid). Secondly, larger validation does help improve test accuracy (compare solid lines). The training behavior is in line with the fact that large capacity networks can perfectly fit and achieve 100% training accuracy [18, 56, 29]. This also means that different accuracy metrics or fairness constraints can be perfectly satisfied. To truly find a model that lies on the Pareto-front of the (accuracy, fairness) tradeoff, the optimization procedure should (approximately) evaluate on the population loss. Thus, as in Figure 4, the validation phase provides this crucial test-proxy in the overpa-
rameterized setting where training error is vacuous. Following the model selection literature [32, 33], the intuition is that, as the dimensionality of the hyper-parameter α is typically smaller than the validation size nV , validation loss will not overfit and will be indicative of the test even if the training loss is zero. Our search space design in Section 2.2 also helps to this end by increasing the oversampling ratio nV/dim(α) via class clustering. In the appendix, we formalize these intuitions for multi-objective problems (e.g., accuracy + fairness). Under mild assumptions, we show that a small amount of validation data is sufficient to ensure that the Pareto-front of the validation risk uniformly approximates that of the test risk. Concretely, for two objectives (L1,L2), uniformly over all λ, the hyperparameter α minimizing the validation risk LSVfair (f) = (1− λ)L SV 1 + λL SV 2 in (2.3) also approximately minimizes the test risk Lfair(f).

heading: 3 Evaluations for Imbalanced Classes
text: In this section, we present our experiments on various datasets (CIFAR-10, CIFAR-100, iNaturalist2018 and ImageNet) when the classes are imbalanced. The goal is to understand whether our bilevel optimization can design effective loss functions that improve balanced error Ebal on the test set. The setup is as follows. Ebal is the test objective. The validation loss Lfair is the balanced cross-entropy
CEbal. We consider various designs for `train such as individually tuning w, l,∆ and augmentation. We report the average result of 3 random experiments under this setup.
Datasets. We follow previous works [53, 13, 8] to construct long-tailed versions of the datasets. Specifically, for a K-class dataset, we create a long-tailed dataset by reducing the number of examples per class according to the exponential function n′i = niµ
i, where ni is the original number of examples for class i, n′i is the new number of examples per class, and µ < 1 is a scaling factor. Then, we define the imbalance factor ρ = n′0/n ′ K , which is the ratio of the number of examples in the largest class (n′0) to the smallest class (n ′ K). For the CIFAR10-LT and CIFAR100-LT dataset, we construct long-tailed versions of the datasets with imbalance factor ρ = 100. ImageNet-LT contains 115,846 training examples and 1,000 classes, with imbalance factor ρ = 256. iNaturalist-2018 contains 435,713 images from 8,142 classes, and the imbalance factor is ρ = 500. These choices follow that of [53]. For all datasets, we split the long-tailed training set into 80% training and 20% validation during the search phase (Figure 1b). Implementation. In both CIFAR datasets, the lower-level optimization trains a ResNet-32 model with standard mini-batch stochastic gradient decent (SGD) using learning rate 0.1, momentum 0.9, and weight decay 1e− 4, over 300 epochs. The learning rate decays at epochs 220 and 260 with a factor 0.1. The upper-level hyper-parameter optimization computes the hyper-gradients via implicit differentiation. Because the hyper-gradient is mostly meaningful when the network achieves near zero loss (Thm 1 of [46]), we start the validation optimization after 120 epochs of the training optimization, using SGD with initial learning rate 0.05, momentum 0.9, and weight decay 1e− 4, we follow the same learning rate decay at epoch 220 and 260. For CIFAR10-LT, 20 hyper-parameters are trained, corresponding to ly and ∆y of each 10 classes. For CIFAR100-LT, ImageNet-LT, and iNaturalist, we reduce the search space with cluster sizes of 10, 20, and 40 as visualized in Figure 3(a) and as described in Section 2.2. For ImageNet-LT and iNaturalist, following previous work [53], we use ResNet-50 and SGD for the lower and upper optimizations, For the learning rate scheduling, we use cosine scheduling starting with learning rate 0.05, and batch size 128. In searching phase, we conduct 150 epoch training with 40 epoch warm-up before the loss function design starts. For the retraining phase, we train for 90 epochs, which is the same as [53] but only due to the lack of training resources we change the batch size to 128 and adjust initial learning rate accordingly as suggested by [22]. Personalized Data Augmentation (PDA). For PDA, we utilize the AutoAugment [12] policy space and apply a bilevel search for the augmentation policy. Our approach follows existing differentiable augmentation strategies (e.g, [24]); however, we train separate policies for each class cluster to ensure that the resulting policies can adjust to class frequencies. Due to space limitations, please see supplementary materials for further details. Results and discussion. We compared our methods with the state-of-the-art long-tail learning methods. Table 1 shows the results of our experiments where the design space is parametric CE (2.1). In the first part of the table, we conduct experiments for three baseline methods: normal CE, LDAM [8] and Logit Adjustment loss with temperature parameter τ = 1 [53]. The latter choice guarantees Fisher consistency. In the second part of the Table 1, we study Algo. 1 with design spaces l, ∆, and l&∆. The first version of Algo. 1 in Table 1 tunes the LA loss parameter τ where l is parameterized by a single scalar τ as ly = τ log(πy). The next three versions of Algo. 1 consider tuning l, ∆, l&∆ respectively (Figure 3b-d shows the evolution of the l and ∆ parameters during the optimization). Finally, in last version of Algo. 1, the loss design is initialized with LA loss with τ = 1 (rather than balanced CE). The takeaway from these results is that our approach consistently leads to
a superior balanced accuracy objective. That said, tuning the LA loss alone is highly competitive with optimizing ∆ and l alone (in fact, strictly better for CIFAR10-LT, indicating Algo. 1 does not always converge to the optimal design). Importantly, when combining l&∆, our algorithm is able to design a better loss function and outperform all rows across all benchmarks. Finally, when the algorithm further is initialized with LA loss, the performance further improves accuracy, demonstrating that warm-starting with good designs improves performance.
In Table 2, we study the benefits of data augmentation, following our intuitions from Lemma 2. We compare to the differentiable augmentation baseline of MADAO [24] which trains a single policy for the full dataset. PDA is a personalized variation of MADAO and leads to noticeable improvement
across all benchmarks (most noticeably in CIFAR10-LT). More importantly, the last two lines of the table demonstrates that PDA can be synergistically combined with the parametric CE (2.1) which leads to further improvements, however, we observe that most of the improvement can be attributed to (2.1).

heading: 4 Approaches and Evaluations for Imbalanced Groups
text: While Section 3 focuses on the fundamental challenge of balanced error minimization, a more ambitious goal is optimizing generic fairness-seeking objectives. In this section, we study accuracyfairness tradeoffs by examining the group-imbalanced setting. X Setting: Imbalanced groups. For the setting with G groups, dataset is given by S = (xi, yi, gi) n i=1 where gi ∈ [G] is the group-membership. In the fairness literature, groups represent sensitive or protected attributes. For (x, y, g) ∼ D, define the group and (class, group) frequencies as
π̄j = PD(g = j), and πk,j = PD(y = k, g = j), for (k, j) ∈ [K]× [G]. The group-imbalance occurs when group or (class, group) frequencies differ, i.e., maxj∈[G] π̄j minj∈[G] π̄j or max(k,j) πk,j min(k,j) πk,j . A typical goal is ensuring that the prediction of the model is independent of these attributes. While many fairness metrics exist, in this work, we focus on the Difference of Equal Opportunity (DEO) [23, 17]. Our evaluations also focus on binary classification (with labels denoted via ±) and two groups (K = G = 2). With this setup, the DEO risk is defined as Ldeo(f) = |L+,1(f) − L+,2(f)|. Here Lk,j(f) is the (class, group)-conditional risk evaluated on the conditional distribution of “Class k & Group j”. When both classes are equally relevant (rather than y = +1 implying a semantically positive outcome), we use the symmetric DEO:
Ldeo(f) = |L+,1(f)− L+,2(f)|+ |L−,1(f)− L−,2(f)|. (4.1)
We will study the pareto-frontiers of the DEO (4.1), group-balanced error, and standard error. Here, group-balanced risk is defined as LGbal(f) = 1 KG ∑K k=1 ∑G j=1 Lk,j(f). Note that this definition treats each (class, group) pair as its own (sub)group. Throughout, we explicitly set the validation loss to cross-entropy for clarity, thus we use CE, CEGbal, CEdeo to refer to L, L G bal, Ldeo. Validation (upper-level) loss function. In Algo. 1, we set Lfair = (1 − λval) · CE + λval · CEdeo for varying 0 ≤ λval ≤ 1. The parameter λval enables a trade-off between accuracy and fairness. Within λval, we use the subscript “val” to highlight the fact that we regularize the validation objective rather than the training objective. Group-sensitive training loss design. As first proposed in [38], the parametric cross-entropy (CE) can be extended to (class, group) imbalance by extending hyper-parameter α to [K]× [G] variables w, l,∆ ∈ R[K]×[G] generalizing (2.1), (2.2). This leads us to the following parametric loss function for group-sensitive classification 1
`train(y, g, f(x);α) = −wyg log
( eσ(∆yg)fy(x)+lyg∑
k∈[K] e σ(∆kg)fk(x)+lkg
) . (4.2)
Here, wyg applies weighted CE, while ∆yg and lyg are logit adjustments for different (class, groups). This loss function is used throughout the imbalanced groups experiments.
1This is a generalization to multiple classes of the proposal in [38] for binary group-sensitive classificaiton.
Baselines. We will compare Algo. 1 with training loss functions parameterized via (1−λ)·CE+λ·Lreg. Here Lreg is a fairness-promoting regularization. Specifically, as displayed in Table 3 and Figure 5, we will set Lreg to be CEGbal, CEdeo and Group LA. “Group LA” is a natural generalization of the LA loss to group-sensitive setting; it chooses weights wg = 1/π̄g to balance group frequencies and then applies logit-adjustment with τ = 1 over the classes conditioned on the group-membership. Datasets. We experiment with the modified Waterbird dataset [58]. The goal is to correctly classify the bird type despite the spurious correlations due to the image background. The distribution of the original data is as follows. The binary classes k ∈ {−,+} correspond to {waterbird, landbird}, and the groups [G] = {1, 2} correspond to {land background,water background}. The fraction of data in each (class, group) pair is π−,2 = 0.22, π−,1 = 0.012, π+,2 = 0.038, and π+,1 = 0.73. The landbird on the water background ({+, 2}) and the waterbird on the land background ({−, 1}) are minority sub-groups within their respective classes. The test set, following [58], has equally allocated bird types on different backgrounds, i.e., π±,j = 0.25. As the test dataset is balanced, the standard classification error E(f) is defined to be the weighted error E(f) = πy,gEy,g(f). Implementation. We follow the feature extraction method from [58], where xi are 512-dimensional ResNet18 features. When using Algo. 1, we split the original training data into 50% training and 50% validation. The search phase uses 150 epochs of warm up followed by 350 epochs of bilevel optimization. The remaining implementation details are similar to Section 3.
5 6 7 8 910.0
12.5
15.0
17.5
20.0
22.5
25.0
Group-LA CEDEO CEbal Algo.1:CEDEO B al an nc
ed er
ro r
Standard classification error 5 6 7 8 90
5 10 15 20 25 30 35
Group-LA CEDEO CEbal Algo.1:CEDEO
D E
O
Standard classification error
achieving lower DEO or balanced error for the same standard error. This demonstrates the value of automatic loss function design for a rich class of fairness-seeking objectives.
Next, in Table 3, we solely focus on optimizing the fairness objectives (rather than standard error). Thus, we compare against CEdeo, CEGbal, Group-LA as the baseline approaches as well as blending CE with λ = 0.1. We also compare to the DRO approach of [58]. Finally, we display the outcome of Algo. 1 with λval = 0.1. While DRO is competitive, similar to Figure 5, our approach outperforms all baselines for all metrics. The improvement is particularly significant when it comes to DEO. Finally, we remark that [38] further proposed combining DRO with the group-adjusted VS-loss for improved performance. We leave the evaluation of AutoBalance for such combinations to future.

heading: 5 Related Work
text: Our work relates to imbalanced classification, fairness, bilevel optimization, and data augmentation. Below we focus on the former three and defer the extended discussion to the supplementary. Long-tailed learning. Learning with long-tailed data has received substantial interest historically, with classical methods focusing on designing sampling strategies, such as over- or under-sampling [40, 60, 9, 41, 1, 71, 57, 68, 5, 50]. Several loss re-weighting schemes [54, 52, 25, 13, 6, 34] have been proposed to adjust weights of different classes or samples during training. Another line of work [70, 37, 53] focuses on post-hoc correction. More recently, several works [44, 35, 16, 34, 8, 13, 53, 67] develop more refined class-balanced loss functions (e.g. (2.1)) that better adapt to the training data. In addition, several works [30, 74] point out that separating the representation learning and class balancing can lead to improvements. In this work, our approach is in the vein of class-balanced losses; however, rather than fixing a balanced loss function (e.g. based on the class probabilities in the training dataset), we employ our Algorithm 1 to automatically guide the loss design. Group-sensitive and Fair Learning. Group-sensitive learning aims to ensure fairness in the presence of under-represented groups (e.g., gender, race). [7, 23, 69, 64] propose several fairness metrics as well as insightful methodologies. A line of research [4, 19] optimize the worst-case loss over the test distribution and further applications motivate (label, group) metrics such as equality of opportunity [23, 17] (also recall DEO (4.1)). [58] discusses group-sensitive learning in an over-parameterized regime and proposes that strong regularization ensures fairness. Closer to our work, [58, 38] also study Waterbirds dataset. Compared to the regularization-based approach of [58], we explore a parametric loss design (inspired by [38]) to optimize fairness-risk over validation. [17] proposes methods and statistical guarantees for fair empirical risk minimization. A key observation of our work is that, such guarantees based on training-only optimization can be vacuous in the overparameterized regime. Thus, using train-validation split (e.g. our Algo 1) is critical for optimizing fairness metrics more reliably. This is verified by the effectiveness of our approach in the evaluations of Section 4. Bilevel Optimization. Classical approaches [61] for hyper-parameter optimization are typically based on derivative-free schemes, including random search [65] and reinforcement learning [76, 3, 66, 72]. Recently, a growing line of works focus on differentiable algorithms that are often faster and can scale up to millions of parameters [46, 49, 59, 28, 48]. These techniques [45, 36, 75, 47] have shown significant success in neural architecture search, learning rate scheduling, regularization, etc. They are typically formulated as a bilevel optimization problem: the upper and lower optimizations minimize the validation and training losses, respectively. Some theoretical guarantees (albeit restrictive) are also available [11, 21, 2, 55]. Different from these, our work focuses on principled design of training loss function to optimize fairness-seeking objectives for imbalanced data. Here, a key algorithmic distinction (e.g. compared to architecture search) is that, our loss function design is only used during optimization and not during inference. This leads to a more sophisticated hyper-gradient and necessitates additional measures to ensure stability of our approach (see Algo 1).

heading: 6 Conclusions and Future Directions
text: This work provides an optimization-based approach to automatically design loss functions to address imbalanced learning problems. Our algorithm consistently outperforms, or is at least competitive with, the state-of-the-art approaches for optimizing balanced accuracy. Importantly, our approach is not restricted to imbalanced classes or specific objectives, and can achieve good tradeoffs between (accuracy, fairness) on the Pareto frontier. We also provide theoretical insights on certain algorithmic aspects including loss function design, data augmentation, and train-validation split. Potential Limitations, Negative Societal Impacts, & Precautions: Our algorithmic approach can be considered within the realm of automated machine learning literature (AutoML) [27]. AutoML algorithms often optimize the model performance, thus reducing the need for engineering expertise at the expense of increased computational cost and increased carbon footprint. For instance, our procedure is computationally more intensive compared to the theory-inspired loss function prescriptions of [53, 8]. A related limitation is that Algo. 1 can be brittle in extremely imbalanced scenarios with very few samples per class. We took the several steps to help mitigate such issues: first, our algorithm is initialized with a Bayes consistent loss function to provide a warm-start (such as the proposal of [53]). Second, to improve generalization and avoid overfitting to validation, we reduce the hyperparameter search space by grouping the classes with similar frequencies. Finally, evaluations show that the designed loss functions are interpretable (Fig. 3(b,c,d)) and are consistent with theoretical intuitions.
Acknowledments This work is supported in part by the National Science Foundation under grants CCF-2046816, CNS-1932254, CCF-2009030, HDR-1934641, CNS-1942700 and by the Army Research Office under grant W911NF-21-1-0312.

Review summaries:
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.

This paper proposes a method for long-tailed classification. The key idea is to use a generalized loss function parameterized by hyper-parameters \alpha and to (automatically) tune \alpha on a validation set. Results are presented on CIFAR-10-LT, CIFAR-100-LT, ImageNet-LT and iNaturalist datasets.

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.