S YSTEMATIC GENERALISATION WITH GROUP INVARIANT PREDICTIONS
A BSTRACT
We consider situations where the presence of dominant simpler correlations with the target variable in a training set can cause an SGD-trained neural network to be less reliant on more persistently correlating complex features.
When the nonpersistent, simpler correlations correspond to non-semantic background factors, a neural network trained on this data can exhibit dramatic failure upon encountering systematic distributional shift, where the correlating background features are recombined with different objects.
We perform an empirical study on three synthetic datasets, showing that group invariance methods across inferred partitionings of the training set can lead to signiﬁcant improvements at such test-time situations.
We also suggest a simple invariance penalty, showing with experiments on our setups that it can perform better than alternatives.
We ﬁnd that even without assuming access to any systematically shifted validation sets, one can still ﬁnd improvements over an ERM-trained reference model.
I NTRODUCTION
If a training set is biased such that an easier-to-learn feature correlates with the target variable throughout the training set, a modern neural network trained with SGD will use that factor to perform predictions, ignoring co-occurring harder-to-learn complex predictive features (Shah et al., 2020).
Without any other criteria, this is arguably desirable behaviour, reﬂecting Occam’s razor.
We consider the situation where although such a simpler correlation is a dominant bias in the training set, a minority group exists within the dataset where the bias does not manifest.
In such cases, relying on more complex predictive features which more pervasively explain the data can be preferable to simpler ones that only explain most of it.
For example, if all chairs are red, redness ought to be a predictive rule for chairhood (without any other criteria for predictions).
However, if some chairs are not red, and all chairs have backs and legs, then one can infer that redness is less relevant.
In this paper, we will study object recognition tasks, where the objects correlate strongly with simpler non-semantic background information for a majority of the images, but not for a minority group.
There is evidence in the literature that modern CNNs tend to ﬁxate on simpler features such as texture (Geirhos et al., 2019; Brendel & Bethge, 2019), canonical pose (Alcorn et al., 2019), or contextual background cues (Beery et al., 2018).
We are assuming that semantic features in a classiﬁcation context (ones that humans would agree contribute to their labelling of objects) are more likely to persistently correlate with the target variable, while simpler non-semantic background biases are more likely to exhibit non-persistent correlations in real-life data collection processes.
Based on this assumption, we will use combinations of objects and backgrounds to compare test-time performances corresponding to particular distributional shifts.
Consider coloured MNIST digits such that there is a dominant, but not universal, correlation between colour and digit identity for a majority of the images.
In the situation we are considering, if the biasing colours in the majority group are not recombined with different digits in the minority group, then there is no signal for the model to disregard these biasing factors, which are retained as important predictive rules.
This can lead to poor performance at systematic generalisation (Lake & Baroni, 2018), where an object occurs with another object’s biasing factor, and at semantic anomaly detection (Ahmed & Courville, 2020), where a novel object appears with one of the biasing factors.
In our example
[Table]
with coloured colour the minority digits with the colours used to bias (different) digits in the majority group, we ﬁnd a marked improvement at systematically shifted tests over the case when the colours in the minority group are different colours altogether (see Table 1).
We investigate the role of encouraging robust predictive behaviour across such groups in terms of improved performance at tasks with such distributional shifts.
Our experiments suggest that training with cross-group invariance penalties can result in models that have learned to be more reliant on persistent complex correlations without being overwhelmed by simpler, yet less stable features, as indicated by improved performance at systematic generalisation and semantic anomaly detection on our synthetic setups.
We ﬁnd that a recently proposed method (Creager et al., 2020) can be effective at inferring the majority and minority groups along a learned feature-bias, and we use this inferred partition to provide us with groups in the training set in our comparative study.
We also suggest a new method for encouraging predictions that rely on persistent correlations across such groups, with the intuition that similar predictive behaviour across the groups should be promoted throughout training.
With experiments on three synthetic datasets, we compare the performance of recently proposed invariance penalties and methods, and ﬁnd that our variant can often perform better at tasks involving such test-time distributional shifts.
S YSTEMATIC AND NON - SYSTEMATIC GENERALISATION
If we assume that data x is generated via a composition C of semantic factors h s and non-semantic factors h n , we can use this decomposition, x = C
( h s , h n ) , to generate test datasets to capture different scenarios.
While h n is actually independent of y , we shall have the independence property p D ( h n | y )
=
p D ( h n ) to not hold when there is bias in the dataset D due to h n – y correlations.
We can evaluate, for a particular target y and our system’s prediction of the target ˆ y ( x ) , the average accuracy E (cid:2) 1 { ˆ y ( C ( h s , h n )) =
y } (cid:3) , as a measure of generalisation for the following different cases. [Figure]
In-distribution generalisation h s ∼ p ( h s | y ) and h n ∼ p ( h n | y ) :
The validation and test sets are assumed to possess the same biases as the training set, in that the class-conditional distribution of non-semantic features in the test set match that of the training set, p ( h n | y ) .
Generalisation under non-systematic-shift h s ∼ p ( h s | y ) and h n (cid:54)∼ p ( h n ) 1 : This estimates a form of generalisation under distributional shift, where the non-semantic factors are sampled from outside the marginal distribution of h n as present in the training set.
Generalisation under systematic-shift h s ∼ p ( h s | y ) and h n ∼ p ( h n | y (cid:48) )
where y (cid:48) ∼ p
( y ) s.t.
y (cid:48) (cid:54) =
y :
This estimates another form of generalisation under distributional shift but one where non-semantic factors are sampled with intent to confuse: non-semantic factors for x are sampled from the marginal distribution of a randomly picked different target, y (cid:48)
(cid:54) =
y .
Although systematicity, as discussed in Fodor & Pylyshyn (1988), and systematic generalisation, as discussed in the NLP literature (Lake & Baroni, 2018; Bahdanau et al., 2019) consider recombinations of intra-semantic factors as well, here, in the context of background-agnostic object recognition tasks, we only consider h s − h n recombinations.
Semantic anomaly detection h s (cid:54)∼ p ( h s ) and h n ∼ p ( h n ) : Such a datapoint should not be conﬁdently categorised as a known y , even if non-semantic features are shared (Ahmed & Courville, 2020).
We can use these x to evaluate anomaly detection, as indicated by decreased predictive conﬁdence, and measured by the area under the precision-recall curve (Hendrycks & Gimpel, 2017).
COLOURED
MNIST: Consider an illustrative dataset with coloured MNIST digits.
For the training set, T r , MNIST digits are coloured with a set of digit-correlated “biasing” colours 80% of the time, and with ten random colours that are different from the biasing colours the remaining 20% of the time.
One digit is held out, for testing semantic anomaly detection.
See Figure 1 for examples of the four test sets corresponding to this setting, and also Appendix A for more details on the construction.
Improving performance for such scenarios involving distributional shift might come at a cost for indistribution performance, since more robust features might be harder to learn than simpler dominant correlations that hold in-distribution.
In real-world deployments where one is likely to encounter unexpected situations, such as in a self-driving car, it can often be preferable to ﬁnd appropriate trade-offs such that classiﬁers can indicate reduced conﬁdence upon encountering anomalous objects, or continue to operate in changing environments, while continuing to achieve a desirable degree of in-distribution predictive performance.
P REDICTIVE GROUP INVARIANCE ACROSS INFERRED SPLITS
In general, we do not expect to have direct knowledge of majority and minority groups corresponding to the biasing non-semantic features in a dataset.
We will later show how one might infer such groups from the data, but we ﬁrst describe an invariance penalty assuming we have access to the groups.
Learning features that are group invariant would require us to match the (class-conditioned) distribution of features from the majority and minority groups (Ganin et al., 2016; Li et al., 2018a).
In terms of predictive performance, we can alternatively ask for the class-conditioned distributions of features to match in the sense that they lead to the same softmax distributions on average as training progresses, without modifying the last linear layer.
This implementation has the advantage of doing away with an adversarial network, and the issues that tend to accompany the training of such models.
We shall refer to this objective as predictive group invariance (PGI).
Intuitively, encouraging matched predictive distributions across the groups with a ﬁxed last layer pushes for over-emphasis on minority-group features in the representation, thus acting as an implicit re-weighting of features in both groups (leading to demoting the relevance of colour in the MNIST case, for example).
When a persistent feature does exist in both groups, using that feature can lead to equal training rates in regularised networks, satisfying the penalty.
Consider a classiﬁer that extracts a feature vector f θ ( x ) , where θ are the parameters of a convolutional neural network for example, with a linear layer w on top.
The predictive distribution is then
[Equation]
where σ is a softmax, and predictions are made by performing an argmax .
Given a partition scheme for splitting the images x in our dataset D such that every i -th image x ( i ) is associated with a partition-label α ( i ) , we deﬁne distributions P c , Q c for the subsets in class c :
[Equation]
We want to minimize empirical risk under the constraint that our feature extractor causes similar predictive distributions on average for pictures of the same object in both partitions.
Formally, we want to optimise
[Equation]
where (cid:96) is the standard loss function for ERM training, for example, the categorical cross-entropy.
A softened objective for stochastic optimisation can be approximated as
[Equation]
Since we are comparing distributions, we make the simplest natural choice of d to be the KLdivergence,
[Equation]
We use this particular ordering of Q (cid:107) P because with our grouping, P consists of examples that are “easy” due to a particular bias, and so the mean predictive distribution for P tends to be correct and low-entropy, while that for Q is more high-entropy and inaccurate.
We take advantage of the zero-forcing property of this KL divergence, encouraging the mean predictive distribution for Q to closely match that of P .
It is likely that different choices for d would be better suited for different settings.
Partitioning the dataset Recently, Creager et al.
(2020) have considered the question of ﬁnding worst-case partitions for invariant learning given a collection of data.
The key intuition is that an invariant learning objective, as formulated by IRM (Arjovsky et al., 2019), is maximally violated by splitting along a spurious correlation when predictions rely exclusively on it in a reference model (see Theorem 1 in Creager et al.
(2020) for details).
In our case, this would consist of partitioning into the majority and minority groups given our ERM-trained model early on in training as reference.
A soft-partition predicting network is used, g ( x, y ) , conditioned on the input and the target, to maximise the IRMv1 penalty (Arjovsky et al., 2019), which gives us soft partition-predictions, ˆ β , for the examples,
[Equation] [Equation]
x ( i ) )) , y ( i ) )
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) 2 , (8) where Φ( x i ) =
w (cid:62) f
θ ( x ) are the logits from the reference model, e ∈ { 0 , 1 } indexes the partition, β ( i ) ( e ) ∈ [0 , 1] signiﬁes the predicted probability for the i -th example being in partition e , such that β ( i ) ( e = 0)
+ β ( i ) (
e = 1)
= 1 , and γ is a hyper-parameter.
We can then compute the partition α ( i ) =
argmax e β ( i ) ( e ) .
In our implementation, we condition the partition predicting network g on the features f θ ( x ) instead of the input x , and use separate networks for each category, i.e. β ( i ) =
g y
( i ) ( f θ ( x ( i ) )) .
We ﬁnd this to perform better in preliminary experiments, improving training and enabling more light-weight g networks.
This also ensures that the same features as the ones used by our ERM-trained reference model are used to predict partitions, resulting in partitions corresponding to more consistent learned-feature biases.
We provide more details in Appendix B.3.
R ELATED WORK
The dominant perspective towards the issue of unreliable behaviour in novel domains has consisted of treating the problem as that of domain generalisation (Blanchard et al., 2011).
One hopes to recover stable features by encouraging invariance across data sampled from different domains, so that performance at test-time out-of-distribution (OoD) scenarios is less likely to be unstable.
Approaches along such lines typically resemble a cross-domain distribution-matching penalty applied to the features being learned, augmenting the usual ERM term (Ganin et al., 2016; Sun & Saenko, 2016; Heinze-Deml & Meinshausen, 2017; Li et al., 2018; Li et al., 2018a;b), and evaluated on datasets that consist of data in different modalities (Li et al., 2017; Peng et al., 2019; Venkateswara et al., 2017), or collected through different means (Fang et al., 2013), or in different contexts (Beery et al., 2018).
Works with the perspective of distributionally robust optimisation (DRO) have generally considered using uncertainty sets around training data (Ben-Tal et al., 2013; Duchi & Namkoong, 2018) to minimise worst-case losses, which can often have a regularising effect by effectively up-weighting harder examples.
More relevant to our discussion, group DRO methods have considered uncertainty sets in terms of different groups of data, for example with different cross-group distributions of labels (Hu et al., 2018), or groups collected differently (Oren et al., 2019), similarly to domain generalisation datasets.
More recently, methods promoting the learning of stable features across data from different environments , or sources, have been proposed by using gradient penalties (Arjovsky et al., 2019), risk-based extrapolation (Krueger et al., 2020), and masking gradients with opposing signs (Parascandolo et al., 2020).
The typical datasets in such existing works are not curated with testing performance under systematic distributional shift in mind, most often not characterising the speciﬁc shift in distribution.
In recent times, a commonly adopted synthetic dataset is the coloured MNIST variant used in Arjovsky et al.
(2019) – since this particular dataset uses ﬂipped colours for the minority group, which is less of a problem with ERM-training, the true digit labels were ﬂipped at a sufﬁciently high frequency to incapacitate ERM performance by forcing reliance on colour.
We believe setups such as ours can be better synthetic testbeds for developing ideas, where it is not necessary to alter ground truth labels to expose a failure mode.
In general, using better models of dataset bias implies a narrower disconnect with realistic settings, with higher chances of the conclusions carrying over.
E XPERIMENTS
We compare performance with our four test sets - in-distribution, non-systematically shifted, systematically shifted, semantic anomalies - for a range of recently proposed methods for a set of three synthesised datasets.
Appendix B describes architectural details and training choices.
5. M ETHODS
We compare recent methods aimed at robust predictions across groups, and which do not require changes to network capacity or additional adversaries to impose invariance penalties.
We also do not include methods based on advances in self-supervised feature learning, such as Carlucci et al.
(2019), since such methods are developed with prior knowledge of the desired invariances, and are thus limited in their generality.
Baseline : This is our reference model, trained via ordinary (regularised) empirical risk minimisation (ERM) without any invariance penalties added.
The choices for architecture and regularisers were made to conform to the way modern networks are typically trained with in-distribution performance in mind (details in Appendix B).
IRMv1 , REx , GroupDRO : IRMv1 (Arjovsky et al., 2019) and REx (Krueger et al., 2020) are two methods that augment the standard ERM term with invariance penalties across data from different sources.
GroupDRO (Sagawa et al., 2020) is an algorithm for distributional robustness, which works by weighting groups of data as a function of their relative losses.
See Appendix C for more details about these methods.
cIRMv1 , cREx , cGroupDRO : We implement label-conditional variants of the above algorithms, which, to our knowledge, has not been explored.
In the context of multi-class classiﬁcation it is reasonable to expect that performances might have multi-modal distributions along different categories earlier in training, which suggests stratiﬁcation by class might improve performance.
Reweight : We weight the losses in the biased group down.
This is a heuristic form of re-balancing the dataset, while choosing a hyper-parameter for the weight using the validation set, with the weight serving to downweight the losses for the biased group.
In preliminary experiments we found this re-weighting variant (King & Zeng, 2001) to signiﬁcantly outperform oversampling the minority group, as suggested in Buda et al.
(2018), or weighting the grouped losses using their population ratios, as performed for imbalanced classes in Cui et al.
cMMD :
Following Li et al.
(2018), we match the MMD (Gretton et al., 2012) of the distribution of features.
In preliminary experiments, we ﬁnd a conditional version (as done with adversarial models in (Li et al., 2018a)) to perform signiﬁcantly better, so we only report cMMD results here.
5. D ATASETS
Evaluating performance in an unambiguous manner for the speciﬁc kinds of generalisation that we aim to study necessitates controlled test-beds.
In order to model these tasks, we use 3 synthetic datasets of progressively higher complexity, approaching photo-realism.
C OLOURED MNIST:
This is the simplest setting, where the
background information
exists as part of the object.
COCO- ON -C OLOURS : We superimpose 10 segmented
COCO
(Lin et al.,
2014) objects
on
coloured backgrounds
The training set has
800 images
per
category,
with
nine
in -distribution categories and one held -out category for anomaly detection.
Validation and test sets have 100 each images per category.
See Figure 2 (left).
This is the most extreme dataset in our experiments in terms of the contrast in complexity between the non-semantic correlating factor (background colour) vs. stable features (objects).
COCO- ON -P LACES :
Here we
superimpose the same
COCO objects on scenes
from
the
P LACES dataset (Zhou et al , 2017), with the place-scenes acting as the bias (
ﬁgure right).
See Appendix A for more details about how these datasets are constructed.
While the backgrounds in this dataset are more complex than colour, they still act as biasing factors, as indicated in the relatively poorer performance at systematic generalisation, and were selected due to visually obvious and distinct colour or texture.
5. R ESULTS
In all cases, we have used the partition predictor to infer the two groups.
The partition accuracies for the three datasets at the end of one epoch of training the base models are in the table below.
We tested a more naïve approach by applying K-Means clustering to the losses, but found it to under-perform, possibly because it cannot account for a consistent feature bias learned by our reference model.
[Table] [Table]
[Table]
In Tables 2,3,
[Table]
we ﬁnd that signiﬁcant improvements can be achieved using group invariance methods.
All hyper-parameters for the results in this set are picked on a validation set consisting of a subset of colours or backgrounds that are different from both the training and test sets, and an equally sized subset of systematically varying colours or backgrounds from the biased majority group.
In all cases, the split is learned after one epoch of training [Table] and the various penalties dropped in at this point with a linearly ramped-in penalty co-efﬁcient.
Details about hyper-parameter selection are in Appendix C.
While conditional variants perform better at systematic generalisation for COLOURED MNIST, perhaps owing to our hyper-parameter selection procedure of using a mixed-shift validation set, performance at systematic shift appears to be traded off with non-systematic shift in some cases for the more complex datasets.
All aggregates are over 5 trials.
5. P RACTICAL CONSIDERATIONS FOR HYPER - PARAMETER SELECTION
While we ﬁnd that with the use of group invariance penalties it is possible to encourage reliance upon complex persistent correlations in the presence of dominant simple biases, this can sometimes come at a cost to in-distribution performance when picking hyper-parameters using validation sets with speciﬁc distributional shift.
One might reasonably expect that this can be mis-aligned with real-life situations: in practice, one typically does not have access to data corresponding exactly to unexpected scenarios, besides not expecting to encounter situations outside the training distribution nearly as often as situations for which a model has been trained and deployed.
A practitioner might wish to aim for a clearer trade-off with such situations, with prior knowledge of how often they might arise compared to in-distribution situations, and with a surrogate validation set to model distributional shift.
Here, we will simply show that picking hyper-parameters without assuming access to validation sets consisting of systematic distributional shift can still provide improvements over the baseline reference model.
We consider three cases.
[Table] [Table]
NS: Hyper-parameters are picked using only the validation set for non-systematic distributional shift (which consists of backgrounds that are different from those in the training set and test sets).
This models the situation where we have access to some data that is different from our training data, and is also considered somewhat representative of any shifts we might encounter.
NS + ID: Hyper-parameters are picked using an (equally-weighted) average of the NS and the in-distribution validation sets.
If we have prior knowledge of the likelihood of encountering data from out-distributions in the wild, we could use this prior to use an appropriately sampled validation set for hyper-parameter optimisation.
ID ONLY : Hyper-parameters are picked using only the in-distribution validation set.
We show results for the different schemes for our method in Tables 5, 6, 7.
While the accuracies under distributional shift are, as expected, less strong than in the previous set of results (NS+S in the tables), we still ﬁnd improvements over the reference model, indicating that one can still achieve an improved classiﬁer.
In Appendix D, we show similar results with all methods, and include only the best performing method for both generalisation under systematic and non-systematic shift corresponding to the different validation strategies in the tables in this section.
[Table]
[Table] [Table]
C ONCLUSION
Our experiments investigate the potential usefulness of invariance penalties and methods at improving performance under distributional shift, such as systematic generalisation and semantic anomaly detection.
While our exploratory experiments are conducted in disambiguated synthetic setups
[Table]
next steps would involve investigating the potential for extending these approaches to real datasets used in the ﬁeld.
Since such methods cannot work when spurious correlations are completely pervasive
[Table]
it is important to include sufﬁcient diversity of data sources and curation in order to be able to reap the advantages such techniques can afford us in real world applications.
We note that peculiarities in datasets and problems might give rise to different potential failings at robustness, calling for more targeted invariance methods.
We ﬁnd that our method of learning features that result in matched predictive behaviour throughout training appears to hold promise at handling certain distributional shifts, although it does not always perform best across different validation schemes.
A practical line of inquiry would be the question of how to make trade-offs in performance between in-distribution and unexpected situations. A CKNOWLEDGMENTS
We thank Ishaan Gulrajani, Tim Cooijmans, David Krueger, and Phong Nguyen for useful discussions, and anonymous reviewers for constructive feedback.
We acknowledge the computational resources provided by Compute Canada and Mila, and the ﬁnancial support of Hitachi, Microsoft Research,
CIFAR, and the Natural Sciences and Engineering Research Council of Canada (NSERC Discovery
Grant).
A D ATASET DETAILS
In this section, we provide more details about how we constructed our synthetic datasets.
A. C OLOURED MNIST
The training set, T r , is constructed with an 80% colour-digit correlation per digit with nine RGBcolours (with the zero digit held out, for testing semantic anomaly detection).
[Table]
The ten colours for the minority group were picked such that their L2 distance is at least 50 units away from the biasing colours.
Prior to colouring, the digits were binarised to avoid grayscale tones potentially resulting in unintentionally similar colours.
For the non-systematic validation and test sets, ten colours each were chosen such that they were at least 50 units away from all other colours.
A. COCO- ON -C OLOURS
We use the following nine categories for in-distribution objects: boat , airplane , truck , dog , zebra , horse , bird , train , and bus .
We hold out motorcycle for anomaly detection experiments.
For background colours, we use the same colours from the coloured MNIST experiments, and also use an 80/20 split for the majority and minority groups.
In case of multiple instances of the same object in an image, we pick the largest one, and ﬁlter our dataset by mask area, such that only images with objects occupying at least 10K pixels are retained.
All images are ﬁnally resized to 64 × 64 .
The training set uses 800 such pictures per category, and the validation and test sets use 100 each.
The colour backgrounds for the minority group, non-systematically shifted validation and test sets are picked using the same strategy as with the COLOURED MNIST dataset.
A. COCO- ON -P LACES
This dataset follows the same procedure as COCO- ON -C OLOURS , except using scenes from the
Places dataset.
In Table 9 we list the backgrounds from the corresponding scenes for the different categories.
B N ETWORK ARCHITECTURES AND TRAINING DETAILS B. C OLOURED MNIST
We use a 4-layer CNN with the ﬁrst three layers being convolutional and the last layer linear.
The convolutional layers have feature dimensions of 64 − 128 − 256 , and are all followed by a MAX
POOL , BATCH NORM layer, and RELU activation.
Before being fed into the ﬁnal linear layer, there is a spatial mean-pooling operation.
An L2 weight decay is added to all parameters with a co-efﬁcient of 1 e − 4 .
[Table]
Training is conducted for 30 epochs, with SGD + Momentum (0.9), using batch sizes of 512.
The learning rate is cut by 10 from its initial value of 0.1 at epochs 9, 18, and 24.
B. COCO- ON - BACKGROUNDS
For both COCO datasets, we use an architecture based off of Wide Resnet 28-10 (Zagoruyko &
Komodakis, 2016).
Since our images are 64 × 64 , we append an extra group of 4 residual blocks with the same layer widths as in the previous group, and use a smaller widening factor of 4 instead of 10 to avoid memory overﬂow (starting base dimension = 64 ).
An L2 weight decay regulariser is applied on all parameters with a coefﬁcient of 5 e − 4 .
We train for 200 epochs with SGD + Momentum (0.9), using batch sizes of 384, with an initial learning rate of 0.1 which is cut by 10 at the 120th, 160th, 180th, and 190th epochs.
We use the initially large learning rate for longer following prior works such as Li et al.
(2019) that have suggested annealing schedules with longer periods of higher learning rates can improve generalisation, which we do ﬁnd to help the base network.
In both cases, we apply data augmentation of random crops (after symmetric padding) and random horizontal reﬂections.
B. P ARTITIONING NETWORK
We use the same MLP with three hidden layers for all our partitioning networks, with dimensions 64 − 32 − 16 .
We use L AYER N ORM (Ba et al., 2016) and R E LU activations after each layer.
To avoid merely memorising hard examples, it is necessary to regularise this network, so we also apply spectral normalisation (Yoshida & Miyato, 2017); this involves spectrally normalising every linear layer, and excluding the scaling term in the layer normalisation transforms, as in Miyato et al.
We use a separate network for each class, training for 100 iterations each, with the same batch size as used for training the rest of the model.
We use the Adam optimiser (Kingma & Ba, 2014) with a learning rate of 1 e − 4 .
In preliminary experiments we found a shared network for all categories to also work, using conditional layer normalisation (Ba et al., 2016; de Vries et al., 2017).
We didn’t investigate it further for all datasets, since in general a larger number of classes in a dataset might require larger capacity in the partition predictor to account for more features, and as the number of classes go up, a number of smaller matrices can have a lower footprint than one very large matrix.
Network architecture design for the partitioning network was done only on the COLOURED MNIST dataset, with access to true oracle group labels for a smaller set of in-distribution validation images (20 per category).
The same network architecture was applied for the two COCO datasets.
The γ hyper-parameter was learned separately for all datasets, using the smaller sets of validation images.
We ﬁnd, in preliminary experiments, that using random partitionings lead to much worse performance.
This suggests that, although not performed for our present study, in more realistic situations one could potentially tune these hyper-parameters by validating over classiﬁcation accuracy as for the invariance penalties.
B. I NVARIANCE PENALTIES
In all cases, we pause training of the base network after 1 epoch of training, and learn a partitioning of the training set.
This learned partition is used to drop in the invariance penalties as training proceeds, and as in prior work (Arjovsky et al., 2019; Krueger et al., 2020), we ﬁnd ramping in the penalty co-efﬁcient over a number of epochs to be useful for stable training.
For IRM and REx (and conditional variants), we ﬁnd it helpful to scale the ERM term down by the penalty co-efﬁcient when the optimal validation co-efﬁcient is greater than 1, as implemented by Arjovsky et al.
and Krueger et al.
C R EVIEW OF BASELINES AND CONDITIONAL VARIANTS
We brieﬂy review the group invariance methods we compared.
C. IRM V
In Arjovsky et al.
(2019), a risk regularisation method is described in order to encourage reliance on features that obey stable correlations with the target variable across data from different environments.
The regularisation consists of a gradient penalty wrt a dummy multiplier on the logits, with the intuition that scaling up or shrinking the logits in different environments can only result in local improvements within each environment if the classiﬁer uses features that correlate at different levels in the different environments.
The objective function is
[Equation]
(9) Φ comprises the predictor, which in our case is w (cid:62) f θ ( x ) .
µ is a dummy multiplier, ﬁxed at 1, and R e is the environment risk, corresponding to the average loss for data in a particular environment when using Φ .
For our conditional variant (cIRMv1), we stratify the gradient penalty over classes, so that the penalty is applied separately per class in each environment.
The hyper-parameters we search over for this method include the penalty co-efﬁcient λ and the number of epochs of training over which to linearly ramp up λ to its full value.
C. RE X
Krueger et al.
(2020) proposed a risk regularisation method that aims to directly match training risks across environments, by imposing a penalty that minimises the variance of risks across environments (V-REx).
[Equation]
For our conditional variant (cREx), we apply the variance penalty stratiﬁed by class.
The hyper-parameters we search over for this method include the penalty co-efﬁcient λ and the number of epochs of training over which to linearly ramp up λ to its full value.
C. G ROUP DRO
Sagawa et al.
(2020) suggest an online algorithm for group-based distributionally robust optimisation, which effectively re-weights group losses as a function of their evolving magnitudes, therefore putting more emphasis on groups that fare worse through training.
For our conditional variant (cGroupDRO), we compute the group weights per class, by using the losses belonging to the classes separately in each group.
The hyper-parameters we search over for this method include the learning rate for the online groupweights and the two group adjustment hyper-parameters.
Additionally, we sample equally from both groups for this method, as suggested, ﬁnding it to improve results in preliminary experiments.
C. R EWEIGHT
We learn a hyper-parameter λ on validation, such that every example in the majority group is weighted with 1 / ( λ + 1) (because we only want to weight the majority group down).
The hyper-parameters we search over for this method include the penalty co-efﬁcient λ and the number of epochs of training over which to linearly ramp up λ to its full value.
C. MMD FEATURE MATCHING
Maximum mean discrepancy based distributional matching of features across domains has been shown to be effective for domain generalisation (Li et al., 2018), and conditional matching of distributions (usually with adversaries, for example, in Li et al.
(2018a)) tends to work better.
We found in preliminary experiments that conditional MMD signiﬁcantly outperformed the unconditional variant, so we only ran full experiments and reported results using cMMD.
The group invariance penalty looks as follows
[Equation]
where φ induces a kernel function K , which in our implementation is a mixture of 3 Gaussians with bandwidths [1 , 5 , 10] , which are the recommended set of bandwidths in Li et al.
Adding sharper or ﬂatter bandwidths appeared to hurt performance in preliminary experiments.
The hyper-parameters we search over for this method include the penalty co-efﬁcient λ and the number of epochs of training over which to linearly ramp up λ to its full value.
C. H YPER - PARAMETER GRID SEARCH RANGES
In all cases, λ is searched over a range of { 1 e − 4 , 1 e − 3 , 1 e − 2 , 0 . 1 , 0 . 5 , 1 , 2 , 5 , 10 , 20 , 50 , 100 , 200 , 500 , 1000 , 10000 , 100000 } , and the number of epochs over which to linearly ramp up λ is searched over { 1 , 5 , 30 } for MNIST and { 1 , 10 , 200 } for COCO.
For the GroupDRO methods, we search over { 0 . 001 . 0 .
01 , 0 . 1 , 1 . 0 , 10 } for the learning rate of the group-weights, and over { 0 , 1 , 2 , 3 , 4 , 5 } for the group-adjustment hyperparameters, as recommended in Sagawa et al.
We also average the losses group-wise as already done in IRMv1 and REx for cMMD and PGI, except for COCO- ON -P LACES , where we ﬁnd this choice to hurt performance.
D D IFFERENT VALIDATION SETS
In this section, we report results for all the methods we compare, when picking hyper-parameters using different validation sets, as discussed in Section 5.4.
We note that contrary to what one would typically do in a real-world deployment, we do not augment the training sets with the validation sets for evaluating test time performance. This is because the presence of data with systematic distributional shift at training time improves performance signiﬁcantly (as observed in Table 1), and our goal here is to perform an illustrative study about the potential effectiveness of invariance methods at learning to generalise systematically. While we could have augmented the training set with validation data when we are not using validation sets with systematic distributional shift, we follow the same protocol in these cases of not augmenting the training set, in order to keep the numbers comparable with each other across different validation schemes.
[Table] [Table]
[Table] [Table]
[Table] [Table] [Table]
[Table] [Table]
E M EASURING SEMANTIC ANOMALY DETECTION
We use the test set with non-systematic distributional shift as the normal data, and the held-out class data combined systematically with the biasing colours or backgrounds as the anomalous data.
For MNIST, this means there is 9 times more normal data than anomalous data, which reﬂects the typical situation of anomalies being rarer.
Our choice of normal data makes this a harder task than usual, since we are assessing for higher (than the anomalies) predictive conﬁdences for non-semantic shift with semantic factors kept the same, and reduced predictive conﬁdence for semantic shift with non-semantic factors from the seen data.
For the COCO datasets, we only sample 100 images from the held-out class to resemble the MNIST experimental setup.
Anomaly detection is measured using average precision, treating the anomalous class as positive, with the negative of the predictive softmax conﬁdence as the score (Hendrycks & Gimpel, 2017).
F A LGORITHM
Algorithm 1: Algorithm for PGI
Initialise all classiﬁer parameters θ, w and partition-predicting networks, g c , ∀ c ∈
[ C ] ; for one epoch do for mini-batches D b ∈ D do grad θ := ∇ θ (cid:96) ( θ, w |D b ) ; grad w := ∇ w (cid:96) ( θ, w |D b ) ; θ, w := optimizer ( grad θ , grad w ) ; end endfor all classes c ∈
[ C ] do
Learn a partition for images in D with labels c , (Eq.
8) endfor T − 1 epochs do for mini-batches D b ∈ D do grad θ := ∇ θ ( (cid:96) ( θ, w |D b )
+ λ.penalty )
(Eq. 6,7) ; grad w := ∇ w (cid:96) ( θ, w |D b ) ; θ, w :=
optimizer
grad θ , grad w ) ; end end