# Contextual Vision Transformers for Robust Representation Learning

arXiv: https://arxiv.org/abs/2305.19402 · HTML: https://mlsanity.com/paper/2305.19402/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 64 references that are in the map.
- First posted on arXiv 30 May 2023.

## Builds on (most cited first, 30 of 42)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Improved Regularization of Convolutional Neural Networks with Cutout](https://mlsanity.com/paper/1708.04552/) (2017)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Transformers in Vision: A Survey](https://mlsanity.com/paper/2101.01169/) (2021)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Albumentations: fast and flexible image augmentations](https://mlsanity.com/paper/1809.06839/) (2018)
- [Improving robustness against common corruptions by covariate shift adaptation](https://mlsanity.com/paper/2006.16971/) (2020)
- [MEMO: Test Time Robustness via Adaptation and Augmentation](https://mlsanity.com/paper/2110.09506/) (2021)
- [Functional Map of the World](https://mlsanity.com/paper/1711.07846/) (2017)
- [Masked Siamese Networks for Label-Efficient Learning](https://mlsanity.com/paper/2204.07141/) (2022)
- [Revisiting Batch Normalization For Practical Domain Adaptation](https://mlsanity.com/paper/1603.04779/) (2016)
- [Evaluating Prediction-Time Batch Normalization for Robustness under Covariate Shift](https://mlsanity.com/paper/2006.10963/) (2020)
- [Efficient Object Localization Using Convolutional Networks](https://mlsanity.com/paper/1411.4280/) (2014)
- [Adaptive Risk Minimization: Learning to Adapt to Domain Shift](https://mlsanity.com/paper/2007.02931/) (2020)
- [On the Robustness of Vision Transformers to Adversarial Examples](https://mlsanity.com/paper/2104.02610/) (2021)
- [Finetune like you pretrain: Improved finetuning of zero-shot vision models](https://mlsanity.com/paper/2212.00638/) (2022)
- [Multi-class Token Transformer for Weakly Supervised Semantic Segmentation](https://mlsanity.com/paper/2203.02891/) (2022)

## Built on by (most cited first, 10 of 10)

- [How to Build the Virtual Cell with Artificial Intelligence: Priorities and Opportunities](https://mlsanity.com/paper/2409.11654/) (2024)
- [Channel Vision Transformers: An Image Is Worth 1 x 16 x 16 Words](https://mlsanity.com/paper/2309.16108/) (2023)
- [ContriMix: Scalable stain color augmentation for domain generalization without domain labels in digital pathology](https://mlsanity.com/paper/2306.04527/) (2023)
- [Morphological Profiling for Drug Discovery in the Era of Deep Learning](https://mlsanity.com/paper/2312.07899/) (2023)
- [C3R: Channel Conditioned Cell Representations for unified evaluation in microscopy imaging](https://mlsanity.com/paper/2505.18745/) (2025)
- [Context is Environment](https://mlsanity.com/paper/2309.09888/) (2023)
- [Towards Context-Aware Domain Generalization: Understanding the Benefits and Limits of Marginal Transfer Learning](https://mlsanity.com/paper/2312.10107/) (2023)
- [EDADepth: Enhanced Data Augmentation for Monocular Depth Estimation](https://mlsanity.com/paper/2409.06183/) (2024)
- [Out-of-distribution evaluations of channel agnostic masked autoencoders in fluorescence microscopy](https://mlsanity.com/paper/2503.19149/) (2025)
- [Who Needs Labels? Adapting Vision Foundation Models With the Metadata You Already Have](https://mlsanity.com/paper/2606.05107/) (2026)

## Abstract (arXiv)

We introduce Contextual Vision Transformers (ContextViT), a method designed to generate robust image representations for datasets experiencing shifts in latent factors across various groups. Derived from the concept of in-context learning, ContextViT incorporates an additional context token to encapsulate group-specific information. This integration allows the model to adjust the image representation in accordance with the group-specific context. Specifically, for a given input image, ContextViT maps images with identical group membership into this context token, which is appended to the input image tokens. Additionally, we introduce a context inference network to predict such tokens on-the-fly, given a batch of samples from the group. This enables ContextViT to adapt to new testing distributions during inference time. We demonstrate the efficacy of ContextViT across a wide range of applications. In supervised fine-tuning, we show that augmenting pre-trained ViTs with our proposed context conditioning mechanism results in consistent improvements in out-of-distribution generalization on iWildCam and FMoW. We also investigate self-supervised representation learning with ContextViT. Our experiments on the Camelyon17 pathology imaging benchmark and the JUMP-CP microscopy imaging benchmark demonstrate that ContextViT excels in learning stable image featurizations amidst distribution shift, consistently outperforming its ViT counterpart.