In-N-Out: Pre-Training and Self-Training using Auxiliary Information for Out-of-Distribution Robustness

Abstract: Consider a prediction setting with few in-distribution labeled examples and many unlabeled examples both in- and out-of-distribution (OOD). The goal is to learn a model which performs well both in-distribution and OOD. In these settings, auxiliary information is often cheaply available for every input. How should we best leverage this auxiliary information for the prediction task? Empirically across three image and time-series datasets, and theoretically in a multi-task linear regression setting, we show that (i) using auxiliary information as input features improves in-distribution error but can hurt OOD error; but (ii) using auxiliary information as outputs of auxiliary pre-training tasks improves OOD error. To get the best of both worlds, we introduce In-N-Out, which first trains a model with auxiliary inputs and uses it to pseudolabel all the in-distribution inputs, then pre-trains a model on OOD auxiliary outputs and fine-tunes this model with the pseudolabels (self-training). We show both theoretically and empirically that In-N-Out outperforms auxiliary inputs or outputs alone on both in-distribution and OOD error.

Reviews and author responses:

This paper introduces a new method for leveraging auxiliary information and unlabelled data to improve out-of-distribution model performance. Theoretically, in a linear model with latent variables, they demonstrate using auxiliary data as inputs helps in-distribution test-error, but can hurt out-of-distribution error, while using auxiliary data to pretrain a "good" representation always improve out-of-distribution error. The proposed method uses the auxiliary data to learn an initial model, which generates psuedolabels to fine-tune the pretrained model.

Pros:
- At a high level, this paper address a question of great interest to the ML community: out-of-distribution generalization.
- The theoretical model shows, albeit in a potentially simple linear setting, that pretraining a low-dimensional shared representation generically improves out-of-distribution accuracy.  I'm not intimately familiar with all of the papers in this area, but I think this emphasis (as opposed to transfer learning) is new. This may be of interest more broadly.
- Through experiments and a concrete example, the paper demonstrates the potential danger of using auxiliary features as input to models evaluated out of distribution.
- The paper reports experimental numbers on two real remote sensing datasets rather than solely evaluating on synthetic data.

Cons:
- Experimental results: My primary complaint with the paper is that, for the tasks considered, In-N-Out does not appear to work much better than the pretraining aux-outputs baseline. For out-of-distribution accuracy, across all of the datasets, the effect sizes are very small and the confidence intervals overlap. For in-distribution accuracy, there's only a large difference for the Landcover dataset. This makes me uncertain about the generality of the method and the potential size of the effects, though it's possible there's a more nuanced story that I'm missing. 
- Clarity: I found the description of the method confusing after several reads (section 2.1 and section 4), and the model sections are very notation heavy without necessarily providing much clarity. The graphical model, however, was very enlightening. 

Question:
- Difference between the theoretical and experimental In-N-Out models: How come the experimental procedure differs from the one that is analyzed, e.g. fine-tuning on h_out(x)? Is the performance in practice worse? More difficult to implement? I don't mind the gap, but some explanation and, if available, associated experiments explaining this would be enlightening.
- Generality of the theoretical model: How universal are the phenomenon captured by the linear model presented in Theorems 1 and 2? It's not obvious if the conclusions are "representative" or generalize beyond the ones explicitly analyzed.

==============
Update after rebuttal:

Thank you for clarifying that aux-outputs is itself a contribution and not simply a baseline for comparison. I also appreciated the additional experiment showing examples where In-N-Out can outperform aux-outputs. I'm raising my score from a 6 to a 7 accordingly.


We thank R1 for the detailed review and questions. R1 appreciates that the “question (is) of great interest to the ML community”, the theory which “may be of interest more broadly”, and that we show the “potential danger of using auxiliary features as input”. They also like that we use “real remote sensing datasets” and not just benchmark datasets like CelebA and CIFAR. We address their questions below.

> R1 notes that In-N-Out does not appear to work much better than the pretraining aux-outputs.

The standard approach to using auxiliary data is to treat them as inputs (aux-inputs), for example in Wang et al 2020, Yeh et al 2020. **One of our key contributions is the aux-outputs approach and showing that this works better out of distribution, and this is a key part of In-N-Out.** On all datasets, aux-outputs does better than the baseline and aux-inputs out of distribution. We view aux-outputs more as an ablation than a baseline.

The main goal of the self-training is to patch up the in-distribution performance of the aux-outputs model (Section 4 intro). In-distribution, on land cover aux-outputs gets 72.5% and In-N-Out gets 77.4%. The oracle model that uses 150k labeled examples gets 80.5% accuracy so we close most (60%) of the gap. On cropland In-N-Out is better than aux-outputs, and this is statistically significant. Aux-outputs gets 95.1%, In-N-Out gets 95.5%, and the oracle gets 95.6%, so **In-N-Out closes a large fraction (60%-80%) of the gap between aux-outputs and oracle on both landcover and cropland, in-distribution, and is statistically significant.**

To further address the reviewer concerns we ran an **additional new experiment**, which we have added to the paper. Our original results in table 1 are from one round of self-training. Following Xie et al 2020, we ran an additional round of self-training, and over 5 runs on landcover we get ID test accuracy of 77.1+-0.3% and OOD test accuracy of 62.6+-0.6% (4.3% improvement over baseline, 7.8% over aux-inputs, 1.6% over aux-outputs), and on cropland we get ID test accuracy of 95.5+-0.2% and OOD accuracy of 92.2+-0.4% which are **statistically significant improvements over aux-outputs, out of distribution, on both landcover and cropland.**

> R1 asks why we fine-tune h_out in practice, but hold h_out fixed in theory.

Practice: We fine-tune the entire model in our experiments because in practice fine tuning works better than linear probing (freezing the lower layers h_out and only tuning the top linear layer). See Kornblith et al 2019  (e.g. Figure 2, logistic regression on features vs fine-tuned). Theory: Unfortunately, the machine learning community does not have the theoretical tools to analyze fine tuning yet, and this is a major open problem.

> R1 asks how universal are the phenomena captured by the linear model.

This is a good question. Our experiments suggest these phenomena are more general, and likely hold for deep neural networks as well under reasonable assumptions, especially for Proposition 1, Example 1, and Theorem 1. In particular, the aux-inputs model often does worse than the baseline, but the aux-outputs model does better out of distribution on all 3 datasets.

S. Kornblith, J. Shlens, and Q. V. Le. Do Better ImageNet Models Transfer Better? CVPR 2019.


We thank all the reviewers for their thorough reviews. The reviewers thought the work tackles an **"especially important problem"**, provides a **”nice description of pitfalls and ways to leverage auxiliary data from OOD evaluation”**, the experiments include evaluations on **”two real-world remote sensing datasets”**, **”empirical validations consistently show...improvement in out-of-distribution (OOD) samples.”**, and the **theory is presented nicely and improves understanding**.

We addressed all the concerns in the individual comments and revised the draft accordingly. We list a summary of the changes here:
* We clarified a misunderstanding from R3 about the **conceptual difference between In-N-Out and the aux-outputs model**. In particular, **In-N-Out uses the aux-inputs model (which uses both x and auxiliary information z as input) to pseudo-label** the in-distribution unlabeled data, whereas **the aux-outputs model does not use z**.  Intuitively, the aux-inputs model (x, z → y) also uses z and therefore produces more accurate pseudolabels in-distribution, so the pseudolabels effectively increase the amount of data In-N-Out model (x → y) is trained on. z is misleading / not robust out-of-distribution which is why the final In-N-Out model does not take z directly as inputs to make predictions.
* R1 asked about the **difference between the proposed In-N-Out model and aux-outputs.** The standard approach to using auxiliary data is using it as additional inputs (aux-inputs). **One of our key contributions is to show that aux-outputs works better OOD**, and this is an integral part of In-N-Out. However, **our empirical results caution that aux-outputs could be suboptimal in-domain, and In-N-Out fixes this issue** while improving both in-distribution and OOD performance. During the rebuttal period, **we added an additional round of self-training on top of In-N-Out, which improved the results across the board (Table 1).** This gives us substantial additional improvements over aux-outputs on both cropland and landcover, in-distribution and OOD.
* R2 and R1 had suggestions for improving clarity of the theory/method, and we revised the appropriate sections in the paper, particularly clarifying the proof of Proposition 1. 


This paper investigates how to use auxiliary information to improve classification performance when few labeled examples are available. As the introduction makes clear, this is an especially important problem area for remote sensing applications, where labels are scarce for many inputs (e.g. satellite photos from countries/regions without much annotation).

The authors present three intuitively plausible baselines/ablations, two of which use auxiliary information, and explain the benefits and downsides of each. For instance, regarding the aux-inputs baseline: "the relationship between the aux-inputs and the target can shift significantly OOD, worsening the OOD error". These claims are later supported with theory using linear models. The theory is presented nicely, improves understanding, and is believable.

Their proposed method is very similar to the aux-out baseline/ablation. The only difference is that they fine-tune on pseudo-labeled in-distribution examples in their method. Seeing as this baseline could also be thought of as an ablation, and taking into account the improved performance of the full In-N-Out method, it is not too worrying. However, I am concerned about Remark 1 in Section 4.1, which says "We train an aux-inputs model gˆin from w,z to y on finite labeled data—since the noise σ 2 = E[ 2 ] is small this model is very accurate." If In-N-Out only improves performance when aux-in is a nearly perfect generator of pseudo-labels for in-distribution data, then doesn't this imply that aux-out would learn just as much from the GT labeled in-distribution examples? How are the pseudo-labels actually helping?

Pros:
- Underexplored, important problem area
- Good clarity of writing and paper structure, including theoretical sections
- Instructive choice of baselines/ablations
- Empirical improvements from the proposed method
- Not just vision datasets; they use a time series dataset as well
- Theory in the case of linear models to improve understanding
- The related work seems appropriate

Cons:
- The proposed method is very similar to one of the baselines/ablations, and I am not certain that there is a meaningful difference between them
- There is no comparison to previously published work that uses auxiliary information for classification. Perhaps there are no suitable baselines from prior work, but it is not clear from the paper that this is the case.

======================================================

Update after rebuttal:

The authors have addressed my concerns. Contrary to my initial understanding, the paper builds off of prior work in a methodical way, and the pseudolabeling stage of In-N-Out makes more sense now. I have raised my score from 6 to 7.

Theoretical and empirical study on how to combine two approaches using auxiliary information for out-of-distribution samples

**Quality:**

_Pros._ The authors propose the method to combine two representative methods (Aux-Inputs and Aux-Outputs) to exploit auxiliary information which is usually available in real-world scenarios. Beginning with theoretical analysis on Aux-Inputs and Aux-Outputs models, they show that the proposed method, In-N-Out is effective in minimizing the risk under a distributional change in a linear regression setting. Empirical validations consistently show the preference of the proposed method highlighting the improvement in out-of-distribution (OOD) samples.

**Clarity:**

_Pros._ The description of the problem setting, approach, and intuition is well-written and persuasive based on their observations in the example in Introduction. The intuitions for the theoretical findings are nicely addressed.
_Cons._ Some missing definitions (e.g., z in Alg. 1 as the output of `\hat{h}_{out}`)  and lacking rigorousness in the text make difficulty in following (e.g., the function `\hat{f}_{in}` takes two inputs, x and z, AND takes one input, `x^{id}_i`, in the 2nd and 4th lines in Alg. 1).  Figure 1 has ambiguities to understand which portion of the model is transferred and how to handle the change of the number of inputs (purely presumably, zero-filling as in the baseline in Sec 2.)

**Originality:**

_Cons._ Aside from their theoretical analysis, the combination of Aux-Inputs and Aux-Outputs is a simple model exploiting pseudo-labels (Xie et al., 2018), used for unsupervised domain adaptation tasks. In terms of novelty, the proposed method has a weak contribution. Isn't possible to borrow a sophisticated model from the works on unsupervised domain adaptation in your experiments?

**What expected in rebuttal:**

(1) Please explain the applicability of a sophisticated model from the literature in unsupervised domain adaptation in your experiments. (As an extension of the related work section.)

(2) In Sec 3., they described, "the aux-outputs model has better risk since w is lower dimensional than x. In particular, the in-domain risk only depends on the dimension but not on the conditioning of the data." However, Aux-Inputs or even a baseline (linear) model can have a lower-dimensional hidden representation in a low-rank linear model (e.g., two-layer perceptron, `f(x) := W_2 W_1 x`). So, the explanation is insufficient for the matter.

(3) The proof of Proposition 1 is incorrect. Where can we find Remark 10? There is a typo missing reference after "We use Theorem 1 in `?`." In Eqn. 23 and 24, if `1+cd/n = 2`, the left-most term in Eqn. 24 is `<= 2 \sigma^2`, but not guarantee the Eqn. 24 if `\sigma_u^2 < \sigma^2`. To satisfy Eqn. 24, `cd/n \sigma^2 < \sigma_u^2` should be true.

(4) In Lemma 8, the square is omitted inside of expectation in LHS of Eqn. 87. And, the dimension of R should be k x (k+m), not k x (k+T). I believe T is accidentally misplaced in this context.

**Minor comments:**

(5) Before Eqn. 26, a missing period right before "For the input model ..."

(6) y' and x' instead of y and x in Eqn. 32.

Thank you for the elaboration in the proof of Proposition 1. Now, the updated appendix seems to be clear to apprehend. 

This is good feedback, and we apologize for the lack of clarity. We have updated the proof (in the supplementary version) with a detailed explanation to your question. Please let us know if this addresses your question or if you have any other feedback.

Here we sketch a response to your question---we focus on the high probability bound, which follows from Hsu et al 2011 (https://arxiv.org/abs/1106.2363). Happy to elaborate further on any point.

The error in linear regression can be decomposed into the minimum possible error (E[ (y - {\theta^*}^T x)^2 ] = \sigma^2) plus the excess error, EE = E[ (\theta^Tx - {\theta^*}^T x)^2 ].

In Theorem 1, Hsu et al give a bound for the excess error EE. Their notation is different, they use \hat{\beta} for the learned linear model, \beta for the optimal linear model, and denote the excess error by || \hat{\beta} - \beta ||_{\Sigma}.

Given fixed delta, we want the theorem to hold with probability at least 1- delta. So we choose t = \log{\frac{3}{\delta}}. Then e.g. Remark 9, equation 11 bounds the excess error, EE since in our case “the linear model is correct”. Since t is fixed, by choosing n large enough, we get that EE <= \sigma^2 (cd / n). \*

So with probability 1 - delta, the total error is upper bounded by EE + \sigma^2 = \sigma^2 (1 + cd/n). 

\* This is precisely what Remark 10 is saying as well (they give simpler forms for the bias and variance terms, and Prop 2 upper bounds the excess error by the sum of these)


Thank you for your explanations to address the raised questions, especially for (2) & (3). 
I guess that Remark 10 is from Hsu et al. (2011), but I recommend clarifying in the Appendix for readers. 

Although I see what I misunderstood in Equation 24 (Proposition 1), I cannot follow Equation 23, specifically the result of "<= \sigma^2 (1 + cd/n)", even after consulting Theorem 1 in Hsu et al. (2011) and Theorem 11.3 in Györfi (2002), which does not directly indicate the Equation 23 is true. Could you elaborate on Equation 23 for clarity? 

> (4) In Lemma 8, the square is omitted inside of expectation in LHS of Eqn. 87. And, the dimension of R should be k x (k+m), not k x (k+T).

We have fixed these, thank you for bringing them up. We have also fixed the minor comments.

> (1) The reviewer asked about the applicability of a sophisticated model from the literature in unsupervised domain adaptation in our experiments

Over the last year, self-training / pseudolabeling has emerged as one of the top performing methods for semi-supervised learning even under domain shifts. For example, Xie et al 2020 self-train on the JFT dataset, and see large improvements on ImageNet-C. Self-training has also achieved state of the art in domain adaptation (e.g. Shu et al 2018, Zou et al 2019, Kumar et al 2020). We believe that these gains can be further boosted with alternative methods such as domain adversarial training (Ganin et al 2015). 

> Ambiguities in Figure 1 and Algorithm 1.

We apologize for the ambiguities in the Figure and Algorithm box. We formally defined the aux-inputs, aux-outputs methods in Section 2, and In-N-Out in Section 4.1, but wanted to have an illustrative figure and algorithm pseudocode. We have clarified that the Figure and Algorithm are informal, and the formal details are all in the text. We have fixed the typos identified.

References:

D. Hsu, S. M. Kakade, T. Zhang. Random design analysis of ridge regression. Arxiv 2011 (published version in COLT 2012, but references to theorems and examples are from Arxiv).

Percy Liang. Statistical learning theory. https://web.stanford.edu/class/cs229t/notes.pdf, 2016.

László Györfi. A Distribution-Free Theory of Nonparametric Regression. Springer, 2002.

N. Tripuraneni, C. Jin, M. I. Jordan. Provable Meta-Learning of Linear Representations. Arxiv 2020.

Y. Zou, Z. Yu, X. Liu, B. V. K. V. Kumar, J. Wang. Confidence Regularized Self-Training. ICCV 2019.


We thank R2 for the positive and detailed review, and suggestions for improvement. R2 appreciates that the "problem setting, approach, and intuition is... persuasive" and the "empirical validations... highlighting the improvement in OOD samples". We address their questions and suggestions below.

> (2) R2 asked what happens to the Aux-Inputs or even a baseline (linear) model if it has a lower-dimensional hidden representation in a low-rank linear model (f(x) := W_2 W_1 x), in relation to the text below Theorem 1: “Without distribution shift... the aux-outputs model has better risk since w is lower dimensional than x”

This is a good question, a low-rank linear model does not help the baseline or aux-inputs because we would still need to estimate W1 and W2 and there are at least d numbers to estimate, which requires O(d) samples. In contrast, the aux-outputs method leverages the low rank structure by learning the low dimensional representation W_1 from pre-training (using multiple auxiliary outputs), so only needs O(k) samples to learn W_2, where k << d is the dimension of the hidden representation.

More precisely, **low rank linear models only help when we have multiple outputs but not for standard linear regression with a single output.** Indeed, any linear regression parameter theta can be written as a low rank linear model, with a low dimensional representation, because we can write theta^T x = W_2 W_1 x, where W_2 = 1 and W_1 = theta^T. Here, the hidden representation is dimension 1, but W_1 is a 1-by-d matrix so the error of linear regression is still proportional to the input dimension d. This would correspond to Theorem 1 in Tripuraneni et al 2020 with t=1.

For more formal details, see e.g. the derivation in Section 2.7 of Liang 2016, leading up to Equation 81: for standard linear regression with a single output, the number of examples needed is proportional to d.

> (3) The reviewer said the proof of Proposition 1 is incorrect. In Eqn. 23 and 24, if 1+cd/n = 2, the left-most term in Eqn. 24 is <= 2 \sigma^2, but not guarantee the Eqn. 24 if \sigma_u^2 < \sigma^2. To satisfy Eqn. 24, cd/n \sigma^2 < \sigma_u^2 should be true.

**We believe this is a misunderstanding. We did not choose n such that 1 + cd/n = 2**. The statement of Proposition 1 says that there exists N, such that if n >= N, Equation 16 holds. In our proof we said that for large enough n, \sigma^2 (1 + cd/n) < \sigma^2 + \sigma_u^2. Indeed, it suffices to choose N > cd \sigma^2 / \sigma_u^2. We have clarified the choice of N in the revised version of the paper. We have also added the missing reference to Theorem 1 in Hsu et al 2011 for the high probability bound, an expectation bound can be shown from Theorem 11.3 in Gyorfi 2002.

> The reviewer asks what the original contributions of our paper are, besides pseudolabeling.

We believe there are thee original contributions that are not mentioned in the review:
- The standard approach to use auxiliary data is to treat them as inputs (aux-inputs), for example in Wang et al 2020, Yeh et al 2020. We show a **cautionary story: while auxiliary information helps in-distribution, it can substantially hurt for under-resourced countries (out of distribution).** We hope this encourages practitioners to be more careful when using auxiliary information.
- We instead **propose using the auxiliary data as prediction targets for pre-training and find that this helps out-of-distribution.** We show in Section 5.4 that OOD unlabeled data is important for this gain.
- The way in which to combine pre-training and self-training is not clear, for example which model should use the auxiliary information z? If the final model uses the auxiliary information z as input, then it will have worse performance OOD even after self-training. In-N-Out leverages the fact that the auxiliary inputs can hurt OOD, but helps in distribution, so z is used for pseudolabeling in-distribution examples but not in the final model. This is key to Theorem 2. In other words, **In-N-Out pseudolabeling is different from standard pseudolabeling**, because it requires careful considerations of when to use auxiliary information.

We thank R3 for the positive and detailed review. R3 appreciates the “important problem area”, “empirical improvements”, and that the “theory… improves understanding”. We address their questions about why pseudolabels help and comparisons to prior work below.

> Aux-outputs vs In-N-Out: R3 asks if there is any meaningful difference between In-N-Out and aux-outputs, and how the pseudolabels actually help. In particular, in the case where aux-in is a nearly perfect generator of pseudolabels, would the ground truth labeled examples be enough, or do the pseudolabels help?

This is a great question.  As a reminder, for inputs x, auxiliary z, label y:
- aux-outputs: x → y [weak in distribution because it doesn’t use z]
- aux-inputs: x, z → y [better in distribution, but bad OOD because z can be misleading OOD]
- In-N-Out: x → y [uses pseudolabels from aux-inputs (stronger model) in-distribution to transfer in-distribution accuracy to x → y]

There is in fact a **meaningful difference because we use the aux-inputs model to do the pseudolabeling** - the final model we fine-tune only uses x to make predictions, whereas the aux-inputs model uses both the inputs x and auxiliary z. Intuitively, the aux-inputs model (x, z → y) also uses z and therefore produces more accurate pseudolabels in-distribution, so the pseudolabels effectively increase the amount of data In-N-Out model (x → y) is trained on. z is misleading / not robust out-of-distribution which is why the In-N-Out model does not use z to make predictions. Theorem 2 proves that using these pseudolabels gets much better loss than only the ground truth labels.

We now give more detailed intuition for why the pseudolabels actually help. For simplicity, consider the extreme example where aux-in is a perfect generator of pseudolabels in-distribution. This means that in-distribution, we can predict y perfectly if we have x and the auxiliary information z. However, if we use x alone, we cannot predict y perfectly, and there is some noise \sigma_u^2.

In-N-Out aims to learn a function \hat{f} that predicts y from x alone (and not z). If we only use n ground truth labeled examples to learn the map from x → y, we will incur excess error proportional to 1/n - intuitively, the more labeled examples we have, the better we perform. Instead, we label lots of unlabeled examples using aux-inputs, which takes in both x and auxiliary information z, and outputs the exact label y. This effectively increases the amount of data we have when training our x -> y classifier (without auxiliary inputs). This is why the excess error for In-N-Out is lower. Our actual argument doesn’t require aux-in to be a perfect classifier, but only requires that z gives additional information that is useful for lowering in-distribution error - the full proof is in Appendix A.

> Comparisons to prior work: R3 asks if we compare with “previously published work that uses auxiliary information for classification”

**Prior work (e.g. Wang et al 2020, Yeh et al 2020) uses auxiliary information as inputs to the model, which we compare to in our paper.** This is the aux-inputs model. Prior work shows that aux-inputs help, but show this on in-distribution accuracy. We show that aux-inputs can often do worse out-of-distribution, but an alternative approach, aux-outputs, does better out-of-distribution, and that In-N-Out can further improve on this. So the baseline from prior work is aux-inputs.


Meta-review: The paper addresses the problem of improving generalization when few annotated data is available by leveraging available auxiliary information. The authors consider the respective merits of two alternatives: using auxiliary information as complementary inputs or as additional outputs in a multi-task or transfer setting.  For linear regression, they show theoretically that the former can help improve in distribution error but may hurt OOD error, while the latter may help improve OOD error. They propose a framework for combining the two alternatives and show empirically that it does so on three different datasets. 


All the reviewers agree on the novelty, interest and impact of the method. The rebuttal clarified the reviewers’ questions. I propose an accept.


Decision: poster-presentations