Just Mix Once: Mixing Samples with Implicit Group Distribution
Abstract
Recent work has unveiled how average generalization frequently relies on superficial patterns in data.
The consequences are brittle models with poor performance in the presence of domain shift in group distribution at test time.
When the subgroups in the training data are known, we can use tools from robust optimization to tackle the problem.
However, group annotation and identification are timeconsuming tasks, especially on large datasets.
A recent line of research [12] is trying to solve this problem with implicit group distribution at training time, leveraging self-supervision and oversampling to improve generalization on minority groups.
Following such ideas, we propose a new class-conditional variant of MixUp [21] for worst-group generalization, augmenting the training distribution with a continuous distribution of groups.
Our method, called Just Mix Once (JM1), is domain-agnostic, computationally efficient, and performs on par or better than the state-of-the-art on worst-group generalization.
Introduction
In supervised learning the goal is to fit a model on a training set to maximize a relevant global metric at test time.
However, the optimization process can exploit spurious correlations between the target y and superficial input patterns in the data. We are interested when undesirable patterns can be categorized in groups g and study the generalization performance in the presence of group distribution shift at test time.
For example, in the case of colored MNIST (Fig. 1), a model can use information about the color (spurious correlation with the target class) instead of the shape to classify a digit (almost all 6s are blue at train time) during training.
Then the model uses this bias at test time with a considerable decrease in generalization performance with a significant shift in the group distribution (for example, 6s are green and blue with equal proportion at test time).
It is well-known that a powerful DNN model [7] tends to exploit easy superficial correlations, like texture, color, background, to solve a task if not constrained in some way [20, 8, 19, 2, 14, 5].
Similarly, in the presence of minority groups in a dataset, the model will tend to disregard such groups, relying on patterns that are frequent in majority groups.
Recent work has shown how to tackle this problem in a supervised [16] and self-supervised [13] way.
In this work, we focus on improving worst-group generalization with implicit group distribution.
Self-supervision and oversampling have successfully been employed to deal with implicit group distribution in Just Train Twice (JTT)
Contribution.
Following this line of work, our contributions are: I)
We propose a simple mechanism, called Just Mix Once (JM1), based on a class-conditional variant of MixUp [21] to improve generalization on minority groups in the data.
II)
Our method does not rely on computationally expensive oversampling or the need to tune the oversampling rate.
III)
We perform extensive experiments with different levels of group annotation, demonstrating that JM1 outperforms the state-of-the-art (SOTA) on vision and language datasets.
( y :6, c : b ) ( y :6, c : b ) ( y :6, c : b ) ( y :6, c : b ) ( y :6, c : b ) ( y :6, c : b ) ( y :6, c : b ) ( y :6, c : b ) ( y :6, c : b ) ( y :6, c : r ) =0.
=0.
=0. =0.
=0. =0.
=0. =0. =0.
( y :2, c : g ) ( y :2, c : g ) ( y :2, c : g ) ( y :2, c : g ) ( y :2, c : g ) ( y :2, c : g ) ( y :2, c : g ) ( y :2, c : g ) ( y :2, c : g ) ( y :2, c : y )
=0. =0. =0.
=0. =0.
=0.
=0. =0. =0.
Just Train Twice
Our work tackles the problem of worst-group generalization with implicit group distribution during training.
This is different from the standard case, where we measure the average error among groups.
In this section we discuss a closely related work, Just Train Twice (JTT) [12]. JTT is a SOTA approach that solves such a problem in a simple two-stage approach:
I) Groups Identification.
The goal is to partition the data into two clusters: one with majority groups (frequent superficial patterns); and one with minority groups (uncommon patterns of interest). The assumption is that samples from minority patterns are difficult to model and frequently misclassified in early-stage of training.
JTT uses the misclassified samples in the early stage of training as a signal to partition the data. Specifically, a supervised learner, parameterized by θ , is trained on the data D = { ( x i , y i ) } ni =1 using ERM loss (up to constant) J ERM ( θ ) : = (cid:80) ( x,y ) ∈D l ( x, y ; θ ) with early stopping. To avoid overfitting, a small validation set with group
annotation is employed to select the best identification epoch T in order to create appropriate partitions. Then misclassified samples (samples with with high loss [13]) are saved in a buffer B := { ( x b , y b ) s.t.
ˆ f T ( x b ) ̸ = y b } .
II) Groups Re-weighting. Once the partitions are identified, the same learner is trained, for a second time, on a re-weighted version of the data, where samples in B are oversampled λ ≫times.
The hope is that if the examples in the error set B come from challenging groups, such as those where the spurious correlation does not hold, then up-weighting them will lead to better worst-group performance [12].
Specifically, the loss for phase II) can be written as (up to constant):
J JTT ( θ ) : = λ (cid:80) ( x b ,y b ) ∈B l ( x b , y b ; θ ) + (cid:80) ( x ¯ b ,y ¯ b ) / ∈B l ( x ¯ b , y ¯ b ; θ ) , where λ is the up-sampling rate for samples
in B identified from phase I
Limitations.
For JTT to work, the authors emphasize that the two hyper-parameters are crucial to guarantee the success of JTT: the number of identification epochs T in phase I) and re-weighting λ in phase II). This requires careful hyper-parameter tuning using a small annotated validation set.
Another caveat is that naively oversampling a portion of the training data λ ≫ 1 times can make training JTT much slower than ERM.
Just Mix Once
We address some of the limitations of JTT by proposing Just Mix Once (JM1).
Remind that our goal is to improve the classification accuracy of minority groups, (with or) without explicit annotation on the groups at training time.
I) (Optional) Groups Identification.
If the explicit group annotation is unavailable on the training set, similarly to phase I) of JTT, we assume that we have annotated groups on a small validation set and resort to using the same self-supervised signal [13, 1] based on the misclassification rate (and the loss magnitude) in the early stage of training like JTT.
Fig. 1 illustrates how relevant samples from the minority groups are identified.
II) Class-conditional Mixing.
To generalize on minority groups at test time, a model should not rely on spurious correlations during training (e.g., texture or color) but on the signal of interest (e.g., shape).
Unlike the oversampling approach taken by JTT, we resort to a better augmentation (Fig. 2) to improve generalization for low-density groups.
Our mixing strategy is inspired by the MixUp strategy [17, 21], and we propose a novel class-conditional variant that achieves robust worst-group generalization.
Specifically, for any two samples ( h ig , h j ¯ g ) in the input (or representation) space from the two partitions g, ¯ g (i.e., difficult/misclassified/minority group g vs the other majority group ¯ g ) with the same class label y ig = y j ¯ g
= y , we mix them up using a convex interpolation [17, 21]:
[Equation]
where α is the mixing parameter (details of how to choose α are deferred to Sec. 4).
Notably, we empirically show that naively implementing the standard MixUp without the proposed two-stage approach fails to generalize in terms of worst-group performance (Fig. 3).
How JM1 works.
We briefly discuss how our mixing strategy enables JM1 to build a "continuous" spectrum of groups in data by sampling α . Assume the training data D = { x i , y i } ni =1 are partitioned into (non-overlapping) groups G 1 , . . . , G m , through sub-populations represented by tuples ( y, c ) of label y and confounding factor(s) c , we can define the following per-group and group-average loss:
[Equation] Notice that in case m = n , each group contains a single data point, and the group-average loss collapses to the standard ERM loss J ERM . Note that, if we define a uniform per-group loss weight p k
= 1 /m, k = 1 , . . , m , we can re-write the group-average loss as: [Equation]
Distributional Robust Optimization (DRO) [16] aims to optimize: J DRO = max k J ( θ, G k ) , which corresponds to a pointwise Dirac distribution p k = 1 if k = arg max k J ( θ ; G k ) else 0 in Eq.
Using MixUp with α mixing rate, we generate new samples drawn from a continuous mixture (Eq. 1), which has as limiting case the training distribution.
In this view, we can interpret each mixed-up sample as drawn from a "mixed-up group" G α parametrized by α .
This means, assuming an oracle partitioning in phase I), JM1 mixes a majority and a minority group in the data at each iteration while generating a continuous group distribution. We denote by J ( θ ; G α ) our per-group loss for a "mixed-up" group
G α , and write the group-average loss by marginalizing out the group distribution:
[Equation]
Note that, first sampling α and then applying MixUp with rate α (Eq. 1) is equivalent to drawing samples directly from a mixed-up group G α drawn implicitly from p ( G ; α ) . Therefore, we can approximate Eq. 4 by a simple MC sampling process in practice: sample α , apply MixUp [21, 17] with rate α , compute the loss using the mixed-up samples; repeat for each mini-batch in SGD updates.
This simple mechanism enables to augment data by marginalizing group information, hence the training cannot rely on spurious group patterns because each sample belongs to a "slightly different" group.
For intuitions behind the class-conditional mixing of JM1’s phase II) in the presence of minority groups, refer to the visualization in Appendix C. We leave thorough theoretical study of JM1 to future work.
Experiments
Benchmark Datasets.
Our focus for all the experiments is to improve worst-group performance for classification problems across vision and language domains (with or) without group annotation at train time.
We report experiments on the benchmark proposed in [16].
The benchmark consists of two vision and one language dataset (Appendix B).
In these datasets, the model can easily exploit superficial, majority group patterns (background in CUB, sex in CelebA, negation in MultiNLI) to solve the classification tasks at train time. However, when tested on a set with a different group distribution (same groups but represented in significantly different proportions), the performance drops, showing a lack of generalization capacity.
JTT [12]) is SOTA on this benchmark, and we primarily compare JM1 to JTT. JM1 Implementation Details.
For JM1, we initially found that mixing groups by sampling α ∼ U (0 , 1) uniformly, even though improving worst-group accuracy over ERM, underperforms compared to JTT.
Therefore, we use a mixing strategy with a slight emphasis on the region closer to the minority groups and propose a coupled strategy: we sample α ∼ U (0 , 1) for half of the epochs, and α ∼ Beta (2 , 5) for the other half.
This simple heuristic ensures that the full mixing domain is spanned, with a focus on the region closer to the misclassified samples in phase I). We use this coupled sampling strategy for CUB and CelebA. For MultiNLI, we sample α only from the uniform prior because the mixing phase runs for a very small number of epochs.
More training details of JM1 are described in Appendix E. It is worth noting that: I) like JTT, JM1 is domain-agnostic and can be applied to vision and language datasets with minimal architectural changes; II) unlike JTT’s oversampling approach, JM1 is a highly scalable approach, which removes the need for tuning the oversampling rate and only introduces an interpolation operation between samples with little computational overhead. Note that we can apply JM1 to mix samples at different levels of abstraction: either the input [21] or the representation space [17]. As the comparison of different MixUp levels is not the focus of our study, we refer the readers to Appendix D for additional ablations.
Worst-Group Generalization.
The first set of experiments evaluate JM1 vs JTT in the setting of implicit group distribution on the training set.
We assume that the models have access to a small validation set with group annotation to guide group identification in phase I).
To ensure a fair comparison between methods, we use the same identified set for both JTT and JM1 in all our experiments. Results in Table 1 show that our approach is competitive with the SOTA on improving worst-group performance on all three datasets with comparable average accuracy.
In Table 2 we report the per-group accuracy on CUB for JTT and JM1 under the same setting.
We see that, compared to JTT, not only does JM1 improve worst-group accuracy, but it improves accuracy on both of the minority groups with little to no loss in performance on the other two majority groups.
Overall, JM is competitive or better than SOTA in terms of worst-group performance on this benchmark
[Table]
[Table]
Oracle Groups. The second set of experiments compare JM1 vs JTT in a different scenario: we assume to have access to the ground-truth minority group assignment on the training set. Note that this does not mean the full group annotation: we only need to know which sample has a groupconflicting confounder against its label (see a comparison of JM1-inspired GroupJM1 and GroupDRO in Appendix D in a setting given full group annotation on training data).
Now we can use such oracle majority vs minority groups in place of the identified set from phase I).
As noted in Appendix C.
[Table] [Table] [Table]
in [12], the performance of JTT drops if not using identified error sets from phase I), since JTT can excessively oversample the minority group.
Given its inner working of mixing groups rather than oversampling, we expect JM1 to outperform JTT in this oracle configuration.
Table 3 validates our intuition.
In fact, JM1 (oracle) outperforms JTT (oracle) on both datasets.
Class-conditional MixUp.
We compare the proposed JM1 with standard (unconditional) MixUp and a class-conditional MixUp baseline, where we mix samples from random groups (i.e., JM without group identification phase).
Note that we use an oracle configuration for simplicity in these experiments.
Fig.
3 shows that, while standard MixUp fails to improve worst-group performance, the proposed mixing strategy performs significantly better than both MixUp baselines by a large margin, despite that all three methods perform very similarly in terms of average performance.
[Figure]
[Table] [Table] Robustness. While JTT and JM1 do not use full group annotation on the training set, they both rely on group annotation on a small validation set to guide group identification (i e., choosing early stopping epoch).
A natural question to investigate is if we can remove the need for group annotation at all, and JM1 can still perform reasonably well.
To that end, we test the robustness of JM1 to different identification epochs/sets in phase I) using the CUB dataset. Specifically, we checkpoint three different identification scenarios (20, 40, 60 epochs) in phase I) for both JM1 and JTT.
Then, we select for JM1 the best phase II) model in terms of average accuracy, while we tune the bestperforming JTT using group annotation on a validation set in terms of worst-group accuracy.
Note that JM1 is now completely free from group annotations.
Table 4 shows that JM1 without group annotation at all loses only around 5% compared to the full-suite JTT with access to group annotation on a validation set.
This is a preliminary result that indicates the potential of training group-robust models in a fully self-supervised manner. We can hope to achieve worst-group robust generalization with no more data requirement than what we need to train typical ML models
[Table]
A Related Work
Robust Classification and Implicit Group Distribution.
Optimizing average accuracy with overparametrized neural networks tends to get rid of minority group information.
This fact is related to how the model is trained and the way the model exploits superficial correlations in the data [7].
Recently approaches to preserve group information have been proposed [16] when full group annotation is available.
However assuming sub-population annotation in large diverse datasets is unrealistic both in terms of time effort and most important because the process is prone to errors.
Learning from Failure [13] proposes an annotation-free and domain agnostic procedure to improve worst-group performance.
The authors propose a self-supervised approach to identify bias-aligned (majority) vs bias-contrastive (minority) groups in the data.
The self-supervised signal consists of the per-sample loss magnitude in the early stage of training; Then the dataset is re-weighted based on the loss magnitude and the procedure applied multiple times during training.
Inspired by this approach Just Train Twice [12] proposes a similar solution, where the identification phase is performed one time, and then a reweighting phase is trained.
Just Train Twice is effective but relies on an annotated and noise-free validation set for the groups. Invariant Learning [2, 5] is another approach that can be used to improve worst-group performance. The goal is to obtain a learner that performs well in different environments, where in each environment we have the same classes but different confounders.
Equality of Opportunities in Machine Learning.
Mixing information from different groups generates a mixture of distributions with a continuous amount of groups.
Having an infinite collection of groups means that the model cannot rely on group information during training and samples in different groups should be considered similarly.
In this setting, we assume to have full annotation for the groups, and robust optimization methods are employed.
The supervised scenario is closely related to the fairness one [14]: wherewith distribution shift we want to improve the worst-group generalization, in fairness we want to learn features and distributions that are invariant to different groups, using explicit constraints [19, 18, 2], auxiliary tasks [6, 9], and causal structure [15, 14]. In the causality literature [15] the group information is interpreted as a confounder. If we have access to the confounder (group annotation) we can control for it using intervention.
The mixing procedure is augmenting the training data with plausible scenarios not present in the data. This approach generates counterfactual scenarios interpolating groups in the data and exploits these imaginative scenarios to account for minority groups. In the case of implicit groups, confounder effects are unknown. We can still use our approach exploiting loss magnitude and pattern formation during early training.
Self-Supervised Learning. Recently a large corpus of works dealt with self-supervised learning.
With the amount and complexity of the data increasing exponentially, the availability of annotation is decreasing fast.
Self-supervised learning tries to exploit structure in the data to cluster data and improve generalization.
Loss magnitude [13, 7, 12, 1] can be used to cluster data.
MixUp [21] is frequently employed to improve generalization and robustness to noise.
The idea is to increase the variety in the data interpolating the training distribution in sample or representation space [17].
MixUp based methods have been extensively used for semi-supervised learning [11, 3, 1, 10] and improve generalization [22, 4]. Notice that in this work the self-supervision is wrt the group annotation and not label annotation as usual. B Benchmark We report experiments on the benchmark proposed in [16].
The benchmark consists of two vision and one language dataset.
In all three datasets, there are factors or confounders (background, sex, hair color) that strongly correlate with the label information. In particular, for the image datasets: in Waterbirds (CUB) the task is to classify land birds vs water birds. The class information is spuriously correlated with the background (land vs water) generating 4 relevant groups in the training set; in
CelebA the task is to classify blond vs not-blond hairs.
The class information is spuriously correlated with the sex information (male vs female) generating 4 relevant groups.
For the language dataset,
MultiNLI, the task is to classify a reasoning process (two sentences) as entailment, contradiction, or neutral.
The class information is spuriously correlated with the presence of negation in the sentences.
In all these datasets the model can easily exploit superficial, majority patterns (background, sex, negation) to solve the tasks at train time.
However, when tested on a set with a different group distribution (same groups but represented in different proportions), the performance drops, showing a lack of generalization capacity.
Waterbirds (CUB).
The task is to classify land birds y l and water birds y w .
The label information is correlated with a confounding factor: land background c l and water background c w .
The majority groups are (land bird, land background) - ( y l , c l ) and (water bird, water background) - ( y w , c w ) .
The minority groups are (land bird, water background) - ( y l , c w ) and (water bird, land background) - ( y w , c l ) .
During training, the minority groups are present in 5% of the data.
At test time the groups are evenly distributed in the data.
CelebA.
The task is to classify blond person y b and not blond person y nb .
The label information is correlated with a confounding factor: male c m and female c f .
The majority groups are (not blond, male) - ( y nb , c m ) and (blond, female) - ( y b , c f ) .
The minority groups are (not blond, female) - ( y nb , c f ) and (blond, male) - ( y b , c m ) .
MultiNLI.
The task is to classify an inference procedure as entailment y e , contradiction y c and neutral y n .
The label information is correlated with a confounding factor: the presence of a negation c neg and absence of a negation c nneg in the reasoning.
C Modelling the Augmented Group Distribution
In this section, we present a toy visualization and a simple idea to better understand the mixing procedure versus an oversampling strategy. Given two classes (blue and brown points) some of the groups are underrepresented on the train set (Fig. 5). These underrepresented groups are the minority groups. On the test set, all the groups are equally represented (Fig. 6).
JTT oversamples the minority groups (Fig. 7) and trains on the augmented distribution.
JM1 mixed minority and majority groups (Fig.
8) and trains on the augmented distribution.
We see that both methods generate a training distribution closer to the target one in Fig.
Mixing Groups.
Consider a multiplicative interaction between group information and samples.
Given a sample ¯ h from the train set, h o = ¯ h g o is a sample from a majority group (over-represented) and h u = ¯
h g u is a sample from a minority group (under-represented). Then a sample from the augmented distribution is obtained using:
[Equation]
We can see that, given a continuous distribution for α , g mix is a sample from a continuous distributions over groups P aug ( g ) =
P mix ( g ) .
It is also easy to see that a generic group sampled from the mixing distribution will be closer to the target than the source domain because g mix is a convex combination of the two groups:
[Equation]
D Additional Experiments
In this section, we report additional ablation experiments. In Section 4 we have considered only the minority groups in the oracle setting, but potentially we can have access to full group annotation on the train set.
Consequently, methods like GroupDRO [16] can be employed.
We compare a variant of JM1, called GroupJM1, and GroupDRO.
In Table 5 we consider the configuration with full group annotation on the train set and we evaluate a variant of
JM1, GroupJM1, with GroupDRO [16].
GroupJM1 is competitive in this scenario, corroborating the idea that the mixing procedure is a general approach to improve worst-group accuracy.
In
Figure 9 we investigate the performance of JM1 with and without class-conditional constraint. We see that empirically the conditional JM1 performs better than the unconditional one. The worst-group accuracy decreases increasing the dataset complexity (in terms of the number of samples, classes, and groups) for both models. In Table 6 we consider the performance of JM1 mixing samples at different levels of abstraction, where the layer indicates mixing in input space, inside the encoder, in output space, or choosing a layer at random.
We also test different sampling distributions.
In Figure 10 and Table 10 we study the behaviour of the models varying the number of identification epochs in phase I).
[Table]
[Figure]
layer
[Table] JM1+ early
U 92.
JM1+ random U/B 94.
[Figure]
A v e r a g e a n d W o r s t - G r o u p A cc u r a c y
[Figure]
JM1-avg JM1-robust avg worst avg worst 10id 86.
20id 92. 30id 93. 40id 93. 50id 94.
60id 93. 70id 93.
80id 92. 90id 93. 100id 91. E Experimental Details
[Table]