Paper: Contextual Vision Transformers for Robust Representation Learning (ICLR-2024)

Reviewer comment (Reproducibility):
How to chose and define the "in-context" prompt is unclear.

Author response:
**Context Choice and Definition:** We carefully select the group membership variable, as depicted in Figure 2, ensuring datasets share a conditional dependency structure with their environment. For instance, pathology images can vary based on the staining procedures or machines used by different hospitals. To address this, we chose the hospital identity indicator as the context variable, which has shown to significantly improve model performance in our experiments.

Reviewer comment (Novelty):
While the paper is well-structured and well-written, it would be beneficial to include more detailed comparisons with related work to highlight the novelty of the proposed approach.

Author response:
**Comparisons with Related Work:** Different from previous work, ContextViT explores token conditioning based on the group membership of the data. ContextViT introduces two novel technical enhancements:
+ Context inference network that maps example sets from the same group into a context token;
+ Layer-wise context conditioning that allows the ViT to integrate context at various network depths, enhancing generalization. 
+ We also establish a principled mathematical link from ContextViT to in-context learning for distribution shifts. 

We have updated our related work section to better highlight these differences.

Reviewer comment (Writing):
In the "Out-of-Distribution Generalization (Pathology Images)" section, it's not entirely clear what "linear probing accuracy" means and how it relates to out-of-distribution generalization. A more in-depth explanation of this metric would improve the clarity of the paper.

Author response:
**Linear probing:** Linear probing, a common protocol for evaluating self-supervised learning [1], involves training a linear classifier on top of fixed pre-trained ViT and ContextViT backbones to predict a target label, such as the presence of cancer. In the Cameylon17 dataset, the classifier is trained on images from three training hospitals and tested on a fourth, unseen hospital. Successful generalization to the held-out hospital indicates aligned feature distributions. We have expanded our discussion on this evaluation metric to provide further clarity.

Reviewer comment (Evaluation):
Are there any specific use cases or domains where ContextViT is particularly well-suited, and are there any limitations or scenarios where it may not perform as effectively?

Author response:
**Use Cases and Limitations:**
+ ContextViT is particularly well-suited for imaging applications where data are grouped or batched. 
+ However, if the group membership variables are mis-specified and fail to capture the environmental context, ContextViT's performance may not exceed that of a regular ViT. For example, a randomly assigned context variable would not improve the model's performance. The effectiveness of ContextViT relies on the user's understanding of the dataset to accurately specify group memberships. 
+ We appreciate the reviewer's insight on this matter and have included a discussion on the potential complexities and dependencies that could affect generalization in our paper.

Reviewer comment (Evaluation):
Could the authors provide more insights into how ContextViT's approach to handling structured variations and distribution shifts could be applied in practical applications outside of the ones discussed in the paper?

Author response:
**Practical Applications and Structured Variations:**
+ This links very well to the previous question. In practice, users want to consider invariant factors of variation per group that systematically shift across groups. For example, shared environments over datasets or global variables in the plated context (see Figure 2 for a graphical model interpretation).
+ We also note that oftentimes this can be an exploratory process where people need to train specific models with different contexts and see whether the resulting models explain away the distributions shifts when there is no obvious choice of the context variable.
+ Lastly, the question interfaces with our desire for a more formal framework in future work based on statistical tests on top of ViT embeddings to detect biases over various covariates and suggest useful context,  which would be a useful tool and would link this work more to causal inference.

Reviewer comment (Novelty):
The novelty of this work is limited.

Author response:
**Novelty:** In addition to our unique perspective on token conditioning based on the data groups, ContextViT is complemented by two novel technical enhancements:
+ Context inference network maps sets of examples from the same group into a context token, enabling ContextViT to generalize to unseen groups during training.
+ Layer-wise context conditioning enables the ViT to integrate context at various depths within the network, not just at the input layer, which significantly improves generalization.
+ Finally, we provide a principled mathematical linking from ContextViT to in-context learning for distribution shift.

Reviewer comment (Novelty):
The intrinsic difference between this work and visual prompting [1] is unclear. It seems that visual prompting can also fit this OOD scenario.

Author response:
**Related work:** We are grateful for the additional related work highlighted by the reviewer. We have incorporated these references into our paper to provide a more comprehensive context for our contributions.
+ Visual prompt tuning: Unlike visual prompt tuning, which learns extra tokens for each downstream task, our work defines the extra token based on the group membership of the data. This distinction is crucial for enabling our model to adapt to new distributions without task-specific tokens.
+ Domain prompt learning: Similar to visual prompt tuning, domain prompt learning generates a consistent output prompt for each task, as the text label set remains unchanged. In our work, we don’t consider text inputs as they are not meaningful in cross-distribution generalization. For instance, the CLIP text encoder would not produce meaningful representations using the group name such as  “plate BR00116991” and “plate BR00116993” or “hospital A” and “hospital B”. Moreover, ContextViT's extra token is dynamically derived from the group membership, allowing for generalization beyond the scope of fixed tasks.

Reviewer comment (Novelty):
The key idea of this work is similar to [2], which also uses a network to predict the context/domain tokens.

Author response:
**Related work:** We are grateful for the additional related work highlighted by the reviewer. We have incorporated these references into our paper to provide a more comprehensive context for our contributions.
+ Visual prompt tuning: Unlike visual prompt tuning, which learns extra tokens for each downstream task, our work defines the extra token based on the group membership of the data. This distinction is crucial for enabling our model to adapt to new distributions without task-specific tokens.
+ Domain prompt learning: Similar to visual prompt tuning, domain prompt learning generates a consistent output prompt for each task, as the text label set remains unchanged. In our work, we don’t consider text inputs as they are not meaningful in cross-distribution generalization. For instance, the CLIP text encoder would not produce meaningful representations using the group name such as  “plate BR00116991” and “plate BR00116993” or “hospital A” and “hospital B”. Moreover, ContextViT's extra token is dynamically derived from the group membership, allowing for generalization beyond the scope of fixed tasks.

Reviewer comment (Experiments):
The comparison in experiment section is insufficient.

Author response:
**Comparison in the experiment section:**
+ ContextViT has been evaluated as a plug-and-play enhancement for three established pre-trained ViTs (DINO, SWAG, CLIP) and has consistently outperformed recent fine-tuning baselines in generalization tasks.
+ In self-supervised representation learning, ContextViT has set a new state-of-the-art on the Camelyon17-WILDS dataset, outperforming 26 other baselines. The full comparison is available on the WILDS leaderboard (https://wilds.stanford.edu/leaderboard).

Reviewer comment (Presentation):
Lack of visualization of the learned context token, which shows the difference of context tokens of different groups.

Author response:
**Analysis of the model:**
+ We have conducted a thorough comparison of ContextViT with its non-context-conditioned ViT counterparts (Tables 1, 2, and 3).
+ Our ablation studies and runtime comparisons (Table 4) carefully examine the impact of layer-wise context conditioning and various implementations of the context inference network.
+ Figures 5 (in the appendix)  explores ContextViT's effectiveness with limited data for context token inference.
+ Figure 7 (in the appendix)  provides a visualization of the context token. We observe that the similarity between the context tokens are closely related to the similarity of the input images.
+ Figure 6 (in the appendix) provides a visual representation of the attention heatmap produced by ContextViT, illustrating its interpretability.

Reviewer comment (Novelty):
The paper is simple and effective, but its novelty is unfortunately limited, and analysis for the insight of this approach is absent.

Author response:
**Novelty:** In addition to our unique perspective on token conditioning based on the data groups, ContextViT is complemented by two novel technical enhancements:
+ Context inference network maps sets of examples from the same group into a context token, enabling ContextViT to generalize to unseen groups during training.
+ Layer-wise context conditioning enables the ViT to integrate context at various depths within the network, not just at the input layer, which significantly improves generalization.
+ Finally, we provide a principled mathematical linking from ContextViT to in-context learning for distribution shift.

**Analysis of the model:**
+ We have conducted a thorough comparison of ContextViT with its non-context-conditioned ViT counterparts (Tables 1, 2, and 3).
+ Our ablation studies and runtime comparisons (Table 4) carefully examine the impact of layer-wise context conditioning and various implementations of the context inference network.
+ Figures 5 (in the appendix)  explores ContextViT's effectiveness with limited data for context token inference.
+ Figure 7 (in the appendix)  provides a visualization of the context token. We observe that the similarity between the context tokens are closely related to the similarity of the input images.
+ Figure 6 (in the appendix) provides a visual representation of the attention heatmap produced by ContextViT, illustrating its interpretability.

Reviewer comment (Theory):
The paper mentioned that the proposed method applies the concept of in-context learning in vision transformer. However, in my opinion, in-context learning is a kind of few-shot learning, which predicts based on the (data, label) pair of a few samples, unlike the usage of all the dataset-c data (or a batch of the data) in this paper. The method looks like a summarization of the dataset information and then makes the prediction based on that summarization.

Author response:
**Concept of in-context learning:** We appreciate your input on this matter. To clarify, we have updated our methods section to distinguish between data conditioning and data+label conditioning. In-context learning is a wide catch term for various types of data conditioning. We consider the case of in-context learning for generalization under group membership of the data and derive our method under that assumption. This is for instance when we condition the model on different environments, a very common assumption. This consideration is also driven by the reality that labels for test distributions are often unavailable, whereas data is plentiful. The context inference network is designed to summarize dataset information, enabling the vision transformer to adapt its representation of the input image accordingly.

Reviewer comment (Evaluation):
The method requires a lot of distribution-c data at the inference stage and increases the inference overhead.

Author response:
**Requires a lot of distribution-c data at the inference stage:**
+ While in an ideal world we would need all the distribution-c data to derive the context token, we empirically tested the data hunger of the model. In Figure 5, we demonstrate ContextViT's performance across a range of testing batch sizes, from 1 to 512. The model's ability to infer context improves with more examples; however, it still achieves state-of-the-art performance with as few as 8 examples, outperforming all existing benchmarks on the WILDS leaderboard (https://wilds.stanford.edu/leaderboard).
+ We also highlight that in many real-world settings, such as new hospitals, devices, users, or biological experiments, testing data is naturally batched. ContextViT is well-suited for these practical applications.

Reviewer comment (Novelty):
The oracle-context model is very similar to some prompt tuning works, like Visual Prompt Tuning & Prompt Learning for Vision-Language Models, but these works are not discussed in the paper.

Author response:
**Related work:** We are grateful for the identification of missing references. These have now been incorporated into our revised paper. Our approach is distinct in its use of group membership to define the context token, coupled with two innovative technical contributions:
+ Context inference network maps sets of examples from the same group into a context token, allowing ContextViT to generalize to new groups unseen during training.
+ Layer-wise context conditioning enables the ViT to integrate context at various network depths, not just at the input layer, which significantly improves generalization.
+ Finally, we provide a principled mathematical linking from ContextViT to in-context learning for distribution shift.

Reviewer comment (Experiments):
The view and impact of this paper are limited, as the method seems to focus on improving the performance of datasets that contain several distinct groups. Although improvements are demonstrated on some specific datasets, the improvement in general image tasks is still unclear.

Author response:
**Scope of the paper:** We agree that ContextViT is tailored for scenarios with grouped datasets, which is a common occurrence in real-world applications where data often arrives in batches or groups. The focus of this paper is on enhancing robustness across distribution shifts. We show that instead of having to learn one model per distribution, we can share the parameters across the tasks and utilize the context inference network to inform the model of the specific distribution. We separately show that by doing so and learning a context inference network we can also learn to generalize to new test distributions effectively. 

**Improvement in general image tasks:**
+ In this paper, we have included three widely-recognized imaging benchmarks from WILDS and a recently released cell imaging benchmark released by the Broad Institute.
+ The reviewer asks us to focus on general image tasks. We highlight that instead of going for artificial benchmarks, we went for real world datasets that exhibit specific problem structure that is not tackled by more general machinery. We would ask the reviewer to see the value of overcoming unsolved problems rather than iterating on known problems.
+ To be more specific, in many practical applications (especially scientific domains such as biological discovery), being able to deal with batch effects is a key hindrance to the successful application of representation learning. This is exemplified by our results on the recent JUMP-CP cell painting dataset released by the Broad institute of Havard and MIT, for which ContextViT consistently improves the generalization across different data distributions. We also extensively tested ContextViT on pathology imaging such as Camelyon17-WILD. This is a key problem as real world applications like medical imaging often suffer from distribution shifts in batches due to the employment of different staining procedures or machines across hospitals. We hope we have demonstrated the importance of the task that we tackle to the reviewer.

Reviewer comment (Experiments):
It is suggested to widely evaluate the framework on other popular datasets and tasks or extend related techniques to improve the capability of transfer learning from one task to other tasks.

Author response:
**Improvement in general image tasks:**
+ In this paper, we have included three widely-recognized imaging benchmarks from WILDS and a recently released cell imaging benchmark released by the Broad Institute.
+ The reviewer asks us to focus on general image tasks. We highlight that instead of going for artificial benchmarks, we went for real world datasets that exhibit specific problem structure that is not tackled by more general machinery. We would ask the reviewer to see the value of overcoming unsolved problems rather than iterating on known problems.
+ To be more specific, in many practical applications (especially scientific domains such as biological discovery), being able to deal with batch effects is a key hindrance to the successful application of representation learning. This is exemplified by our results on the recent JUMP-CP cell painting dataset released by the Broad institute of Havard and MIT, for which ContextViT consistently improves the generalization across different data distributions. We also extensively tested ContextViT on pathology imaging such as Camelyon17-WILD. This is a key problem as real world applications like medical imaging often suffer from distribution shifts in batches due to the employment of different staining procedures or machines across hospitals. We hope we have demonstrated the importance of the task that we tackle to the reviewer.

Reviewer comment (Experiments):
The method should be compared with more related works.

Author response:
**Evaluation and baselines:**
+ ContextViT has been tested as a plug-and-play enhancement for three well-established pre-trained ViTs (DINO, SWAG, CLIP) in a supervised learning context. It has consistently outperformed two recent fine-tuning baselines in terms of generalization.
+ In the realm of self-supervised representation learning, we have considered both the in-distribution generalization and out-of-distribution generalization. In the first, ContextViT consistently outperformed ViT across all testing configurations. In the second, ContextViT has achieved a new state-of-the-art on the Camelyon17-WILDS dataset, surpassing 26 other baselines. Due to space constraints, we only included the top baselines in Table 3. The complete list of baselines is available on the WILDS leaderboard (https://wilds.stanford.edu/leaderboard).

Reviewer comment (Novelty):
Despite the proposed contextual learning paradigm, the technical contributions in this paper are limited and not novel enough.

Author response:
**Technical contributions:** We recognize prior work on conditioning ViTs with extra tokens in our related work section. Our approach uniquely leverages group membership to define the context token and introduces two novel technical enhancements:
+ The context inference network effectively maps sets of examples from the same group into a context token. This enables ContextViT to generalize on new groups that it hasn’t seen during training, as demonstrated by our empirical studies. 
+ Layer-wise context conditioning allows the ViT to utilize context at multiple levels of the network, rather than just the input layer, enhancing generalization capabilities.
+ Finally, we provide a principled mathematical linking from ContextViT to in-context learning for distribution shift.

Reviewer comment (Presentation):
Figure 1 is unclear and somehow misleading. The source of the context (where those images come from) and the function (input, output) of the inference model should be labeled. It is strongly suggested to redo this figure.

Author response:
**Clarity:**
  + Based on your suggestions, we have revised the figure to enhance its clarity and readability.

Reviewer comment (Presentation):
The end of page 5 is missing.

Author response:
**Clarity:**
  + The end of page 5 continues on page 6 (under Table 1).

Reviewer comment (Experiments):
What would the performance be if this framework is applied to a large dataset that was combined with several small datasets?

Author response:
**Application to large datasets with small sub datasets:** The JUMP-CP benchmark, comprising multiple data plates from various perturbations, serves as an example of combining "small datasets".

Reviewer comment (Theory):
If the subgroups of the data are unknown, is there any way to benefit from the proposed framework?

Author response:
**Without knowledge of the group membership, is there any way to benefit from the proposed framework?**
+ For this paper, we assume that the indicator function of the group membership is known during training, a realistic assumption given that data is often generated in group fashion in the real world. In future work, one might strive to infer latent group membership from the data, but we don’t tackle that in this submission.
+ In Figure 5 (in the appendix), we explored the performance of ContextViT with a testing batch size of one. This would be equivalent to the situations where we don’t assume knowledge of the group membership at test time. We have demonstrated that this self-conditioning still improves the performance on top of the standard ViT.