Project and Probe: Sample-Efficient Adaptation by Interpolating Orthogonal Features
Authors: Annie S Chen
Yoonho Lee
Amrith Setlur
Sergey Levine
Chelsea Finn
Venue: ICLR 2024 spotlight (2024)
Keywords: distribution-shift robustness
fine-tuning
adaptation
transfer learning
Primary area: transfer learning, meta learning, and lifelong learning
TL;DR: We propose Project and Probe, a lightweight, sample-efficient approach that learns a diverse set of predictive features and adapts to a target distribution by interpolating among them with a small target dataset.

Abstract: Transfer learning with a small amount of target data is an effective and common approach to adapting a pre-trained model to distribution shifts. In some situations, target data labels may be expensive to obtain, so we may only have access to a limited number of target data points. To make the most of a very small target dataset, we propose a lightweight, sample-efficient approach that learns a diverse set of features and adapts to a target distribution by interpolating these features. Our approach, Project and Probe (Pro$^2$), first learns a linear projection that maps a pre-trained embedding onto orthogonal directions while being predictive of labels in the source dataset. The goal of this step is to learn a variety of predictive features, so that at least some of them remain useful after distribution shift. Pro$^2$ then learns a linear classifier on top of these projected features using a small target dataset. Theoretically, we find that Pro$^2$ results in more sample-efficient generalization by inducing a favorable bias-variance tradeoff. Our experiments on four datasets, with multiple distribution shift settings for each, show that Pro$^2$ improves performance by 5-15% when given limited target data compared to prior methods such as standard linear probing.

Published as a conference paper at ICLR 2024
PROJECT AND PROBE: SAMPLE-EFFICIENT ADAPTA- TION BY INTERPOLATING ORTHOGONAL FEATURES
Annie S. Chen"!, Yoonho Lee Amrith Setlur?, Sergey Levine Chelsea Finn'
Stanford University', Carnegie Mellon University?, UC Berkeley'
asc8@stanford.co yoonho@stannoore
ABSTRACT
Transfer learning with small amount of target data is an effective and common approach to adapting pre-trained model distribution shifts. In some situations, target data labels may be expensive to obtain, so we may only have access to a limited number of target data points. To make the most of very small target dataset, we propose a lightweight, sample-efficient approach that learns a diverse set of features and adapts to a target distribution by interpolating these features. Our approach, PROJECT AND PROBE (PRo²), first learns a linear projection that maps a pre-trained embedding onto orthogonal directions while being predictive of labels in the source dataset. The goal of this step is to learn a variety of predic- tive features, so that at least some of them remain useful after distribution shift. PRO² then learns linear classifier top of these projected features using small target dataset. Theoretically, we find that PRO² results in more sample-efficient generalization by inducing a favorable bias-variance tradeoff. Our experiments on four datasets, with multiple distribution shift settings for each, show that PRO² improves performance by 5-15% when given limited target data compared prior methods such as standard linear probing.
INTRODUCTION
Machine learning models can face significant challenges when there is a distribution shift between training and evaluation data. A model trained on specific source dataset may not perform well when deployed on a target domain with distribution of inputs that differs significantly from the source domain. One common and reliable approach for adapting to distribution shifts is fine-tuning a trained model on small amount of labeled data from the new target domain. However, in some situations, target data labels may be expensive to obtain, which limits the number of available labeled datapoints for fine-tuning. As an example, a hospital may have imaging software that slightly differs from what used for dataset collection, but they may not be able to acquire many samples. In such conditions, conventional fine-tuning approaches may overfit to the small target dataset and distort the information learned during initial training. Therefore, require a method that can reliably extract information from the new target domain with less overfitting.
labeled
Recent works have demonstrated the effectiveness of re-training final linear head using target data for adapting to distribution shifts due to spurious correlations or domain shift (Rosenfeld et al., 2022 Kirichenko al. 2022, Mehta et al. 2022). However, it is unclear whether this standard approach of re-training a linear layer is the most data-efficient method to adapt pre-trained features to various target distributions. While versatile, feature embeddings may not necessarily contain the most suitable set of features for adapting to target distributions: they may also contain redundant, nonppredec, or noisy information. Our primary insight is that the key to more sample-efficient adaptation to target domains lies in starting with compact and diverse set of useful features. Each feature in this set should not only be predictive, but also hold unique information distinct from others inside the set. We leverage source data, which is substantially more abundant than target data, in performing this selection of features for target adaptation.
We propose PROJECT AND PROBE (PRO²). simple and sample-efficient method for adapting to unknown target distributions. PRO2 first learns projection of pre-trained embedding vectors. which is optimized extract diverse set of features that are each predictive of labels. More specifically,

---

Published as a conference paper at ICLR 2024
Distribution 1: Ideal Conditions
Curb
Car
Distribution Harsh Weather
Embeddings
Pre-Trained Backbone
Project with Large Source Dataset
Probe with Small Target Dataset
Figure 1: The Project and Probe (PRO²) framework for adapting to different target distribu- tions. (a) We first use large source dataset to project pre-trained feature embeddings onto a set of predictive features while enforcing orthogonality. (b) For new target distribution, we learn linear layer on top of the projected features. This step adaptively chooses features in a data-efficient manner.
first source dataset project pre-trained feature embeddings dictive features. We enforce pairwise orthogonality among all features, thereby ensuring that each projected dimension carries unique information not present in others. We expect this learned feature space to compactly contain diverse set of predictive features while discarding information that is redundant or not predictive on the task. PRO² then uses the reduced set of features as basis space for adaptation. Specifically, we fit linear head on top of the projected embedding using labeled target data. Both the linear projection and the linear head require minimal computational overhead, making PRO² practical method for adapting to new target distributions. Fig. 1 shows a visual summary of ?RO?.
smaller of pre-
To support approach, we provide theoretical analysis, in both general setting with minimal distribution assumptions as well as the more specific setting of a shifted homoscedastic Gaussian model, showing how ?RO? learns projection matrix that results in better generalization due to favorable bias-variance tradeoff. From this analysis, PRO² improves sample efficiency because can learn useful, diverse features that is more likely better recover the important directions for adaptation with a smaller projection dimension, allowing us to combat the variance introduced by very small target dataset while maintaining low bias. We conduct experiments on variety of distribution shift settings across 4 datasets. We find that standard linear probing, which is the default method used by prior works, not the most data-efficient adaptation approach. Using PRO?, i.e. projecting with source data onto an informative feature-space basis and probing with target data, improves performance by 5.15% in few-shot adaptation new target distributions.
RELATED WORK
Robustness and zero-shot generalization. Many prior works aim to improve robustness to various
distribution shifts TTzng et al. 2014 Ganin et al. 2016
Arjovsky et al. 2019 Sagawa al. 2020
Nam et al.) 2020 Creager et al.
2022 Liu et al.
2021
Zhang and Ré 2022). Additionally, prior
works have studied how to adapt pre- praited features
target distribution via fine-tuning Oquab
jet al. (2014): Yosinski et al. (2014); /hariff Razavian et al. (2014). Such fine-tuning works ty typically frame robustness distribution shift as zero-shot generalization problem Kornblith et al. (2018)
Zhai et (219); Wortsman et al. (2022) Kumar et al. (2022), where the model is trained on source and evaluated on target. Both of the above classes of approaches fundamentally cannot handle the problem settings we consider, where single function is insufficient for achieving good performance on different distributions, In this paper, we evaluate on a variety of test distributions, some of which are mutually exclusive, and it is therefore crucial to perform adaptation on the target distribution.
Adapting to distribution shifts. Recent works have proposed various methods for adapting models at test time with some labeled target data (Sun et al. (2020): Varsavsky et al. (2020); Iwasawa and Matsuo (202I): Wang et al. (2020)] Zhang et al. 2202 ); Gandelsman et al. (2022); Lee et al. (2022a). In particular, given a feature embedding produced by pretrained network with sufficient expressiv- ity, training a final linear head, also known as linear probing, suffices for adapting to datasets with spurious correlations Kirichenko al. (202)); Mehta et al. (2022): Izmailov et al. (2022) well as in the setting of domain generalization Rosenteld et al. (2022). As detailed further in Sec. 3) we specifically focus on scenarios in which we have very little target data (only 4 ~ 256 datapoints). We
2

---

Published as a conference paper at ICLR 2024
find that in this setting, training final linear head in the default manner is not the most data-efficient way to adapt. PRO?, which breaks this training down into 2 steps, is able to more effectively extract useful features and interpolate between them for varying target distributions, leading to improved sample efficiency with limited target data.
Learning diverse features for spurious datasets. Prior works have explored learning diverse or
orthogonal features in standard supervised learning settings (Bansal et al.) 2018 Xie et al. 2017b a
Laakom et al. 2023b; Cogswell et al.) 2015) Laakom et al. [2023]] Zbontar et al. 2021). and we show diversification can lead to more sample efficient adaptation to distribution shifts. Neural net-
works tend be biased towards learning simple functions that rely on shortcut features (Arpit et al.,
2017
Gunasekar et al. 2018] (Shah et al.] 2020] Geirhos et al.] 2020 Pezeshki et al.] [202]] et al.,
important to consider the entire set of functions that are predictive on the training data (Fisher et al., 2019 [Semenova et al.) 2019
2022
Lubana et al.) 2022). To better handle novel distributions,
Xu et al. 2002]] Recent diversification methods for adaptation discover such set (Teney et al,
2022 Pagliardini et al.) 2022; Lee et al. 2022b). The latter two methods use additional assump-
tions such as unlabeled data. With a simular motivation to ours, TTney et al. (2022) penalizes the similarity between different features, but does so with additional loss term instead of explicitly
enforcing orthogonality. We observe that this implementation detail matters in Sec. 6 where PRO* outperforms [Teney al. (2022). A concurrent work (Morwani al.) 2023) also proposes an orthog- onal projection method to learn diverse classifiers. However, the Probe step of PRO² additionally interpolates between the orthogonal features, and we provide theoretical and empirical analysis of how distribution shift severity affects sample efficiency during probing.
Compression & feature selection. In aiming to extract important features and discarding repetitive information, 2RO² is related to work on compression May et al. (2019) and information bottle- neeks Tishby et al. (2000; ; Alemi et al. (2016). Our method is also closely related to methods that learn projections such as principal component analysis (PCA) and linear discriminant analysis (LDA)). Beyond these representative methods, there an immense body of work on feature selec- tion (Dash and Liu 1997 Liu and Motoda 2007: Chandrashekar and Sahin, 2014 Li et al.] (017) and dimensionality reduction (Lee et al. [007]]Sorzan et al.] 2011:/20mi11nghhii Cunningham and Ghahramani 2015). Among all projection-based methods, LDA is the most related to ours, but it only learns the single most discriminative direction. In Corollary 9 we show that PRO² with dimensionality d = 1 provably recovers the LDA direction in shifted homoscedastic Gaussian model, and that using higher values of d is critical in adapting to higher degrees of distribution shift. Generally, most methods (including LDA) operate in the setting without distribution shift.
3 ADAPTATION TO DISTRIBUTION SHIFT
We now describe our problem setting, where the goal is to adapt model so to provide accurate decision boundary under distribution shift given a limited amount of target distribution information. We consider a source distribution ps y) and multiple target distributions p (2., y), p7 y).... The source dataset D3 (X x 2)N is sampled from the source distribution p.-. We evaluate adaptation to each target distribution ph given a small set of labeled target data Dj- E (X' x ))M, where M < N so the model must learn from both the source and target data for best performance. We measure the post-adaptation average accuracy of the model on a held-out target dataset from the same distribution pr.
We note that this setting differs from the setting studied in prior works on spurious correla- tions (Sagawa et al., 2020), which train model only on source data Ds and evaluate the model's performance on the hardest target distribution (i.e., worst-group accuracy). This is also different from the setting used in fine-tuning methods for zero-shot generalization (Wortsman et al. 2022] Kumar et al., 2022): such methods fine-tune pretrained model on source data Dg and directly evaluate performance on target data Dj without any exposure to labeled target data. Compared to these zero-shot evaluation settings, we argue that small amount of target data may realistically be required to handle the arbitrary distribution shifts that arise in the real world, Target data can be an effective point of leverage because it can be available or easy to collect, and we find that even a small dataset can reveal a lot about what features are effective in the target distribution. Our problem setting of adapting with target data has been used in some recent works (Kirichenko et al., 2022 Rosenfeld et al. 2022 Izmailov et al. 2022; Lee et al. 2022a), but we specifically focus on the setting in which only have access to very small target dataset, i.e., M < N.
3

---

Published as a conference paper at ICLR 2024
Algorithm Project and Probe
Input: Source data Ds. Target data Dr, Backbone f ->
Initialize 11 ,d for i in 1...d d do
#Project with source
11; f- arg min (L, subject I1 II; for all j
[a) Two orthogonal features [b) Interpolated classifier
end for
Visualization of PRO?: (a) orthog- onal decision boundaries learned during the Project stage, and (b) the interpolated classi- fier learned during the Probe stage.
Figure 2:
Initialize g: . y
#Probe with target
f arg min L (T()(I((())),
PROJECT AND PROBE
We now describe PRO?, a framework for few-shot adaptation to distribution shifts. PRO² is com- posed of two steps: (1) learn projection II that maps pre-trained embeddings onto orthogonal directions, and (2) learn a classifier using projected embeddings.
Before Step (1): dimensional embeddings. This backbone model extracts meaningful features from the raw inputs, resulting in low-dimensional embedding space, for example 224 x 224 x 3 images to D 1024- dimensional embeddings.
pre-trained backbone model f X' RD to map the datapoints D-
Project with source. Recall that we operate in the few-shot setting, where we may have fewer target datapoints than even embedding dimensions (M < D). Intuitively, we would like to select suitable decision boundary from a set of decision boundaries that worked well in the source domain. If this set is discrete, that might correspond to training some sort of diversified ensemble of linear classifiers on top of the features, strategy adopted some prior works (Teney et al., 2021 Lee et al. 2022b Pagliardini et al. 2022).
Step 1:
However, in general, we might need the expressive power of a continuous set of decision boundaries to adapt to the target domain, and we can construct this set by interpolating over basis of decision boundaries. Mathematically, this is identical to selecting a set of linear features. Thus, the question must answer is: which set of linear features of the D-dimensional feature space should we retain? First, it should be clear that the features should form an orthogonal basis, as otherwise they will be redundant. Second, the features should be discriminative, in the sense that they sufficient to solve the desired prediction task. Lastly, there should not be too many of them, since the more features we include (i.e., the larger the rank of the basis learn), the more samples we'll lll need from the target domain to find the best decision boundary in the corresponding set.
To learn feature space that satisfies these desiderata, we parameterize linear projection II: RD Rd that maps the embeddings to a reduced space (d < D). Specifically, we use the source data to learn complete orthonormal basis for the embedding space II,, II2.... ... ILa e", by learning each basis vector with the constraint that it is orthogonal to all vectors before it:
11, = arg min (x.y)) (fIL) s.t.
for all j <.
(PRo²)
Note that this induces natural ranking among the basis vectors. This collection of orthogonal vectors constitute the rows of our projection matrix IT. In our implementation, we do projected gradient descent, enforcing orthogonality using QR decomposition on the projection matrix after every gradient step. See Appendix B for a short PyTorch implementation.
Empirically and theoretically, we find that it is particularly beneficial to use a small d & D, even d = 1, in when adapting to small distribution shifts and use larger d for more severe distribution shifts.
Step 2: Probe with target. After learning 11, we learn a classifier g Rd + y that maps the projected embeddings to the target labels:
g arg min Clog))), v).
4

---

Published as a conference paper at ICLR 2024
Since the projection II was optimized to diverse set of the most discriminative features for the source data, we expect the initial projected features to be particularly predictive when the distribution shift is relatively small.
In summary, PRO? is a simple and lightweight framework that addresses the problem of few-shot adaptation in the presence of distribution shifts. We summarize its overall structure in Algorithm 1 and show a simplified 3D visualization in Fig. 2 In our implementation, we use cached embeddings for all source and target datapoints, such that feeding raw inputs through is amortized over epochs and experiments, making our framework scalable and efficient. As an orthogonal improvement to our work, one could additionally fine-tune the backbone network on source data. In Sec.. 5] we theoretically analyze the properties of the projection and classifier learned by PRO?. We then empirically evaluate PRO² on a variety of distribution shifts and publicly available backbone networks in Sec. 6
one-time cost that
5 ANALYSIS
In this section, we present a theoretical analysis of PRO², aiming to understand how our proposed orthogonal feature selection procedure can lead to sample-efficient adaptation under distribution shifts. Intuitively, the more shift we can expect, the more features we should need to adapt to it, which in turn requires more samples during adaptation (to fit the features accurately). However, the choice of how we extract features determines how the sample complexity grows under distribution shift: while large shifts may still require many features, if the features prioritized well, then smaller shifts might require only a very small number of features, and thus require fewer samples.
Given M target samples, in our analysis we first show that using (fewer) d features leads to lower variance, which scales as (d)/M), as opposed to ODD/M) But this improvement comes at a cost in bias, which in some cases scales as O(vD d. (T. (PT)). (p)))))rr) Note that this term grows with the amount of shift between the source ps and target pr distributions. In Sec. 5.1 we first analyze the specific features learned by PRO² with minimal distributional assumptions. Then, in Sec.. 5.2, we apply our general results to a shifted homoscedastic Gaussian (SHOG) model, where the bias and variance terms take more intuitive forms. We also empirically verify our results using synthetic SHOG data. Additional theoretical results and proofs can be found Appendix A
5.1
TRADEOFFS FOR GENERAL SHIFTS.
In this section, analyze the properties of II (learned projection) on the target distribution to understand why sample efficiency during adaptation can be improved by first extracting a set of diverse features that are predictive on the source distribution.
Probing on the target distribution. We first introduce some additional notation specific to the and Ma denote the projection matrix for Sa,
target distribution. Let Sa denote span([
Ma = (I1 T 11) '11T where, II II,,,,,II
(1)
AE We denote the target error for classifier as (w(()) A Epp ((w, x), y), and the bias incurred by probing over the projected features span((1l; (d as:
ba
A min Lr(w') min fr(w).
w'ES.
WEW
Finally, the d-dimensional weight vector learned by PRO² on the M projected target samples is:
M
Σ
u(( y()).
arg
min
WES
|| 1
In Proposition[] I show a bound on the bias by incurred by only using features in Sa to fit the target predictor, as opposed to using all possible d-dimensional linear features. Note that the upper bound on ba reduces to 0 as we add more features d + D. The rate at which bu - 0 is determined by how the optimal linear classifier w (on target) interplays with the projection matrix Ma learned on the souree data. On one extreme, when there is distribution shift, we know that for the projection II, returned by PRO?, II, x w/, and thus (Ip - M2 ) = O, the bias by - 0 with just one direction. On the other extreme, if I consists of random orthogonal features, then the bias ba decreases at
5

---

Published as a conference paper at ICLR 2024
Fer 000
1.0
For O3D hear OOD D
90
Nokepcce Nerm
10-
2.6
80
80
70
70
60
50-
20
12 128 512 2010 0369
2
0
12
330 512 2010 $192
2
0
12 138 517 2008 8192
Number Source Sirectons
Target Datue Siee
Trrpet Datent
Terper Detent
Evaluation of 2RO2 on shifted homoscedastic Gaussian data. (left) The x- and y-axes denote the rank of Ma and the nullspace norm [[[p - M2) respectively. The nullspace norm drops slowly for more severe distribution shifts, indicating that ba would be low enough only at very high values of d for severe shifts. (right) For less severe distribution shifts (ID and Near OOD), low- dimensional projections suffer from less bias, resulting in higher accuracy in the low-data regime. For the Far OOD distribution, using all 20-dimensional features is best, as bias drops more slowly.
Figure 3:
rate O d) even when there is distribution shift. This indicates that the rate at which the bias reduces as we increase d should be controlled by degree of distribution shift, and how informative the source features (in IIa) remain under this shift. This argument is captured more formally in the following Proposition.
Proposition 1 (bias induced by distribution shift). For some wi that is the optimal linear predictor on distribution pr over the full feature space, and an L-Lipschitz smooth convex loss I, the bias ba S T. [[ - Ma)wrrl2 When Ma is random rank d projection matrix with columns drawn uniformly over the sphere SD-1, then baLv/
||
In Theorem 2
describe the full bias-variance tradeoff where that the variance term is also controlled by the number features but unlike the bias is independent of the nature shift between source and the target.
Theorem 2 (bias-variance tradeoff). When the conditions in Lemma 7 hold and when ||x||oo (L), for B-bounded loss I, w.h.p. 1- 8, the excess risk for the solution wa of PRO" that uses features is Lr(Wa) minwew CT(w)
I(ID TLD-MMaie Ma) Vd + Bd+BBuu B//loo(1//)))
VM
where the first term controls the bias and the second controls the variance.
This result provides insights on what factors affect generalization when probing on target data. Tighter compression of the original representation, i.e., using a smaller d, increases bias while de- creasing variance. The rate of bias increase is determined by the degree of distribution shift, where more severe shifts correspond to steeper increase in bias. However, this bias can be mitigated as long as the important features needed for prediction on the target domain are covered by the com- pressed representation. Thus, ?RO² induces favorable bias- variance tradeoff, as the features tracted are predictive and diverse and hence are more likely to cover the important features needed for the target domain, allowing compression to smaller d while still maintaining low bias. The distribution shift has no effect on variance, and variance can only be decreased by using a low- dimensional represent (at the cost of bias) learning from larger target dataset (higher M).
5.2 BIAS-VARIANE VARIANCE TRADEOFF IN A SIMPLIFIED GAUSSIAN MODEL.
In this subsection, we consider a simplified setting of shifted homoscedastic Gaussian (SHOG). Within this model, we show that the more general statement in Theorem 2 can be simplified further to provide a more intuitive relationship between the factors that affect excess risk on target. We also empirically demonstrate the behavior predicted by our bounds on synthetic SHOG data.
Shifted homoseedastic Gaussian (SHOG) model of distribution shift. We model the source dis- tribution as a Bernoulli mixture model of data in which binary labels are balanced (y ~ B0rn().5)) and the class conditional distributions are homoscedastic multivariate Gaussians:
x y ~ (My, Es) for E (0, 1},
where Ho: #r E are mean vectors and Σg E DxD is the shared covariance matrix. The target distri- bution has the same label distribution and Gaussian means, but different covariance matrix given

---

Published as a conference paper at ICLR 2024
by Ey. Under this model, the degree of shift is fully realized by the distance between Σs, En under some reasonable metric. This directly impacts the bias term by when IIa is either returned by PRO* or random projection matrix with columns drawn uniformly over the sphere S We specialize the more general bias-variance tradeoff result to a shifted homoscedastic Gaussian (SHOG) model in Corollary 3 where we derive a simpler bound distilling the effect of d, and the degree of distribution shift on the target excess risk.
Corollary 3 (tradeoff under SHOG). Under SHOG model, and conditions for random projection Ma in Proposition the target excess risk (T( Lr(w)) d (K(PpsllTT) + M
Recall from Proposition I the bias scales directly with the norm along the null space of Ma- In Fig. left)), we plot the nullspace norm for different d in three target distributions of varying dis- tribution shift severity under the SHOG model. We see that the more severe shifts have higher nullspace norm even at large values of d. This indicates that the excess risk on target (OOD) may suffer from high bias severe shifts. Fig. 3) right), see that the source distribution (ID) suffers from virtually no bias, since d=1 achieves optimal target accuracy for even small target datasets. In contrast, the "Near OOD" and "Far OOD" distributions suffer from high bias since even with large datasets they fall short of optimal target performance (by >40%), and higher projection dimension d is needed for adaptation. This also validates our findings in Theorem 2 and Corollary 3
6 EXPERIMENTS
In this section, aim empirically answer the following questions: (1) Can PRO² identify feature-space basis for rapid adaptation, and how does it compare to other methods for extracting features? (2) How does the dimensionality of the feature-space basis affect sample efficiency in different distribution shift conditions? We provide additional empirical results and analyses, such as showing that the adaptation performance of PRO? improves with better pre-trained backbones, in Appendix C Details on pre-trate pre-trained models and training details are in Appendix B
6.1 EXPERIMENTAL SETUP
Datasets. We run experiments on six datasets with distribution shifts: 4-way collages (Teney et al. 2021 Waterbirds (Sagawa et al. 2020), CelebA (Liu et al., (015). Camelyon (Bandi et al. 2018)) Li (Santurkar et al. (020). and FMoW (Koh et al. 2022)) datasets. Each of these datasets have a source distribution that we use for training. For the first four datasets, we construct multiple target distributions for evaluation, representative of a range of potential test distributions. For the latter two datasets, which are larger datasets representing shifts that may occur in the wild, we evaluate on the given test set. For all settings, we use the original source datasets, which each contain thousands of datapoints. For target data, we subsample very small aabel-balanced datasets for adaptation, with {2, 8, 32, 128} images per label for the first four datasets and {1, 2, 5) images per label for the latter two datasets. The remaining target distribution datapoints are used for evaluation. Due to space constraints, we describe the different target distributions in Appendix B
Computational efficiency. Similarly to Mehta et al, (2022), we use feature embeddings from a pre- trained backbone without fine-tuning. Our aim is to develop methods that can leverage pretrained models out-of-the-box with minimal computational requirements: our training involves at most two linear layers on top of cached feature vectors. For all comparisons, we hyperparameter tune over different learning rates 10.1, 0.01, and 0.001) as well as 3 different L2 regularization weights (0.1, 0.01, 0.001). In our main experiments in /ecc. 6.2 we also sweep over 6 different projection dimen- sions (d = 1, 4, 16, 64, 256, 1024) and report results over 10 runs. For hyperparameter tuning, we adopt the typical practice of using a target validation set, which is common prior work in similar transfer learning settings (Kirichenko et al., 2022. Mehta et al., 2022. Lee et al., (022a). The chal- lenge of hyperparameter tuning for a target domain without additional domain-specille information remains an open problem that we hope can be addressed in future work. As a demonstration of the computational efficiency of PRO?, after caching pre-trained embeddings, we can collectively run all experiments in Sec. 6.2] which is nearly 30k runs due to hyperparameter tuning, within 24 hours hypers using four standard CPUs and no GPUs. We find that PRO" is robust to learning rate, which is expected as the optimization problem is linear.
7

---

Published as a conference paper at ICLR 2024
Waterbirds
CelebA
Camelyon 17
Collages
Spurious
Spurious
Test
Mollages-MNIST
1.0
1.0
0.9
0.8
0.8
0.7-
0.6
0.6
Balanced
Balanced
0OD Test
Col agees-
1.0
1.0
1.0
1.0
0.9
0.9
Test Accuracy
0.3-
0.0-
0.7-
0.8
0.6-
0.7
0.7-
0.8
0.5
Minority
Minority
Collages Fashion
1.0
1.0
0.9
0-
0
0.8
0.7
0.7
0.6
0.8
0.6
0.5
Average
Average
Average
Average
1.0
1.0
0.9
0.9
0.9
0.8
0.8
0.8
0.7-
0.7
0.8
0.6
120
32
120
2
8
32
128
8
32
128
Target Train Data Size
Random Projection
DFR (Kirichenko et al.)
Teney
Pro3 (ours)
Figure 4: Main results. We compare 4 different methods for learning features to adapt to a target distribution: targge.distribution (1) Random Projection, (2) DFR Kirichenko et al. (2022). (3) Teney al. (2021). and (4) PRO?. We report average target accuracies after probing with different target dataset sizes ranging from 2 to 128 datapoints per
label; error bars indicate standard error across 10 runs. PRO* is the best performing or tied for best performing method acroxs each of these 4 datasets with any amount of target data. PRO* substantially outperforms Random Projection and DFR in the low-data regime on all four datasets. PRO² also outperforms Tency ct al. 2021 1) on average 3 of the 4 datascts particularly when given more target data.
Living17
FMoW
Target Train Data Size (per label)
2
5
Random Projection
85.7 (0.6)
92.7 (1.0)
99.2 (0.1)
16.3 (0.7)
23.6
33,3 (0.6)
DFR (Kirichenko al.)
87.1 (0.9)
95.0 (0.9)
98.8 (0.3)
17.5 (0.8)
24.0 (0.6)
35.1 (0.6)
PRO?
91.2 (0.4)
95.7 (0.7)
99.2 (0.06)
20.2 (0.9)
28.7 (1.1)
37.2 (0.8)
Table 1: Additional main results. We run additional experiments on the Living17 dataset from the Breeds benchmark (Santurkar al. 2020) and FMoW (Koh et al. 202T). reporting adaptation accuracy and standard error across T0 runs. Both of these datasets are challenging multi-class distribution shift tasks and are represen- tative of worldworld scenarios. We find that similar to the other datasets, PRO* is the best performing or tied for best performing method on these datasets when given a limited amount of target data.
6.2 COMPARISON TO PRIOR PROJECTION METHODS
We investigate whether PRO can extract features that can facilitate adaptation to different distri- bution shifts, and how it compares other feature extraction methods. We perform comprehensive experimental evaluation on the six datasets, comparing PRO² against four other projection methods: (1) Random Projection, (2) DFR Kirichenko et al. (2022). which uses standard linear probing, and (3) Teney et al. (002T). which aims to learn multiple predictive patterns by minimizing the alignment of input gradients over pairs of features. Experiments Fig. 4|ddd Tab. iindicate that across all dif- ferent target distributions six datasets, PRO significantly outperforms Random Projection and DFR, especially in the low-data regime. In particular, these results show that DFR or standard linear prob- ing, the strategy adopted by several additional prior works by default Mehta et al. (2022) y Izmailov et al. (2022), is not the most data-efficias way to utilize pre- pre-trained embeddings when given limited target data. This is because such embeddings contain redundant or non con-predietive information, and including these features during adaptation leads to higher variance without decreasing bias, which in turn means that need more labeled samples. In contrast, PRO² improves sample efficiency
8

---

Published as a conference paper at ICLR 2024
Waterbirds
CelebA
Camelyon17 17
Collages
Prof tours) on Spurious
Prof lourst on Spurious
Pro (ours) on ID Test
Pro" lours] on Co lages lags HNIST
1424
07.3
20.0
91.9
25.4
1024
53.5
25.2
D
91.7
1624
80.3
85.3
88.9
22.4
1004
84.8
25.3
58.0
256
09.0
30.8
92.
31.3
256
50.
15.9
56.5
57.6
256
79.0
87.0
08.8
33.7
256
75.4
57.4
55.3
57.7
66-
81.7
27.
90.7
23.3
54
90.1
94.3
95.6
94.7
95.2
56.6
95.9
16
90.5
$3.3
99.0
25.7
57.0
21.3
4
57.
25.3
19.8
97.0
81.1
10.0
75.8
99.8
93.2
51.9
00.0
23.3
93.4
23.5
Number Features
on Balanced
[nurs]
Baianced
Prof [o.ar) on GOD Test
(ours) on Collages-CIFAR
1036
6
46.6
90.2
95.6
1024
632
76.4
16.6
03.
08.4
F62
1024
75.3
83.5
05.5
50.5
256
78.8
87.4
60.0
53.4
158
81.9
$7.6
91A
66.7
75.3
80.7
85.1
64
63.0
46.0
90.8
92.3
02.4
33.4
66.2
16
66
49.6
62.4
92.
08.9
76.0
60
16
23.
62.5
43.
636
863
40.3
92.2
22.0
91.8
53.2
$1.3
12.4
61.9
90.9
91.9
90.2
93.2
71.3
11.9
Prof [eurst Minority
Pro" (ours) on Minorty
Pro loursi Co inges- MNIST
1824
1024
2004
62.5
92.3
94.3
256-
89.5
91.0
93.2
94.5
76.3
63.3
66
85.1
26.3
35.7
66-
51.9
91.
50
15.1
64
35.5
630
78.1
16.5
16
04.3
16.9
05.6
85.0
663
042
60
52.3
06.6
87.5
26.2
50.3
26.3
1
135
75.1
70.8
75.0
31.4
61.
33.4
51.2
32
128
Target Train Data Size
Figure 5: Feature-space dimensionality of PRO2 and severity of distribution shift. We vary the feature- space dimensions d (y-axis) of PRO? and report held-out accuracy after training on target datasets of different size (x-axis) our 4 datasets. Higher accuracies are in blue and lower accuracies are in red. We see that smaller feature-space dimensions suffice for target distributions with milder distribution shifts while higher dimensions are required for more severe shifts. For example, on the spurious test distribution (small dist, shift) of Waterbirds/CelebA, the bottom row, which uses d = 1 is bluest, while the blue is concentrated in the top right squares (which use more features and more data) for more difficult distrihution shifts such as Minority for Waterbirds/CelebA and the collages test sets.
by first extracting a predictive feature-space basis from the source distribution, removing redundant information. TTnny et al. (202T) T) is sufficient in some scenarios with milder distribution shift, where a diverse range of features are not needed for adaptation. However, it fails to achieve high accuracy given a large target dataset more severe distribution shifts, such as the Minority distributions on Waterbirds and CelebA or the Fashion-MNIST and CIFAR distributions in 4-y Collages. This indicates that the feature diversity from the orthogonality constraint gives PRO2 better coverage of different features, enabling better adaptation to severe distribution shifts given enough target data. These results demonstrate the effectiveness of PRO compared to existing methods in the few-shot adaptation problem setting.
6.3 PROJECTION DIMENSION AND SHIFT SEVERITY
In this subsection, we investigate how the feature-space dimension d affects the sample efficiency of PRO?, for different degrees of distribution shift. Experiments in Fig. 5 5/show that when the distribu- tion shift is less severe, such as the Spurious test distributions on Waterbirds and CelebA, it helpful to reduce the number of features used. This scenario is analogous to the ID setting in Fig. 3 In such scenarios, the top-ranked features from the source data are also predictive on the target distribution, and incorporating additional features worsens generalization because it increases variance without sufficiently decreasing bias. However, when the distribution shift is more severe, such as the Minor- ity distributions on Waterbirds and CelebA or Dollages-Fassion MNIST and Collages-CIFAR, it is helpful to increase the number of features used. This scenario is analogous to the Far OOD setting in Fig. 3) These empirical results are supported formally by our theoretical results in SSc.. 5, which show that the optimal number of features to use increases with distribution shift severity.
CONCLUSION
In this paper, we propose PRO², lightweight framework consisting of 2 steps: (1) projection step that extracts diverse and predictive feature-space basis and (2) a probing step that interpolates between the projected features to efficiently adapt varying target distributions. Our theoretical and empirical analyses reveal a number of interesting novel insights: (i) standard linear probing is not the best approach for few-shot adaptation; (ii) Retaining a diverse range of potentially useful features that different target distributions may require improves sample efficiency, (iii) we can trade off how much adapt (size of the feature-space basis) vs number of samples, picking the best basis adapt
9

---

Published as a conference paper at ICLR 2024
for each level of shift. These insights open up range of exciting paths for future work. First, framework may be extended to other problem settings, such as the active learning setting, in which the model can adaptively request target labels. Another interesting direction would be selecting which features to use an unsupervised fashion, without any labeled target data. For limitations, more complex target distributions generally require Our analysis on the bias-variance tradeoff' expressed in Theorem 2 depends on the dataset size, so PRO² will have smaller benefits in settings with large amount of target data.
directions and therefore more data fit.
ACKNOWLEDGMENTS
We thank members of the IRIS and RAIL labs for helpful discussions this project. This work was supported by NSF, KFAS, Apple, Juniper, and ONR grant N00014-20-1-2675.
REFERENCES
Alemi, A. A., Fischer, I., Dillon, J. V., and Murphy, K. (2016). Deep variational information bottle- neck. arXiv preprini arXiv:/612.004/ 16/2.00410.
Arjovsky, M., Bottou, L., Gulrajani, I,, and Lopez-Paz, D. (2019). Invariant risk minimization. arXiv preprint arxiv: 1K07.0997702893
Arpit, D., Jastrzebski, S., Ballas, N., Krueger, D., Bengio, E., Kanwal, M. S., Maharaj, T., Fischer, A., Courville, A., Bengio, Y., et al. (2017). A closer look memorization in deep networks. In International Conference on Machine Learning.
Bandi, P., Geessink, O., Manson, Q., Van Dijk, M., Balkenhol, M., Hermsen, M., Bejnordi, B. E., Lee, B., Paeng, K., Zhong, A., al. (2018). From detection of individual metastases classifi- cation of lymph node status at the patient level: the camelyon17 challenge. IEEE Transactions on Medical Imaging.
Bansal, N., Chen, X., and Wang. Z. (2018). Can gain more from orthogonality regularizations in training deep networks? Advances in Neural Information Processing Systems, 31.
Bartlett, P. L. and Mendelson, S. (2002). Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463-482.
Caron, M., Misra, I., Mairal, J., Goyal, P., Bojanowski, P., and Joulin, A. (2020). Unsupervised
learning of visual features by contrasting cluster assignments. Advances in Neural Information Processing Systems, 33:9912-9924.
Chandrashekar, G. and Sahin, F. (2014). survey feature selection methods. Computers & Electrical Engineering, 40(1))16-28.
Cogswell, M., Ahmed, F., Girshick, R., Zitmick, L., and Batra, D. (2015). Reducing overfitting in deep networks by decorrelating representations. arXiv preprini 111.v015/1.0666.
Creager, E., Jacobsen, J.-H., and Zemel, R. (2021). Environment inference for invariant learning. In International Conference on Machine Learning.
Cunningham, J. P. and Ghahramani, Z. (2015). Linear dimensionality reduction: Survey, insights, and generalizations. The Journal of Machine Learning Research, 16(1):2859-200..
Dash, M. and Liu, H. (1997). Feature selection for classification. Intelligent data analysis, 1(1- 4):131-156.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., al. (2020). An image is worth 16x16 words: Trans- formers for image recognition at scale. arXiv preprint arxiv:2010.11922
Fisher, A., Rudin, C., and Dominici, (2019). All models are wrong, but many are useful: Learning variable's importance by studying entire class of prediction models simultaneously. J. Mach. Learn. Res., 20(177):1-81.
Gandelsman, Y., Sun, Y., Chen, X., and Efros, A. A. (2022). Test-time training with masked autoen- coders. arXiv preprint arXiv:2209.0752..
10

---

Published as a conference paper at ICLR 2024
Ganin, Y., Ustinova, E., Ajakan, H., Germain, P., Larochelle, H., Laviolette, F., Marchand, M., and Lempitsky, V. (2016). Domain-adversaan training of neural networks. The journal of machine learning research, 17(1):2096-2033 2030.
Geirhos, R., Jacobsen, J.-H., Michaelis, C., Zemel, R., Brendel, W., Bethge, M., and Wichmann, F. A. (2020). Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(1 1):655- 673.
Gunasekar, S., Lee, J. D., Soudry, D., and Srebro, N. (2018). Implicit bias of gradient descent on linear convolutional networks. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa- Bianchi, N., and Garnett, R., editors, Advances in Neural Information Processing Systems.
He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778.
Iwasawa, Y. and Matsuo, Y. (2021). Test-time classifier adjustment module for model- domain generalization. Advances Neural Information Processing Systems, 34:2427-2440.
Izmailov, P. Kirichenko, P., Gruver, N., and Wilson, G. (2022). On feature learning in the presence of spurious correlations. urXiv preprint arXiv:2210.1 133.9
Kakade, S. M., Sridharan, K., and Tewari, A. (2008). On the complexity of linear prediction: Risk bounds, margin bounds, and regularization. Advances in neural information processing systems, 21.
Kirichenko, P., Izmailov, P., and Wilson, A. G. (2022). Last layer re-training is sufficient for robust- ness to spurious correlations. arXiv preprint arXiv:2204.02937
Koh, P. W., Sagawa, S., Xie, S. M., Zhang, M., Balsubramani, A., Hu, W., Yasunaga, M., Phillips, R. L., Gao, I., Lee, T., al. (2021). Wilds: A benchmark of in-the-will wild distribution shifts. In International Conference on Machine Learning, pages 5637-5664. PMLR.
Korbliihh S., Shlens, J., and Le, Q. (2018). Do better imagenet models transfer better? arxiv 2018. arXiv preprint arXiv: 0005 00974..74.
Kumar, A., Raghunathan, A., Jones, R., Ma, T., and Liang, P. (2022). Fine-tuning distort pretrained features and underperform -isttribiiiiiioon..... arXiv preprint 1rXiv22202.10054
Laakom, F., Raitoharju, J., Iosifidis, A., and Gabbouj, M. (2023a). Learning distinct features helps, provably. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 206-222. Springer.
Laakom, F., Raitoharju, J., Iosifidis, A.. and Gabbouj. M. (2023b). Wid-reg: A data-dependent lithin-layer diversity regularizer. arXiv preprint arXiv:2301.01352.
Lee, J. A., Verleysen, M., et (2007). Nonlinear dimensionality reduction, volume 1. Springer.
Lee, Y., Chen, S., Tajwar, F., Kumar, A., Yao, H., Liang, P., and Finn, C. (2022a). Surgical fine-tuning improves adaptation to distribution shifts. arXiv preprint 2rXiv:2210. 1466..
Lee, Y., Yao, H., and Finn, C. (2022b). Diversify and disambiguate: Learning from underspecified data. arXiv preprint 182.4.220..03418.
Li, J., Cheng, K., Wang, S., Morstatter, F., Trevino, R. P., Tang, J., and Liu, (2017). Feature selection: A data perspective. ACM computing surveys (CSUR), 50(6):1-45.
Li. Z., Evtimov, I., Gordo, A., Hazirbas, C., Hassner, T.,, Ferrer, C. C., Xu, C., and Ibrahim, M. (2022). A whac-a-mole dilemma: Shortcuts come multiples where mitigating one amplifies others.
Liu, E. Z., Haghgoo, B., Chen, A. S., Raghunathan, A., Koh, P. W. Sagawa, S., Liang, P., and Finn, C. (2021). Just train twice: Improving group robustness without training group information. In International Conference Machine Learning, pages 6781-6792. PMLR.
Liu, H. and Motoda, H. (2007). Computational methods of feature selection. CRC press.

---

Published as a conference paper at ICLR 2024
Liu, Z., Luo, P., Wang, X., and Tang, X. (2015). Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV).
Loshchilov, I. and Hutter, F. (2017). Decoupled weight decay regularization. arXiv preprint 1rXiv:1111.05100
Lubana, E. S., Bigelow, E. J., Dick, R. P., Krueger, D., and Tanaka, H. (2022). Mechanistic mode connectivity. arXiv preprint arXiv:2211.08-27
May, A., Zhang. J., Dao, T., and Ré, C. (2019). the downstream performance of compressed word embeddings. In Wallach, H.,, Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.
Mehta, R., Albiero, V., Chen, L., Evtimov, I., Glaser, T., Li, Z.,, and Hassner, T. (2022). You only need a good embeddings extractor to fix spurious correlations.
Morwani, D., Batra, J., Jain, P., and Netrapalli, P. (2023). Simplicity bias in 1-hidden layer neural networks. arXiv preprint 202.02230..00457
Nam, J., Cha, H., Ahn, S., Lee, J., and Shin, J. (2020). Learning from failure: Training debiased classifier from biased classifier. Conference on Neural Information Processing Systems.
Oquab, M., Bottou, L., Laptev, I., and Sivic, J. (2014). Learning and transferring mid-level image representations using convolutional neural networks. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 1717-1724.
Pagliardini, M., Jaggi, M., Fleuret, F,, and Karimireddy, S. P. (2022). Agree to disagree: Diversity through disagreement for better transferability. arXiv preprint 220i::2202.04414..
Petridis, S. and Perantonis, S. J. (2004). On the relation between discriminant analysis and mutual information for supervised linear feature extraction. Pattern Recognition, 37(5)8857-774
Pezeshki, M., Kaba, S.-O., Bengio, Y., Courville, A., Precup, D., and Lajoie, G. (2021). Gradient starvation: A learning proclivity in neural networks. In Beygelzimer, A.. Dauphin, Y., Liang, P., and Vaughan, J. W., editors, Advances in Neural Information Processing Systems.
Rosenfeld, E., Ravikumar, P., and Risteski, A. (2022). Domain-adjusted regression or: Erm may already learn features sufficient for out-of-distribution generalization. arXiv preprint arXiv:2202.06856.
Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang. (2020). Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. International Conference on Learning Representations.
Santurkar, S., Tsipras, D., and Madry, A. (2020). Breeds: Benchmarks for subpopulation shift. arXiv preprint ar arXiv:2008.04833
Semenova, L., Rudin, C., and Parr, R. (2019). A study in rashomon curves and volumes: A new perspective on generalization and model simplicity in machine learning. arXiv preprins arXiv: /00
Shah, H., Tamuly, K., Raghumathan, A., Jain, P., and Netrapalli, P. (2020). The pitfalls of simplicity bias in neural networks. Conference on Neural Information Processing Systems.
Sharif Razavian, A., Azizpour, H., Sullivan, J., and Carlsson, S. (2014). Cnn features off-the-shelf: astounding baseline for recognition. In Proceedings the IEEE conference computer vision and pattern recognition workshops, pages 806-813.
the
Sorzano, C. O. S., Vargas, J., and Montano, A. P. (2014). A survey of dimensionality reduction techniques. arXiv preprint arXiv: 403.2877.
Sun. Y., Wang, X.. Liu, Z., Miller, J., Efros, A., and Hardt. M. (2020). Test-time training with self- supervision for generalization under distribution shifts. In International conference on machine learning, pages -229.9248. PMLR.
12

---

Published as a conference paper at ICLR 2024
Teney, D., Abbasnejad, E., Lucey, S., and Hengel, A. v. d. (2021). Evading the simplicity bias: Training a diverse set of models discovers solutions with superior ood generalization. arXiv preprint ar 1rX1:2200..0612.
Teney, D., Abbasnejad, E., Lucey, S., and van den Hengel, A. (2022). Evading the simplicity bias: Training a diverse set of models discovers solutions with superior ood generalization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 16761-16772.
Tishby, N., Pereira, F. C., and Bialek, W. (2000). The information bottleneck method. arXiv preprint physics/0004057.
Tzeng, E., Hoffman, J., Zhang, N., Saenko, K., and Darrell, T. (2014). Deep domain confusion: Maximizing for domain invariance. arXiv preprint arXiv: 141..3444
Varsavsky, T., Orbes-Arteaga, M., Sudre, C. H., Graham, M. S., Nachev, P., and Cardoso, M. J. (2020). Test-time unsupervised domain adaptation. In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 428-4346 438.436.75 Springer.
Wainwright, M. J. (2019). High-dimensional statistics: A non-asymptotic viewpoint, volume 48. Cambridge university press.
Wang, D., Shelhamer, E., Liu, S., Olshausen, B., and Darrell, T. (2020). Tent: Fully test-time adaptation by entropy minimization. arXiv preprint arXiv:2006. 10726.
Wortsman, M., Ilharco, G., Kim, J. W., Li, M., Kornblith, S., Roelofs, R., Lopes, R. G., Hajishirzi, H., Farhadi, A., Namkoong, H., and Schmidt, L. (2022). Robust fine-tuning zero-shot mod- els. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 7959-7971.
Xie, B., Liang, Y., and Song, L. (2017a). Diverse neural network learns true target functions. In Artificial Intelligence and Statistics, pages 1216-1224. PMLR.
Xie, P., Singh, A., and Xing, E. P. (2017b). Uncorrelation and evenness: new
regularizer. In International Conference on Machine Learming, pages 3811-3881 1-3820. PMLR.
Xu, Y., He, H., Shen, T., and Jaakkola, T. (2022). Controlling directions orthogonal to classifier. arXiv preprint arXiv:2201.11255
Yosinski, J., Clune, J., Bengio, Y., and Lipson, H. (2014). How transferable are features in deep neural networks? Advances in neural information processing systems, 27.
Zbontar, J., Jing, L., Misra, I. LeCun, Y., and Deny, S. (2021). Barlow twins: Self-supervised learning via redundancy reduction. In International Conference on Machine Learning, pages 12310-12320. PMLR.
Zhai, X., Puigcerver, J., Kolesnikov, A., Ruyssen, P., Riquelme, C., Lucie, M., Djolonga, J., Pinto, A. S., Neumann, M., Dosovitskiy, A., et al. (2019). A large-scale study of representation learning with the visual task adaptation benchmark. arXiv preprint urXiv: 1910::990.004 04867.
Zhang, M., Marklund, H., Dhawan, N., Gupta, A., Levine, S., and Finn, C. (2021). Adaptive risk minimization: Learning to adapt to domain shift. Advances in Neural Information Processing Systems, 34:23664 2226666633678. 23678.
Zhang, M. and Ré, C. (2022). Contrastive adapters for foundation model group robustness. arXiv preprint arXiv:2207.07188
13

---

Published as a conference paper at ICLR 2024
A PROOFS FOR THEORETICAL ANALYSIS
We present proofs for our theoretical analysis in Sec. along with some additional statements, As in the main paper, we denote the dimensionality of the feature-space basis learned by PRO* as d, the original dimension of the representations given by the feature backbone f as D, source and target distributions ps and pr, and the number of source and target datapoints N and M. We let Ma denote the projection matrix for span( {11)) M [[][[I Hss] T. If the target error for the feature w is (ww( := ED, ((W,[((w)))))))))))( x), y), then the bias incurred by probing the subspace Ma consisting of source features is: ba
min w' (Sspa((([.))))
Cr((w'
min Cr(w),
wEW
and we denote the feature-space basis of dimensionality d learned by PRO² as follows:
M
arg min
Σ ([Wwxx!)) y()).
(3)
wespan (T, H
From the original D-dimensional feature representations given by our feature backbone want our learned linear projections 11 _ to retain as much information as possible that is relevant in predicting the label y. In other words, we want maximize the mutual information between the projected features II(x) and the labels y. In Theorem 6 We first formally characterize the solution found by projection step in PRO as maximizing this mutual information amongst all rank matrices with orthogonal columns, using the following two lemmas.
Lemma 4 (Entropy of sub-gaussian under low rank projection). For a- suub-gaussian random variable x. and rank r orthonormal matrix A edx the differential entropy H (Ax) for projection Ax is O(1/d) when log o O(./d)).
Proof. Let A; denote the 7th row of A, then:
i=d
i=d
H(Ax) = x
x
Σ
H(4Tx
Σ x(A
t=2
x=1
Since, x is o sub- -abbsagassin, using standard bounds on differential entropy and the inequality above can argue that:
Ax S E Vi,t
H(A x) S log((Tee7) Vi
HAA x) (1/d)) Vi (since logo = O(1/d)) H(Ax) ((1/d)
Lemma 5 (Entropy for mixture of o- suub-aasssians. For a d dimiinnsio random variable x that is mixture of two a o-sub- gaussian bounded Jensen-Shannon divergence JS(r" T x1|l0' x2) S 3 and mixture proportion a E [0, 1] the density function p(x) = px + (1 - a)) px(x,, then the entropy H for ||l|| 1, is at most Ollog o + -).
o-sub-gaussian random variables x1 and with overlapping supports,
Proof.
Using Jensen inequality and the definition of KL divergence, the differential entropy can be broken down as follows:
H(v x) = f (a- (a. Pi.pii" (x)
a) p2("" x)) log(a Pi T Tx) (1 ax) p2(w'' T xx)) da x) + +1
I HHu Tx x1) (1-a)) a1 Tx2)
log aa(1 aa) / p1d x)) log pz (u T x) da a(1 ai) p2(00 x:) log pı (v) xx) dx
< aHH x1) + (1- a) xx,) + 2a -),3
log(o) +
where the last step follows from the first arguments made in the proof for Lemma
14

---

Published as a conference paper at ICLR 2024
Theorem 6 (Information in projected input). When the distributions p(x y) are exp ( sub- gaussian for each y and the Jensen-Shamonn divergence JS(p(v's x y = 0) p( x = 1)) = O1/d). the solution {II,) returned by PRO? maximizes tight lower bound (difference bounded by an O(1) constant) on the mutual information criterion T AAx; y) among all x D row-orthonormal matrices A. (See end of the proof for discussion on assumptions and remark tightness).
Proof. We use an inductive argument on d. Consider the following maximization problem where By is the set of all row orthonormal matrices of row rank d < D:
max yAx;y).
(4)
A
Let d >>> 1. Then, we can re-write the above as:
max y(Ax:)) AEBOXD
max
(A'x, :y) where, E NNulSpace(A") (A'), |/0122 1.
A'EBa-I,VED
(5)
Now, we can decompose this expression using the conditional mutual information identities:
I A'xx U : = I(A'x;; ()) + I(v x: - I(v x: A'x) + [(v] x; A'x y)
((A'ayy) + [( xxy) - (I(oTx: A'x) I( A'x v))
(6)
Now, we upper bound the drop in information when condition y: (I(vTx; A'x)
I(v x; A'x y)) using Lemma and Lemma 5
I(v x: A'x) - [(v] x; A'x 3) = H([ x; A'x] T 1y) - H(vTxly) - (A'x y) + 1(0 x; A'x)
-
< H( x; A x T 1y) H(Txx) = H( y) + H( x)
3 (1/d).
where the last statement applies Lemma 4 to bound [AA y) (since Ax y 0 and Ax y= 1 are a -ubb-gaausiann) and Lemma to bound the entropy on the mixture of sub-gaussians H( x).
Also, note that the conditional distributions differ in Jensen-Shannon divergence by (l/dd) and the sub -oo--gussiin assumption gives us log o = (1)dd).
Using equation 6
equation we have:
max y(Ax;y) 2 AEBa
max A 2 uaNull (A) =1l==1
T (44x; + T(o T(xTxx)) x:y)
O(1d))
(8)
Let A, denote the tth row of A, then. Then, applying the above inductively for all d:
i=d
Σn(T ( 0(11) -
max yAx;y) ACR.
max ACR,
(9)
Let A be the solution of the right hand side and 8" arg max 2=1 I(v y) Next, we note that 3i such that A = v² It is easy to prove this by contradiction. Consider the case where z such that A; = v². Then, we can construct a solution {(In v²v" )A; order them by mutual information ((AA)) (ID e )x; y), take the top d - 1 entries and append to this set v". The new solution would have higher value of the objective on the right hand side of equation 9 since the new solution retains optimal directions perpendicular to 8" while adding v" to the set. Thus, we arrive at a contradiction and it is clear that v² belongs to the solution A" for the objective on the right side of equation 9
Knowing that x² has to be part of A+, we can now write the right side of equation as the following:
max 7(Ax;y)= 2 max AERaXD
(v] x;y)
EIDDD
I ( [/2 [I [I-vi 01 x:y)
+ max
E.ERD
T (1 0i0y T) (1 viie x: + +..+++.0000 0(1,, (10)
+ max PER
Ug3
15

---

Published as a conference paper at ICLR 2024
where vi. ......: 2 denote the solutions to each subsequent max term. This sequence of solu- tions is the same as that returned by solving the following iterative optimization problem because maximizing mutual information with label for linear projection of the input is the same find- ing a direction that minimizes Bayes error of the linear projection (Petridis and Perantonis (2004) connects mutual information cross entropy loss and Bayes error):
1. vं arg s x(),,x),))
2. Project data in the null space of oi: (T oiol
3. Re-solve (1.) to get next o7 and so on.
Finally, it easy to see that solution returned by the above iterative optimization is the same as that returned by the project step of PRO?.
Discussion on assumptions: Following are some remarks and intuitions behind the assumptions we make:
Sub-gaussianity: We need sub-gaussianity bound the entropy of linear projections, which is easily satisfied for inputs with bounded support. Note that the sub-gaussian pa- rameter need only satisfy log = (1/dd)), where d < D which is the input dimension of the data.
Bounded JS-divergence: The main intuition behind why need the class conditional distributions to not differ too much (bounded JS-divergence) along linear projections is that if they are very different from each other it is possible that even with sub-gaussianity assumptions there may exist linear projections that have a high mutual information over the mixture of conditionals (which is the marginal imput distribution p(x) i.e., I(v" x; A'x) is high) but not when we condition on the label (i.e., /(v> x; A x y) is low). Now, since we iteratively search for linear projections, our project step is oblivious and we may recover both of these directions (see equation and Lemma [5. But only one may be present the information theoretically optimal linear projection.
these interactions
Remark on tightness of our lower bound: We show that we maximize a lower bound in equation 9 But, in the special case when the class conditionals log-concave (e.g., multivariate Gaussian) can also show something much tighter: max yAA:y)) = max (Ei (ATx;y)
(11) This is because our upper bounds on the entropy terms have matching lower bounds for log- concave distributions, which then be applied lower bound the negative terms in the first step of equation
We now provide proof' of the generalization bounds in Section Showing the varianee tradcoff. Lemma 7 (generalization bound for probing projected features). For an L--Lipshitz, B-bounded
loss I, with probability 2 1- 8, wa in equation has generalization error va+ B //og(110) when ||x||0 O().
VM
Proof. For this proof, we invoke the following two lemmas.
Lemma 1 (generalization bound for linear functions Bartlett and Mendelson (0022). For an L-
Lipshitz B-bounded loss the generalization error for predictor wa. contained in the class of lz norm bounded linear predictors W is bounded with probability 2 1 1- -8:
M
Tog (og(1/8)
1((wa. x),
Σ (((ww, M,x(),, y(0) S 2LR. BV 2=1
2M
where R., is the empirical Rademacher complexity of l2 norm bounded linear predictors.
16

---

Published as a conference paper at ICLR 2024
Lemma 2 (W) bound for linear functions Kakade al. 2008)). Let W be a convex set inducing W) for input space X, bounded in norm || by some value R > 0. If there exists a mapping h W + that is k-strongly convex with
the set of linear functions F(W) & {(w, x) X +
respect the dual norm I| I| and some subset W' C W takes bounded values h(-) {h(w) K
w E W for some K > 0. then the empirical Rademacher complexity of the subset W is bounded
by R., (F(W") & R 2k
Let |11122 be the function h W + in Lemma 2
know that || [12 is 2-strongly convex in l2 norm. Further, take the standard 12 norm as the norm over X. So, the dual norm . is also given by la 2. We also know that W bounded in ||2 by 1, based
norm. Thus, Thus, K =
setup definition.
Further, note that ||x||> O(). We apply Cauchy-Schwand and use the fact that Ma 1 to bound the norm of the projected vector:
Max S Ma|xx|1a & Mavd|xx/lo Vd.
(11)
By Lemma we get the empirical Rademacher complexity RM (W) and plugging this into Lemma I yields the main result in Lemma
Theorem (bias-varianee variance tradeoff, Theorem 2} When the conditions in Lemma T Thold and when
||x O(1), for B-bounded loss I, w.h.p. 8, the excess risk for the solution of PRO that uses features
Vd+ 12/log(1)
CT(Wa)
min CT(W) - Ma) +
(12)
WEW
VM
where the first term of the RHS controls the bias and the second controls the variance.
Proof. The excess risk for Wa is
Lr(Wa) - min Lr(w) WEW
Cr(Wa)
min (T.))
(w)) + min WEspan[TT)]
Cr(w) - min (w())
wEW
min (Es
Cr(w)
min Cr(w) (r min wEW
Cr(w)
Vd+ + B Aoollo(//)
(ID - Mal
(13)
VM
where the first term the bias (bounded using Lemma T]. and the second term the generalization enror or the variance (bounded using Lemma
Corollary 9. Under the SHOG model, II, recovers the linear discriminant analysis (LDA) solution, II - =-==-1 E (1))/(/E- (p2 - -11)l22).
Proof. Since the LDA solution is Bayes optimal under the HOG model. it is exactly characterized by the top eigen vector of Σ- (1) T. Thus, the Bayes optimal solution target w x Σ- (122 #1), and since II, returns the Bayes optimal linear predictor, following Theorem the above corollary is proven.
Lemma 10 (bias under SHOG). Let Ma be the projection returned by PRO2 The bias ba term Σ=
under our SHOG (Ip U wrll) Here, us
a3 A12 and VT
E Further, when
|ES lop is bounded, and Ma is random rank d projection matrix, ba = 0 1 KL(psllPT).
17

---

Published as a conference paper at ICLR 2024
Proof. From Corollary 9 know that M1 is exactly the rank-1 projection matrix given by the
direction Egg (H2 - f1 /1) [15 (#2 (1)ll)). Therefore
ba 2 (I 2al & (Lp Mi)w = (1 -
(14)
This gives us the first result for us, vr.
For the second result, we note that the KL divergence between multivariate Gaussian distributions convex.
KL(ps] (PT) = (L(p())ps(x ())lp(s))r(x yy)
- <L yp)(p y))
= 0.5 (MI Es)l ΣT)) + 0.5 KL(N(22, Es)l Er))
LIr(ET -
Σ log Af + Σ log XT D.
(15)
Refer to Wainwright (2019) for the final step, where X and X? are the eigenvalues of source and target covariance matrices, respectively. The final term in the above derivation is (tt((E,)) when Eg 0(1). From Lemma we know that under random projections onto dimensions,
ba S L. /1 (d/D)ll
- (w2 p))l(tr(E7")
(16)
where we use Corollary 9 Thus from (66] and (15)
we get our desired bound:
bay D
d
(L(psllpT).
Corollary 11 (tradeoff under SHOG, Corollary Under our SHOG model shift, and conditions for random projection Ma in Lemma 10 the target error (r(wa) (L(ps||pr)+ M when Σx 01).
d
Proof. Direct application the variance result in Lemma and bias result in Lemma TO using the same technique used to prove Theorem
B EXPERIMENTAL DETAILS
B.1 PYTORCH PSEUDOCODE FOR THE PROJECTION STEP
Below, we provide PyTorch pseudocode for the projection step PRO" for binary classification datasets.
def Learn
_earn_ feature _eeatue_spaaeee____e____e_e space_pasis (x, y, num features)
projection worch.nn. Linear (x. shape/ num features) opt corch.opting AdamW (pro paw.meterrrarrrsseeeamao (), 1r-0.01. 01,
weight_decay=0.001
max_steps 100
for i in range (max saa_st)
legits projection (x)
logits (logits, y, reduction=" lcss f.binayycrrosseettrry..bbbb none") mean
opt zer.ger)_grad()
loss backward ste.step() knforce orthogonality;
performing projected gradient descent
Q, R torch. linalg.. (projection... wwwghw.igtt...... detach () .T) projection. weettionnwwigtt.....m wwwght. (Q * torch.diag (R)
feature space projection. weight detach () return feature_space
18

---

Published as a conference paper at ICLR 2024
B.2 ADDITIONAL DATASET DETAILS
4-Way Collages Teney et al. 2021). This binary classification dataset consists of 4- way collages of four images per datapoint, one from each of (1) CIFAR, (2) MNIST, (3) Fushion-MNISS and (4) SVHN. All four image features are completely correlated in the source data, and we consider four target distributions, where only one of the image features are predictive of the label in each target distribution.
Waterbirds (Sagawa et al. 2020). This dataset tasks the model with classifying images of birds either waterbird or landbird. The label is spurious correlated with the background of the image, which is either water or land. There are 4,795 training samples, of which 95% of the data follows the spurious correlation. We use the original training set as the source data and evaluate on different target distributions constructed from the original test dataset: (1) Minority, which contains the test data points that do not follow the spurious correlation, (2) Spurious, containing the points that do, and (3) Balanced, which contains an equal number of points from each of the 4 (bird, background) groups.
CelebA (Liu et al. 2015). Similar to Waterbirds, we use the original training set as source data and evaluate on (1) Minority, (2) Spurious, and (3) Balanced target distributions. In target distributions corresponding the spurious correlation typically used for evaluation (spurious attribute- attribute-gender with label-hair color). Below, in Appendix Cinclude additional results on 4 other variants following the settings used in Lee et al. (2022b): (1) 1 uses slightly open mouth as the label and wearing lipstick as the spurious attribute, (2) CelebA-2 uses attractive as the label and smiling as the spurious attribute, (3) CelebA-3 uses wavy hair as the label and high checkbones as the spurious attribute, and finally (4) CelebA-4 uses heavy makeup as the label and big lips as the spurious attribute.
our main experiments in Sec. 6
Camelyon17 (Bandi et al. 2018). This dataset is part of the WILDS benchmark Koh et al. (2021) and contains medical images where variations in data collection from different hos- pitals induce naturally occurring distribution shifts. We evaluate on target distributions: (1) ID-Test: held out test set of images from the source distribution, and (2) OOD-Test the actual test distribution with a distribution shift due to evaluating data from a different hospital.
Living17 (Santurkar et al. 2020 The task is to classify images into one of 17 animal categorie. This dataset presents a subpopulation shift, in that while the ID and OOD dis- tributions have the same overall classes, they contain different subpopulations. We test on the given test set.
FMoW (Koh et al. (021). This dataset contains satellite images from 5 geographic regions, and the task is the classify the image as one of 62 building or land use types. For the target distribution, we use the subset of the OOD test data belonging to the Africa region.
Pre-trained models and additional training details. We extract penultimate embeddings of all source and target datapoints from a pre- pre-trained backbone. We preprocess all datapoints according to the augmentation used during pre- pre-training, and obtain feature embeddings with eval-mode batch normalization. We cache all embeddings for (backbone, dataset) pair to single file and train our linear models from the cached file. We use VT.IP-ViT--/16 Dooovviskki et al. (2020) in our main experiments, and additionally experiment with ResNet18 He et al. (2016), ResNet50, ResNet50- SWav CCronn et al. (2020), CLIP- models in Appendix C.3 All pretrained models are publicly available online. We train all models using the Adam W optimizer Loshchilov and Hutter 2017) with weight decay 0.01. For all experiments, we perform early stopping with accuracy on held-out target data and report mean and standard deviation across 10 runs.
C ADDITIONAL EXPERIMENTAL RESULTS
C.I ADDITIONAL VISUALIZATIONS FOR SYNTHETIC GAUSSIAN EXPERIMENT
In Fig. 6) we approximate the bias and variance in the synthetic HOG experiment studied in Fig. 3 On the left, for each test distribution (ID, Near OOD, and Far OOD), we plot the relationship be- tween approximate bias (using error at the largest target dataset size) and nullspace norm and find that they have a roughly linear relationship. Thus, this plot empirically supports the connection sup- ported in the theory between bias and the number of features used, as the nullspace norm decreases as the dimension of the feature-space basis increases.
19

---

Published as a conference paper at ICLR 2024
40
Far OOD
Near OOD
-
2
30
ID
Excess Error
30
10
20
20
10
10
010
0
0.0
0.2
0.4
0.6
0.8
1.0
2
8
32
128
512
2048 8192
Nullspace Norm
Target Dataset Size
Visualization of bias and variance in the synthetic homoscedastic Gaussian experi- ment ig. 3 (Left) We approximate bias by the error at the largest target dataset size, and compare to the nullspace norm. The two quantities have roughly linear relationship. (Right) We approximate variance by the difference between the error at each dataset size and the error at the largest. We report the average across the three test distributions. Note on the left plot, ID is easily learned and so the corresponding line is therefore clustered near (0, 0), as the nullspace norm and bias are both near 0.
Figure 6:
C.2 EMPIRICAL ANALYSIS PROJECTED FEATURE SPACE
We begin by observing the empirical properties of the projected feature space learned during the first projection phase of PRO2. The Waterbirds dataset consists of "spurious" groups where the background type (land or water) correlates with the bird type (land water), on which using shortcut feature that relies on background type will perform optimally, as well as "minority" groups in which the correlation does not hold and requires robust feature that focuses the bird itself. On this dataset, we first extract oracle shortcut and robust features by minimizing loss on spurious and minority groups on target data, respectively. These two directions serve as proxies for the optimal classifier on two different target distributions. In addition to PRO?, we also evaluate a random feature extraction method, which simply samples random orthonormal basis for the original embedding space. We plot the nullspace norm of these two features in the subspace spanned by the first k directions, for S k < by ?RO² are more similar to the shortcut feature than the robust feature. Because the orthogonality constraint forces the features to be different from each other, the nullspace norm reduces zero at the highest value & 1024. This experiment shows that the basis learned by PRO² contains both the robust and shorteut features for this dataset, and that the robust and shortcut features emerge even for very low-rank bases (i.e., for small values of d). In contrast, random orthogonal basis only captures these two predictive features when the rank is larger. This indicates that our orthogonal projection approach quickly picks up the most important directions in feature space, which in this correspond to the shortcut feature representing the background and the robust feature representing the type of bird, as discussed in prior work (Sagawa et al.) 2020).
1024 in ??. As expected, we see that the earlier features learned
C.3 USING VARIOUS PRETRAINED BACKBONES
Finally, as PRO² relies on using a pre-trained backbone model that is not fine-tuned to initially extract features, study how different backbones affect performance. In Fig. 7 we plot the accuracy of PRO" using 5 pre-trained backbone models that achieve a range of Image-Net accuracies. We find that PRO² improves significantly with better pre-trained backbones. These experiments demonstrate the promise of the PRO² framework. The quality of pre-trained feature extractors will continue to improve with future datasets and architectures, and PRO² leverages such pre-trate pre-trained backbone models for distribution-shi adaptation in a computationally efficient manner.
C.4 ABLATION ON THE IMPORTANCE OF ENFORCING ORTHOGONALITY
For the purposes of our empirical analysis, we additionally consider a simpler variant that optimizes the projection matrix II with No Constraint orthogonality:
II, arg min ((LL))
(PRo2-NC)
20

---

Published as a conference paper at ICLR 2024
Waterbirds
CelebA
Minerlyy
0.57
0.75
6:6
Test Accuracy
Baanced
Pro-2 on Balanced
0.35
0.9
0.50
0.85
0.35
2r on Spur ous
Pro on Spurious
2.00
0.35
Target Train Data Size
ResNet18 v1
ResNet50_vc v2
ResNet50 SWWV SWAV
VIFB-15,5MAG-EEE_11
ME-L-16-SWAG L2L_vl
Figure 7: Different backbones. We show the accuracy of PRO where we use various pretrained backbones, which are not fine-tuned. PRO² is able to leverage improvements in the backbone with minimal computational overhead.
Waterbirds
CelebA
Camelyon 17
Collages
Test Accuracy
1.00
Average Acress Bists
1.0
Average Across Bists
Average
1.0
Average Across Dists
05
0.90
0.8
05
0.6
12
0.8
12'
10
0.-
12'
12
Target Train Data Size
Pro2-NC
Pro² (ours)
Figure 8: Importance of orthogonality. We show the adaptation accuracy of PRO? compared to PRO2-NC, variant without orthogonality enforced, averaged across the varying target distributions for each dataset.
21

---

Published as a conference paper at ICLR 2024
CelebA-1
CelebA-2
CelebA-3
CelebA-4
Spurious
Sourious
Spurtous
Spurious
1.0
1.0
0.9
9
0.9
0.9-
0.8
0.7
0.7
0.6
0.6
0.6
0.6
Balanced
Balanced
anced
Dalanced
1.0
1.0
1.0
0.9
0.9
0.9
0.9
0.8
0.8-8
Test Accuracy
0.7
0.7
0.7
0.6
0.6
0.6
0.6
0.5
Minority
Minority
Minortty
Minoriyy
1.0
1.0
1.0
1.0
0.9-
0.9
0.9
0.9
0.8
0.8
0.7
0.7
0.6
0.6
0.6
0.8
Average
Average
Average
Average
1.0
1.0
1.0
0.9 0
0.8
0.8
0.8-
8
0.8
0.7
0.7
0.7
0.7
0.6
32
120
32
120
128
Target Train Data Size
Random Projection
DFR (Kirichenko al.)
+ Teney
+ Pro3 (ours)
Figure 9: Main results on additional CelebA variants. We compare 4 different methods for learn- ing features to adapt to target distribution: (1) Random Projection, (2) DFR Kirichenko et al. (022), i.c. standard linear probing, (3) Teney et al. (2021). and (4) PRO?. We report target accura- cies after probing with different target dataset sizes; we report mean and standard deviation across 10 runs. Similar to the trends seen in Fig. 4 PRO² achieves high accuracy in the low-data regime, substantially outperforming both random orthogonal projection and no projection most target dis- tributions on all four datasets.
We compare PRO to PRO--NC in Fig. While PRO2-NC sufficient some scenarios with milder distribution shift, where the shorteut feature continues to be informative, it fails to learn diverse set of predictive features and often only learns shortcut features, often failing on more severe distribution shifts.
C.5 EVALUATION ON ADDITIONAL CELEBA VARIANTS
Finally, in Fig. 9 we supplement our main results in Fig. 4 with additional results from 4 additional variants of CclebA. The takeaways from these results line up with those from Fig. 4 In the few-shot adaptation problem setting, PRO is consistently the most effective, compared to Random Projection, DER Kirichenko et al. (2022), which uses standard linear probing, and Teney et al. (2021).
C.6 ADDITIONAL COMPARISONS
We provide an additional comparison replacing the projection step with PCA (see results below), which performs significantly worse than Pro" on both Waterbirds and CelebA using range of target data sizes. This is expected since PCA an unsupervised approach that does not consider labels, whereas Pro² leverages the given labels from the source data in its projection step.
Additionally, we evaluate against two additional points comparison that use both small amount of labeled target domain data like Pro?. and additionally use unlabeled data: D-BAT (Pagliardini et al. 2022) and DiyDis (Lee et al., 2022b), shown in Table 3 While both methods perform well, Pro" at least matches the performance with any target data size and sometimes even outperforms the other two methods, especially with few datapoints (2) on Waterbirds and more datapoints (32) on CelebA, despite not using any additional unlabeled data, which may be difficult to obtain.
22

---

Published as a conference paper at ICLR 2024
Waterbirds (Balanced)
CelebA (Balanced)
Method
Target Train Data Size (per label)
Method
Target Train Data Size (per label)
2
8
32
2
8
32
PCA
81.5 (1.0)
85.4 (0.5)
90.7 (0.3)
PCA
61.4 (1.8)
65.6 (0.8)
75.2 (0.6)
Pro²
91.8 (0.2)
91.7 (0.4)
92.7 (0.3)
Pro²
76.5 (1.4)
78.3 (0.9)
80.8 (0.9)
Table 2: PCA comparison. PRO² significantly outperforms PCA by leveraging the source data in its projection step.
Waterbirds (Balanced)
CelebA (Balanced)
Method
Target Train Data Size (per label)
Method
Target Train Data Size (per label)
2
32
2
32
D-BAT
90.1 (0.5)
91.5 (0,4)
92.3 (0.2)
D-BAT
76.7 (0.5)
78,4 (0.7)
78.5 (0.7)
DivDis
89.5 (0.8)
91.1 (0.5)
93.0 (0.2)
DivDis
76.2 (0.7)
77.8 (0.5)
77.7 (1.0)
Pro 91.8 (0.2)
91.7 (0.4)
92.7 (0.3)
Pro²
76.5 (1.4)
78.3 (0.9)
80.8 (0.9)
Table 3: Comparison against methods that additionally require unlabeled data. PRO matches and outperforms two state-of-the-arr methods that additionally require unlabeled data.
Finally, in Table 4 we provide the zero-shot transfer performance using the same pre-trained back- bone (ViT-L-16_SW Wwt..SSAG_FBR.. 1 trained on the source distribution. We see that using a small amount of target data, PRO² results in significant gains, particularly OOD distributions, and thus produces consistently strong performance across a range test distributions.
Waterbirds
CelebA
Spurious
Minority
Balanced
Spurious
Balanced
Minority
No Target Data
97.6 (0.1)
75.1 (0.4)
86.1 (0.1)
No Target Data Pro²
97.8 (0.1)
59.9 (0.4)
22.6 (0.5)
Pro
97.9 (0.2)
97.0 (0.3)
94.3 (0.2)
98.0 (0.1)
83.7 (0.3)
96.0 (0.2)
Camelyon17
Collages
ID Test
OOD Test
MNIST
CIFAR
Fashion MNIST
No Target Data
93.7 (0.1)
93.1 (0.0)
No Target Data
75.9 (0.1) 71.8 (0.2)
61.3 (0.1)
Pro?
94.5 (0.
93.5 (0.2)
Pro?
98.0 (0.1)
90.9 (0.3)
85.0 (0.3)
Table 4: Comparison against zero-shot transfer. PRO2 significantly outperforms the zero-shot transfer performance after training the source distribution.
C.7 NUMERICAL VALUES OF EXPERIMENT RESULTS
We include numerical values of experimental results in Figure 4
23

---

Published as a conference paper at ICLR 2024
Train Data Size
Method
Balanced
Spurious
Minority
DFR
82.4 (0.61)
88.3 (0.88)
83.2 (0.34)
Pro"2
94.1 (0.16)
97.9 (0.12)
91.0 (0.30)
Random Projection
82.9 (0.92)
88.4 (0.70)
88.1 (0.61)
8
DFR
86.8 (0.53)
92.3 (0.15)
91.5 (0.19)
Pro"2
94.2 (0.11)
(8.19)
91.2 (0.37)
8
Random Projection
87.8 (0.54)
90.0
90.7 (0.54)
32
DFR
90.4 (0.57)
90.1())13)
94.9 (0.29)
32
Pro"2
94.1 (0.11)
98.2
(0.24)
32
Random Projection
92.9 (0.33)
(0.50)
94.4 (0.44)
128
DFR
94.0 (0.15)
97.3 (0.11)
96.8 (0.32)
128
Pro"2
94.4 (0.17)
98.2 (0.12)
96.9 (0.24)
128
Random Projection
93.8 (0.15)
96.9 (0.23)
96.2 (0.43)
Table 5: Numerical values of Waterbirds results.
Train Data Size
Method
Balanced
Spurious
Minority
2
DFR
60.1 (0.80)
94.5 (0.36)
84.8 (0.46)
Pro"2
77.4 (0.65)
97.8 (0.23)
89.1 (0.51)
Random Projection
59.2 (0.22)
92.8 (0.43)
84.4 (0.49)
8
DFR
66.3 (0.38)
96.8 (0.51)
89.0 (0.15)
8
Pro"2
78.2 (0.17)
98.1 (0.71)
91.3 (0.52)
8
Random Projection
65.6 (0.45)
95.7 (0.58)
89.3 (0.46)
32
DFR
72.3 (0.32)
97.1 (0.33)
94.5 (0.21)
32
Pro"2
79.4 (0.18)
(0.49)
(0.34)
32
Random Projection
72.9 (0.31)
97.3 (0.54)
94.3 (0.42)
128
DFR
78.4 (0.35)
97.9 (0.13)
95.9 (0.68)
128
Pro"2
81.8 (0.73)
98.8 (0.24)
96.4 (0.16)
128
Random Projection
79.2 (0.25)
98.2 (0.17)
97.0 (0.33)
Table 6: Numerical values of CelebA results.
24


Reviews:
[{"review_id": "zk3dDAgeBk", "reviewer": "Reviewer_sNr1", "summary": "This paper deals with the problem of transfer learning with a small amount of target data. It proposes Project and Probe, which first learns a liner projection that maps the pre-trained embedding onto orthogonal directions and then learns a linear classifier on top of the projected features on the small target dataset. The proposed method outperforms prior methods on transfer learning when given very limited target data.", "questions": "See weaknesses.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is clearly written and organized.\n\n- The enhanced sample-efficient generalization of the proposed method is supported by theoretical analysis.", "weaknesses": "- A critical comparison is missing from the experiments: How does the proposed method perform compared to zero-shot transfer learning methods (i.e., no target training data), as cited in related work?\n\n- For reproducibility, it is necessary to include the numerical values of the experimental results in addition to the line charts.", "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": "", "extra_scores": {}}, {"review_id": "m9H5pZUhrs", "reviewer": "Reviewer_Ft2e", "summary": "The paper introduces Project and Probe (PRO^2), a transfer learning method designed for scenarios with limited target data due to distribution shifts. PRO^2 is based on a two-step approach: first, it projects pre-trained embeddings from the source dataset onto orthogonal directions to derive a diverse, non-redundant set of predictive features; next, it trains a linear classifier on these projected features using the target data. Theoretical analyses emphasize the method's favorable bias-variance tradeoff, and experimental results across four datasets demonstrate an improved performance by 5-15% compared to traditional linear probing methods.", "questions": "1) Regarding the empirical assessment, was a consistent hyper-parameter search strategy employed across all the evaluated baselines, or was it exclusively used for the proposed model?\n\n2) There are approaches that exploits Linear Discriminant Analysis (LDA) and Quadratic Discriminant Analysis (QDA) for the adaptation of the head of a pretrained model with success (Shysheya et al. 2022). Can the authors comment on the differences between their method and these methods? While the paper covers the relation with respect to LDA, it does not seem to mention the relation with QDA. Adding those baselines to the empirical evaluation may be beneficial.\n\n3) The authors briefly mention the relation with the previous work of Morwani et al. (2023) in the related work section. This section appears somewhat limited and would benefit from a more in-depth exploration. Could the authors elaborate on the parallels and distinctions between their work and Morwani et al. (2023)?\n\n4) Would the authors be able to provide a detailed analysis of the complexity for the proposed method (e.g. FLOPs and/or MACs)? Understanding complexity is crucial as it essentially represents the computational budget. How does the method's time complexity stand in comparison to leading fine-tuning approaches, such as BiT by Kolesnikov et al. (2020), which adapt the entire model body or FiT (Shysheya et al. 2022), which adapt a subset of the body parameters?\n\n\nReferences\n-----------\n\nMorwani, D., Batra, J., Jain, P., & Netrapalli, P. (2023). Simplicity bias in 1-hidden layer neural networks. arXiv preprint arXiv:2302.00457.\n\nKolesnikov, A., Beyer, L., Zhai, X., Puigcerver, J., Yung, J., Gelly, S., & Houlsby, N. (2020). Big transfer (bit): General visual representation learning. In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part V 16 (pp. 491-507). Springer International Publishing.\n\nShysheya, A., Bronskill, J., Patacchiola, M., Nowozin, S., & Turner, R. E. (2022). Fit: Parameter efficient few-shot transfer learning for personalized and federated image classification. arXiv preprint arXiv:2206.08671.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper stands out in terms of clarity, organization, and overall presentation. It also offers an extensive appendix that provides in-depth coverage of related topics, adding value for the reader.\n\n- The authors have presented a robust theoretical framework that substantiates their approach. They effectively highlight the method's capability to achieve a desirable balance between bias and variance.\n\n- The empirical experiments are detailed and present a wide range of scenarios. While there are certain reservations (addressed below), the breadth and depth of this section are commendable.", "weaknesses": "- The study by Morwani et al. (2023) has previously explored orthogonal projections as a remedy for feature collapse and simplicity bias. This prior exploration somewhat diminishes the uniqueness of the approach presented in this paper.\n\n- Some aspects of the empirical evaluation are unclear and require further details from the authors. See the \"Questions\" section for more details.\n\n- The paper presents a rather limited set of baselines. Given that the experimental setup seems relatively straightforward, it would be advantageous to have a more comprehensive range of baselines. Specifically, a comparative analysis involving methods anchored in LDA and QDA (as pointed out in Shysheya et al., 2022) could enrich the paper.", "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": "", "extra_scores": {}}, {"review_id": "9oOawGtEyF", "reviewer": "Reviewer_fuyr", "summary": "In this paper, the authors propose PROJECT AND PROBE, a lightweight framework consisting of 2 steps: (1) a projection step that extracts a diverse and predictive feature-space basis and (2) a probing step that interpolates between the projected features to efficiently adapt varying target distributions. The core idea is to ensure orthogonality among each component of the feature vector. Each component is utilized by an identical predictor but contains distinct information. Subsequently, these components are employed for predicting the target data. The proposed approach is supported by a theoretical analysis showing that the proposed approach improves sample efficiency.", "questions": "See Section above.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and well-organized.\n- Enforcing feature orthogonality is intuitive and seems suitable for learning features that remain invariant to distribution shifts.\n- The proposed algorithm outperforms the baselines, especially when the sample size of the target data is relatively small.", "weaknesses": "- One major limitation of this work is the fact that the project and probe processes are considered solely in the linear model regime due to the pre-trained feature extraction. Also, the theoretical analysis was conducted with a linear model. how would it  extended to large-scale problems?\n\n- The authors only compare with projection-based baselines. I think comparisons with recent general unsupervised domain adaptation methods are needed. \n\n- While learning diverse/orthogonal features is novel in the context of domain adaptation. There is an active line of research that explores this idea in the standard supervised learning setting, such as [1-7]. I think these methods should be discussed in the related work. \n\n\n[1] Bansal, N., Chen, X., & Wang, Z. Can we gain more from orthogonality regularizations in training deep networks?. Neurips (2018)\n\n[2] Xie, P.; Singh, A.; and Xing, E. P. . Uncorrelation and evenness: a new diversity-promoting regularizer. ICML (2017)\n\n[3] Xie, B.; Liang, Y.; and Song, L.. Diverse neural network learns true target functions. In Artificial Intelligence and Statistics (2017)\n\n[4] Laakom, F., Raitoharju, J., Iosifidis, A., & Gabbouj, M. WLD-Reg: A Data-dependent Within-layer Diversity Regularizer. AAAI (2023)\n\n[5] Cogswell, M.; Ahmed, F.; Girshick, R. B.; Zitnick, L.; and Batra, D. Reducing Overfitting in Deep Networks by Decorrelating Representations. ICLR (2016)\n\n[6] LLaakom, F., Raitoharju, J., Iosifidis, A., & Gabbouj, . Learning distinct features helps, provably. ECML (2023).\n\n[7] Zbontar, J., Jing, L., Misra, I., LeCun, Y., & Deny, S. (2021, July). Barlow twins: Self-supervised learning via redundancy reduction. ICML (2021)", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)