# Generalizing to any diverse distribution: uniformity, gentle finetuning and rebalancing

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 63 references that are in the map.
- First posted on arXiv 8 Oct 2024.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [Deep Bayesian Active Learning with Image Data](https://mlsanity.com/paper/1703.02910/) (2017)
- [Generalizing to Unseen Domains: A Survey on Domain Generalization](https://mlsanity.com/paper/2103.03097/) (2021)
- [Domain Generalization: A Survey](https://mlsanity.com/paper/2103.02503/) (2021)
- [CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation](https://mlsanity.com/paper/2102.04664/) (2021)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [A Brief Review of Domain Adaptation](https://mlsanity.com/paper/2010.03978/) (2020)
- [Recent Advances in Adversarial Training for Adversarial Robustness](https://mlsanity.com/paper/2102.01356/) (2021)
- [The Risks of Invariant Risk Minimization](https://mlsanity.com/paper/2010.05761/) (2020)
- [User-friendly introduction to PAC-Bayes bounds](https://mlsanity.com/paper/2110.11216/) (2021)
- [Bayesian SegNet: Model Uncertainty in Deep Convolutional Encoder-Decoder Architectures for Scene Understanding](https://mlsanity.com/paper/1511.02680/) (2015)
- [Simple data balancing achieves competitive worst-group-accuracy](https://mlsanity.com/paper/2110.14503/) (2021)
- [Wasserstein Distributionally Robust Optimization: Theory and Applications in Machine Learning](https://mlsanity.com/paper/1908.08729/) (2019)
- [Are Labels Required for Improving Adversarial Robustness?](https://mlsanity.com/paper/1905.13725/) (2019)
- [Diverse Weight Averaging for Out-of-Distribution Generalization](https://mlsanity.com/paper/2205.09739/) (2022)
- [Leveraging Unlabeled Data to Predict Out-of-Distribution Performance](https://mlsanity.com/paper/2201.04234/) (2022)
- [Mixout: Effective Regularization to Finetune Large-scale Pretrained Language Models](https://mlsanity.com/paper/1909.11299/) (2019)
- [Assessing Generalization of SGD via Disagreement](https://mlsanity.com/paper/2106.13799/) (2021)
- [Bias and Unfairness in Information Retrieval Systems: New Challenges in the LLM Era](https://mlsanity.com/paper/2404.11457/) (2024)

## Built on by (most cited first, 2 of 2)

- [Conformation-Aware Structure Prediction of Antigen-Recognizing Immune Proteins](https://mlsanity.com/paper/2507.09054/) (2025)
- [Composable and adaptive design of machine learning interatomic potentials guided by Fisher-information analysis](https://mlsanity.com/paper/2504.19372/) (2025)

## Abstract (arXiv)

As training datasets grow larger, we aspire to develop models that generalize well to any diverse test distribution, even if the latter deviates significantly from the training data. Various approaches like domain adaptation, domain generalization, and robust optimization attempt to address the out-of-distribution challenge by posing assumptions about the relation between training and test distribution. Differently, we adopt a more conservative perspective by accounting for the worst-case error across all sufficiently diverse test distributions within a known domain. Our first finding is that training on a uniform distribution over this domain is optimal. We also interrogate practical remedies when uniform samples are unavailable by considering methods for mitigating non-uniformity through finetuning and rebalancing. Our theory provides a mathematical grounding for previous observations on the role of entropy and rebalancing for o.o.d. generalization and foundation model training. We also provide new empirical evidence across tasks involving o.o.d. shifts which illustrate the broad applicability of our perspective.