Multi-Task Distribution Learning
Authors: Connor Shorten
Venue: ICLR 2022 Submitted (2022)
Keywords: Data Augmentation
Distribution Shift
Multi-Task Learning

Abstract: Multi-Task Learning describes training on multiple tasks simultaneously to leverage the shared information between tasks. Tasks are typically defined as alternative ways to label data. Given an image of a face, a model could either classify the presence of sunglasses, or the presence of facial hair. This example highlights how the same input image can be posed as two separate binary classification problems. We present Multi-Task Distribution Learning, highlighting the similarities between Multi-Task Learning and preparing for Distribution Shift. Even with rapid advances in large-scale models, a Multi-Task Learner that is trained with object detection will outperform zero-shot inference on object detection. Similarly, we show how training with a data distribution aids with performance on that data distribution. We begin our experiments with a pairing of distribution tasks. We then show that this scales to optimizing 10 distribution tasks simultaneously. We further perform a task grouping analysis to see which augmentations train well together and which do not. Multi-Task Distribution Learning highlights the similarities between Distribution Shift and Zero-Shot task inference. These experiments will continue to improve with advances in generative modeling that enables simulating more interesting distribution shifts outside of standard augmentations. In addition, we discuss how the WILDS benchmark of Domain Generalizations and Subpopulation Shifts will aid in future work. Utilizing the prior knowledge of data augmentation and understanding multi-task interference is a promising direction to understand the phenomenon of Distribution Shift. To facilitate reproduction, we are open-sourcing code, leaderboards, and experimental data upon publication.

Under review as a conference paper at ICLR 2022
LEARNING MULTIPLE DATA DISTRIBUTIONS
Anonymous authors
Paper under double-blind review
ABSTRACT
Vulnerability to Distribution Shift is major challenge for real-world applications of Deep Learning. Distribution Shift occurs when the test data distribution is not identical to the training distribution. We present Multi-Distribution Learning, ex- ploring the effectivvness of Data Augmentation to prepare for Distribution Shift. We use common Data Augmentations simulate multiple data distributions and test how many of these distributions can be learned with single model. We first illustrate the impact of adding the test distribution to the training distribution, such that it is no longer Distribution Shift. We find success with pairs of distributions, however, the model limited when attempting to learn 10 distributions simulta- neously. In order to search for subset of distributions that train well together, we use a Lookahead search derived from research in Multi-Task Learning. From this analysis, select distributions with high affinities to train together and I negative distribution to use for performance comparison. These experiments con- firm that the accuracy of models evaluated test set included in the training distribution is much higher compared to test distributions not included in training. We evaluate our distribution grouping with Zero-Shot Distribution Inference a held-out positive distribution and the held-out negative distribution. We compare this Zero-Shot Distribution Inference to strong baseline of the RandAugment training scheme, as well as to a weak baseline of a model trained without Data Augmentation. We conclude by discussing directions for future work in using Data Augmentation to combat Distribution Shift.
INTRODUCTION
Distribution Shift, defined as the performance on test distributions not identical to the training set, are the Achilles heel of Deep Learning. We propose simple solution to use Data Augmentation for targeted Generalization, alleviating the problem of Distribution Shift. In training, we aim to cover as many data distributions as possible, such that Distribution Shift is not completely novel to our model. We measure our models using the Average Distribution Accuracy and propose a new framework of Distribution Generalization Metrics. These performance scores are weighted average of the model loss or accuracy, the number of distributions it can cover, and the diversity of these distributions.
We begin with two important questions for thinking about Distribution Shift. What are different data distributions? What does Distribution Shift in the real-world look like? A data distribution can be characterized based on the frequency of class labels. For example. the training set may contain 80% dogs and 20% cats, but the test distribution shifts 40% dogs and 60% cats. In addition to distributions of labels, we can think about distributions of the inputs themselves. This is significantly harder to characterize due to the nigh-dimensional nature of raw inputs, such as pixel grids or word embeddings, or even intermediate vector representations from these sequential processing models.
The WILDS benchmark Koh et al. (2020) has been a very useful dataset to understand Distribution Shift in the real-world. Koh et al. collected examples of Distribution Shifts across applications in Wildlife Monitoring, Molecular Engineering, and Python Code Completion, to give a few ex- amples. The authors characterize Distribution Shift into two categories, Subpopulation Shift and Domain Generalization. Subpopulation Shift is well illustrated by our previous example of a shift- ing frequency of cats and dogs. However, Subpopulation Shift can also refer to more high-level distributions, such as an increase in certain features of the input. Domain Generalization describes settings where the train and test domains are completely disjoint. For example, sentiment classifier

---

Under review as a conference paper at ICLR 2022
Theming
Tel
Die
Taking
Duyy Burlen
DusButbo
Figure 1: The Multi-Dise Distribution Learning framework evaluates how training with data distribu- tion helps with test performance on that distribution. The framework tests the limits of how many distributions we can learn with single model.
trained on movie reviews and then tested restaurant reviews Gururangan et al. (2020), or image classifier trained on photorealistic images and tested on paintings Peng et al. (2019).
Another technique for studying Distribution Shift is through corruption tests. Datasets such as are constructed by applying common Data Augmentations to the original ImageNet test set. Models that achieve high perfor- mance on ImageNet typically do not generalize to ImageNet-C C as well. There are several other ex- amples of these datasets, such smaller-scale CIFAR-C and MNIST-C tests Mu & Gilmer (209)), ImageNet-P with more subtle corruptions, and Stylized- Nmaieed- Geirhos et al. (2019) which uses style transfer algorithm to disentangle content from style in images. Corruption testing, achieved by applying Data Augmentation to an original test set, is an important lens for viewing the phe- nomenon Distribution Shift. We explore how we can similarly use Data Augmentation to prepare for Distribution Shift and the limits of this technique.
ImageNet-C Henddyyck & Dietterich (2019),
Our objective use Data Augmentation to target the Generalization we want the model to achieve. This is a simple, but promising strategy to inject prior knowledge about how the distribution will shift. We believe these assumptions on the test data distribution are necessary to make progress, inspired by the following quote from Arjovsky (2019). "If the test data is arbitrary or unrelated to the training data, then generalization is obviously futile." " This quote supports the idea that we should make assumptions and bake in prior knowledge about the test distributions.
We begin our experiments by confirming that models which achieve high performance on an original test set will fail generalize to the same test set corrupted with common data augmentations. This holds even for models trained with RandAugment Cubuk et al. (2020). a strong regularization strat- egy based Data Augmentation. Our training strategy achieve generalization to these test sets is illustrated in Figure 1. We apply N corruptions to the original data to form N training distributions. These N corruptions are additionally applied the original test data form an equivalent N test distributions. We report the average accuracy across the N test distributions. We begin by showing how this strategy increases the performance using N=2 augmentations.
We then scale our strategy up to learning N=10 distributions at once. We find that this requires much longer training times achieve rivaling performance to single augmentation specializations trained with N=1. We borrow inspiration from research in Multi-Tast Task Learning and use a Task Groupings algorithm Fifty al.] (202T) to search for distributions that train well together. We refer to these distributions as having high affinity with each other. We use a Lookahead search to compare the performance of a proposed gradient update after training on a given distribution. Distributions are determined to train well together if the proposed update improves the performance across both test sets. We further attempt to disentangle the impact of repeating data in our distribution construction strategy by partitioning the CIFAR-10 set into 5,000 images each, to remove the confounding factor of repeating data. We find similar distribution groupings with and without repeating the original data.
2

---

Under review as a conference paper at ICLR 2022
We select 4 augmentations from our Lookahead search to train together. We additionally select augmentation with a negative affinity for the sake comparison. We find that the grouping achieves much higher Average Distribution Score than the N=10 model trained with similar computation. We then generalize these experiments to more realistic evaluation of Distribution Shift. We pose this as problem of Zero-Shot Distribution Inference. We propose that groupings with high affinity to the novel test distribution will perform better. We use a cross-validation analysis to iteratively pair 3 out of the augmented distributions for training, evaluating the held-out distribution Zero-Shot. We find much better performance with this cluster compared to strong RandAugment baseline, which is not completely Zero-Shot, as well as a weak baseline of a model trained without any Data Augmentation.
Our experiments utilize Data Augmentation to simulate diverse data distributions. We are currently limited to geometric transformations such as Rotations Horizontal Flipping. However, we note that new frontiers in Data Augmentation such as CycleGAN Zhu al. (2017) translations between domains or Texl-te-Image generation Ramesh et al. (2021) should have massive impact on this framework. We believe that Distribution Shift should be viewed similarly to Zero-Shot task infer- commonly used in probing large models Liu al. ( An interesting characteristic of Zero-Shot inference the difference between massive models typically trained with less supervi- sion than smaller, expert models. We view training on specialized distributions as analogous to supervised learning with particular tasks. Finally, we propose directions to extend our Distribution Generalization Score. Rather than the naive average accuracy, which we have reported in these experiments, we want metric that rewards covering many distributions and the diversity of these distributions. We propose some ideas for measuring the diversity of these distributions.
In summary, we propose framework of using Data Augmentation to anticipate and alleviate the challenges of Distribution Shift. Our contributions are as follows:
We propose a new framework for Learning and Evaluating Multiple Data Distributions.
We present a new metric for Distributional Generalization, the Average Distribution Score.
We adapt a Task Grouping algorithm from Multi-Task Learning to study Multiple Distri- bution Learning.
We present an analysis of Zero-Shot Distribution Inference, inspired by research on Zero- Shot Task Inference.
RELATED WORK
2.1 DISTRIBUTION SHIFT
Understanding the nature of Generalization is one of the core goals of Artificial Intelligence Chol- let (2019). Many recent works have characterized types of Generalization. These include Do- main Transfer Peng al. (2019). Compositional and Systematic Generalization Johnson al. 2016). Concept Drift, or Robustness, to give few examples. Further characterizations include In-Distribution versus lii--offDDtributiion Generalization, viewing Generalization gaps in a similar lens as fine- and coarse-grained classification. From practical perspective, Koh et al. (2020) have published the WILDS benchmark characterizing the generalization needed for many real-worls ap- plications. In this work, propose simple framework for studying the relationship between train- ing and test distributions. Andreassen et al. (202T) and Wortsman et al. (002 have also explored the relationship between Zero-Shot robustness and fine-tuned models. Regardless of the particular type of generalization, we can likely form Data Augmentations that communicate expected Distribution
Shifts.
2.2 CORRUPTION TESTING
Many works have found that Deep Learning systems vulnerable to corruption tests. These tests are generally head-scratching cases because they very closely resemble the original input. This de- fies natural intuition about the data and requires thinking closely about the distributions of inputs. Corruption tests range from injecting adversarial noise maps to applying Data Augmentations, as we have done in this work. [Hendrycks & Dietterich (009) presented ImageNet-C, dataset constructed
3

---

Under review as a conference paper at ICLR 2022
from of 75 Data Augmentations and ImageNet-P constructed from adversarial injections. Ova- dia et al. (2019) used corruption sets to evaluate Uncertainty techniques in Deep Learning. ? utilized style randomization to understand intrinsic properties of Convolutional Neural Networks, finding texture bias.
2.3 DATA AUGMENTATION
Multi-Distribbution Learning a technique building on Data Augmentation for improving the Gen- eralization of Deep Neural Networks. Many works have covered different Data Augmentation tech- niques across many domains. Some examples include Kernel Filters, Geometric Transformations, Random Erasing, Color Space Transformations, Mixing Images, Adversarial Training, Neural Style Transfer, and GAN Data Augmentation [Shorten & Khoshgoltaar (2019). In addition to individ- ual augmentations, recent works have explored new strategies for utilizing Data Augmentation. This includes application in contrastive losses self-supervised learning such as SimCLR Chen al. a (2020) well as consistency losses in Unsupervised Data Augmentation Xie al. (2020). Sinha et al. (2021) have additionally explored utilizing Data Augmentations that intentionally create Out-of-Distribution data. Our work is similar in that we are looking for new ways to utilize Data Augmentation in Deep Learning.
3 EXPERIMENTS
Our experiments report training and evaluation strategies for the ResNet152V He et al. (016) on the CIFAR-10 10 dataset. The experiments are implemented using the Keras framework. The experi- ments are run on an NVIDIA A100 GPU in Google Colab Pro+ runtime, with a maximum training time of roughly 20 hours. We experiment with up to common Data Augmentations used for image data: Rotation, Translation, Crop, Cutout, Bright, Dark, Gamma Contrast, Horizontal Flip, Vertical Flip, and Jigsaw. These augmentations are implemented using the imgaug open-source library. Following is quick description of how these augmentations transform the original data distribution.
The Rotation augmentation rotates an image along polar coordinates between 0 to 360 degrees. The Translation augmentation shifts an image along the x and y axes. This is done by moving pixels over/up by pixels and adding Os to displaced values. The Crop augmentation selects a subset of the original image and then resizes the subset to the original resolution with bilinear interpolation. The Cutout augmentation selects patches of the image and replaces the patch with 0 values. The Bright augmentation adds k to each pixel value in the image, essentially pushing it closer to 255 (white). The Dark augmentation inversely subtracts pushing it closer to 0 (black). The Gamma Contrast augmentation alters the distribution of pixel values a smaller window between 0 and 255. The Horizontal Flip (LR Flip) augmentation mirrors the original image along the y-axis. The Vertical Flip (UD Flip) augmentation inversely mirrors the original image along the x-axis. The Jigsaw augmentation divides the image into n pixel windows and shuffles their location in the image.
each pixel value in the image, essentially
As a quick summary of our experiments, Training with 2 Distributions characterizes the problem of Distribution Shift and the simple solution of adding the test distribution to the training set. The model is not training with the test set itself, just the underlying distribution that will be encoun- tered in-the-wild. We simulate the different distributions through Data Augmentations. We then scale this up to N=10 distributions in training. Lookahead Search to Group Training Distributions utilizes Lookahead analysis to find distributions that complement each other during training. Re- sults of Grouping Training Distributions compares the Average Distribution Score with the grouped distributions N=1 and N=10 distribution training strategies. Table illustrates the Zero-Shot Dis- tribution Inference of the grouping with a strong Rand/ugment baseline and weak baseline of training with No Data Augmentation.
3.1 TRAINING WITH 2 DISTRIBUTIONS
We begin by highlighting the problem of Distribution Shift, reproducing it in experimental setup, and conducting simple experiments within the framework described in Figure 1. As shown in Table 1, when trained without any Data Augmentation the model performs much worse the Rotation
4

---

Under review as a conference paper at ICLR 2022
Training Technique
Original Test
Rotation Test
Jigsaw Test
Average Distribution Score
No Augmentation
73.9%
41.6%
61.3%
51.5%
RandAugment
83.8% %
55.0%
69.6%
62.3%
Rotation
75.9%
75.2%
63.2%
69.2%
Jigsaw
75.9%
42.7%
71.6%
57.2%
Seq-RJ
78.2%
66.5%
73.0%
69.8%
Seq-JR 75.9%
71.5%
69.4%
70.5%
Alt-RJ
80.0%
73.7%
74.8%
74.3%
Alt-SH-RJ
80.2%
73.3%
75.0% %
74.2%
Table 1: The efficacy of training with specialized distribution. The bolded values indicate the maximum accuracy achieved on each data distribution. The bolded values indicate the maximumm accuracy for each data distribution.
Epoch
10
30
50
200
350
1000
Original Test
65.5%
68.4%
72.5%
77.0%
79.8%
83.4%
Average Aug Score 57.4%
58.7%
61.8%
68.9%
71.6%
75.3%
Max Aug Score
78.5%
77.1%
78.0%
79.2%
79.5%
83.2 %
Min Aug Score
50.2%
49.3%
49.4%
62.7%
65.1%
70.6 %
Max Aug Name
UD Flip
UD Flip
UD Flip
UD Flip
LR Flip
LR Flip Crop
Min Aug Name
Rotate
Translate
Crop
Crop
Crop
Table 2: Learning curves for Multi-Distributicn Learning of 10 Augmentations. The schematic of
this training strategy is shown in Figure
and Jigsaw distributions compared to the original CIFAR-10 data. Even though the Rand Augment training scheme dramatically improves performance the original set, it is still susceptible to these augmentation derived distribution shifts. We then report the efficacy of specialized training on the distribution used in testing, such that it is no longer Distribution Shift. We shift our aim to learning to specialize on multiple distributions in training. We consider two high-level schema for doing this. The first of which is Sequential Training (abbreviated Seq in Table 1). Sequential Training refers to first learning the Rotation distribution and then learning the Jigsaw distribution. We encounter an interesting case of Catastrophic Forgetting where performance on the previously learned distribution quickly drops as it shifts focus to another data distribution. We find better result with Alternating Training (abbreviated as Alt in Table 1). In Alternating Training we switch between each training distribution every epoch, rather than after convergence. We additionally test sharing a base feature extractor between each distribution and using separate classification heads model each distribution (abbreviated as Alt-S in Table 1). We do not find much difference with separate heads compared to fine-tuning the entire model. We continue with Alternating Training throughout the rest of the experiments.
3.2 TRAINING WITH 10 DISTRIBUTIONS
Following the success of learning 2 distributions in training, we our sights on increasing the number of learned data distributions. Table 2 illustrates the learning curve of 10 data distributions. We find that it takes much longer for this strategy to achieve strong Average Augmentation Score. However, after 1,000 epochs (roughly 20 hours of training with our computing setup described previously), the model achieves a fairly strong Average Augmentation Score with small variance and no significant minimum or maximum outliers in the accuracy scores. Throughout training the Rotation, Translation, and Crop Augmentations have the lowest performance, whereas Vertical and Horizontal Flipping (shown UD Flip and LR Flip) have the highest score. At the end of training the multi-distribution learning model achieves nearly identical score on the Horizontal Flip as the Original Test it could be interesting future work this distribution shift in the prediction and representation spaces. We find more interesting details about the Vertical Flip augmentation in analysis of Lookahead Search.
if the model has achieved invariance to
5

---

Under review as a conference paper at ICLR 2022
Repeating Data
0.21
0.031
0.017
0.02
-0.024
0.0049
-0.03
-0.0048
-0.019
0.0092
0.0043
0.19
0
0.042
0 046
0.063
054
0.086
0.052
0.019
0.029
00017
0.029
0.022
+0 0.0013
-0.036
0..11
-0.02
-0 00.029
-0.019
-0.004
0.062
-0.039
-0.665
0.054
-0.04
-0.0
038
-0.046
-0.007
00002
0.029
-0.553
060666
0.073
00.00
0.0075
0.065
0.15
-0.
0.022
-0 02025
0.023
0.011
-02
00.11
-0.018
-0.032
0.0096
-0.068
0.048
0.026
000025
-0.04
0.012
-0.3
-0.021
0.0067
00.014
0.0049
-0.017
0.0024
0.11
-0.0099
-0.031
0.0003
-044
0.011
0.016
-0 0.006
0.032
0000881
013 0 0002
0.014
-0.014
0.027
-0.5
0..9
-0.33
-0.6
-0.55
-0.55
0.64
-0.54
-0.66
-0.15
0053
0.0035
0.018
0
0.04
0.0053
0.02
0.022
0.013
0.015
0.081
-0.6
2
3
6
7
8
Figure 2: This table shows the performance difference by taking the proposed gradient update on the y-axis on the test distributions on the x-axis. In this index: Rotate = O, Translate = 1, Crop = 2, Cutout 3, Bright = Dark 5, Gamma Contrast 6, LR Flip 7, UD Flip 8, and Jigsaw 9.
3.3 LOOKAHEAD SEARCH TO GROUP TRAINING DISTRIBUTIONS
Due to the long training time of learning 10 distributions at once, we aim to find a subset of dis- tributions that have a faster learning curve and higher performance. We adapt the Task Groupings algorithm Fifty et al.] (2021) to search for this subset. This algorithm uses Lookahead search to compute what the change in performance on other tasks, or distributions in our example, will change. For example, in Multi-Task Learning the Task Groupings algorithm computes the change in performance on Semantic Segmentation after taking a gradient update on Object Detection. We similarly compute the change in performance the Translation distribution after taking gradient step on the Rotated training set.
The results of the Distribution Grouping analysis shown in Figure 2. The x-axis is the evaluated distribution and the y-axis is the distribution that proposes a gradient update. The augmentations are indexed according their position the following list: Rotate 0. Translate
Crop 2, Cutout = 3, Bright = 4, Dark = Gamma Contrast = 6, LR Flip 7, UD Flip 8, and Jigsaw We find a strong outlier result where proposed gradient update on Vertical Flipping (UD Flip) causes significant drop in performance across the other augmentations. This likely explains the optimization difficulty behind learning all 10 distributions simultaneously. The second strongest negative result is the impact of updating on Bright and evaluating the Dark test distribution. We leave to future work explore if the inverse augmentations of increased and decreased Brightness have opposite gradient directions.
From Figure 2, we propose to group Rotation, Translation, Dark, and Jigsaw to form our subset. In order to further understand the impact of gradient update on one distribution to the performance of another, we avoid repeating data. For example, the rotation distribution still has roughly the same image the translation distribution in Figure 2. The right hand side of Figure illustrates the lookahead updates when we split the data such that no distribution has the same original CIFAR-10 10 image. We find similar grouping analysis with and without repeating data.

---

Under review as a conference paper at ICLR 2022
Repeating Data
No Repeating Data
0.1
0:33
017
0.02
024 0049 039 0048 019 0092
0.16
0.088
0.19
0.026
0.04
0.29
0.945
0.018
0.013
0.03
0.004
0.9
0 042
046 063 4 054 086 0.052 0819 0.29
01
0.066
0.95
0,027
4
0.068
0.071
0.11
0.064
0.41
0.07
0.017
0.29
0022
4 013 036
4 016 0.02 4 929 6 0.019 4 004
00
0.047
0.099
0.64
0.001
0.017
0013
4 001
008
0.42
0002
00
0.062
0.039
0.045
0054
0 044
4 038
4 046
4 007 6 022 0.29
0.02
0.033
0.06
0.010
4 022
0.022
0.2
4 001
0.24
01
0.053
0.066
0.073
8.0075
0065
0.15
0.022
0 025
6 0.023 0.01
0.033
0.044
0.044
0 001
0011
008
0 013
0.04
0.04
007
--01
0.01
0:018
0.022
6.096
0.000
0048
0626
0.0025
0.04
0012
0.048
0.854
0.049
0.055
0.25
0.06
0.045
0012
0.022
0.001
-03
00049
0.017
011
0000
0.03
0.03
0041
0:01
0.25
0011
0016
0081
0002
0.056
0:022
0.08
010
001
--03
--85
-06
--04
Figure 3: A Side-by-Side view of the impact of repeating data for Lookahead update analysis. Please see the caption of Figure 2 for more details on the x and y axes and the distribution indexing from O to 9.
3.4 RESULTS OF GROUPING TRAINING DISTRIBUTIONS
We analyze the performance of our distribution grouping on learning the training distributions and zero-shot inference novel distributions. Table 3 reports enumerating through 3 of the 4 selected distributions for training. We encountered an optimization difficulty with the Rotate, Translate, Jigsaw (abbreviated as RIJ Table3) that achieves low performance. However, the other 3 clusters outperform the 10 Augmentation scheme other than the Jigsaw test. We find an interesting division between distributions that benefit from specialization (Rotation and Translation) and others that benefit from the ensemble (Darker and Jigsaw). This could be an interesting property for the design of ensembles and specialist distribution models that we leave for future work.
The ultimate goal of learning multiple data distributions is to lessen the decrease in accuracy when faced with Distribution Shift. We test how the grouped distributions help with Zero-Shot Distribu- tion Inference through cross-validation analysis within the group. For example, we hold-out the Rotation distribution and train with Translate, Darker, and Jigsaw. The entires for 3 Augmentations in Table reflect the Zero-Shot accuracy without training that augmented distribution. For Rota- tion and Translation, we find a very strong Zero-Shot generalization with the group. As mentioned in Table 3, ran into optimization difficulties with our run leaving out the Darker distribution and it achieves a very low accuracy. We found similar performance on Zero-Shot Jigsaw generalization. We held-out Gamma Contrast an example of distribution with negative affinity to the group. We report the best Zero-Shot accuracy on Gamma Contrast, supporting our analysis that the negative augmentation does not benefit from this particular cluster of training distributions. We compare this with the RandAugment model reported in Table 1, which achieves 83.8% accuracy on the original CIFAR-10 test set. Although RandAugment has never seen any of these distributions in isolation, it has seen compositions of them such as rotation, then translation, then darker applied to a single image. Thus, the RandAugment baseline is not truly Zero-Shot. The No Augmentation baseline highlights the poor performance when faced with Distribution Shift. In both RandAugment and No Augmentation we see the lowest performance decrease with the weakest augmentation, Darker. This highlights the need to develop distribution distance scores to weight the performances.
4 DISCUSSION
4.1 FUTURE DIRECTIONS FOR SIMULATING DATA DISTRIBUTIONS
Our experiments tested simple data augmentations to simulate distribution shifts. We have used the terms "augmentations" and "distributions" interchangeably to propose their similarities. We

---

Under review as a conference paper at ICLR 2022
Training Distributions
Rotate Test
Translate Test
Darker Test
Jigsaw Test
Average Distribution Score
Augs (TDJ)
N/A
72.4%
82.2%
67.4%
74.0%
Augs (RDJ)
72.7%
N/A
81.0%
66.3%
73.3%
Augs (RTJ)
37.3%
33.2%
N/A
42.2%
37.6%
3 Augs (RTD)
71.9%
71.0%
81.4%
N/A
74.8%
10 Augs, 100 Epochs
54.3%
49.3%
62.7%
57.6%
56.0%
10 Augs, 1000 Epochs
70.9%
74.8%
76,3%
76.0%
74.5% %
Aug
75.2%
79.4%
74.9%
71.6%
N/A
Table 3: Comparison of Generalization to Test Distributions included in the Training Distribution.
Training Distributions
Rotate Test
Translate Test
Darker Test
Jigsaw Test
Gamma Contrast Test
3 Augmentations
72.4%
71.9%
45.5%
65.2%
62.6%
RandAugment
57.4%
49.7%
82.2 %
69.5%
72.6 %
No Augmentation 46.5%
35.6%
68.1%
62.5%
54.2%
Table 4: Zero-Shot Distribution Inference.
believe techniques such as CycleGAN Zhu al. (2017) and DALL-E Ramesh et al. (2021) will allow us to control Domain Generalization with Data Augmentation--ttlee interfaces. For example, the Wiidlife-monitoriie monitoring camera at test time. CycleGAN may be able to map the existing training data into this new location, further guided with text interfaces from DALL-E. We propose that we can train with this inferred novel distribution to perform better in the Distribution Shift.
camera in the WILDS benchmark encounters new location for wildelife
4.2 ZERO-SHOT DISTRIBUTION INFERENCE
Few-Shot Learning describes learning a new task with a few examples, hopefully achieved by lever- aging information from previously learned tasks. Zero-Shot Learning refers to more extreme case in which the model performs a new task, but no additional learning is allowed. We can similarly view Distribution Shifts as having few examples from the new distribution having to perform Zero-Shot task inference. Another interesting trend in Few- and Zero-Shot learning is the efficacy of larger-scale models. Models such as GPT-3 can be guided with prompts to perform novel language tasks. It remains to be seen whether collections of expert models assigned to their respective infer- ences will outperform a single monolithic model. Even if it becomes unnecessary to train expert models for collections of tasks, distributions still be useful for evaluation.
have argued, understanding novel tasks will
4.3 DISTRIBUTION GENERALIZATION SCORE
Our experiments report a simple average augmentation accuracy. For future work we intend to explore metrics that better capture the diversity of model's generalization ability. We intend to weight the average accuracies by the number evaluated distributions and the distance between these distributions. The distance between distributions is challenging quantity measure. One solution, similar to ideas in the automated evaluation of generated images such as the Inception Score, could be embed each distribution into the feature space of pre-trained classifier and use these vector representations to compute distances. strong Distribution Generalization Score should also be correlated with performance Zero-Shot Distribution Inference.
CONCLUSION
In Conclusion, we have presented Multi-Distribution Learning, framework to study the impact of training on potential test distributions. This is simple framework to leverage the prior knowledge about how a data distribution is likely to change in deployment. Key to our framework are the number of distributions used in training and the metric used score the diversity of distribution performance. We have presented several experiments that highlight the challenge of learning many distributions once. We have additionally presented technique group together distributions for
8

---

Under review as a conference paper at ICLR 2022
training. Finally, propose framing the problem of Distribution Shift analogous to Zero-Shot Task Inference. We believe that rapid advances in Data Augmentation and Generative Modeling will have a large impact on this research direction.
REFERENCES
Andres Andreassen, Yasaman Bahri, Behnam Neyshabur, and Rebecca Roelofs. The evolution of out-of-distribution robustness throughout fine-tuning. arXiv preprint, 2021. URL arXiv:
2106. 15831
Martin Arjovsky. Phd thesis:
Out of distribution generalization machine learning. New York University, Courant Institute of Mathematical Sciences, 2019.
Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. ICML, 2020.
Francois Chollet. On the measure of intelligence. arXiv preprint, 2019. URL arxiv: 01547
Ekin D. Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V. Le. Randaugment: Practical automated data augmentation with reduced search space. In NeurlPS, 2020.
Christopher Fifty, Ehsan Amid, Zhe Zhao, Tianhe Yu, Rohan Anil, and Chelsea Finn. Efficiently
identifying task groupings for multi-task learning. arXiv preprint, 2021. URL arxiv:2109.
04617
Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A. Wichmann, and Wieland Brendel. Imagenet-traiie trained cnns are biased towards texture; increasing shape bias im- proves accuracy and robustness. ICLR, 2019.
Suchin Gururangan, Ana Marasovic, Swabha Swayamdipta, Kyle Lo, Beltagy, Doug Downey, and Noah A. Smith. Don't 't stop pretraining: Adapt language models to domains and tasks. In ACL, 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In ECCV, 2016.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common cor- ruptions and perturbations. In ICLR, 2019.
Justin Johnson, Bharath Hariharan, Laurenas van der Maaten, Li Fei-Fei, C. Lawrence Zitnick, and Ross Girshick. Clevr: A diagnostic dataset for compositional language and elementary visual reasoning. arXiv preprint, 2016. URL arwiv:1612..1... 06890
Pang W. Koh, Shiori Sagawa, Henrik Marklund, Sang M. Xie. Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard L Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton A. Earnshaw, Imran S. Haque, Sara Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percey Liang. Wilds: A benchmark of in-the-will distribution shifts. arXiV preprint, 2020. URL 2rl.aa2012..000000
Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre- train, prompt, and predict: A systematic survey of prompting methods in natural language pro- cessing. arXiv preprint, 2021. URL
Norman Mu and Justin Gilmer. Mnist-c: A robustness benchmark for computer vision. arXiv
preprint, 2019. URL arxii: 02337
Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D Sculley, Sebastian Nowozin, Joshua V. Dillion, Balaji Laksminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. NeurlPS, 2019.
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang. Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In ICCV, 2019.

---

Under review as a conference paper at ICLR 2022
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alee Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. arXiv preprint, 2021. URL arXiv:
2102.
Connor Shorten and Taghi M. Khoshgoftaar. survey image data augmentation for deep learn- ing. Springer Journal of Big Data, 2019.
Abhishek Sinha, Kumar Ayush, Jiaming Song, Burak Uzkent, Hongxia Jin, and Stefano Ermon. Negative data augmentation. In ICLR, 2021.
Mitchell Wortsman, Gabriel Ilharco, Mike Li, Jong W. Kim, Hannanch Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt. Robust fine-tuning of zero-shot models. arXiv preprint, 2021. URL arxiv: 2119::210.0193
Qizhe Xie, Zihang Dai, Eduard Hovy, Tinh-Thang Luong, and Quoc V. Lc. Unsupervised data augmentation for consistency training. In NeurIPS, 2020.
Yan Zhu, Taesung Park, Phillip Isola, and Alexei A. Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In ICCV, 2017.
10


Reviews:
[{"review_id": "WR-C-DW9mee", "reviewer": "Reviewer_w1qQ", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is more like an undergraduate course project report instead of an academic paper. In this paper, they mainly described how they do, but many important comparisons and studies are totally missing. For example, there is no literature study, no evidence to show the limitation of the literature, and no comparison with state-of-the-art. No motivation is described. Using data augmentation and multi-task learning are widely used techniques in many papers, so they can not be considered a novelty. ", "summary_of_the_paper": "In this work, the authors present a data augmentation method to address the data distribution shift problem in a multi-task learning framework. ", "summary_of_the_review": "Many important comparisons and studies are missing. No technical novelty. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "CJ5EXEtFmMS", "reviewer": "Reviewer_F3K8", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n\nThe expression of this paper is okay.\n\n\nCons:\n\nThe novelty and contribution of this work are marginal.", "summary_of_the_paper": "This paper presents Multi-Distribution Learning, exploring the effectiveness of Data Augmentation to prepare for Distribution Shift.", "summary_of_the_review": "The novelty and contribution of this work are marginal.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "5c0qWqCiS3e", "reviewer": "Reviewer_PK1M", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Overall thoughts:**\n\nPreparing for or adapting to distribution shift is a an interesting, yet in my opinion understudied, challenge for machine learning models. Real world applications would likely benefit the most from research in this domain, and a work which develops models robust to distribution shift, or able to adapt quickly to such a shift, is likely to be highly celebrated and well received by industry. This paper collapses distributional learning and data augmentations into one perspective and investigates how training models on many \"distributions of augmentations\" may lead to models which are robust to significant distribution shifts during inference. \n\n**Method**\n\nThe goal of this work is to improve a model's robustness to test-time distribution shift by training one model on many different dataset distributions (i.e. data augmentations). After training, this model should perform well on all dataset distribution on which it was trained. The authors judge this is a good indication of robustness to distribution shift and cite Martin Arjovsky to support their framing of the problem (i.e. the requirement to inject prior knowledge about the types of distribution shifts into the training dataset). Later, the authors analyze which distributions should be trained together in one model to increase model generalizability. \n\nIt is unsurprising to me that one model can learn multiple dataset augmentations (i.e. \"distribution shifts\") throughout training and perform well on each of them during test time. This also creates a disconnect from real-world applications of this method in that the distribution shift for a live model is typically unknown (both in terms of what the shift will be and when it will occur). Moreover, simply scaling a model's size may to enable it to learn a robust feature extractor which can handle many different distributions of data. \n\n**Weaknesses and Suggestions for improvement**\n\nOverall, it feels as if this paper was a bit rushed and incomplete. I do not recommend it for acceptance to ICLR, but encourage the authors to invest additional time into crafting the story and improving the empirical analysis for later resubmission. As it is presented, this work feels a bit all over the place, and while the ideas are interesting and promising, they are not developed into a strong or cohesive story. \n\n1. In my opinion, much of the value of distributional shift learning is invested in creating a model which can quickly adapt to changes in the test distribution or is robust to these changes. Table 4 begins to analyze this dimension, but significantly more analysis would be needed to support the authors' method in finding dataset augmentations which would train a model that is especially robust to distribution shift. \n2. Similar to point 1, the authors quote Martin Arjovsky to support their framing of the problem (i.e. the requirement to inject prior knowledge about the types of distribution shifts into the training dataset), but this requirement is often unavailable in many real world applications. It  severely limits the scope of this work to the scenarios where a future distribution shift can be anticipated in the present.\n3. I would recommend the authors consider refocusing the story of this paper to decide which dataset augmentations to apply to a model to make it robust to corrupted data and test-time distribution shifts. In other words, given ImageNet1k, which dataset augmentations would likely maximize the performance on Stylized ImageNet or ImageNet-Corrupted? \n\n**Strengths**\n\n1. I think the concept of choosing dataset augmentations from a predefined set to create a model which is robust to distributional shifts is very interesting and promising. Given the number of augmentations available to image data, and the pervasiveness of distributional shift for real-world applications, this work addresses an important challenge in computer vision. \n2. The introduction and related work contextualize the authors’ method with the most recent work in this domain. It is apparent the authors are familiar with the distributional shift landscape and knowledgeable in this domain. \n3. The proposal to develop a \"generalization score\" disparate from test-time metrics like test loss and test accuracy is also very relevant.\n\n**Nitpicks**\n\nTypos:\n 1. Achilles heel => Achilles' heel (should be plural)\n\nA lot of nouns are also capitalized that should not be capitalized. For instance\n1. for targeted Generalization => for targeted generalization\n2. Data Augmentation => data augmentation\n3. Deep Learning => deep learning\n4. Distributional Shift => distributional shift\n5. Average Distributional Accuracy => average distributional accuracy\n6. Distributional Generalization Metrics => distributional generalization metrics\n\nFinally, there are several latex/formatting issues:\n1. **?** utilized style randomization to understand intrinsic properties of Convolutional Neural Networks, finding a texture bias => missing reference in latex.\n2. Figures and tables have width which exceeds the margins.", "summary_of_the_paper": "The authors investigate distribution shift between the data seen during training and the data seen during testing by proposing a framework to analyze distribution shift as well as empirical analysis. In particular, they define a training regime for learning from many distributions within one model as well as a lookahead method to determine which distributions should be chosen for a 0-shot distribution shift benchmark.   ", "summary_of_the_review": "I think the authors investigate a very interesting/relevant challenge that could be applicable to a wide range of machine learning applications, but the work as it is presented now -- the story, writing, and empirical analysis -- needs significant revision to present a compelling case for the utility, correctness, and applicability of their method. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Though some concepts discussed in the submission are interesting, there are many major concerns: there is a lack of literature review, comparison experiments with the state-of-the-art methods are missing, the technical novelty of the proposed method is very limited. 
In the rebuttal, the authors agreed with reviewers' comments and did not provide responses to address reviewers' concerns.

Therefore, based on its current form, this submission does not meet the standard of publication at ICLR.
Decision: Reject