# Rich Feature Construction for the Optimization-Generalization Dilemma

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

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 36 references that are in the map.
- First posted on arXiv 24 Mar 2022.

## Builds on (most cited first, 24 of 24)

- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [Intriguing properties of neural networks](https://mlsanity.com/paper/1312.6199/) (2013)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [On First-Order Meta-Learning Algorithms](https://mlsanity.com/paper/1803.02999/) (2018)
- [Adversarial Examples Are Not Bugs, They Are Features](https://mlsanity.com/paper/1905.02175/) (2019)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Recognition in Terra Incognita](https://mlsanity.com/paper/1807.04975/) (2018)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Self-Challenging Improves Cross-Domain Generalization](https://mlsanity.com/paper/2007.02454/) (2020)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [Empirical Analysis of the Hessian of Over-Parametrized Neural Networks](https://mlsanity.com/paper/1706.04454/) (2017)
- [Gradient Starvation: A Learning Proclivity in Neural Networks](https://mlsanity.com/paper/2011.09468/) (2020)
- [Gradient Matching for Domain Generalization](https://mlsanity.com/paper/2104.09937/) (2021)
- [Fishr: Invariant Gradient Variances for Out-of-Distribution Generalization](https://mlsanity.com/paper/2109.02934/) (2021)
- [On Calibration and Out-of-domain Generalization](https://mlsanity.com/paper/2102.10395/) (2021)
- [Learning from Failure: Training Debiased Classifier from Biased Classifier](https://mlsanity.com/paper/2007.02561/) (2020)
- [Does Invariant Risk Minimization Capture Invariance?](https://mlsanity.com/paper/2101.01134/) (2021)
- [Evading the Simplicity Bias: Training a Diverse Set of Models Discovers Solutions with Superior OOD Generalization](https://mlsanity.com/paper/2105.05612/) (2021)
- [Domain-Adjusted Regression or: ERM May Already Learn Features Sufficient for Out-of-Distribution Generalization](https://mlsanity.com/paper/2202.06856/) (2022)
- [Predict then Interpolate: A Simple Algorithm to Learn Stable Classifiers](https://mlsanity.com/paper/2105.12628/) (2021)
- [Meta-Learned Invariant Risk Minimization](https://mlsanity.com/paper/2103.12947/) (2021)

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

- [On Feature Learning in the Presence of Spurious Correlations](https://mlsanity.com/paper/2210.11369/) (2022)
- [Model Ratatouille: Recycling Diverse Models for Out-of-Distribution Generalization](https://mlsanity.com/paper/2212.10445/) (2022)
- [A Whac-A-Mole Dilemma: Shortcuts Come in Multiples Where Mitigating One Amplifies Others](https://mlsanity.com/paper/2212.04825/) (2022)
- [Probable Domain Generalization via Quantile Risk Minimization](https://mlsanity.com/paper/2207.09944/) (2022)
- [Towards Last-layer Retraining for Group Robustness with Fewer Annotations](https://mlsanity.com/paper/2309.08534/) (2023)
- [Pareto Invariant Risk Minimization: Towards Mitigating the Optimization Dilemma in Out-of-Distribution Generalization](https://mlsanity.com/paper/2206.07766/) (2022)
- [Understanding and Improving Feature Learning for Out-of-Distribution Generalization](https://mlsanity.com/paper/2304.11327/) (2023)
- [Calibrating Multi-modal Representations: A Pursuit of Group Robustness without Annotations](https://mlsanity.com/paper/2403.07241/) (2024)
- [SelecMix: Debiased Learning by Contradicting-pair Sampling](https://mlsanity.com/paper/2211.02291/) (2022)
- [Discovering environments with XRM](https://mlsanity.com/paper/2309.16748/) (2023)
- [Spurious Feature Diversification Improves Out-of-distribution Generalization](https://mlsanity.com/paper/2309.17230/) (2023)
- [No Reason for No Supervision: Improved Generalization in Supervised Models](https://mlsanity.com/paper/2206.15369/) (2022)
- [Mitigating Dataset Bias by Using Per-sample Gradient](https://mlsanity.com/paper/2205.15704/) (2022)
- [Spuriosity Didn't Kill the Classifier: Using Invariant Predictions to Harness Spurious Features](https://mlsanity.com/paper/2307.09933/) (2023)
- [Nuisances via Negativa: Adjusting for Spurious Correlations via Data Augmentation](https://mlsanity.com/paper/2210.01302/) (2022)
- [Safety Reasoning with Guidelines](https://mlsanity.com/paper/2502.04040/) (2025)
- [Learning useful representations for shifting tasks and distributions](https://mlsanity.com/paper/2212.07346/) (2022)
- [FOOGD: Federated Collaboration for Both Out-of-distribution Generalization and Detection](https://mlsanity.com/paper/2410.11397/) (2024)
- [Mitigating Simplicity Bias in Deep Learning for Improved OOD Generalization and Robustness](https://mlsanity.com/paper/2310.06161/) (2023)
- [Sufficient Invariant Learning for Distribution Shift](https://mlsanity.com/paper/2210.13533/) (2022)
- [Malign Overfitting: Interpolation Can Provably Preclude Invariance](https://mlsanity.com/paper/2211.15724/) (2022)
- [Domain Generalization by Rejecting Extreme Augmentations](https://mlsanity.com/paper/2310.06670/) (2023)
- [Robustness to Spurious Correlations Improves Semantic Out-of-Distribution Detection](https://mlsanity.com/paper/2302.04132/) (2023)
- [Learning Diverse Features in Vision Transformers for Improved Generalization](https://mlsanity.com/paper/2308.16274/) (2023)
- [TIF: Learning Temporal Invariance in Android Malware Detectors](https://mlsanity.com/paper/2502.05098/) (2025)
- [Learning an Invertible Output Mapping Can Mitigate Simplicity Bias in Neural Networks](https://mlsanity.com/paper/2210.01360/) (2022)
- [Fine-tuning with Very Large Dropout](https://mlsanity.com/paper/2403.00946/) (2024)
- [Generalizable Prompt Learning of CLIP: A Brief Overview](https://mlsanity.com/paper/2503.01263/) (2025)
- [Improving Group Robustness on Spurious Correlation via Evidential Alignment](https://mlsanity.com/paper/2506.11347/) (2025)
- [Towards Generalizable PDE Dynamics Forecasting via Physics-Guided Invariant Learning](https://mlsanity.com/paper/2509.24332/) (2025)

## Abstract (arXiv)

There often is a dilemma between ease of optimization and robust out-of-distribution (OoD) generalization. For instance, many OoD methods rely on penalty terms whose optimization is challenging. They are either too strong to optimize reliably or too weak to achieve their goals. We propose to initialize the networks with a rich representation containing a palette of potentially useful features, ready to be used by even simple models. On the one hand, a rich representation provides a good initialization for the optimizer. On the other hand, it also provides an inductive bias that helps OoD generalization. Such a representation is constructed with the Rich Feature Construction (RFC) algorithm, also called the Bonsai algorithm, which consists of a succession of training episodes. During discovery episodes, we craft a multi-objective optimization criterion and its associated datasets in a manner that prevents the network from using the features constructed in the previous iterations. During synthesis episodes, we use knowledge distillation to force the network to simultaneously represent all the previously discovered features. Initializing the networks with Bonsai representations consistently helps six OoD methods achieve top performance on ColoredMNIST benchmark. The same technique substantially outperforms comparable results on the Wilds Camelyon17 task, eliminates the high result variance that plagues other methods, and makes hyperparameter tuning and model selection more reliable.