Kiran Krishnamachari, See-Kiong Ng, & Chuan-Sheng Foo (2024). Uniformly Distributed Feature Representations for Fair and Robust Learning. Transactions on Machine Learning Research.
Abstract:
A fundamental challenge in machine learning is training models that generalize well to distributions different from the training distribution. Empirical Risk Minimization (ERM), which is the predominant learning principle, is known to under-perform in minority sub-populations and fail to generalize well in unseen test domains. In this work, we propose a novel learning principle called Uniform Risk Minimization (URM) to alleviate these issues. We first show theoretically that uniform training data distributions and feature representations support robustness to distribution shifts. Motivated by this result, we propose an empirical method that trains deep neural networks to learn a uniformly distributed feature representation in their final activation layer for improved robustness. Our experiments on multiple datasets for sub-population shifts and domain generalization show that URM improves the generalization of deep neural networks without requiring knowledge of groups or domains during training. URM is competitive with the best existing methods designed for these tasks and can also be easily combined with them for improved performance. Our work sheds light on the importance of the distribution of learned feature representations for model robustness and fairness. Code is available at https://github.com/kiranchari/UniformRiskMinimization.
License type:
Attribution 4.0 International (CC BY 4.0)
Funding Info:
There was no specific funding for the research done