Contextual Vision Transformers for Robust Representation Learning

We thank the reviewer for their detailed review and constructive feedback.

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.

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.

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).

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.

We trust that this rebuttal has addressed the points raised by the reviewer. We are thankful for your time and effort and are open to any further questions.

We thank the reviewer for their detailed review and constructive feedback.

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.

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.

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.

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.

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.

We trust that this rebuttal has addressed the points raised by the reviewer. We are thankful for your time and effort and are open to any further questions.

References:

1. Chen, Ting, et al. "A simple framework for contrastive learning of visual representations." International conference on machine learning. PMLR, 2020.

We thank the reviewer for their detailed review and constructive feedback.

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.

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.

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.

We believe this rebuttal has addressed the points raised by the reviewer. We are thankful for your time and effort and are open to any further questions.

We thank the reviewer for their detailed review and constructive feedback.

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.

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).

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.

Clarity:

• Based on your suggestions, we have revised the figure to enhance its clarity and readability.
• The end of page 5 continues on page 6 (under Table 1).

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".

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.

We trust that this rebuttal addresses the concerns raised. We are grateful for your time and effort and welcome any further questions.

Dear Authors,

Thank you for your effort and revisions. However, I remain unconvinced about the broader impact of your method. Please provide more substantial evidence of its applicability to general image recognition tasks. Additionally, I agree with the other reviewers about the need for greater novelty in your approach. As it stands, I maintain my score and believe the paper is currently below the acceptance threshold.

Best regards,

Reviewer