Representation Matters:
Assessing the Importance of Subgroup Allocations in Training Data

Targeted data collection in ML. Recent research evidences that targeted data collection can be an effective way to reduce disparate performance of ML models evaluated across sub-populations (Raji & Buolamwini, 2019). Chen et al. (2018) present a formal argument that the addition of training data can lessen discriminatory outcomes while improving accuracy of learned models, and Abernethy et al. (2020) show that adaptively collecting data from the lowest-performing sub-population can increase the minimum accuracy over groups. At the same time, there are many complexities of gathering data as a solution to disparate performance across groups. Targeted data collection from specific groups can present undue burdens of surveillance or skirt consent (Paullada et al., 2020). When ML systems fail portions of their intended population due to issues of measurement and construct validity, more thorough data collection is unlikely to solve the issue without further intervention (Jacobs & Wallach, 2019).

With these complexities in mind, we study the importance of numerical representation in training datasets in achieving diverse objectives. Optimal allocation of subpopulations in statistical survey designs dates back to at least Neyman (1934), including stratified sampling methods to ensure coverage across sub-populations (Lohr, 2009). For more complex prediction systems, the field of optimal experimental design (Pukelsheim, 2006) studies what inputs are most valuable for reaching a given objective, often focusing on linear prediction functions. We consider a constrained sampling structure and directly model the impact of group allocations on subgroup performance for general model classes.

Valuing data. In economics, allocations indicate a division of goods to various entities (Cole et al., 2013). While we focus on the influence of data allocations on model accuracies across groups, there are many approaches to valuing data. Methods centering on a theory of Shapley valuations (Yona et al., 2019; Ghorbani & Zou, 2019) complement studies of the influence of individual data points on model performance to aid subsampling data (Vodrahalli et al., 2018).

Methods for handling group-imbalanced data. Importance sampling and importance weighting are standard approaches to addressing class imbalance or small groups sizes (Haixiang et al., 2017; Buda et al., 2018), though the effects of importance weighting for deep learning may vary with regularization (Byrd & Lipton, 2019). Other methods specifically address differential performance between groups. Maximizing minimum performance across groups can reduce accuracy disparities (Sagawa et al., 2020) and promote fairer sequential outcomes (Hashimoto et al., 2018). For broader classes of group-aware objectives, techniques exist to mitigate unfairness or disparate performance of black box prediction functions (Dwork et al., 2018; Kim et al., 2019). It might not be clear a priori which subsets need attention; Sohoni et al. (2020) propose a method to identify and account for hidden strata, while other methods are defined for any subsets (Hashimoto et al., 2018; Kim et al., 2019).

In addition to modifying the optimization objective or learning algorithm, one can also modify the input data itself by re-sampling the training data to match the desired population by downsampling or by upsampling with data augmentation techniques (Chawla et al., 2002; Iosifidis & Ntoutsi, 2018).

$\Delta^{k}$ denotes the $k$-dimensional simplex. $\mathbb{Z}^{+}$ denotes non-negative integers and $\mathbb{R}^{+}$ non-negative reals.

In light of the decomposition of the population-level risk as a weighted average over group risks in Eq. (1), we now consider the composition of fixed-size training sets, in terms of how many samples come from each group.

It will be illuminating to consider $\vec{\alpha}$ not only as a property of an existing dataset, but as a parameter governing dataset construction, as captured in the following definition.

The procedure in Definition 2 suggests formalizing data collection as a component of the learning process in the following way: in addition to choosing a loss function and method for minimizing the risk, choose the relative proportions at which to sample the groups in the training set:

In Section 3, we show that when a dataset curator can design dataset allocations in the sense of Definition 2, they have the opportunity to improve accuracy of the trained model. Of course, one does not always have the opportunity to collect new data or modify the composition of an existing dataset. Section 2.2 considers methods for using fixed datasets that have groups with small training set allocation $\alpha_{g}$, relative to $\gamma_{g}$, or high risk for some groups relative to the population.

In classical empirical risk minimization (ERM), one learns a function from class $\mathcal{F}$ that minimizes average prediction loss over the training instances $(x_{i},y_{i},g_{i})\in\mathcal{S}$ (we also abuse notation and write $i\in\mathcal{S}$) with optional regularization $R$:

There are many methods for addressing small group allocations in data (see Section 1.1). Of particular relevance to our work are objective functions that minimize group or population risks. In particular, one approach is to use importance weighting (IW) to re-weight training samples with respect to a target distribution defined by $\vec{\gamma}$:

This empirical risk with instances weighted by ${\gamma_{g}}/{\alpha_{g}}={\gamma_{g}}n/{n_{g}}$ is an unbiased estimate of the population risk, up to regularization. While unbiasedness is often desirable, IW can induce high variance of the estimator when $\gamma_{g}/\alpha_{g}$ is large for some group (Cortes et al., 2010), which happens when group $g$ is severely underrepresented in the training data relative to their population prevalence.

Alternatively, group distributionally robust optimization (GDRO) (Hu et al., 2018; Sagawa et al., 2020) minimizes the maximum empirical risk over all groups:

For losses $\ell$ which are continuous and convex in the parameters of $f$, the optimal GDRO solution corresponds to the minimizer of a group-weighted objective: $\tfrac{1}{n}\sum_{i=1}^{n}w(g_{i})\cdot\ell(f(x_{i}),y_{i})$, though this is not in general true for nonconvex losses (see Prop. 1 of Sagawa et al., 2020, and the remark immediately thereafter).

Given the correspondence of GDRO (for convex loss functions) and IW to the optimization of group-weighted ERM objectives, we now investigate the joint roles of sample allocation and group re-weighting for estimating group-weighted risks. For prediction function $f$, loss function $\ell$, and group weights $w:\mathcal{G}\rightarrow\mathbb{R}^{+}$, let $\hat{L}(w,\alpha,n;f,\ell)$ be the random variable defined by:

where the randomness in $\hat{L}$ comes from the draws of $x_{i},y_{i}$ from $\mathcal{D}_{g_{i}}$ according to procedure $\mathcal{S}(\vec{\alpha},n)$ (Definition 2), as well as any randomness in $f$.

The following proposition shows that group weights and allocations play complementary roles in risk function estimation. In particular, if $w(g)$ depends on the sampling allocations $\alpha_{g}$, then there are alternative group weights $w^{*}$ and allocation $\vec{\alpha}^{*}$ such that the alternative estimator has the same expected value but lower variance.

Since the estimation of risk functions is a key component of learning, Proposition 1 illuminates an interplay between the roles of sampling allocations and group-weighting schemes like IW and GDRO. When allocations and weights are jointly maximized, the optimal allocation accounts for an implicit target distribution $\gamma^{\prime}$ (defined above), which may vary by objective function. The optimal weights account for per-group variability $Var_{(x,y)\sim\mathcal{D}_{g}}[\ell(f(x),y))]$. In Section 4 we find that it can be advantageous to use IW and GDRO when some groups have small $\alpha_{g}/\gamma_{g}$; though the boost in accuracy is less than having an optimally allocated training set to begin with, and diminishes when all groups are appropriately represented in the training set allocation.

We model the impact of allocations on performance with scaling laws that describe per-group risks as a function of the number of data points from their respective group, as well as the total number of training instances.

1 is similar to the scaling law in Chen et al. (2018), but includes a $\tau_{g}^{2}n^{-q}$ term to allow for data from other groups to influence the risk evaluated on group $g$. It additionally requires that the same exponents $p,q$ apply to each group, an assumption that underpins our theoretical results in Section 3. We examine the extent to which 1 holds empirically in Section 4.3, and will modify Eq. (1) to include group-dependent terms $p_{g},q_{g}$ when appropriate. The following examples give intuition into the form of Eq. (1).

It is often advantageous to pool training data to learn a single classifier. In this case, model performance evaluated on group $g$ will depend on both $n_{g}$ and $n$, as the next examples show.

Example 3 suggests that in some settings, we can relate $\sigma_{g}$ and $\tau_{g}$ to ‘group specific’ and ‘group agnostic’ model components that affect performance for group $g$. In general, the relationship between group sizes and group risks can be more nuanced. Data from different groups may be correlated, so that samples from groups similar to or different from $g$ have greater effect on $\mathcal{R}(\hat{f};\mathcal{D}_{g})$ (see Section 4.5). Eq. 1 is meant to capture the dominant effects of training set allocations on group risks and serves as our main structural assumption in the next section, where we study the allocation that minimizes the approximate population risk.

We now study properties of the allocation that minimizes the approximated population risk:

The following proposition lays the foundation for two corollaries which show that:

1. (1)

   when only the population prevalences $\vec{\gamma}$ vary between groups, the allocation that minimizes the approximate population risk up-represents groups with small $\gamma_{g}$;

2. (2)

   for two groups with different scaling parameters $\sigma_{g}$, the optimal allocation of the group with $\gamma_{g}<\tfrac{1}{2}$ is bounded by functions of $\sigma_{A},\sigma_{B}$, and $\vec{\gamma}$.

The proof of Proposition 2 appears in Section A.2. Note that $\vec{\alpha}^{*}$ does not depend on $n$, $\{\tau_{g}\}_{g\in\mathcal{G}}$, or $q$; this will in general not hold if powers $p_{g}$ differ by group.

We now study the form of $\vec{\alpha}^{*}$ under illustrative settings. Corollary 1 shows that when the group scaling parameters $\sigma_{g}$ in Eq. (1) are equal across groups, the allocation that minimizes the approximate population risk allocates samples to minority groups at higher than their population prevalences. The proof of Corollary 1 appears in Section A.3.

This shows that the allocation that minimizes population risk can differ from the actual population prevalences $\vec{\gamma}$. In fact, Corollary 1 asserts that near the allocation $\vec{\alpha}=\vec{\gamma}$, the marginal returns to additional data from group $g$ are largest for groups with small $\alpha_{g}$, enough so as to offset the small weight $\gamma_{g}$ in Eq. (1). This result provides evidence against the idea that small training set allocation to minority groups might comply with minimizing population risk as a result of a small relative contribution to the population risk.

Remark. A counterexample shows that $\alpha_{g}^{*}\leq\gamma_{g}$ does not hold for all $g$ with $\gamma_{g}>{1}/{|\mathcal{G}|}$. Take $\vec{\gamma}=[.68,.30,.01,.01]$ and $p=1$; Eq. (6) gives $\alpha_{2}^{*}>0.3=\gamma_{2}>1/4$. In general, whether group $g$ with $\gamma_{g}\geq{1}/{|\mathcal{G}|}$ gets up- or down-sampled depends on the distribution of $\vec{\gamma}$ across all groups.

Complementing the investigation of the role of the population proportions $\vec{\gamma}$ in Corollary 1, the next corollary shows that the optimal allocation $\vec{\alpha}^{*}$ generally depends on the relative values of $\sigma_{g}$ between groups. Inspecting Eq. (1) shows that $\sigma_{g}$ defines a limit of performance: if $\sigma_{g}^{2}$ is large, the only way to make the approximate risk for group $g$ small is to make $n_{g}$ large. For two groups, we can bound the optimal allocations $\vec{\alpha}^{*}$ in terms of $\{\sigma_{g}\}_{g\in\mathcal{G}}$, and the population proportions $\vec{\gamma}$. We let $A$ be the smaller of the two groups without loss of generality. From Eq. (6), we know that for two groups, $\alpha^{*}_{A}$ is increasing in $\tfrac{\sigma_{A}}{\sigma_{B}}$; Corollary 2 gives upper and lower bounds on $\alpha^{*}_{A}$ in terms of $\sigma_{A}$ and $\sigma_{B}$. Corollary 2 is proved in Section A.4.

Altogether, these results highlight key properties of training set allocations that minimize population risk. Experiments in Section 4 give further insight into the values of weights and allocations for minimizing group and population risks and apply the scaling law model in real data settings.

Here we give a brief description of each dataset we use. For more details, see (Tables 1 and B.1).

Modified CIFAR-4. To create a dataset where we can ensure class balance across groups, we modify the CIFAR-10 dataset (Krizhevsky, 2009) by subsetting to the bird, car, horse, and plane classes. We predict whether the image subject moves primarily by air (plane/bird) or land (car/horse) and group by whether the image contains an animal (bird/horse) or vehicle (car/plane); see Figure B.1. We set $\gamma=0.9$.

ISIC. The International Skin Imaging Collaboration dataset is a set of labeled images of skin lesions designed to aid development and standardization of imaging procedures for automatic detection of melanomas (Codella et al., 2019). For our main analysis, we follow similar preprocessing steps to (Sohoni et al., 2020), removing any images with patches. We predict whether a lesion is benign or malignant, and group instances by the approximate age of the patient of whom each photo was taken.

Goodreads. Given publicly available book reviews compiled from the Goodreads database (Wan & McAuley, 2018), we predict the rating (1-5) corresponding to each review. We group instances by genre.

Mooc. The HarvardX-MITx Person-Course Dataset contains student demographic and activity data from 2013 offerings of online courses (HarvardX, 2014). We predict whether a student earned a certification in the course and we group instances by highest completed level of education.

Adult. The Adult dataset, downloaded from the UCI Machine Learning Repository (Dua & Graff, 2017) and originally taken from the 1994 Census database, has a prediction task of whether an individual’s annual income is over $\$50,000$. We group instances by sex, codified as binary male/female, and exclude features that directly determine groups status.

We first investigate (a) the change in group and population performance at different training set allocations, and (b) the extent to which optimizing the three objective functions defined in Section 2.2 decreases average and group errors.

For each dataset, we vary the training set allocations $\vec{\alpha}$ between $(0,1)$ and $(1,0)$ while fixing the training set size as $n=\min_{g}n_{g}$ (see Table 1) and evaluate the per-group and population losses on subsets of the heldout test sets. For the image classification tasks, we compare group-agnostic empirical risk minimization (ERM) to importance weighting (implemented via importance sampling (IS) batches following the findings of Buda et al. (2018)) and group distributionally robust optimization (GDRO) with group-dependent regularization as described in Sagawa et al. (2020). For the non-image datasets, we implement importance weighting (IW) by weighting instances in the loss function during training, and do not compare to GDRO, as the gradient-based algorithm of Sagawa et al. (2020) is not easily adaptable to the predictors we use for these datasets. We pick hyperparameters for each method based on cross-validation results over a coarse grid of $\vec{\alpha}$; for IS, IW, and GDRO, we allow the hyperparameters to vary with $\vec{\alpha}$; for ERM we choose a single hyperparameter setting for all $\vec{\alpha}$ values.

Figure 1 highlights the importance of at least a minimal representation of each group in order to achieve low population loss (black curves) for all objectives. For CIFAR-4, the population loss increases sharply for $\alpha_{A}<0.1$ and $\alpha_{A}>0.8$, and for ISIC, when $\alpha_{A}<0.2$. While not as crucial for achieving low population losses for the remaining datasets, the optimal allocations $\vec{\alpha}^{*}$ (stars) do require a minimal representation of each group. The $\vec{\alpha}^{*}$ are largely consistent across the training objectives (different star colors). The population losses (black curves) are largely consistent across mid-range values of $\alpha_{A}$ for all training objectives. The relatively shallow slopes of the black curves for $\alpha_{A}$ near $\alpha^{*}_{A}$ (stars) stand in contrast to the per-group losses (blue and orange curves), which can vary considerably as $\vec{\alpha}$ changes. From the perspective of model evaluation, this reinforces a well-documented need for more comprehensive reporting of performance. From the view of dataset design, this exposes an opportunity to choose allocations which optimize diverse evaluation objectives while maintaining low population loss. Experiments in Section 4.4 investigate this further.

Across the CIFAR-4 and ISIC tasks, GDRO (dotted curves) is more effective than IS (dashed curves) at reducing per-group losses. This is expected, as minimizing the largest loss of any group is the explicit objective of GDRO. Figure 1 shows that GDRO can also improve the population loss (see $\alpha_{A}>0.7$ for CIFAR-4 and $\alpha_{A}<0.2$ for ISIC) as a result of minimizing the worst group loss. Importance weighting (dashed curves) has little effect on performance for Mooc and Adult (random forest models), and actually increases the loss for Goodreads (multiclass logistic regression model).

For all the datasets we study, the advantages of using IS or GDRO are greatest when one group has very small training set allocation ($\alpha_{A}$ near $0$ or $1$). When allocations are optimized (stars in Figure 1), the boost that these methods give over ERM diminishes. In light of Proposition 1, these results suggest that in practice, part of the value of such methods is in compensating for sub-optimal allocations. We find, however, that explicitly optimizing the maximum per-group loss with GDRO can reduce population loss more effectively than directly accounting for allocations with IS.

Section C.2 shows that similar phenomena hold for different loss functions and models on the same dataset, though the exact $\vec{\alpha}^{*}$ can differ. In Section C.1, we show that losses of groups with small $\alpha_{g}$ can degrade more severely when group attributes are included in the feature matrix, likely a result of the small number of samples from which to learn group-specific model components (see Example 3).

For each dataset, we combine the results in Figure 1 with extra subsetting runs where we vary both $n_{g}$ and $n$. From the combined results, we use nonlinear least squares to estimate the parameters of modified scaling laws, where exponents can differ by group

The estimated parameters of Eq. (7) given in Table 2 capture different high-level phenomena across the five datasets. For CIFAR-4, $\hat{\tau}_{g}\approx 0$ for both groups, indicating that most of the group performance is explained by $n_{g}$, the number of training instances from that group, whereas the total number of data points $n$, has less influence. For Goodreads, both $n_{g}$ and $n$ have influence in the fitted model, though $\hat{\tau}_{g}$ and $\hat{q}_{g}$ are larger than $\hat{\sigma}_{g}$ and $\hat{p}_{g}$, respectively. For ISIC, $\hat{\tau}_{A}\approx 0$ but $\hat{\tau}_{B}\not\approx 0$, suggesting other-group data has little effect on the first group, but is beneficial to the latter. For the non-image datasets (Goodreads, Mooc, and Adult), $0\!<\!\hat{\sigma}_{g}\!<\!\hat{\tau}_{g}$ and $\hat{p}_{g}\!<\!\hat{q}_{g}$ for all groups.

These results shed light on the applicability of the assumptions made in Section 3. Figure B.2 in Section B.5 shows that the fitted curves capture the overall trends of per-group losses as a function of $n$ and $n_{g}$. However, the assumptions of Proposition 2 and Corollaries 1 and 2 (e.g., equal $p_{g}$ for all $g\in\mathcal{G}$) are not always reflected in the empirical fits. Results in Section 3 use Eq. (1) to describe optimal allocations under different hypothetical settings; we find that allowing the scaling parameters vary by group as in Eq. (7) is more realistic in empirical settings.

The estimated models describe the overall trends (Figure B.2), but the parameter estimates are variable (Table 2), indicating that a range of parameters can fit the data, a well-known phenomenon in fitting power laws to data (Clauset et al., 2009). While we caution against absolute or prescriptive interpretations based on the estimates given in Table 2, if such interpretations are desired (Chen et al., 2018), we suggest evaluating variation due to subsetting patterns and comparing to alternative models such as log-normal and exponential fits (cf. Clauset et al., 2009).

We now study the use of scaling models fitted on a small pilot dataset to inform a subsequent data collection effort. Given the results of Section 4.2, we aim to collect a training set that minimizes the maximum loss on any group. This procedure goes beyond the descriptive use of the estimated scaling models in Section 4.3; important considerations for operationalizing these findings are discussed below.

We perform this experiment with the Goodreads dataset, the largest of the five we study. The pilot sample contains 2,500 instances from each group, drawn at random from the full training set. We estimate the parameters of Eq. (7) using a procedure similar to that described in Section 4.3. For a new training set of size $n_{\textrm{new}}$, we suggest an allocation to minimize the maximum forecasted loss of any group:

For $n_{\textrm{new}}\in[1\times,2\times,4\times,8\times]$, the pilot sample size, we simulate collecting a new training set by drawing $n_{\textrm{new}}$ fresh samples from the training set with allocation $\vec{\alpha}=\hat{\alpha}^{*}_{\textrm{{minmax}}}(n_{\textrm{new}})$. We train a model on this sample (ERM objective) and evaluate on the test set. For comparison, we also sample at $\vec{\alpha}=\vec{\gamma}$ (population proportions) and $\vec{\alpha}=(0.5,0.5)$ (equal allocation to both groups). We repeat the experiment, starting with the random instantiation of the pilot dataset, for ten trials. As a point of comparison, we also compute the results for all $\alpha$ in a grid of resolution $0.01$, and denote the allocation value in this grid that minimizes the average maximum group loss over the ten trials as $\alpha^{*}_{\textrm{grid}}$.

Among the three allocation strategies we compare, $\vec{\alpha}=\hat{\alpha}^{*}_{\textrm{minmax}}$ minimizes the average maximum loss over groups, across $n_{\textrm{new}}$ (Figure 2). Since per-group losses generally decrease with the increased allocations to that group, we expect the best minmax loss over groups to be achieved when the purple and grey bars meet in Figure 2. The allocation strategy $\vec{\alpha}=\hat{\alpha}^{*}_{\textrm{minmax}}$ does not quite achieve this; however, it does not increase the population loss over that of the other allocation strategies. This reinforces the finding of Section 4.2 that different per-group losses can be reached for similar population losses and provides evidence that we can navigate these possible outcomes by leveraging information from a small initial sample.

While the results in Figure 2 are promising, error bars highlight the variation across trials. The variability in performance across trials for allocation baseline $\alpha^{*}_{\textrm{grid}}$ (which is kept constant across the ten trials) is largely consistent with that of the other allocation sampling strategies examined (standard errors in Figure 2). However, the estimation of $\hat{\alpha}^{*}$ in each trial does introduce additional variation: across the ten draws of the pilot data, the range of $\hat{\alpha}^{*}$ values for subsequent dataset size $n_{\textrm{new}}=5000$ is [2e-04,0.04], for $n_{\textrm{new}}=10000$ it is [1e-04,0.05], for $n_{\textrm{new}}=20000$ it is [5e-05,0.14], and for $n_{\textrm{new}}=40000$ it is [2e-05,0.82]. Therefore, the estimated $\hat{\alpha}^{*}$ should be leveraged with caution, especially if the subsequent sample will be much larger than the pilot sample. Further caution should be taken if there may be distribution shifts between the pilot and subsequent samples. We suggest to interpret estimated $\hat{\alpha}^{*}$ values as one signal among many that can inform a dataset design in conjunction with current and emerging practices for ethical data collection (see Section 5).

We now shift the focus of our analysis to explore potential between- and within- group interactions that are more nuanced than the scaling law in Eq. (1) provides for. The results highlight the need for and encourage future work extending our analysis to more complex notions of groups (e.g., intersectional, continuous, or proxy groups).

As discussed in Section 3, data from groups similar to or different from group $g$ may have greater effect on $\mathcal{R}(\hat{f}(\mathcal{S});\mathcal{D}_{g})$ compared to data drawn at random from the entire distribution. We examine this possibility on the ISIC dataset, which is aggregated from different studies (Section B.1). We measure baseline performance of the model trained on data from all of the studies. We then remove one study at a time from the training set, retrain the model, and evaluate the change in performance for all studies in the test set.

Figure 3(a) shows the percent changes in performance due to leaving out studies from the training set. The MSK and UDA studies are comprised of 5 and 2 sub-studies, respectively; Figure 3(b) shows the results of leaving out each sub-study. Rows correspond to the study withheld from the training set and columns correspond to the study used for evaluation. Rows and columns are ordered by $\%$ malignancy. For Figure 3(a) this is the same as ordering by dataset size, SONIC being the largest study.

Consistent with our modelling assumptions and results so far, accuracies evaluated on group $g$ decrease as a result of removing group $g$ from the training set (diagonal entries of Figure 3). However, additional patterns show more nuanced relationships between groups.

Positive values in the upper right regions of Figures 3(a) and 3(b) show that excluding studies with low malignancy rates can raise performance evaluated on studies with high malignancy rates. This could be partially due to differences in label distributions when removing certain studies from the training data. Importantly, this provides a counterpoint to an assumption implicit in 1, that group risks decrease in the total training set size $n$, regardless of the groups these $n$ instances belong to. To study more nuanced interactions between pairs $g^{\prime}\!\neq\!g$, future work could modify Eq. (1) by reparameterizing $r(\cdot)$ to directly account for $n_{g^{\prime}}$.

Grouping by substudies within the UDA and MSK studies reveals that even within well defined groups, interactions between subgroups can arise. Negative off-diagonal entries in Figure 3(b) suggest strong interactions between different groups, underscoring the importance of evaluating results across hierarchies and intersections of groups when feasible.

Of the 16,965 images in the full training set, 7,401 are from the SONIC study. When evaluating on all non-SONIC instances (like the evaluation set from the rest of the paper), withholding the SONIC study from the training set (akin to the training set of the rest of the paper) leads to higher AUROC (.905) than training on all studies (0.890). This demonstrates that more data is not always better, especially if the distributional differences between the additional data and the target populations are not well accounted for.