PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Sang Michael Xie, Ananya Kumar, Robbie Jones, Fereshte Khani, Tengyu Ma, Percy Liang

arXiv:2012.04550v3cs.LGstat.ML

TL;DR

The paper asks how auxiliary information should be used to achieve strong performance under scarce labels and distribution shifts. In-N-Out uses an auxiliary-input model to pseudolabel in-distribution data, then pre-trains on auxiliary outputs and fine-tunes with labeled and pseudolabeled data. Auxiliary inputs can improve in-distribution and OOD performance on curated datasets but hurt OOD performance in real-world datasets, whereas auxiliary outputs improve OOD performance.

  • Problem

    The paper asks how auxiliary information should be used to achieve strong performance under scarce labels and distribution shifts.

    Source: §1 · p. 1

  • Method

    In-N-Out uses an auxiliary-input model to pseudolabel in-distribution data, then pre-trains on auxiliary outputs and fine-tunes with labeled and pseudolabeled data.

    Source: §1 · p. 2

  • Results

    Auxiliary inputs can improve in-distribution and OOD performance on curated datasets but hurt OOD performance in real-world datasets, whereas auxiliary outputs improve OOD performance.

    Source: §8 · p. 9

  • Takeaways & Limitations

    Auxiliary information can support robustness when used as pre-training supervision while still supplying in-distribution signals through pseudolabels.

    Source: §1 · p. 2, §7 · p. 9

    The paper leaves automatic selection of the input-versus-auxiliary-information division under distribution shifts as an open question.

    Source: §7 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper asks how cheaply available auxiliary information should be used when labels are scarce and test distributions may shift. It shows that auxiliary inputs favor in-distribution accuracy, auxiliary outputs favor OOD robustness, and In-N-Out combines these benefits.

Motivation

  • Auxiliary information is valuable but underused in semi-supervised settings with few labeled examples and many in- and out-of-distribution unlabeled examples.The motivating applications include remote sensing, where labels are scarce and geographic or political shifts require OOD extrapolation.Source: §1 · p. 1

Auxiliary inputs

  • Auxiliary inputs often improve in-distribution accuracy but can worsen OOD accuracy by introducing spurious label correlations.The paper reports this pattern on CelebA and real-world satellite datasets.Source: §1 · p. 1

Auxiliary outputs

  • Auxiliary outputs improve OOD accuracy through pre-training but can have worse in-distribution accuracy than auxiliary inputs, revealing a tradeoff.The paper connects this approach to auxiliary supervision and robustness-oriented pre-training.Source: §1 · pp. 1–2

In-N-Out

  • In-N-Out uses auxiliary-input predictions to pseudolabel in-distribution data, pre-trains on auxiliary outputs, and fine-tunes with the expanded labeled set.This sequence is designed to transfer in-distribution performance while retaining the OOD benefits of auxiliary-output pre-training.Source: §1 · p. 2, Figure 1 · p. 2

Empirical evidence

  • Across CelebA and two remote sensing tasks, In-N-Out improves OOD accuracy and matches or exceeds auxiliary inputs and outputs alone on in-distribution accuracy.On remote sensing tasks it improves 1–2% in-distribution and 2–3% OOD over not using auxiliary information; OOD unlabeled data is crucial for OOD gains.Source: §1 · p. 2

2 SETUP

The setup distinguishes models that ignore auxiliary information, use it as input features, or use it as auxiliary prediction outputs during pre-training. The output-based model learns a shared representation before target-task transfer.

2 SETUP

  • The prediction problem uses input x, target y, and cheaply available auxiliary information z under in-distribution and OOD distributions.Auxiliary information may come from additional data sources or be derived from the original input.Source: §2 · p. 2

2.1 MODELS

  • The baseline minimizes labeled-data risk while ignoring z, whereas aux-inputs minimizes labeled-data risk using z directly as features.These models differ only in whether auxiliary information is excluded or supplied to the target predictor.Source: §2.1 · p. 3
  • The aux-outputs model predicts z from x during pre-training, then transfers the learned representation to predict y from labeled data.It seeks a lower-dimensional representation shared by the auxiliary and target tasks.Source: §2.1 · p. 3
  • The aux-outputs model ignores auxiliary information at prediction time, using it only to learn the feature representation.Its final target predictor maps the pre-trained representation of x to y.Source: §2.1 · p. 3

3 THEORETICAL ANALYSIS OF AUX-INPUTS AND AUX-OUTPUTS MODELS

The theoretical analysis uses a linear multi-task setting with latent shared features and distribution shifts to explain why auxiliary inputs help in-distribution but may fail OOD. It proves that auxiliary-output pre-training improves robustness under broad covariate shifts.

3 THEORETICAL ANALYSIS OF AUX-INPUTS AND AUX-OUTPUTS MODELS

  • The theory models x and u as independent latent sources, with shared feature w=B⋆x influencing both auxiliary information z and target y.The linear setting assumes full-rank task matrices, sufficient auxiliary dimension, unlimited ID and OOD unlabeled data, and squared loss.Source: §3 · pp. 3–4

3.1 AUXILIARY INPUTS HELP IN-DISTRIBUTION,

  • Auxiliary inputs improve in-distribution prediction because z helps recover unobserved variation u that the baseline cannot capture from x alone.The aux-inputs model combines x and z to estimate target variation associated with both latent sources.Source: §3.1 · p. 4
  • Auxiliary inputs can be worse than the baseline OOD because shifts in u can increase z variance and magnify estimation error in the auxiliary-feature coefficient.The theory gives a setting where this OOD disadvantage persists for every training-set size.Source: §3.1 · p. 4
  • Auxiliary-output pre-training recovers a lower-dimensional bottleneck representation from x and improves risk under arbitrary covariate shift relative to the baseline.The representation can reduce error directions; in the stated linear setting, lower dimension also reduces expected excess risk without distribution shift.Source: §3.2 · pp. 4–5
  • The theoretical In-N-Out procedure pre-trains on auxiliary outputs using unlabeled data, then trains the target map with labeled and pseudolabeled examples.Its stages combine an aux-inputs predictor for pseudolabel generation with an aux-outputs representation.Source: §3.2 · p. 5

4 IN-N-OUT: COMBINING AUXILIARY INPUTS AND OUTPUTS

In-N-Out combines aux-inputs pseudolabeling with aux-outputs pre-training so predictions avoid relying on auxiliary information that may become misleading under covariate shift. The theory shows this combination can improve over aux-outputs under stated noise and sample-size conditions.

4 IN-N-OUT: COMBINING AUXILIARY INPUTS AND OUTPUTS

  • In-N-Out uses an aux-inputs model to pseudolabel in-distribution unlabeled data, then fine-tunes an aux-outputs model pre-trained on auxiliary prediction.The final model predicts from x without using z at test time, reducing reliance on potentially spurious auxiliary features.Source: §4 · p. 5

4.1 IN-N-OUT IMPROVES OVER PRE-TRAINING UNDER ARBITRARY COVARIATE SHIFT

  • The theoretical analysis studies a modified In-N-Out setup in which all models are linear, with n≥m+d training points and noise restricted to values satisfying a−bσ2 > 0.Source: §4 · p. 5, §4.1 · p. 6
  • When σ2=0, the aux-inputs model can produce perfect pseudolabels, giving In-N-Out a larger correctly labeled training set.This noiseless case provides intuition for why self-training can improve the pre-trained representation.Source: §4.1 · p. 5
  • With nonzero noise, pseudolabel errors propagate, but the analysis shows In-N-Out can reduce dependence on auxiliary noise relative to aux-outputs.The proof compares singular values of the In-N-Out and aux-outputs error matrices and uses a random-matrix lower bound.Source: §4.1 · pp. 5–6

4 IN-N-OUT: COMBINING AUXILIARY INPUTS AND OUTPUTS

  • The experiments use auxiliary information derived from inputs or obtained from other sources across datasets summarized in the experimental setup.The supplied dataset summary identifies CelebA, Cropland, and Landcover as examples of these two auxiliary-information sources.Source: Figure 3 · p. 6

4.1 IN-N-OUT IMPROVES OVER PRE-TRAINING UNDER ARBITRARY COVARIATE SHIFT

  • As σ →0, In-N-Out’s excess-risk ratio relative to aux-outputs goes to 0, making the In-N-Out estimator much better in this limit.The comparison is between In-N-Out risk and aux-outputs risk on the shifted test example.Source: §4.1 · p. 6

5 EXPERIMENTS

Across three datasets, In-N-Out generally improves OOD accuracy while matching or improving ID accuracy, combining auxiliary-input pseudolabeling with auxiliary-output pre-training.

5.1 EXPERIMENTAL SETUP

  • Experiments use few labeled examples, many unlabeled ID and OOD examples, and dataset-specific auxiliary information derived from inputs or external sources.CelebA uses image attributes, Cropland uses vegetation bands and location, and Landcover uses climate data.Source: §5 · p. 6, §5.1 · pp. 6–7

5.2 MAIN RESULTS

  • In-N-Out generally improves both ID and OOD accuracy over auxiliary-input and auxiliary-output models across datasets.Table 1 reports results averaged over five trials with 90% intervals.Source: §5.2 · p. 7, Table 1 · p. 7
  • CelebA auxiliary information improves both ID accuracy by 2–4% and OOD accuracy by 5%, with In-N-Out achieving the best OOD and comparable ID performance.The dataset uses image-derived attributes, and hat-wearing images define the OOD split.Source: §5.1 · p. 6, §5.2 · p. 7
  • On remote-sensing tasks, auxiliary inputs can improve ID accuracy while hurting OOD, whereas auxiliary outputs improve OOD but can reduce ID accuracy.In cropland, auxiliary inputs improve ID nearly 1% but reduce OOD by 6%; in land cover, climate inputs reduce OOD by over 4%, while auxiliary outputs reduce ID by 3%.Source: §5.2 · p. 7
  • In-N-Out raises auxiliary-output ID accuracy from 72.5% to 77.4% in land cover and from 95.1% to 95.5% in cropland.It closes 62% and 80% of the respective gaps to oracle models using substantially more labeled ID data.Source: §5.2 · p. 8

5.3 CHOICE OF AUXILIARY INPUTS MATTERS

  • The choice of auxiliary inputs materially changes the ID–OOD tradeoff, with location coordinates helping ID interpolation but hindering extrapolation to new locations.Using vegetation data without locations improves OOD accuracy by 1.5% while leaving ID accuracy similar to the RGB baseline.Source: §5.3 · p. 8

5.4 OOD UNLABELED DATA IS IMPORTANT FOR PRE-TRAINING

  • OOD unlabeled examples are important for pre-training: on Landcover, ID-only unlabeled data does not improve OOD accuracy, whereas OOD-only data performs well on both metrics.The comparison standardizes unlabeled sample size across pre-training variants.Source: §5.4 · p. 8, Table 2 · p. 8

6 RELATED WORK

Related work connects the paper to multi-task learning, pre-training, self-training, domain adaptation, and robustness, while emphasizing its broader shift and task assumptions.

Multi-task learning and weak supervision

  • Unlike prior multi-task and weak-supervision analyses, this work allows arbitrary auxiliary–target alignment beyond a shared latent variable and considers distribution shifts.Related multi-task approaches often require closely aligned auxiliary and target parameters or correlated weak labels.Source: §6 · pp. 8–9

Transfer learning, pre-training, and self-supervision

  • The paper extends pre-training theory by addressing OOD robustness and combining auxiliary inputs with auxiliary-output pre-training.Earlier linear-regression analyses focus on in-distribution generalization bounds.Source: §6 · p. 9

Self-training for robustness

  • Compared with robust self-training, In-N-Out studies multi-task linear regression, pre-training, arbitrary covariate shifts, and their combined robustness benefits.The paper also relates to work on repeated self-training, spurious-feature removal, and combined self-training and pre-training.Source: §6 · p. 9

Domain adaptation

  • In-N-Out can provide a better starting point for domain-adaptation methods that benefit from target-domain models or source-model inductive bias.The connection is presented as complementary to self-training, entropy minimization, and domain-adversarial adaptation.Source: §6 · p. 9

7 DISCUSSION

The discussion frames In-N-Out as using auxiliary information flexibly: potentially spurious features support ID prediction while auxiliary outputs and unlabeled data support robustness.

Using spurious features for robustness

  • In-N-Out uses potentially spurious auxiliary features as useful ID signals rather than removing them, while improving OOD robustness.This contrasts with approaches that explicitly remove spurious features.Source: §7 · p. 9

General robustness with unlabeled data

  • The method instantiates a broader robustness paradigm that learns representations from unlabeled data across the input space before labeled training.Its distinctive extension is using some collected features as both inputs and outputs for robustness to arbitrary shifts.Source: §7 · p. 9

Leveraging metadata and unused features in applications

  • Metadata such as locations or timestamps can join additional auxiliary data sources for use in In-N-Out, including information otherwise discarded.The discussion presents this as applicable to improving both ID and OOD accuracy.Source: §7 · p. 9

Division between input features and auxiliary information

  • The method applies to any input division, but automatically choosing the split between ordinary inputs and auxiliary information under distribution shift remains open.This is the discussion’s explicit scope boundary.Source: §7 · p. 9

8 CONCLUSION

The paper concludes that auxiliary inputs can help in-distribution performance yet hurt OOD performance, whereas auxiliary outputs improve OOD performance; In-N-Out combines both for improvements on both axes.

8 CONCLUSION

  • Auxiliary inputs can improve in-distribution and OOD performance on curated datasets but hurt OOD performance on real-world datasets.Source: §8 · p. 9, A PROOF FOR SECTIONS 3 AND 4 · p. 14
  • Pretraining with auxiliary information as outputs improves OOD performance, while auxiliary inputs can worsen OOD error under distribution shift.Source: §8 · p. 9
  • In-N-Out combines auxiliary inputs and outputs to improve both in-distribution and OOD performance.Source: §8 · p. 9

A PROOF FOR SECTIONS 3 AND 4

A.1 MODELS AND EVALUATION

The appendix defines baseline, auxiliary-input, auxiliary-output, and In-N-Out estimators, then states theoretical conditions under which auxiliary inputs improve in-distribution risk.

A.1 MODELS AND EVALUATION

  • The baseline uses ordinary least squares with x alone, while the auxiliary-input estimator uses both x and z as predictors.Source: §A.1 · p. 14
  • Auxiliary outputs pretrain a low-dimensional mapping from x to z, then regress the target y on the learned embedding w.Source: §A.1 · p. 14
  • In-N-Out learns an output representation, pseudolabels in-domain unlabeled examples with an input model, and self-trains without z on those pseudolabels.Source: §A.1 · p. 15
  • Under bounded inputs, sub-Gaussian noise, and T=m, the auxiliary-input model eventually improves in-distribution risk over the baseline with probability at least 1−δ.Source: §A.2 · p. 15
  • The proof bounds baseline risk by an irreducible-noise term and uses well-specified linear regression to upper-bound auxiliary-input excess error.Source: §A.2 · pp. 15–16
  • The excess-error bound relies on ordinary least squares concentration under the stated covariance, boundedness, and noise regularity conditions.Source: §A.2 · pp. 16–17

A.3 AUXILIARY INPUTS CAN HURT OUT-OF-DISTRIBUTION

A constructed linear example shows that auxiliary inputs can worsen OOD risk arbitrarily under a shifted test distribution, despite their in-distribution benefit.

A.3 AUXILIARY INPUTS CAN HURT OUT-OF-DISTRIBUTION

  • There exists a problem setting and shifted test distribution for which the auxiliary-input model has greater OOD risk than the baseline.Source: §A.3 · p. 17
  • The construction uses x as the target-relevant feature, while auxiliary variables include one correlated component and one irrelevant component whose test variance can change.Source: §A.3 · p. 17
  • The example’s linear design uses z1=u1+w and z2=u2, with y=x+u1+ϵ and bounded training distributions.Source: §A.3 · p. 17
  • By increasing the shifted distribution’s variance for the irrelevant latent component, the auxiliary-input model’s OOD lower bound can become arbitrarily large without changing baseline risk.Source: §A.3 · p. 17

A.4 PRE-TRAINING IMPROVES RISK UNDER ARBITRARY COVARIATE SHIFT

The theory shows that auxiliary-output pretraining recovers the target representation up to its row space and yields lower expected risk than the baseline under arbitrary covariate shift.

A.4 PRE-TRAINING IMPROVES RISK UNDER ARBITRARY COVARIATE SHIFT

  • For every problem setting, noise distribution, and test distribution satisfying the theorem’s assumptions, auxiliary outputs achieve lower risk than the baseline.Source: §A.4 · pp. 18–19
  • Pretraining learns a low-rank map from x to z, recovers the unobserved feature representation up to identifiability, and then regresses y on that representation.Source: §A.4 · p. 18
  • Population-risk minimizers recover the row space of the true input-to-representation map when the auxiliary regression is rank-k and E[xx⊤] is invertible.Source: §A.4 · p. 18
  • The proof compares instance-wise expected risks over training labels and establishes the auxiliary-output advantage through a positive-semidefinite matrix difference.Source: §A.4 · pp. 18–19

A.5 IN-N-OUT IMPROVES RISK UNDER ARBITRARY COVARIATE SHIFT

The theory formalizes In-N-Out as pseudolabeling with an auxiliary-input model followed by self-training from an auxiliary-output representation, and proves improved risk under arbitrary covariate shift. The proof establishes the required estimator equivalence and compares risk bounds for In-N-Out and auxiliary pre-training.

A.5 IN-N-OUT IMPROVES RISK UNDER ARBITRARY COVARIATE SHIFT

  • Under absolute-continuity assumptions, random feature matrices have minimum singular values bounded below with high probability.This lemma supplies the conditioning needed for the subsequent risk bounds.Source: §A.5 · p. 20
  • In-N-Out first fits y from w and z, pseudolabels unlabeled data, and self-trains a model from w to y.The linear analysis represents this procedure through an estimator based on the learned representation W and latent matrix U.Source: §A.5 · pp. 20–21
  • The In-N-Out estimator is equivalent to an alternative optimization whose minimum equals the original joint optimization minimum.The argument proves both inequalities, yielding p*=q* and the corresponding estimator form.Source: §A.5 · pp. 21–22
  • The proof lower-bounds auxiliary-output risk and upper-bounds In-N-Out risk with high probability over training data and shifted test inputs.Theorem 2 combines the two bounds to establish the claimed risk comparison under the stated linear setting and assumptions.Source: §A.5 · pp. 22–24

B.1 CELEBA

The experiments evaluate auxiliary-input, auxiliary-output, and In-N-Out strategies across CelebA and two remote-sensing tasks. The setups vary the auxiliary modalities and distribution splits, while the reported CelebA analysis shows a significant correlation between in-distribution and OOD accuracy for auxiliary-input models.

B.1 CELEBA

  • CelebA uses seven auxiliary binary attributes, 2,000 labeled examples, and separate in-distribution and OOD unlabeled and test sets.Models use a ResNet-18 backbone and report averages over five trials with balanced gender ratios across test distributions.Source: §B.1 · p. 24
  • CelebA auxiliary-input models concatenate auxiliary attributes with baseline logits, whereas auxiliary-output models pre-train on seven-way attribute prediction before gender fine-tuning.The experiments also vary the number and identity of auxiliary inputs across model configurations.Source: §B.1 · pp. 24–25
  • CelebA auxiliary-input experiments show a significant correlation between in-distribution and OOD accuracy across 100 sampled auxiliary-input configurations.Each plotted result averages five trials, supporting comparison of the two accuracy measures across auxiliary-information choices.Source: §B.1 · p. 25

B.2 CROPLAND

  • Cropland simulates geographic OOD generalization by assigning Illinois, Iowa, and Missouri to the in-distribution split and Indiana and Kentucky to OOD.Auxiliary information combines latitude-longitude bands with vegetation bands, and auxiliary-output models reconstruct these bands with image-to-image prediction.Source: §B.2 · pp. 25–26

B.3 LANDCOVER

  • Landcover uses MODIS time series with NDVI and ERA5 climate variables, pairing each point with its nearest climate-data neighbor.Auxiliary inputs concatenate resampled ERA5 measurements with MODIS and NDVI, while auxiliary outputs predict six time-averaged climate variables.Source: §B.3 · p. 26
  • For Landcover, In-N-Out initializes from the auxiliary-output model and obtains pseudolabeled in-distribution unlabeled data from the auxiliary-input model.The procedure uses a smaller learning rate during pre-training, with additional dropout regularization for repeated self-training.Source: §B.3 · p. 27

Related papers

Loading 2012.04550v3…