Examining and Combating Spurious Features under Distribution Shift

Many machine learning models that minimize the average training loss via empirical risk minimization (ERM) are trained and evaluated on randomly shuffled and split training and test sets. However, such in-distribution learning setups can hide critical issues: models that achieve high accuracy on average often underperform when the test distribution drifts away from the training one (Hashimoto et al., 2018; Koenecke et al., 2020; Koh et al., 2020). Such models are often “right for the wrong reasons” due to reliance on spurious correlations (or “dataset biases”) (Torralba & Efros, 2011; Goyal et al., 2017; McCoy et al., 2019; Gururangan et al., 2018), heuristics that hold for most training examples but are not inherent to the task of interest, such as strong associations between the presence of green pastures background with the label “cows” in image classification. Naturally, models that use such features will fail when tested on data where the correlation does not hold.

Recent work has investigated how models trained with ERM learn spurious features that do not generalize, from the points of view of causality (Arjovsky et al., 2019), understanding model overparameterization (Sagawa et al., 2020b) and information theory (Lovering et al., 2021). However, these works have not characterized the idea of spurious features mathematically. In this paper, we characterize spurious features from an information-theoretic perspective. We consider prediction of target random variable $Y\in\mathcal{Y}$ from input variable $X\in\mathcal{X}$ and characterize spurious features learned under changes to the input distribution $p(X)$ (i.e. covariate shift).

A central goal of machine learning is to learn true causal relationships between $X$ and $Y$ in a manner robust to spurious factors concerning the variables. We assume that there exists an “ideal” data distribution $p_{\text{ideal}}$ (short for $p_{\text{ideal}}(X,Y)$ below) which contains data from all possible experimental conditions concerning the confounders that cause spurious correlations, both observable and hypothetical (Lewis, 2013; Arjovsky et al., 2019; Bellot & van der Schaar, 2020). For example, consider the problem of classifying images of cows and camels (Beery et al., 2018). Under the ideal conditions, we assume that pictures of cows and camels on any background can be collected, including cows in deserts and camels in green pastures. Therefore, under $p_{\text{ideal}}$ the background of the image $X$ is no longer a spurious factor of the label $Y$. However, such an “ideal” distribution $p_{\text{ideal}}$ is not accessible in practice (Bahng et al., 2020; Koh et al., 2020; McCoy et al., 2019), and our training distribution $p_{\text{train}}$ (often, in practice, an associated empirical distribution) does not match $p_{\text{ideal}}$. ERM-based learning algorithms indiscriminately fit all correlations found in $p_{\text{train}}$, including spurious correlations based on confounders (Tenenbaum, 2018; Lopez-Paz, 2016).

To investigate the spurious features learned under the distribution shift from $p_{\text{ideal}}$ to $p_{\text{train}}$, we first characterize those features of $X$ which most efficiently capture all possible information needed to predict $Y$. We define these robust features using the notion of minimal sufficient statistic (MSS) (Dynkin, 2000; Cvitkovic & Koliander, 2019) under $p_{\text{ideal}}$. We then examine whether the features learned under $p_{\text{train}}$ contain spurious features compared to the MSS learned under $p_{\text{ideal}}$. Through our analysis, we find that even only with covariate shift, the features learned on $p_{\text{train}}$ can contain spurious features or miss robust features of $p_{\text{ideal}}$.

Models that fit spurious correlations in $p_{\text{train}}$ can be vulnerable to groups (subpopulations of $p_{\text{ideal}}$/$p_{\text{test}}$) where the correlation does not hold. A common approach to avoid learning a model that suffers high worst group errors is group distributionally robust optimization (group DRO), a training procedure that efficiently minimizes the worst expected loss over a set of groups in the training data (Oren et al., 2019; Sagawa et al., 2020a). The partition of groups can be defined in several ways, such as by presence of manually identified potentially spurious features (Sagawa et al., 2020a), data domains (Koh et al., 2020), or topics of text (Oren et al., 2019). In a typical setup, the groups of interest in the test set align with those used to partition the training data. Under such setups, group DRO usually outperforms ERM with respect to the worst-group accuracy. We contend that this is because it promotes learning robust features that perform uniformly well across all groups. However, in many tasks, we can not collect clean group membership of training examples due to expensive annotation cost or privacy concerns regarding e.g. demographic identities of users or other sensitive information.

Inspired by our analysis of spurious features, we demonstrate that group DRO can fail under “imperfect” partitions of training data that are not consistent with the test set, especially when reducing spurious correlation in one group could exacerbate the spurious correlations in another (§4.2), as shown in Fig. 1. This is because group DRO treats each training group as a unit, preventing it from adjusting learning weights differently for subgroups within each group. Recent work has proposed to use sophisticated unsupervised clustering algorithm to search for meaningful subclasses (Sohoni et al., 2020) and execute group DRO on the found subclasses. To learn robust models under noisy protected groups, Wang et al. (2020) designs robust approaches that is based on an estimate of a noise model between the clean and noisy groups. Instead of relying on good partitions of groups or a not readily available noise model, we propose group-conditional DRO (GC-DRO) that defines the uncertainty set over the joint distribution of groups and their instances (i.e. $q(G)q(X,Y|G)$). Every training example is reweighted by both its group weight and the instance-level weight, which offers a more flexible uncertainty set compared to group DRO. Through extensive experiments on three tasks — facial attribute classification, natural language inference, and toxicity detection, we show that GC-DRO significantly outperforms both ERM and group DRO in various partitions of training data and demonstrate the robustness of GC-DRO against various group partitions.

To study spurious features, we need to formally define which features or properties of the data describe spurious correlations, and which features are robust features relevant to the task at hand. In supervised learning we are interested in finding a good representation $T(X)$ of the input $X$¹¹1We assume that $T(X)$ is a deterministic mapping of $X$ given neural network parameters. that is useful to predict a target label $Y$. What characterizes the optimal representations of $X$ w.r.t. $Y$ is much debated, but a common assertion is that $T(X)$ should be a minimal sufficient statistic (MSS) of $X$ for $Y$ (Adragni & Cook, 2009; Schwartz-Ziv & Tishby, 2017; Achille & Soatto, 2018; Cvitkovic & Koliander, 2019), which is:

(i) $T(X)$ should be sufficient for $Y$, i.e. $\forall x\in\mathcal{X},t\in\mathcal{T},y\in\mathcal{Y},p(x|t,y)=p(x|t)$, which is equivalent to $p(y|t,x)=p(y|t)$. This means given the value of $T(X)$, the distribution of $X$ does not depend on the value of $Y$.

(ii) Given that $T(X)$ is sufficient, it should be minimal w.r.t. $X$, i.e. for any sufficient statistic $S$, there exists a deterministic function $f$ such that $T=f(S)$ almost everywhere w.r.t. $X$. This means for any measurable, non-invertible function $g$, $g(T)$ is no longer sufficient for $Y$.

In other words, the minimal sufficient statistics most efficiently capture all information useful for predicting $Y$. The notion of MSS has been connected to Shannon’s information theory (Kullback & Leibler, 1951; Cover, 1999) and extended to any joint distribution $P(X,Y)$ of $X$ and $Y$ in the information bottleneck (IB) framework (Tishby et al., 2000; Shamir et al., 2010; Kolchinsky et al., 2019), which provides a principled way to characterize the extraction of relevant information from $X$ for predicting $Y$. Loosely speaking, learning a MSS $T$ is equivalent to maximizing $I(T(X);Y)$ (sufficiency) and minimizing $I(X;T(X))$ (minimality).

The training data is often marred by various abnormalities, such as selection biases (Buolamwini & Gebru, 2018) and confounding factors (Gururangan et al., 2018). We ask if the MSS learned under $p_{\text{train}}$ are robust features under $p_{\text{ideal}}$. Note that we do not study how to learn MSS via ERM in this paper, on the other hand, considering that MSS provides a good candidate for robust representations, we want to study if the MSS learned under $p_{\text{train}}$ contains spurious features with respect to the MSS learned under $p_{\text{ideal}}$, which are universal robust features against various spurious factors.

We consider the distribution shift in $p(X)$,³³3It is often assumed that $p(Y|X)$ is invariant in supervised learning problems (Arjovsky et al., 2019). also known as covariate shift (David et al., 2010), and we show that the entropy of MSS learned under $p_{\text{train}}$ conditioned on the robust features is zero in Theorem 1 with proofs in §A.

Theorem 1 tells us that $T_{train}(X)$ is deterministic given $T_{ideal}(X)$ under ${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}p_{\text{train}}}$ (shown in blue to distinguish from Eq. 1). However, this does not imply $H_{ideal}(T_{train}(X)|T_{ideal}(X))=0$ under ${\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}p_{\text{ideal}}}$. Thus, we cannot conclude that $T_{train}(X)$ contains no spurious features. We further discuss the implications with two cases based on the relationship between the support of input $\mathcal{X}_{train}$ and that of $\mathcal{X}_{ideal}$: (1) $\mathcal{X}_{train}=\mathcal{X}_{ideal}$ and (2) $\mathcal{X}_{train}\subset\mathcal{X}_{ideal}$. When the input support of $p_{\text{train}}$ is equal to that of $p_{\text{ideal}}$, we have the following corollary:

Corollary 1 corroborates the findings in Wen et al. (2014) that the (unweighted) solution learned by ERM is also the robust solution when only covariate shift exists and $\mathcal{X}_{train}=\mathcal{X}_{ideal}$. In practice, however, this assumption does not hold (because we only have datasets with limited support) and thus the representation $T_{train}(X)$ learned by ERM is not necessarily equivalent to $T_{ideal}(X)$. By Theorem 1, $T_{train}(X)$ is deterministic given $T_{ideal}(X)$ under $p_{\text{train}}$, which implies that the information contained in $T_{train}(X)$ is equal to or less than that contained in $T_{ideal}(X)$. In the former case, $T_{train}(X)$ can be equivalent in representation to $T_{ideal}(X)$ but can also contain spurious features that co-occur with the robust features in the training data. In the latter case, $T_{train}(X)$ can miss robust features in $T_{ideal}(X)$. We demonstrate these two cases with synthetic experiments in Appendix D due to space limit.

The discussions in §3 suggest that under covariate shift, directly learning from the empirical data distribution can result in learning the spurious correlations satisfied by the majority of the training data. When the spurious factors are known, we can apply group distributionally robust optimization (group DRO), which reweights the losses of different groups associated with spurious factors to alleviate covariate shift and learn robust features that generalize to both minority and majority groups. In this section, we first review group DRO and discuss under which cases it can fail.

Since group DRO can be problematic with imperfect partitions, we propose a more flexible uncertainty set over the joint distribution of $(x,y,g)$, i.e. $q(g)q(x,y|g)$, using fine-grained weights over instances within each group instead of treating the entire group as a whole. We extend the $\alpha$-covered distribution to both the group-level ($q(g)$) and conditional instance-level ($q(x,y|g)$) distributions to define the uncertainty set $\mathcal{Q}$. At training time, a sample is weighted by both its group weight induced from $q(g)$ as well as the instance-level weight induced from $q(x,y|g)$. Specifically, the new uncertainty set is

where $N$ is the number of training examples and $\alpha,\beta\in(0,1]$. Denote $n_{i}$ the number of samples in group $i$, then $p^{train}(x,y|g=i)=\frac{1}{n_{i}}$. The second constraint of Eq. 6 can be rewritten as $\frac{1}{N}\leq q(x,y|g)\leq\frac{1}{\beta n_{i}}$. Compared with the $\beta$-covered distribution, we add a lower bound $q(x,y|g)\geq\frac{1}{N}$ to compensate for imbalanced group sizes. With a plain $\beta$-covered distribution for $q(x,y|g)$, the DRO objective roughly upweights a $\beta$-fraction of instance losses of each group. However, we only want to emphasize a small subset of examples that perform badly in the majority groups. Thus, we add this lower bound to $q(x,y|g)$ in Eq. 6 to directly “punish” larger groups. To see this, the percentage of examples that are upweighted by $\frac{1}{\beta}$ in group $i$ is roughly $\frac{N-n_{i}}{N-n_{i}\beta}\beta$, which is monotonically decreasing function w.r.t. $n_{i}$. Therefore, the larger the group size $n_{i}$ is, the smaller fraction of instances in group $i$ are upweighted.

In this section, we evaluate the proposed group-conditional DRO on one image classification task and two language tasks — natural language inference and toxicity detection. To demonstrate the effectiveness of our method under various partitions of data, we first introduce the clean (group number $m=|\mathcal{A}|\times|\mathcal{Y}|$) and imperfect data partitions of each task. As we discussed at the end of §4.2, there are various cases where the partitions of training data are imperfect such that each group is not purely associated with examples from one pair of $(a,y)$. In this section, we inspect several cases reflecting diverse properties of partitions to evaluate our method. First, on the image and NLI tasks, we manually design adversarial partitions of data such that there are explicit conflicts between groups and purely reweighing over groups cannot eliminate covariate shift (§6.1). Second, we use the attributes provided by a supervised classifier to create the imperfect partitions of the toxicity data set (§6.1). Third, we also perform unsupervised clustering on the toxicity data set to obtain imperfect partitions in §6.4.

Through a mathematical characterization of features used in prediction, we have demonstrated that under covariate shift ERM models can pick up spurious features or miss robust features. The GC-DRO algorithm resulting from this analysis allows for a more flexible uncertainty set that performs consistently well in the worst test group under different partitions. This new understanding of features opens up new avenues in both redesigning our distributionally robust algorithms, and further characterizing possible spurious factors that may influence model robustness, for example through unsupervised learning.

This work in this paper was supported in part by a Facebook SRA Award and the NSF/Amazon Fairness in AI program under grant number 2040926. We thank the anonymous reviewers for useful suggestions.