Improving training data selection to boost model robustness and fairness
Mind the Tail: Data Selection Through Sparse Concepts
Computer Vision and Pattern Recognition
Summary
Training AI models on images involves choosing which examples to use carefully. The authors show that picking examples based on how well they represent rare or less frequent features, called tail concepts, helps models perform better, especially when facing tricky or unusual images later. They use a way to find and count these features in images, then select examples that fill gaps instead of just focusing on common features. This approach leads to models that are both more reliable and less prone to errors caused by unusual backgrounds or confusing details.
What this means in practice
- •For machine learning engineers: Improve training datasets by including underrepresented visual features to build more robust image classifiers.
- •For computer vision teams: Enhance model reliability against rare or confusing image contexts by targeting data selection toward sparse concepts.
Authors
Ziheng Zhang, Neha Hulkund, Colin Lee, Zheda Mai, Hui Shen, Ying Ding, Tanya Berger-Wolf, Sara Beery, Wei-Lun Chao, Jianyang Gu
Abstract
Data selection reallocates training support across visual concepts, including objects, attributes, color, texture, and backgrounds. Different allocation strategies may produce varied downstream model behavior, especially for tail concepts supported by only a few samples. Discarding them can hurt test-time robustness, while overemphasizing samples lying on category boundaries may introduce unstable training signals. This work explicitly investigates this relationship across fine-grained, spurious-correlation, and general image classification, as well as multimodal pretraining. Specifically, we use sparse autoencoders (SAEs) to extract concept activations from individual images. During selection, we accumulate the dataset-level concept support for each added sample. Applying diminishing returns to this distribution down-weights well-represented concepts and shifts the focus to the distribution tail. We then incorporate local support to find samples that reliably strengthen under-supported concepts. Empirical results show consistent improvements over strong data selection baselines across all benchmarks. We also reveal the connection between concepts our method enhances and the induced behavioral changes, including improved robustness against out-of-distribution and spurious contexts. This work establishes concept support, beyond individual sample importance, as a useful principle for data-efficient and robust learning.