# AutoBalance: Optimized Loss Functions for Imbalanced Data

arXiv: https://arxiv.org/abs/2201.01212 · HTML: https://mlsanity.com/paper/2201.01212/

## Key facts

- Cited by 49 papers in the map (data through 14 Oct 2026).
- Builds on 60 of its 86 references that are in the map.
- First posted on arXiv 4 Jan 2022.

## Builds on (most cited first, 30 of 60)

- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [Equality of Opportunity in Supervised Learning](https://mlsanity.com/paper/1610.02413/) (2016)
- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [Improved Regularization of Convolutional Neural Networks with Cutout](https://mlsanity.com/paper/1708.04552/) (2017)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Class-Balanced Loss Based on Effective Number of Samples](https://mlsanity.com/paper/1901.05555/) (2019)
- [Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss](https://mlsanity.com/paper/1906.07413/) (2019)
- [Decoupling Representation and Classifier for Long-Tailed Recognition](https://mlsanity.com/paper/1910.09217/) (2019)
- [AutoAugment: Learning Augmentation Policies from Data](https://mlsanity.com/paper/1805.09501/) (2018)
- [Long-tail learning via logit adjustment](https://mlsanity.com/paper/2007.07314/) (2020)
- [Exploring the Limits of Weakly Supervised Pretraining](https://mlsanity.com/paper/1805.00932/) (2018)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [A systematic study of the class imbalance problem in convolutional neural networks](https://mlsanity.com/paper/1710.05381/) (2017)
- [Does Learning Require Memorization? A Short Tale about a Long Tail](https://mlsanity.com/paper/1906.05271/) (2019)
- [Bilevel Programming for Hyperparameter Optimization and Meta-Learning](https://mlsanity.com/paper/1806.04910/) (2018)
- [BBN: Bilateral-Branch Network with Cumulative Learning for Long-Tailed Visual Recognition](https://mlsanity.com/paper/1912.02413/) (2019)
- [Fairness Beyond Disparate Treatment & Disparate Impact: Learning Classification without Disparate Mistreatment](https://mlsanity.com/paper/1610.08452/) (2016)
- [Gradient-based Hyperparameter Optimization through Reversible Learning](https://mlsanity.com/paper/1502.03492/) (2015)
- [Optimizing Millions of Hyperparameters by Implicit Differentiation](https://mlsanity.com/paper/1911.02590/) (2019)
- [SNAS: Stochastic Neural Architecture Search](https://mlsanity.com/paper/1812.09926/) (2018)
- [Learning from Simulated and Unsupervised Images through Adversarial Training](https://mlsanity.com/paper/1612.07828/) (2016)
- [Designing Neural Network Architectures using Reinforcement Learning](https://mlsanity.com/paper/1611.02167/) (2016)
- [What is the Effect of Importance Weighting in Deep Learning?](https://mlsanity.com/paper/1812.03372/) (2018)
- [Fast AutoAugment](https://mlsanity.com/paper/1905.00397/) (2019)
- [Empirical Risk Minimization under Fairness Constraints](https://mlsanity.com/paper/1802.08626/) (2018)

## Built on by (most cited first, 30 of 49)

- [Imbalance Trouble: Revisiting Neural-Collapse Geometry](https://mlsanity.com/paper/2208.05512/) (2022)
- [FedNest: Federated Bilevel, Minimax, and Compositional Optimization](https://mlsanity.com/paper/2205.02215/) (2022)
- [Multi-stage Neural Networks: Function Approximator of Machine Precision](https://mlsanity.com/paper/2307.08934/) (2023)
- [SoLar: Sinkhorn Label Refinery for Imbalanced Partial-Label Learning](https://mlsanity.com/paper/2209.10365/) (2022)
- [Understanding Why Generalized Reweighting Does Not Improve Over ERM](https://mlsanity.com/paper/2201.12293/) (2022)
- [A Systematic Review on Long-Tailed Learning](https://mlsanity.com/paper/2408.00483/) (2024)
- [Balancing Logit Variation for Long-tailed Semantic Segmentation](https://mlsanity.com/paper/2306.02061/) (2023)
- [On the Implicit Geometry of Cross-Entropy Parameterizations for Label-Imbalanced Data](https://mlsanity.com/paper/2303.07608/) (2023)
- [CBA: Improving Online Continual Learning via Continual Bias Adaptor](https://mlsanity.com/paper/2308.06925/) (2023)
- [Constrained Bi-Level Optimization: Proximal Lagrangian Value function Approach and Hessian-free Algorithm](https://mlsanity.com/paper/2401.16164/) (2024)
- [Selective Attention: Enhancing Transformer through Principled Context Control](https://mlsanity.com/paper/2411.12892/) (2024)
- [Online Bilevel Optimization: Regret Analysis of Online Alternating Gradient Methods](https://mlsanity.com/paper/2207.02829/) (2022)
- [Multi-Objective Bayesian Optimization with Active Preference Learning](https://mlsanity.com/paper/2311.13460/) (2023)
- [Alternating Implicit Projected SGD and Its Efficient Variants for Equality-constrained Bilevel Optimization](https://mlsanity.com/paper/2211.07096/) (2022)
- [Towards Automated Imbalanced Learning with Deep Hierarchical Reinforcement Learning](https://mlsanity.com/paper/2208.12433/) (2022)
- [Improved Balanced Classification with Theoretically Grounded Loss Functions](https://mlsanity.com/paper/2512.23947/) (2025)
- [Nystrom Method for Accurate and Scalable Implicit Differentiation](https://mlsanity.com/paper/2302.09726/) (2023)
- [Toward Operationalizing Pipeline-aware ML Fairness: A Research Agenda for Developing Practical Guidelines and Tools](https://mlsanity.com/paper/2309.17337/) (2023)
- [Data Optimization in Deep Learning: A Survey](https://mlsanity.com/paper/2310.16499/) (2023)
- [On how to avoid exacerbating spurious correlations when models are overparameterized](https://mlsanity.com/paper/2206.12739/) (2022)
- [Generalized Categories Discovery for Long-tailed Recognition](https://mlsanity.com/paper/2401.05352/) (2023)
- [Class-Aware Universum Inspired Re-Balance Learning for Long-Tailed Recognition](https://mlsanity.com/paper/2207.12808/) (2022)
- [Fair Canonical Correlation Analysis](https://mlsanity.com/paper/2309.15809/) (2023)
- [ImbaGCD: Imbalanced Generalized Category Discovery](https://mlsanity.com/paper/2401.05353/) (2023)
- [Sharp analysis of out-of-distribution error for "importance-weighted" estimators in the overparameterized regime](https://mlsanity.com/paper/2405.06546/) (2024)
- [Understanding Representation of Deep Equilibrium Models from Neural Collapse Perspective](https://mlsanity.com/paper/2410.23391/) (2024)
- [Persian Emotion Detection using ParsBERT and Imbalanced Data Handling Approaches](https://mlsanity.com/paper/2211.08029/) (2022)
- [Class-attribute Priors: Adapting Optimization to Heterogeneity and Fairness Objective](https://mlsanity.com/paper/2401.14343/) (2024)
- [Dual-CBA: Improving Online Continual Learning via Dual Continual Bias Adaptors from a Bi-level Optimization Perspective](https://mlsanity.com/paper/2408.13991/) (2024)
- [ReFine: Boosting Time Series Prediction of Extreme Events by Reweighting and Fine-tuning](https://mlsanity.com/paper/2409.14232/) (2024)

## Abstract (arXiv)

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.