Active Learning of Robust Models

Mizu Nishikawa-Toomey; Andreas Kirsch; Laurent Charlin; Dhanya Sridhar

cs.LGNov 6, 2026

Active Learning of Robust Models

Authors: Mizu Nishikawa-Toomey, Andreas Kirsch, Laurent Charlin, Dhanya Sridhar

Organizations: Mila – Quebec AI Institute · Université de Montréal · HEC Montréal

Abstract

In many domains from species identification to cancer diagnosis, labelling data to train predictive models remains costly. In these settings, it is advantageous to intelligently select data points to label to achieve the maximum information gain, this is the motivation behind the study of active learning. In addition to the high costs of acquiring data in such domains, there is the added complexity of data distribution shifts that can occur between training and deployment. For example, the features of cancerous skin cells taken in one country may be different in another region of the world. These two challenges - efficient labelling of data, and learning models robust to distribution shifts have been studied separately but not together, even though these challenges often co-occur in real-life settings. We identify failure modes of standard active learning approaches when the distribution of data at deployment is different from the distribution in which the data is sampled for labelling. We propose a new strategy for efficiently learning classifiers that generalise under specific forms of distribution shift, namely subpopulation shift. Drawing on work in learning from heterogeneous data, we aim to train models that perform well across subpopulations of the data, through active learning. To handle this challenging setting, we exploit the fact that many datasets include metadata about data collection that can serve as a noisy proxy for each sample's underlying subpopulation. We demonstrate how to leverage these noisy labels during active learning selection, and demonstrate empirically that our approach improves classifier accuracy across subpopulations.

Figures & tables

Panel of Figure 1 — Figure 1: A causal graph that reflects the invariance properties we consider. The variable $e$ denotes the latent environment which influences the distribution of the observed data, and the variable $s$ is the source label which has a noisy relationship to $e$ , the source label will be further explained in Section 3.3 . $\mathbf{a}_{\mathrm{inv}}$ and $\mathbf{a}_{\mathrm{sp}}$ denote the invariant and spurious features of the data.

Panel of Figure 2 — Figure 2: Binary classification of a two dimensional input, using a max margin classifier. The green line indicates the learned decision boundary, the underlying true decision boundary is on the x axis (y=0 line). Red shaded regions indicate the regions in which the model will incorrectly classify points. Gray shaded regions show regions of highest uncertainty for the classifier. When the data is spuriously correlated, there are large regions of confidently incorrect points.

Panel of Figure 3 — Figure 3: Active learning on semi-synthetic datasets with spurious correlations.

Panel of Figure 6 — Figure 6: WGA reported on active learning on the Waterbirds (top), CelebA (middle) and ColoredMNIST (bottom) datasets with increasing numbers of sources $|S|=8\text{ (left)},|S|=12\text{ (right)}$ and $|S_{0}|=1$ . The higher, the better.

Figure 7: Changing the mixture of data from $e=0$ and $e=1$ in each source for the ColoredMNIST dataset. Each source contains a mixture of data from both environments, containing 1%, 5%, or 10% from one environment and the rest from another. Going from left to right we have increasing amount of noise.

Figure 8: Changing the difficulty of the invariant feature for the ColoredMNIST dataset. Dotted lines show the original dataset with all 10 digits, solid lines show experiments with an easier invariant feature with only two digits in the dataset.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 11: Uncertainty quantification using a Bayesian logistic regression model of unseen data from a test set. Red hues (Red, yellow, orange) are y=0 predictions, blue hues (blue, green) are y=1 predictions, with a higher uncertainty for those classes represented by lighter colours - yellow and green, low uncertainty with darker colours - blue and red. The data on which the Bayesian logistic regression model is trained on is indicated by black crosses. There are regions where the model is highly confident and wrong as well as regions of high uncertainty and the model is correct.

Hyper-parameter	ColoredMNIST	Waterbirds	CelebA	Camelyon
num_epochs	50	30	60	50
batch_size	130	2	2	10
learning rate	$1\times 10^{-4}$	$1\times 10^{-6}$	$1\times 10^{-6}$	$1\times 10^{-5}$
AL batch size	10	30	90	5
Dropout probability	0.5	0.5	0.5	0.5

Appendix

Table 7: Hyper-parameters used for each dataset.

Figure 12: Changing hyper-parameter K for the ColoredMNIST dataset.

Figure 13: Changing the hyper-parameter K for the Waterbirds dataset.

Figure 14: Changing the hyper-parameter K for the Camelyon dataset.

Figure 15: The number of points acquired for each acquisition strategy from environment $e=0$ on active learning on the ColoredMNIST dataset. Our methods successfully up-sample the proportion of the less common and more unusual environment $e=0$ . Random sampling indicates the naturally occurring base proportion of the environment in the dataset. It was observed that random and Random g-DRO samples the same number of points from $e=0$ , in this dataset, so the results for Random g-DRO are omitted from the plot.

Panel of Figure 16 — Figure 16: Number of points queried from the minority environment $e=0$ for the Waterbirds dataset, left and the CelebA dataset, right, both with $|S|=4,|S_{0}|=1$ . Our methods are able to up-sample this group above the baseline rates.

Figure 17: Comparing standard active learning acquisition functions that apply scores at a per point level, compared to our methods SoftMax-SoftRank and Top1 SoftRank on the ColoredMNIST dataset.

Figure 18: Comparing standard active learning acquisition functions that apply scores at a per point level, compared to our methods SoftMax-SoftRank and Top1 SoftRank on the Waterbirds dataset. Our methods that apply scores to determine the source budget outperform baseline methods that only apply the score at a point level.

Figure 19: Comparing standard active learning acquisition functions that apply scores at a per point level, compared to our methods SoftMax, SoftRank and Top1 SoftRank on the Camelyon17 dataset.

Figure 20: Changing the mixture of data from $e=0$ and $e=1$ in each source for the ColoredMNIST dataset. Each source contains a mixture of data from both environments, containing 1%, 5%, or 10% from one environment and the rest from another. Going down the rows we have increasing number of total sources $|S|$ , and going across the columns we have increasing amount of noise. More mixing results in less information contained in the source identifier $s$ , resulting in the advantage of our method declining.

Figure 21: The difference in per-group accuracies across groups reported on active learning on the Waterbirds dataset, the smaller the better. Our methods SoftMax-SoftRank and Top1-SoftRank achieve the lowest range of group-accuracies.

Figure 22: The difference in per group accuracies across groups reported on active learning on the CelebA dataset for $|S|=4$ .

Panel of Figure 24 — Figure 24: For $e{=}0$ we have two Gaussians centered at $\pm|w_{\mathrm{inv}}-w_{\mathrm{sp}}|$ (close to zero), and for $e{=}1$ they are centered at $\pm|w_{\mathrm{inv}}+w_{\mathrm{sp}}|$ (further from zero). The region $[-c,c]$ is the region where the point will be selected for querying. The area shaded in orange shows the region $q_{0}(c)$ , where points from $e=0$ will be selected for querying, and the area shaded in purple shows the region $q_{1}(c)$ where points from $e=1$ will be selected for querying. The ratio $R=q_{0}(c)/q_{1}(c)$ can be changed by the noise of the two Gaussians, with increasing noise resulting in smaller $R$ and the size of $w_{\mathrm{sp}}$ , driving apart the two Gaussians, resulting in larger $R$ . As a result, the proportion of minority samples is highly affected by the parameters of the model.

Active Learning

Active learning selects the examples, labels, features, or experiments most useful for improving a model under limited budgets.

See all 315 papers

Active Learning with Imperfect Labels: Optimal Labeler Assignment and Sample Selection
Dec 14, 2025Pouya Ahadi, Blair Winograd, Camille Zaug +3Human-in-the-Loop Annotation Active Learning
School of Industrial and Systems Engineering Georgia Institute of Technology Atlanta, GA 30332, USA · Ford Motor Company
Guiding Data Allocation for Robust Subpopulation Generalization
Nov 9, 2026Zhaoying Pan, Yipei Wang, Shenyu Lu +1Subgroup Robustness Learning with Distribution Shift
Elmore Family School of Electrical and Computer Engineering, Purdue University
Active Learning under Simulated Class Imbalance and Label Shift in Binary Semantic Segmentation
Jan 8, 2026Julien Combes, Jean-François Coeurjolly, Alexandre DervilleActive Learning Label Shift
SVH, Michelin · SVH · Michelin

Dec 14, 2025J/K move · Enter open · S save

Active Learning with Imperfect Labels: Optimal Labeler Assignment and Sample Selection

Pouya Ahadi, Blair Winograd, Camille Zaug, Karunesh Arora+2

School of Industrial and Systems Engineering Georgia Institute of Technology Atlanta, GA 30332, USA · Ford Motor Company

Human-in-the-Loop Annotation Active Learning Learning with Noisy Labels

Open paper AskPDF arXiv

Active Learning (AL) is commonly used in applications where labeling data is expensive or time-consuming. In practice, however, labels are often noisy due to varying labeler expertise and annotation uncertainty, especially for complex or ambiguous samples. Learning from such imperfectly labeled data can degrade classifier performance. We propose an AL framework that explicitly accounts for label noise by optimally assigning labelers and selecting samples to minimize labeling error. Our approach, called OLAS (Optimal Labeler Assignment and Sampling), uses a noise model that depends on both labeler accuracy and model uncertainty to guide these decisions. We develop two tractable optimization formulations: one for assigning samples to labelers to minimize worst-case noise, and another for selecting samples while controlling overall label noise. Theoretical results provide closed-form solutions under mild conditions. Empirical evaluations on benchmark datasets and a real-world warranty claim classification problem show that OLAS achieves the highest or near-highest classification accuracy among existing AL strategies across most settings, using only a single label per sample.

Abstract

Active Learning

Explore similar work

Active Learning with Imperfect Labels: Optimal Labeler Assignment and Sample Selection

Guiding Data Allocation for Robust Subpopulation Generalization

Active Learning under Simulated Class Imbalance and Label Shift in Binary Semantic Segmentation