Contextual Vision Transformers for Robust Representation Learning

TL;DR

This paper presents ContextViT, a novel vision transformer architecture that incorporates group-specific context tokens and inference mechanisms to enhance robust image representation learning under distribution shifts.

Contribution

Introduction of ContextViT with context tokens and inference network, enabling improved out-of-distribution generalization and stable feature learning across diverse applications.

Findings

Improved OOD generalization in supervised fine-tuning on iWildCam and FMoW.

Outperforms standard ViT in pathology and microscopy imaging benchmarks.

Effective in self-supervised learning for stable feature extraction.

Abstract

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.