2023/05/30 by Yujia Bao, Bao, Yujia, Theofanis Karaletsos +1 · 1 citation
Biochemistry, Genetics and Molecular Biology · Computer Science · #AI in cancer detection #Artificial Intelligence (cs.AI) #Cell Image Analysis Techniques #Computation and Language (cs.CL) #Computer Vision and Pattern Recognition (cs.CV) #Digital Imaging for Blood Diseases #FOS: Computer and information sciences
paper · pdf · doi:10.48550/arxiv.2305.19402
openalex publication_date 2023/05/30 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28
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.