Paper: Feature Accompaniment: Is It Feasible to Learn Out-of-Distribution Generalizable Representations with In-Distribution Data? (ICLR-2024)

Reviewer comment (Evaluation):
I am unconvinced that the empirical results of section 2 imply the existence of a failure mode related to the feature learning process. The empirical result may not necessarily be due to nonlinear feature-learning dynamics in this experiment, but rather just that pre-trained CLIP models contain features covering a much larger data distribution than is captured by models distilled on ImageNet. In particular, if you distill a model from CLIP on the ImageNet training set, are some CLIP features that are not represented in ImageNet not likely to be left out? These features could still be helpful in OOD classification. For example, the OOD bird and car could have core features that look different from the core features seen in-distribution. Pretrained models may contain these features, while distilled models may not have learned them if they do not appear in the training set.

Author response:
Thank you for your detailed and insightful review! Your feedback has led to some fruitful discussions among the authors and is helpful for us to improve our paper. In general, there are three major concerns raised in your review:

- Connections between feature accompaniment and the empirically observed OOD generalization failure and the significance of our results.

- Can the empirical observations in Sec. 2 be explained by the fact that CLIP may have learned some OOD core features that are absent in the ID data?

- Discussion on how feature accompaniment relates to the simplicity bias of neural networks.

For the second concern, we believe that **our experimental results cannot be solely explained by CLIP having the ability to extract OOD core features due to large-scale pre-training**. This is due to our linear probe evaluation protocol: first train the linear probe on top of given (CLIP/distilled) representations _**on the training set (ID data)**_, then test its accuracy on both ID and OOD test sets. In other words, _this linear probe still needs to solve an OOD generalization task_---although with inputs as CLIP/distilled representations instead of raw images.

- Having this in mind, now assume that the distribution shift robustness of CLIP is attributed to it having the ability to extract OOD core features from OOD images. If this is indeed the case, then the linear probe based on CLIP representations should **not** exhibit very good OOD generalization performance. This is because those OOD features are **absent** in the ID images (and hence also absent in the CLIP representations of ID images as well) and thus the linear probe trained using only ID data is **not** able to learn to use those OOD features to achieve OOD generalization.

- As a piece of empirical evidence that supports the above explanation, we have visualized the CLIP representations for both ID and OOD images (see Fig. 9 in Sec. G for details). The visualization shows that ID and OOD representations exhibit the geometry where the distribution shift on the features is roughly in a direction **orthogonal** to the classification boundary, so it is not likely that the CLIP's advantage in robustness is due to learning "more" OOD features that provide additional information for prediction.

Reviewer comment (Theory):
Could there be discussion on how feature accompaniment relates to the previous studies on simplicity bias and gradient starvation [1,2], which find that networks rely on simple features and ignore more complex features? In particular, work on Gradient Starvation [2] suggests that an increase in strength of a simpler feature inhibits the learning of other more complex features. Are these results contradictory to those suggested by feature accompaniment?

Author response:
Thank you for your detailed and insightful review! Your feedback has led to some fruitful discussions among the authors and is helpful for us to improve our paper. In general, there are three major concerns raised in your review:

- Connections between feature accompaniment and the empirically observed OOD generalization failure and the significance of our results.

- Can the empirical observations in Sec. 2 be explained by the fact that CLIP may have learned some OOD core features that are absent in the ID data?

- Discussion on how feature accompaniment relates to the simplicity bias of neural networks.

For the third concern, we believe that feature accompaniment and simplicity bias are two **complementary** feature learning biases of neural networks that may manifest in different data distributions: both works [1] and [2] consider a data model involving multiple features that are **equally predictive** for the task; the network may then only learn some of them due to the simplicity bias. By contrast, the background features in our data model are assumed to be **not predictive at all** for the task, yet we theoretically show that the network may still learn them due to feature accompaniment---note that _this does not contradict the simplicity bias_ since the resulting solution can still be "simple", e.g., learning sparse core/background features. Instead, **we believe that feature accompaniment as a novel inductive bias of neural networks helps extend our general understanding of the feature learning biases of neural networks**. Technically, we would also like to note that the theory in [1] only considers a specific dataset (LSN dataset with one linear coordinate and one 3-slab coordinate) while we consider a more general data model, and the theory in [2] is under the neural tangent kernel regime, in which the network is essentially linearized and the features are **not** learned (they stay at random initialization).

Reviewer comment (Theory):
I think the experiment part in Section 2 is a bit disconnected from the theoretical part in Section 4. They consider different settings and different learning objective. I don't think theory in Section 4 can explain the experimental results in Section 2. I know that Section 2 is probably just a starting point of studying whether OOD-generalizable representations are learnable, so this's okay. But I think it could be better to connect them more in the writing.

Author response:
Thank you for your positive review and your approval of our work! We are very glad to know that you found our theoretical results novel, foundamental and may benefit understanding the learning process of neural networks in other scenarios. In general, your review has raised two major concerns:

- The experiment part and the theoretical part seem a bit disconnected.

- What will happen if we have **multiple** training domains (with different background feature distributions) instead of only one training domain?

Since the first concern is common among all reviewers, we provide a general response to it in the "General response to all reviewers" panel, along with a discussion on the significance of our contributions. Please see our general response for more details.

Reviewer comment (Theory):
In the empirical part, this work treats "good OOD linear probing performance" as "good generalization representation" (Figure 1). The principle here is "rich-representation"[1][2]. I suggest the author clarify the principle.

Author response:
- Your review mentioned that _"In the empirical part, this work treats 'good **OOD linear probing** performance' as 'good generalization representation'. The principle here is 'rich-representation' [1,2]."_ We are aware of the insightful works by Zhang et al. [1,2] and would like to clarify that **our evaluation protocol is _not OOD linear probing_** (i.e., training and testing the linear probe on both OOD data, similar to the linear probing protocol in [2]), **but _OOD testing of the linear probe_** (i.e., training the linear probe on **ID** data and testing it on OOD data). Note that this is a **fundamental distinction**: in OOD linear probing, the linear probe essentially solves an **ID generalization** task since its training and testing are on the **same** distribution---in this case, the richness of the representation indeed plays a key role (under the theoretical framework of [2]). However, in our experiments, the linear probe still needs to solve an **OOD generalization** task, albeit with inputs as some pre-trained representations rather than raw images. In this setting, _better linear probing performance does not necessarily mean a 'richer' representation_ (see below for more discussion).

Reviewer comment (Theory):
The author treats -- a supervised pretrained model on Imagenet dataset with more objective information (i.e. prediction the representation of a CLIP model) -- as the **oracle** objective function. By comparing this model (pretrained on Imagenet) with CLIP (pretrained on a large dataset), the author concludes that *"OOD generalizable representations may not be learnable using only ID data without explicitly taking into account the inductive biases of the model or the task."*

Author response:
- Given the above clarification, "an OOD-generalizable representation (oracle representation)" in our experiments should be viewed as "a representation based on which an ID-trained (linear) classifier can generalize OOD". In this regard, **the CLIP representation is indeed a proper approximation of the "oracle representation"** since it has been known for exhibiting strong robustness to distribution shifts [3]. Our experiments thus **do support our conclusion** that "OOD-generalizable representations may not be learnable using only ID data".

Reviewer comment (Experiments):
The comparison between the Imagenet pretrained model and the CLIP model does not support the stated conclusion. From the rich-representation's principle (which is used in the linear probing experiment), OOD linear probing benefits from a representation that contains diverse and simple features. CLIP, pretrained on a large dataset, contains rich features, but also uses more data. It is possible that the model pretrained on Imagenet is already the best possible Imagenet-pretrained model in terms of OOD linear probing performance. Therefore, the CLIP model should not be assumed to be the achievable upper bound for other Imagenet-pretrained models.

Author response:
- **On the rich-representation interpretation:** as mentioned in your review, a potential explanation of the generalizability of the CLIP representation is its richness (in the sense that it _"contains a rich set of features that could be redundant in-distribution but crucial out-of-distribution"_, taken from your review). We **do agree** that this representation richness contributes to generalizability and hypothesize that this may be one of the reasons for the distilled model outperforming standard models. However, _if the generalizability of CLIP can be solely attributed to its richness, then the distilled model _should also achieve good OOD generalization_ since it **directly learns**_ this rich representation during training (instead of learning the representation by a _proxy_ classification task, which may fail to capture _all_ useful features, as shown by [1,2]). However, our experiments suggest that this is **not** the case. Instead, we believe that **feature accompaniment can explain the remaining large gap between distilled models and CLIP**; please see our general response for more details.

Reviewer comment (Experiments):
The Imagenet pretrained model is not an **oracle** in terms of rich-representation. Compared with Imagenet's 1k target categories, it does contain more supervision information (with the help of CLIP and CLIP's pretraining dataset). However, this does not address how it would compare to a model trained on 22k target categories, for example.

Author response:
- **On the "oracle" objective:** Thank you for your clarification on your comment. Yes, we believe that our representation distillation objective **is** "oracle" in terms of representation learning (if we do not have additional prior knowledge of the model/the task and thus treat the model as a black box). We can actually argue this mathematically: consider each model $h$ to be composed of an encoder $\Phi$ and a (linear) classifier $f$. Given each input $x$ (e.g., an image), we treat $\Phi(x)$ as **the representation of $x$** and $f$ as the classifier on top of the representations extracted by $\Phi$. Given a training set $(X,Y) = {(x_i,y_i)}_{i=1}^N$, **the aim of an ideal representation learning process is that the learned representation $\Phi(x_i)$ recovers some oracle representation** $\Phi^*({x_i})$ **for every** $x_i \in X$, where we approximate $\Phi^*$ using CLIP. In this regard, our representation distillation objective (i.e., $\mathrm{minimize}\  \lVert \sum_i \Phi(x_i) - \Phi^*(x_i) \rVert_2^2$ ) is indeed consistent with the above idealization since its minimizer does uniquely recovers $\Phi^*(x)$ for every $x$---we thus refer to it as the **oracle objective** in the above sense. **By contrast**, existing objectives perform representation learning using proxy tasks such as a classification objective defined over the final output of $h(x)$, or other auxiliary objectives based on $h(x)$ and/or $\Phi(x)$. However, minimizing such objectives does **not** guarantee the unique recovery of $\Phi(x) = \Phi^*(x)$ as guaranteed by our oracle objective. Of course, even by minimizing the oracle objective, we cannot guarantee $\Phi = \Phi^*$ since we only have $\Phi(x) = \Phi^*(x)$ for a finite training set, which leads to the remaining gap between the distilled model and CLIP---our theoretical results suggest that feature accompaniment plays an important role in this gap when $\Phi$ is instantiated as a non-linear neural network trained by SGD.

   **Empirically**, in our preliminary experiments, we have also tested using both the original labels and the CLIP representations in representation learning, as the ablation mentioned in your review. Our results suggested that this variant indeed does **not** outperform only using CLIP representations. Due to time limitations, we will add those results to the camera-ready version of the paper if it is accepted.

Reviewer comment (Theory):
The theoretical section did not discuss the relationship between works about SGD and features, e.g., [3][4][5].

Author response:
- **Comparison with work on neural network features:** To our knowledge, our discovery of feature accompaniment is novel, and **no** existing theoretical study on SGD-trained neural networks has explored this feature learning bias. As shown by the cited works in your review and **Reviewer jEjy**, one of the major focuses of existing theory on the inductive biases of neural networks is the **simplicity bias**, i.e., the bias of finding solutions that are "simple" or sparse. Feature accompaniment **complements existing observations** by showing that neural networks also have the bias of learning "_irrelevant_ features" in certain tasks (note that this does not contradict with sparsity).