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

1 INTRODUCTION . When models are tested on distributions that are different from the training distribution , they typically suffer large drops in performance ( Blitzer and Pereira , 2007 ; Szegedy et al. , 2014 ; Jia and Liang , 2017 ; AlBadawy et al. , 2018 ; Hendrycks et al. , 2019a ) . For example , in remote sensing , central tasks include predicting poverty , crop type , and land cover from satellite imagery for downstream humanitarian , policy , and environmental applications ( Xie et al. , 2016 ; Jean et al. , 2016 ; Wang et al. , 2020 ; Rußwurm et al. , 2020 ) . In some developing African countries , labels are scarce due to the lack of economic resources to deploy human workers to conduct expensive surveys ( Jean et al. , 2016 ) . To make accurate predictions in these countries , we must extrapolate to out-of-distribution ( OOD ) examples across different geographic terrains and political borders . We consider a semi-supervised setting with few in-distribution labeled examples and many unlabeled examples from both in- and out-of-distribution ( e.g. , global satellite imagery ) . While labels are scarce , auxiliary information is often cheaply available for every input and may provide some signal for the missing labels . Auxiliary information can come from additional data sources ( e.g. , climate data from other satellites ) or derived from the original input ( e.g. , background or non-visible spectrum image channels ) . This auxiliary information is often discarded or not leveraged , and how to best use them is unclear . One way is to use them directly as input features ( aux-inputs ) ; another is to treat them as prediction outputs for an auxiliary task ( aux-outputs ) in pre-training . Which approach leads to better in-distribution or OOD performance ? Aux-inputs provide more features to potentially improve in-distribution performance , and one may hope that this also improves OOD performance . Indeed , previous results on standard datasets show that improvements in in-distribution accuracy correlate with improvements in OOD accuracy ( Recht et al. , 2019 ; Taori et al. , 2020 ; Xie et al. , 2020 ; Santurkar et al. , 2020 ) . However , in this paper we find that aux-inputs can introduce more spurious correlations with the labels : as a result , while aux-inputs often improve in-distribution accuracy , they can worsen OOD accuracy . We give examples of this trend on CelebA ( Liu et al. , 2015 ) and real-world satellite datasets in Sections 5.2 and 5.3 . Conversely , aux-output methods such as pre-training may improve OOD performance through auxiliary supervision ( Caruana , 1997 ; Weiss et al. , 2016 ; Hendrycks et al. , 2019a ) . Hendrycks et al . ∗Equal contribution . 𝑥 𝑧 𝑤 𝑦 𝑢 𝐵∗ 𝐴∗ 𝐶∗ 𝜃 '' 𝜃 # Figure 2 : Graphical model for our theoretical setting : prediction task with input x , target y , and auxiliary information z , which is related to y through the latent variable w and latent noise u . ( 2019a ) show that pre-training on ImageNet can improve adversarial robustness , and Hendrycks et al . ( 2019b ) show that auxiliary self-supervision tasks can improve robustness to synthetic corruptions . In this paper , we find that while aux-outputs improve OOD accuracy , the in-distribution accuracy is worse than with aux-inputs . Thus , we elucidate a tradeoff between in- and out-of-distribution accuracy that occurs when using auxiliary information as inputs or outputs . To theoretically study how to best use auxiliary information , we extend the multi-task linear regression setting ( Du et al. , 2020 ; Tripuraneni et al. , 2020 ) to allow for distribution shifts . We show that auxiliary information helps in-distribution error by providing useful features for predicting the target , but the relationship between the aux-inputs and the target can shift significantly OOD , worsening the OOD error . In contrast , the aux-outputs model first pre-trains on unlabeled data to learn a lower-dimensional representation and then solves the target task in the lower-dimensional space . We prove that the aux-outputs model improves robustness to arbitrary covariate shift compared to not using auxiliary information . Can we do better than using auxiliary information as inputs or outputs alone ? We answer affirmatively by proposing the In-N-Out algorithm to combine the benefits of auxiliary inputs and outputs ( Figure 1 ) . In-N-Out first uses an aux-inputs model , which has good in-distribution accuracy , to pseudolabel in-distribution unlabeled data . It then pre-trains a model using aux-outputs and finally fine-tunes this model on the larger training set consisting of labeled and pseudolabeled data . We prove that In-N-Out , which combines self-training and pre-training , further improves both in-distribution and OOD error over the aux-outputs model . We show empirical results on CelebA and two remote sensing tasks ( land cover and cropland prediction ) that parallel the theory . On all datasets , In-N-Out improves OOD accuracy and has competitive or better in-distribution accuracy over aux-inputs or aux-outputs alone and improves 1–2 % in-distribution , 2–3 % OOD over not using auxiliary information on remote sensing tasks . Ablations of In-N-Out show that In-N-Out achieves similar improvements over pre-training or self-training alone ( up to 5 % in-distribution , 1–2 % OOD on remote sensing tasks ) . We also find that using OOD ( rather than in-distribution ) unlabeled examples for pre-training is crucial for OOD improvements . 2 SETUP . Let x∈Rd be the input ( e.g. , a satellite image ) , y ∈R be the target ( e.g. , crop type ) , and z ∈RT be the cheaply obtained auxiliary information either from additional sources ( e.g. , climate information ) or derived from the original data ( e.g. , background ) . Training data . Let Pid and Pood denote the underlying distribution of ( x , y , z ) triples in-distribution and out-of-distribution , respectively . The training data consists of ( i ) in-distribution labeled data { ( xi , yi , zi ) } ni=1 ∼ Pid , ( ii ) in-distribution unlabeled data { ( xidi , zidi ) } mid i=1 ∼ Pid , and ( iii ) out-of-distribution unlabeled data { ( xoodi , zoodi ) } mood i=1 ∼Pood . Goal and risk metrics . Our goal is to learn a model from input and auxiliary information to the target , f : Rd×RT →R . For a loss function ` , the in-distribution population risk of the model f is Rid ( f ) =Ex , y , z∼Pid [ ` ( f ( x , z ) , y ) ] , and its OOD population risk isRood ( f ) =Ex , y , z∼Pood [ ` ( f ( x , z ) , y ) ] . 2.1 MODELS . We consider three common ways to use the auxiliary information ( z ) to learn a model . Baseline . The baseline minimizes the empirical risk on labeled data while ignoring the auxiliary information ( accomplished by setting z to 0 ) : f̂bs =argmin f 1 n n∑ i=1 ` ( f ( xi,0 ) , yi ) . ( 1 ) Aux-inputs . The aux-inputs model minimizes the empirical risk on labeled data while using the auxiliary information as features : f̂in =argmin f 1 n n∑ i=1 ` ( f ( xi , zi ) , yi ) . ( 2 ) Aux-outputs . The aux-outputs model leverages the auxiliary information z by using it as the prediction target of an auxiliary task , in hopes that there is a low-dimensional feature representation that is common to predicting both z and y . Training the aux-outputs model consists of two steps : In the pre-training step , we use all the unlabeled data to learn a shared feature representation . Let h : Rd→Rk denote a feature map and gz-out : Rk→RT denote a mapping from feature representation to the auxiliary outputs . Let ` aux denote the loss function for the auxiliary information . We define the empirical risk of h and gz-out as : R̂pre ( h , gz-out ) = 1 mid+mood ( mid∑ i=1 ` aux ( gz-out ( h ( x id i ) ) , z id i ) + mood∑ i=1 ` aux ( gz-out ( h ( x ood i ) ) , z ood i ) ) . ( 3 ) The estimate of the feature map is ĥout =argminhmingz-outR̂pre ( h , gz-out ) . In the transfer step , the model uses the pre-trained feature map ĥout and the labeled data to learn the mapping gy-out : Rk→R from feature representation to target y . We define the transfer empirical risk as : R̂trans ( ĥout , gy-out ) = 1 n n∑ i=1 ` ( gy-out ( ĥout ( xi ) ) , yi ) ( 4 ) The estimate of the target mapping is ĝy-out = argmingy-out R̂trans ( ĥout , gy-out ) . The final aux-outputs model is f̂out ( x , z ) = ĝy-out ( ĥout ( x ) ) . ( 5 ) Like the baseline model , the aux-outputs model ignores the auxiliary information for prediction . 3 THEORETICAL ANALYSIS OF AUX-INPUTS AND AUX-OUTPUTS MODELS . We now analyze the baseline , aux-inputs , and aux-outputs models introduced in Section 2 . Our setup extends a linear regression setting commonly used for analyzing multi-task problems ( Du et al. , 2020 ; Tripuraneni et al. , 2020 ) . Setup . See Figure 2 for the graphical model . Letw=B ? x∈Rk be a low-dimensional latent feature ( k≤d ) shared between auxiliary information z and the target y . Let u∈Rm denote unobserved latent variables not captured in x . We assume z and y are linear functions of u andw : y=θ > ww+θ > u u+ , ( 6 ) z=A ? w+C ? u , ( 7 ) where ∼ P denotes noise with mean 0 and variance σ2 . As in Du et al . ( 2020 ) , we assume the dimension of the auxiliary information T is greater than the feature dimension k , that is T ≥k , and thatA ? , B ? andC ? have full rank ( rank k ) . We also assume T ≥m , wherem is the dimension of u . Data . Let Px and Pu denote the distribution of x and u in-distribution ( ID ) , and let P ′x , P ′u denote the distribution x and uOOD . We assume x and u are independent , have distributions with bounded density everywhere , and have invertible covariance matrices . We assume the mean of u is zero in- and out-of-distribution1 . We assume we have n≥m+d in-distribution labeled training examples and unlimited access to unlabeled data both ID and OOD , a common assumption in unsupervised domain adaptation theory ( Sugiyama et al. , 2007 ; Kumar et al. , 2020 ; Raghunathan et al. , 2020 ) . Loss metrics . We use the squared loss for the target and auxiliary losses : ` ( ŷ , y ) = ( y− ŷ ) 2 and ` aux ( z , z ′ ) =‖z−z′‖22 . Models . We assume all model families ( f , h , gz-out , gy-out ) in Section 2 are linear . Let S= ( A ? , B ? , C ? , θw , θu , Px , Pu ) denote a problem setting which satisfies all the above assumptions . 3.1 AUXILIARY INPUTS HELP IN-DISTRIBUTION , BUT CAN HURT OOD . We first show that the aux-inputs model ( 2 ) performs better than the baseline model ( 1 ) in-distribution . Intuitively , the target y depends on both the inputs x ( throughw ) and latent variable u ( Figure 2 ) . The baseline model only uses x to predict y ; thus it can not capture the variation in y due to u . On the other hand , the aux-inputs model uses x and z to predict y . Since z is a function of x ( through w ) and u , u can be recovered from x and z by inverting this relation . Note that u is unobserved but implicitly recovered . The aux-inputs model can then combine u and x to predict y better . Let σ2u=Eu∼Pu [ ( θ > u u ) 2 ] denote the ( in-distribution ) variance of y due to the latent variables u . The following proposition shows that if σ2u > 0 then with enough training examples the aux-inputs model has lower in-distribution population risk than the baseline model.2 Proposition 1 . For all problem settings S , P , assuming regularity conditions ( bounded x , u , sub-Gaussian noise , and T =m ) , and σ2u > 0 , for all δ > 0 , there existsN such that for n≥N number of training points , with probability at least 1−δ over the training examples , the aux-inputs model improves over the baseline : Rid ( f̂in ) < Rid ( f̂bs ) . ( 8 ) Although using z as input leads to better in-distribution performance , we show that the aux-inputs model can perform worse than the baseline model OOD for any number of training examples . Intuitively , the aux-inputs model uses z , which can be unreliable OOD because z depends on u and u can shift OOD . In more detail , the aux-inputs model learns to predict ŷ= θ̂ > x , inx+θ̂ > z , inz , where the true output y=θ > x x+θ > z z , and θ̂z , in is an approximation to the true parameter θz , that has some error . Out-of-distribution u and hence z can have very high variance , which would magnify ( θ̂z , in−θz ) > z and lead to bad predictions . Example 1 . There exists a problem setting S , P , such that for every n , there is some test distribution P ′x , P ′ u with : E [ Rood ( f̂in ) ] > E [ Rood ( f̂bs ) ] ( 9 )

Summary: 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.