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.

§ INTRODUCTION

In recent years, Vision Transformers (ViTs) have emerged as a powerful tool for image representation learning [12]. However, real-world datasets often exhibit structured variations or shifts in latent factors across different groups.
These shifts, which occur when known or unknown factors of variation change during data acquisition, can lead to inaccurate or biased predictions when the model encounters new data.
To address this issue, we introduce Contextual Vision Transformers (ContextViT), a novel method that generates robust feature representations for images, effectively handling variation in underlying latent factors.

Transformer-based models can perform test-time generalization by integrating available information through in-context learning [Brown et al., 2020], where input-label pairs are added to the transformer inputs. In this work, we leverage this principle to address distribution shifts at test time. However, direct application of this strategy presents challenges for vision transformers due to the quadratic scaling of image patch tokens with input resolution.

Derived from the principle of in-context learning, we propose ContextViT, a variant of vision transformer that condenses the relevant context into a single token. This context token is shared across all images with the same underlying distribution and varies across different distributions.
For instance, in medical imaging, this could correspond to learning a single hospital token for each hospital. The context token is appended to the input sequence of the transformer, conditioning the resulting image representation on the context. This approach reduces the need for conditioning the model on extra input examples and enables efficient inference at test time.  Despite these advantages, one limitation remains: it does not allow for generalization to unseen hospitals in absence of inferred context tokens.

To overcome this limitation, we introduce a context inference network that estimates these tokens on-the-fly from the input images. Given a group of images from the same context (e.g., a hospital), the context inference network predicts the token that encapsulates the characteristics of this context. We argue that this procedure should be applied iteratively across all transformer layers, a process we term layer-wise context conditioning. This is because earlier transformer layers capture local image patterns, while later layers capture more abstract concepts [Ghiasi et al., 2022]. Adjusting the covariates at the first layer is not sufficient to accommodate the covariate shift, which can be a concept-level change.

Our empirical analysis begins with supervised fine-tuning, where we integrate our context conditioning mechanism with pre-trained ViTs. Our results show consistent performance gains across three pre-trained ViT models, namely DINO [Caron et al., 2021], SWAG [Li et al., 2016], and CLIP [Radford et al., 2021], on the WILDS benchmark iWildCam [Beery et al., 2021] and FMoW [Christie et al., 2018]. We then explore the integration of self-supervised representation learning with ContextViT. On the microscopy cell imaging benchmark JUMP-CP [Haghighi et al., 2022], we observe an interesting phenomenon: while the performance of ViT decreases due to the batch effect as we introduce more data plates for pre-training, the performance of ContextViT steadily increases with more data. On the histopathology benchmark Camelyon17-WILDS [Bandi et al., 2018, 38], ContextViT significantly improves the performance on unseen hospitals, setting a new state-of-the-art.

The main contributions of this paper are:

* We propose ContextViT, an adaptation of ViT for producing robust image representations using learned context tokens to capture distribution shift. We derive ContextViT from in-context learning under distribution shift.

* We incorporate a context inference mechanism into ContextViT, enhancing its ability to adapt and generalize to previously unseen distributions on the fly.

* We present layer-wise context conditioning for ContextViT, a process that uses per-layer context tokens iteratively across all transformer layers to handle concept-level changes across distributions.

* We explore the integration of self-supervised learning with cell-imaging and histopathology datasets, demonstrating significant performance improvements under distribution shifts and establishing a new state-of-the-art.

Contextual Vision Transformer (ContextViT). Before each Transformer layer, ContextViT applies a context inference model over all patches (from images within the current batch that share the same group membership) to produce a context token.
To simplify the illustration, we present only two Transformer layers and depict a single patch embedding for each image on the left side.

§ VISION TRANSFORMERS (VITS)

Patch token embeddings
Let $x\in \mathbb{R}^{H\times W\times C}$ be an image with $C$ channels and resolution $H$ by $W$. ViTs first partition the image into a sequence of non-overlapping 2D patches $[x_{p_1},\ldots,x_{p_N}]$, each with resolution $(H_p, W_p)$ and represented by $x_{p_i}\in\mathbb{R}^{H_p\times W_p\times C}$. ViTs treat each image patch as a “1D token” for the Transformer, and we obtain patch token embeddings by flattening each patch $x_{p_i}$ into a 1D vector and applying a trainable affine projection. The resulting patch token sequence is denoted as $[t_{p_1},\ldots,t_{p_N}]$ where each $t_{p_i}\in\mathbb{R}^d$.

token and position embeddings
In ViTs, we prepend a trainable  token $t_{\texttt{CLS}}$ to the input sequence. This enables the Transformer encoder to capture global image features by aggregating information across the patch sequence. We retain positional information for each patch token using a trainable 1D position embedding $p_i$. The input sequence for the Transformer encoder is thus given by $[t_{\texttt{CLS}}, t_{p_1} + pos_{p_1}, \ldots, t_{p_N}+pos_{p_N}]$.

The Transformer layer [Vaswani et al., 2017] is a critical component of the ViT architecture. It comprises a self-attention layer and a feed-forward layer, each with residual connections [He et al., 2016] and layer normalization [Ba et al., 2016]. Self-attention enables the model to capture dependencies between patches in the input image by embedding each patch based on its similarity to other patches. ViTs utilize a stack of Transformer layers $T^{(1)},\ldots T^{(L)}$ to encode  the input sequence into a sequence of 1D features:

\[
[y_{\texttt{CLS}}^{(L)}, y_{p_1}^{(L)},\ldots,y_{p_N}^{(L)}]=T^{(L)}\cdots T^{(2)} T^{(1)}([t_{\texttt{CLS}}, t_{p_1} + pos_{p_1}, \ldots, t_{p_N}+pos_{p_N}]).
\]

§ CONTEXTUAL VISION TRANSFORMER (CONTEXTVIT)

Let $\mathcal{D}_c$ denote a collection of samples from a distribution $P_c(\cdot)$ under a context with index $c$, which can point to an unobserved environmental effect, a set of covariates such as experimental conditions, or another factor of variation with a scope over the set $\mathcal{D}_c$ that systematically shapes its distribution. Our goal is as follows: given a collection of such related datasets from distributions with varying factors of variation, captured as varying contexts $c$, we would like to learn a model that can generalize gracefully across these different domains and ideally adapt to new contexts $c^{*}$ during test time.

We assume that we observe data from multiple distributions jointly, each varying by a context $c$, so that our dataset  is comprised by their collection $\mathcal{D}=\{\mathcal{D}_1, ..., \mathcal{D}_c\}$.

We aim to learn a single shared ViT model, parameterized by $\theta$, across all sub-datasets with varying contexts.
We denote a tuple $\{x,y,c\}$ as the data describing the sample, where $x$ denotes an input, the image, and $y$ denotes an output variable such as a label or an unknown embedding, and $c$ denoting group membership for each datum to each sub-distribution $P_{c}(\cdot)$. Over the next sections we will explain our assumptions for ContextViT and will present practical implementations.

§.§ In-Context Model

A popular paradigm for incorporating test-time conditioning in Transformers is given by in-context learning [Brown et al., 2020] and prompt-conditioning [Radford et al., 2021].
To translate that paradigm to our task, we incorporate group membership to the representation learning task by conditioning on the members of the group, assuming the following factorization of the joint distribution over all data:

\begin{align}
\label{eq:in_context}
P(Y|X, C) = \prod_c \prod_{\{x,y\}\in \mathcal{D}_c} P(y|x, \mathcal{D}_c;\,\theta).
\end{align}

Here, we would concatenate all images belonging to context $c$ to the query image $x$. During ViT processing, these images are patchified and represented as patch tokens, resulting in an input to the Transformer consisting of the tokens belonging to the query image and to the context images.
This process is closest to in-context reasoning as commonly applied for LLMs, but ported over to the scenario of representation learning for vision.

ContextViT inference (solid line) vs. the implicit generative model (dotted line)

§.§ Context-Token Model

For the sake of exposition,  let's assume an implicit hierarchical generative model of images with shared global latent variables per group $t_c \in \mathbb{R}^{d}$ which summarizes the shared latent characteristics of $\mathcal{D}_c$.  This variable represents an embedding of the context (which can characterize the environment or implicit covariates determining each distribution, or other factors of variation shifting in $c$) and is treated as a context token, see Fig. <ref>. Such models are common in generative modeling when environmental effects or style-class disentanglement are proposed.
The context token $t_c$ can be characterized by its posterior distribution given the members of the group $P(t_c| \mathcal{D}_c)$ and can then be utilized in an adjusted model $P(y|x,t_c ;\theta)$ to explain the impact of  $\mathcal{D}_c$.

Upon closer inspection, it is evident that the in-context-learning process described in Sec. <ref> can be interpreted as inferring and marginalizing the posterior distribution $P(t_c| \mathcal{D}_c)$ of the context-specific variable $t_c$ as shown in the following:

\begin{align}
\label{eq:token1}
P(y|x, \mathcal{D}_c)= \int P(y|x,t_c;\,\theta) P(t_c| \mathcal{D}_c) \,\mathrm{d}t_c .
\end{align}

Under this viewpoint using Eq.<ref>, we can now see that the context token is shared over all members of a distribution with index $c$ (similar to a variable shared in a plate in graphical models), allowing us to rewrite Eq.<ref> such that we can factorize the dependence of the model on $\mathcal{D}_c$ for each data point given this context token:

\begin{align}
\label{eq:token2}
P(Y|X, C) = \prod_c \int  P(t_c| \mathcal{D}_c) \prod_{\{x,y\}\in \mathcal{D}_c} P(y|x, t_c;\,\theta) \,\mathrm{d}t_c.
\end{align}

We now have established equivalence between the in-context model, and performing inference in an implicit forward generative process which assumes existence of a context token as an additional factor of variation capturing distribution shifts for each group sharing a context.
We note that recent work proposes a deeply related interpretation of in-context learning

in [Xie et al., 2022] supporting our view of an implicit probabilistic model.

§.§ Oracle-Context Model

To simplify the setup for use in a ViT framework, we devise simpler formulations of the model and perform maximum likelihood inference over the context token, simplifying Eq. <ref> to:

\begin{align}
\label{eq:token3}
P(Y|X, C) = \prod_c \prod_{\{x,y\}\in \mathcal{D}_c} P(y|x ; t_c, \theta),
\end{align}

where $t_c$ now is a shared parameter that can be inferred per existing group during training.
We denote this model the Oracle-Context Model, since it only assumes knowledge of the indicator $c$ about which group an image belongs to, and does not require access to the other members of the group.

We can instantiate this oracle model $P(y|x ; t_c, \theta)$ by conditioning the ViTs with a learnable context token $t_c$ and append it to the input sequence of the Transformer: $[t_{\texttt{CLS}} , t_c, t_1 + pos_1, \ldots, t_N+pos_N]$.

The limitation, however, is that such a method cannot be applied during test-time beyond known distributions, since it lacks the ability to infer the token embedding without prior exposure to examples from this group during the training phase.

§.§ Context Inference Model

We overcome this limitation by closer matching the objective and assuming a model that performs amortized inference over the context-token parameters on the fly given observations $\mathcal{D}_c$ by utilizing an inference network $h(\cdot;\phi)$.
Concretely, we will again simplify $p(t_c| \mathcal{D}_c)$ to be given by maximum likelihood but this time also utilize $\mathcal{D}_c$ by posing:

\begin{align}
\label{eq:token4}
P(Y|X, C) = \prod_c \prod_{\{x,y\}\in \mathcal{D}_c} P(y|x, t_c ; \theta),\quad
\text{with}~t_c = h(\mathcal{D}_c;\, \phi).
\end{align}

Here, $h(\cdot;\phi)$ infers $t_c$ on the fly when observing $\mathcal{D}_c$, and can indeed also be used during test time on a previous unknown distribution with context $c^{*}$ time given a set of samples $\mathcal{D}_{c^{*}}$.
We note that the assumption of availability of samples $\mathcal{D}_c$  during testing is highly realistic and common, as datasets $\mathcal{D}_c$ (comprised of observed inputs sharing a context) typically appear in groups for free also during test time, for example one would never image a single cell in an experiment but a collection under fixed conditions, and the expensive step is to assess their unknown mapping to output labels per group.
In our experiments we will show how contexts $c$ can also be persistent objects, such as hospitals having measurement processes for pathology images, that easily allow us to access samples from the set of measurements given context $c$.

There are multiple options available for instantiating the inference model $h(\cdot;\phi)$ (refer to the Appendix for more details). In the case of ContextViT (Figure <ref>), the inference model $h(\cdot;\phi)$ comprises three key components:

Mean pooling
Given an image $x\in \mathcal{D}_c$ with its corresponding patch embeddings $t_{p_i}$, a direct way to obtain the context token $t_c$ is to aggregate the patch embeddings of all images in $\mathcal{D}_c$ through an average pooling operation
$t_c^{\text{mean}} = h_c^{\text{mean}}(\mathcal{D}_c) \coloneqq \frac{1}{|\mathcal{D}_c|} \sum_{x\in \mathcal{D}_c} \frac{1}{N} \sum_{x_{p_i} \in x} t_{p_i}$.
The input sequence for the image $x$ would then be represented as $[t_{\texttt{CLS}} , t_c^{\text{mean}}, t_1 + pos_1, \ldots, t_N+pos_N]$. In practice, instead of pooling over the entire $\mathcal{D}_{c}$, we apply mean pooling over $\mathcal{B}\cap \mathcal{D}_c$ where $\mathcal{B}$ is the current batch.

Linear transformation with gradients detaching
To further enhance the context representation, we introduce a trainable linear transformation to the pooled representation. In order to prevent the patch embeddings from being penalized by the context inference model, we detach their gradients. This results in the expression  $h_c^{\text{linear}}(\mathcal{D}_c;\, b, W) \coloneqq b + W\cdot \texttt{detach}(t_c^{\text{mean}})$.

Layerwise context conditioning
Recent work [Ghiasi et al., 2022] has shown that Transformer features progress from abstract patterns in early layers to concrete objects in later layers. We explore the application of context conditioning beyond the input layer driven by the hypothesis that patch embeddings may not be able to capture higher-level concepts. For the $l$-th Transformer layer, we use $y^{(l)}$ to denote its output and $\mathcal{D}^{(l)}_c$ to denote the collection of hidden patch embeddings of context $c$. We propose layerwise context conditioning which performs amortized inference over context-token parameters for each layer in the ViT instead of just propagating the input-layer token:

\[
P(y^{(l)}|y^{(0)}\coloneqq x; t_c) = \prod_{l=1}^{L} P(y^{(l)} | y^{(l-1)}; t_c^{(l-1)}),\quad
\text{with}~t_c^{(l)} = h(\mathcal{D}^{(l)}_c; \phi^{(l)}).
\]

For the $l$-th Transformer layer, we can express the layerwise context conditioning as

\[
[y_{\texttt{CLS}}^{(l)}, y_c^{(l)},  y_{p_1}^{(l)},\ldots,y_{p_N}^{(l)}]=  T^{(l)}([y_{\texttt{CLS}}^{(l-1)}, t_c^{(l-1)},  y_{p_1}^{(l-1)},\ldots,y_{p_N}^{(l-1)}]).
\]

§ EXPERIMENTS

We demonstrate the utilities of ContextViT across a variety of applications.
We start with the common supervised fine-tuning setting where we augment three existing pre-trained ViTs with our context inference model (Section <ref>). Next we experiment ContextViT with self-supervised representation learning. We demonstrate the importance of context conditioning for both in-distribution generalization <ref> and out-of-distribution generalization <ref>. We present our ablation study and analysis in Section <ref>. Due to the constraints of space, we kindly direct our readers to the Appendix for the pseudo code implementation details. Our code is available at <https://github.com/insitro/ContextViT>.

§.§ Fine-tuning pre-trained ViTs with context conditioning

To evaluate robustness over data shifts, we consider two image classification datasets from the WILDS benchmarks [Koh et al., 2021]: iWildCam [Beery et al., 2021] and FMoW [Christie et al., 2018].
In iWildCam, the task is to classify the animal species (182-way classification) from images taken by different cameeras. We use the camera trap id for context inference and report the testing accuracy on unseen camera traps.
In FMoW, the task is to classify the land use of a satellite image (62-way classification). We use the region id for context inference and report the  worst-region accuracy for image from the testing time period.

We consider three existing pre-trained ViT models, each with a different number of parameters:
1) DINO ViT-S/8 [Caron et al., 2021], which is pre-trained on ImageNet with self-distillation between a teacher and a student network.
2) SWAG ViT-B/16 [Singh et al., 2022], which is pre-trained on IG3.6B using weak supervision (hashtags from Instagram).
3) CLIP ViT-L/14, which is pre-trained on the Multi-modal WebImageText using language-image contrastive learning.
Despite their differences in pre-training objectives and datasets, these models share the same ViT backbone. Therefore we can augment these pre-trained models with our proposed context conditioning mechanism and fine-tune the combined ContextViT jointly with empirical risk minimization.

Table <ref> presents our results for supervised fine-tuning. We note that our direct fine-tuning baseline (using CLIP ViT-L/14) outperforms the number reported by [Wortsman et al., 2022] on iWildCam (81.5 vs. 78.3) and underperforms the number reported by [Kumar et al., 2022] on FMoW (46.6 vs. 49.9). We think this is likely caused by the differences in the implementation details (data augmentation, learning rate scheduling, etc.), and unfortunately we cannot find the configuration online to reproduce the exact numbers. Nevertheless, upon comparing our implementations, we make the following observations: 1) Smaller ViTs exhibit inferior generalization compared to larger ViTs; 2) Incorporating our context inference model consistently enhances the performance of ViTs; 3) Layerwise context inference further enhances generalization.

§.§ In-distribution generalization with self-supervised learning

Microscopy imaging with cell painting has demonstrated its effectiveness in studying the effects of cellular perturbations [Haghighi et al., 2022, Moshkov et al., 2022, 43]. Despite meticulous regulation of experimental parameters like cell density and exposure time, technical artifacts can still confound measurements from these screens across different batches. Learning a robust cell representation that can effectively handle batch variations remains an ongoing challenge.

We consider three cell plates (, , ) from the JUMP-CP dataset released by the JUMP-Cell Painting Consortium [Chandrasekaran et al., 2022].
Each plate consists of 384 wells with perturbations (either via a  chemical compound or a crispr guide) targeted at 160 different genes. Our input data consists of single-cell images: 229228 images for , 226311 images for  and 239347 images for . We note that here each single-cell image has dimension $224 \times 224 \times 8$ (5 fluorescence channels and 3 brightfield channels). For each plate, we split the images randomly into training (40%), validation (10%) and testing (50%). We use the plate id for context inference.

For each model, we use DINO [Caron et al., 2021] to pre-train the ViT and ContextViT from scratch for 100 epochs. Once pre-training is completed, we freeze the backbone parameters and attach an MLP (with one hidden ReLU layer of dimension 384) to predict the targeted gene of the cell. Given the unique nature of each plate, we train an individual MLP classifier for each separate plate using the training split of that specific plate, and subsequently report the corresponding testing accuracy in Table <ref>.

We note that the representation learned by DINO on a single plate () exhibits poor generalization across plates (42.7 for  and 46.8 for ). Moreover, directly combining all three plates for pre-training results in a degradation of the model's performance: -2.9 for , +0.8 for  and -1.2 for . In contrast, by utilizing ContextViT, the model effectively accounts for batch effects through context conditioning during the pre-training stage, resulting in superior performance across all three plates.

§.§ Out-of-distribution generalization with self-supervised learning

In medical applications, models are often trained using data from a limited number of hospitals with the intention of deploying them across other hospitals more broadly [Yala et al., 2019, Yala et al., 2021, Bandi et al., 2018]. However, this presents a challenge for the generalization of out-of-distribution data. In this section, we aim to evaluate the ability of self-supervised learning with ContextViT to achieve better out-of-distribution generalization.

Dataset We consider the Camelyon17-WILDS benchmark [Bandi et al., 2018, 38]. The dataset contains 455,954 labeled and 2,999,307 unlabeled pathology images across five hospitals (3 for training, 1 for validation and 1 for testing). Given a $96\times96\times3$ image, the task is to predict whether the image contains any tumor tissue.
We use the hospital id for context inference.

We utilize DINO to pre-train our ViT and ContextViT models from scratch on unlabeled pathology images. Unlike previous work ([Sagawa et al., 2021]) that incorporates both in-distribution and out-of-distribution unlabeled data (Unlabeled ID&OOD), we exclusively use images from the three training hospitals, denoted as “Unlabeled ID”, during pre-training stage to prevent any potential information leakage from out-of-distribution data. Given the 96 by 96 input resolution of our dataset, we opt for a smaller patch size (8 instead of 16) for the ViT models. Once pre-training is complete, we freeze the ViT parameters and train a linear classifier, shared across all training hospitals, with SGD to predict the target label, the presence of tumors. We report the testing accuracy on the out-of-distribution hospitals.

Table <ref> presents a comparison of our linear probing results (marked with$^\dagger$) with other published results on the Camelyon17 benchmark. Our DINO pre-trained ViT-S/8 baseline outperforms the much larger CLIP model in terms of linear probing (93.8 vs. 92.6), which is not surprising given the ViT-S/8 has been pre-trained on the same pathology domain. Next, we see that by conditioning our ViT-S/8 representation with other images from the same hospital within the testing batch, ContextViT-S/8 achieves an OOD accuracy of 97.5%, significantly outperforming all baselines. Moreover, Figure <ref> demonstrates that the linear classifier built upon ContextViT-S/8 continues to enhance its OOD performance as the training of the linear classifier progresses, while the one built upon ViT-S/8 exhibits signs of over-fitting to the training data.

Ablation study on the context inference model

Due to space constraints, we have deferred additional analyses to the Appendix.
In Table <ref>, we perform an ablation study of our context inference model $h(\cdot;\phi)$. We observe that utilizing the mean patch embeddings as context (94.1) and incorporating an additional linear transformation (94.4) result in improved performance compared to the no context baseline (93.8). Notably, we find that detaching the gradients when calculating the context token from the patch embeddings is crucial, leading to a performance boost from 94.4 to 97.2. Furthermore, the application of layerwise context conditioning further enhance the performance.

We also experimented with other set representation models for context inference, including deep sets [Zaheer et al., 2017] (see Appendix for details). Although deep sets (with detached patch embeddings) outperformed the baseline (94.9 vs. 93.8), it still fell short of the simpler approach of mean pooling with linear transformation. Our hypothesis is that the Transformer block's self-attention mechanism, which utilizes the same key, query, and value transformations across the sequence, results in the mean pooling with linear transformation producing a context embedding that is more comparable to the original patch sequence.

Time efficiency
One limitation of ContextViT is the extra time cost for the context conditioning. In Table <ref>, we present the DINO pre-training time cost for different context inference methods. We observe that adding the non-layerwise context conditioning increases the baseline DINO pre-training time by 5%-9%. Applying layerwise context conditioning further increases the time cost by 29%.

Learning curves of linear probing on Camelyon17. Unlike ViT-S/8, the OOD accuracy on top of ContextViT-S/8 steadily improves as the training of the linear classifier progresses.

On Camelyon17, we observe that even when we use a testing batch size of 1 (computing the context by aggregating 144 8 x 8 patches from a single testing image into a single context token), ContextViT-S/8 still significantly outperforms ViT-S/8.

Sensitivity to testing batch size
During the inference phase, ContextViT-S/8 takes a random batch of testing examples (512 in all previous experiments) and groups them based on their group membership, such as the hospital ID in the case of Camelyon17, to infer the corresponding context tokens. In Figure <ref>, we present a visualization of ContextViT's linear probing performance sensitivity across different testing batch sizes. Remarkably, even with a testing batch size of 1, where ContextViT-S/8 leverages patches from a single testing image (144 patches) to establish its context, it still outperforms our ViT-S/8 baseline by a significant margin (97.0 vs. 93.8). It's worth noting that after DINO pre-training, ContextViT-S/8 acquires the ability to condition its output features with respect to the context. As highlighted by [13], pathology images from different hospitals exhibit distinct color distributions. The exceptional performance of ContextViT-S/8 with a small testing batch size demonstrates the model's capability to automatically estimate the color distribution shift from a few representative images during test time.

Attention maps from multiple heads
In line with the findings of [Caron et al., 2021], which demonstrate that DINO can learn class-specific features leading to unsupervised object segmentations, we visualize the attention heatmaps for different heads learned with ContextViT in the Camelyon17 dataset. Figure <ref> illustrates these attention maps, showcasing that the learned attentions are focused on meaningful aspects for both in-distribution data and out-of-distribution data. Furthermore, different attention heads exhibit distinct preferences for different cell types. For instance, head-2 primarily focuses on the contour of the cells, while head-3 directs its attention to the actual cells themselves.

Visualizing the context tokens
We employ PCA to visualize the learned context tokens for each hospital, and the results are presented in Figure <ref>. One approach to visualizing the context tokens is by inferring them from all examples belonging to each hospital, resulting in five unique context tokens. However, in practice, we infer the context token on-the-fly for the current mini-batch. Using a batch size of 256, we sample 300 batches for each hospital.

Remarkably, the inferred context tokens for each hospital exhibit high similarity, appearing closely clustered together. Additionally, we include example images for each hospital on the right side of Figure <ref>. Notably, the distances between context tokens from different hospitals align well with their visual dissimilarities. For instance, hospital 3 (highlighted in red) is closer to hospital 1 (highlighted in orange) than it is to hospital 4 (highlighted in purple).

Left: PCA visualization of context tokens inferred by ContextViT using pathology images from different hospitals. Right: Example images from different hospitals.

Attention heatmap of ContextViT on patholog images from in-distribution hospitals (top) and OOD hospitals (bottom) in Camelyon17.

§ EXPERIMENTS

We demonstrate the utilities of ContextViT across a variety of applications.
We start with the common supervised fine-tuning setting where we augment three existing pre-trained ViTs with our context inference model (Section <ref>). Next we experiment ContextViT with self-supervised representation learning. We demonstrate the importance of context conditioning for both in-distribution generalization <ref> and out-of-distribution generalization <ref>. We present our ablation study and analysis in Section <ref>. Due to the constraints of space, we kindly direct our readers to the Appendix for the pseudo code implementation details. We are committed to release our code.

§.§ Fine-tuning pre-trained ViTs with context conditioning

To evaluate robustness over data shifts, we consider two image classification datasets from the WILDS benchmarks [Koh et al., 2021]: iWildCam [Beery et al., 2021] and FMoW [Christie et al., 2018].
In iWildCam, the task is to classify the animal species (182-way classification) from images taken by different cameeras. We use the camera trap id for context inference and report the testing accuracy on unseen camera traps.
In FMoW, the task is to classify the land use of a satellite image (62-way classification). We use the region id for context inference and report the  worst-region accuracy for image from the testing time period.

We consider three existing pre-trained ViT models, each with a different number of parameters:
1) DINO ViT-S/8 [Caron et al., 2021], which is pre-trained on ImageNet with self-distillation between a teacher and a student network.
2) SWAG ViT-B/16 [Singh et al., 2022], which is pre-trained on IG3.6B using weak supervision (hashtags from Instagram).
3) CLIP ViT-L/14, which is pre-trained on the Multi-modal WebImageText using language-image contrastive learning.
Despite their differences in pre-training objectives and datasets, these models share the same ViT backbone. Therefore we can augment these pre-trained models with our proposed context conditioning mechanism and fine-tune the combined ContextViT jointly with empirical risk minimization.

Table <ref> presents our results for supervised fine-tuning. We note that our direct fine-tuning baseline (using CLIP ViT-L/14) outperforms the number reported by [Wortsman et al., 2022] on iWildCam (81.5 vs. 78.3) and underperforms the number reported by [Kumar et al., 2022] on FMoW (46.6 vs. 49.9). We think this is likely caused by the differences in the implementation details (data augmentation, learning rate scheduling, etc.), and unfortunately we cannot find the configuration online to reproduce the exact numbers. Nevertheless, upon comparing our implementations, we make the following observations: 1) Smaller ViTs exhibit inferior generalization compared to larger ViTs; 2) Incorporating our context inference model consistently enhances the performance of ViTs; 3) Layerwise context inference further enhances generalization.

§.§ In-distribution generalization with self-supervised learning

Microscopy imaging with cell painting has demonstrated its effectiveness in studying the effects of cellular perturbations [Haghighi et al., 2022, Moshkov et al., 2022, 43]. Despite meticulous regulation of experimental parameters like cell density and exposure time, technical artifacts can still confound measurements from these screens across different batches. Learning a robust cell representation that can effectively handle batch variations remains an ongoing challenge.

We consider three cell plates (, , ) from the JUMP-CP dataset released by the JUMP-Cell Painting Consortium [Chandrasekaran et al., 2022].
Each plate consists of 384 wells with perturbations (either via a  chemical compound or a crispr guide) targeted at 160 different genes. Our input data consists of single-cell images: 229228 images for , 226311 images for  and 239347 images for . We note that here each single-cell image has dimension $224 \times 224 \times 8$ (5 fluorescence channels and 3 brightfield channels). For each plate, we split the images randomly into training (40%), validation (10%) and testing (50%). We use the plate id for context inference.

For each model, we use DINO [Caron et al., 2021] to pre-train the ViT and ContextViT from scratch for 100 epochs. Once pre-training is completed, we freeze the backbone parameters and attach an MLP (with one hidden ReLU layer of dimension 384) to predict the targeted gene of the cell. Given the unique nature of each plate, we train an individual MLP classifier for each separate plate using the training split of that specific plate, and subsequently report the corresponding testing accuracy in Table <ref>.

We note that the representation learned by DINO on a single plate () exhibits poor generalization across plates (42.7 for  and 46.8 for ). Moreover, directly combining all three plates for pre-training results in a degradation of the model's performance: -2.9 for , +0.8 for  and -1.2 for . In contrast, by utilizing ContextViT, the model effectively accounts for batch effects through context conditioning during the pre-training stage, resulting in superior performance across all three plates.

§.§ Out-of-distribution generalization with self-supervised learning

In medical applications, models are often trained using data from a limited number of hospitals with the intention of deploying them across other hospitals more broadly [Yala et al., 2019, Yala et al., 2021, Bandi et al., 2018]. However, this presents a challenge for the generalization of out-of-distribution data. In this section, we aim to evaluate the ability of self-supervised learning with ContextViT to achieve better out-of-distribution generalization.

Dataset We consider the Camelyon17-WILDS benchmark [Bandi et al., 2018, 38]. The dataset contains 455,954 labeled and 2,999,307 unlabeled pathology images across five hospitals (3 for training, 1 for validation and 1 for testing). Given a $96\times96\times3$ image, the task is to predict whether the image contains any tumor tissue.
We use the hospital id for context inference.

We utilize DINO to pre-train our ViT and ContextViT models from scratch on unlabeled pathology images. Unlike previous work ([Sagawa et al., 2021]) that incorporates both in-distribution and out-of-distribution unlabeled data (Unlabeled ID&OOD), we exclusively use images from the three training hospitals, denoted as “Unlabeled ID”, during pre-training stage to prevent any potential information leakage from out-of-distribution data. Given the 96 by 96 input resolution of our dataset, we opt for a smaller patch size (8 instead of 16) for the ViT models. Once pre-training is complete, we freeze the ViT parameters and train a linear classifier, shared across all training hospitals, with SGD to predict the target label, the presence of tumors. We report the testing accuracy on the out-of-distribution hospitals.

Table <ref> presents a comparison of our linear probing results (marked with$^\dagger$) with other published results on the Camelyon17 benchmark. Our DINO pre-trained ViT-S/8 baseline outperforms the much larger CLIP model in terms of linear probing (93.8 vs. 92.6), which is not surprising given the ViT-S/8 has been pre-trained on the same pathology domain. Next, we see that by conditioning our ViT-S/8 representation with other images from the same hospital within the testing batch, ContextViT-S/8 achieves an OOD accuracy of 97.5%, significantly outperforming all baselines. Moreover, Figure <ref> demonstrates that the linear classifier built upon ContextViT-S/8 continues to enhance its OOD performance as the training of the linear classifier progresses, while the one built upon ViT-S/8 exhibits signs of over-fitting to the training data.

Ablation study on the context inference model

Due to space constraints, we have deferred additional analyses to the Appendix.
In Table <ref>, we perform an ablation study of our context inference model $h(\cdot;\phi)$. We observe that utilizing the mean patch embeddings as context (94.1) and incorporating an additional linear transformation (94.4) result in improved performance compared to the no context baseline (93.8). Notably, we find that detaching the gradients when calculating the context token from the patch embeddings is crucial, leading to a performance boost from 94.4 to 97.2. Furthermore, the application of layerwise context conditioning further enhance the performance.

We also experimented with other set representation models for context inference, including deep sets [Zaheer et al., 2017] (see Appendix for details). Although deep sets (with detached patch embeddings) outperformed the baseline (94.9 vs. 93.8), it still fell short of the simpler approach of mean pooling with linear transformation. Our hypothesis is that the Transformer block's self-attention mechanism, which utilizes the same key, query, and value transformations across the sequence, results in the mean pooling with linear transformation producing a context embedding that is more comparable to the original patch sequence.

Time efficiency
One limitation of ContextViT is the extra time cost for the context conditioning. In Table <ref>, we present the DINO pre-training time cost for different context inference methods. We observe that adding the non-layerwise context conditioning increases the baseline DINO pre-training time by 5%-9%. Applying layerwise context conditioning further increases the time cost by 29%.

§ RELATED WORK

In-context learning
Our work draws inspiration from and expands the idea of in-context learning [Brown et al., 2020].
We realize a form of conditioning that can position transformer-based models to perform test-time adaptation to new tasks, but instead of conditioning on a dataset explicitly for each query, we infer context tokens that are shareable across a group.
Like  [Xie et al., 2022], we interpret in-context learning via an implicit generative model. The key distinction lies in our focus on adaptation to distribution shift, and our expansion of the framework to include explicit context tokens.

There is a substantial body of previous work on conditioning transformers with extra tokens. Li and Liang, 2021 introduced prefix tuning as a method to adapt a frozen Transformer model for various downstream tasks. Xu et al., 2022, Li et al., 2021, Mao et al., 2022 employed additional tokens to encode extra domain knowledge, such as multi-modal information from language. We focus on the scenario where we do not assume existence multi-modal measurements and our domain knowledge is structural (knowing group membership). White et al., 2022 explored hierarchical models for transformers in the context of language modeling, highlighting similar group-specific effects and connections to mixed models.
ContextViT sets itself apart by inferring the context information directly from the input data and learning the model parameters jointly throughout training. Similar to [Hao et al., 2022], our approach also achieves test time adaptation without the enormous memory costs of in-context learning since we share inference over the context token across examples using it and only require a single additional token to represent the group information.

Robust learning
Robustness in machine learning can be interpreted in various ways.
[Mahmood et al., 2021] discovered that ViTs are vulnerable to white-box adversarial attacks [Hendrycks and Dietterich, 2019]. Consequently, numerous strategies have been proposed to enhance the adversarial robustness of ViTs ViTs [Mao et al., 2022, Chefer et al., 2022].

In this study, we concentrate on robustness in the face of systematic distribution shifts. Test time adaptation is a widely adopted approach [Sun et al., 2020, Liu et al., 2021, Zhang et al., 2022]. Li et al., 2022 developed a matching network to select the most suitable pretrained models for testing. The use of batch normalization at test time has also been shown to improve the generalization of CNN  across various applications [Zhang et al., 2021, Nado et al., 2020, Li et al., 2016, Schneider et al., 2020, Lin and Lu, 2022, Kaku et al., 2020].
Our approach, however, diverges from directly normalizing feature statistics or training multiple expert models. Instead, we incorporate context information as an additional latent input to the Transformer layer.
Finally, it's worth noting that the representations generated by ContextViT can be seamlessly integrated with robust learning algorithms such as Group Distributionally Robust Optimization [Sagawa et al., 2019] and Invariant Risk Minimization [Arjovsky et al., 2020], should users desire a specific form of invariance.

§ CONCLUSION

We have presented Contextual Vision Transformers, a method that addresses challenges posed by structured variations and covariate shifts in image datasets. By leveraging the concept of in-context learning and introducing context tokens and token inference models, ContextViT enables robust feature representation learning across groups with shared characteristics. Through extensive experimental evaluations across diverse applications, ContextViT  consistently demonstrates its utility compared to standard ViT models in terms of out-of-distribution generalization and resilience to batch effects. This success highlights the power and versatility of ContextViT when handling structured variations in real-world applications, where invariance to such nuisances may be desirable.

[Arjovsky et al., 2020]
M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz.
Invariant risk minimization.
stat, 1050:0 27, 2020.

[Ba et al., 2016]
J. L. Ba, J. R. Kiros, and G. E. Hinton.
Layer normalization.
arXiv preprint arXiv:1607.06450, 2016.

[Bandi et al., 2018]
P. Bandi, O. Geessink, Q. Manson, M. Van Dijk, M. Balkenhol, M. Hermsen, B. E.
Bejnordi, B. Lee, K. Paeng, A. Zhong, et al.
From detection of individual metastases to classification of lymph
node status at the patient level: the camelyon17 challenge.
IEEE transactions on medical imaging, 380
(2):0 550–560, 2018.

[Beery et al., 2021]
S. Beery, A. Agarwal, E. Cole, and V. Birodkar.
The iwildcam 2021 competition dataset.
arXiv preprint arXiv:2105.03494, 2021.

[Brown et al., 2020]
T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal,
A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss,
G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu,
C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess,
J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei.
Language models are few-shot learners, 2020.

[Buslaev et al., 2020]
A. Buslaev, V. I. Iglovikov, E. Khvedchenya, A. Parinov, M. Druzhinin, and
A. A. Kalinin.
Albumentations: Fast and flexible image augmentations.
Information, 110 (2), 2020.
ISSN 2078-2489.
URL <https://www.mdpi.com/2078-2489/11/2/125>.

[Caron et al., 2021]
M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and
A. Joulin.
Emerging properties in self-supervised vision transformers.
In Proceedings of the IEEE/CVF international conference on
computer vision, pages 9650–9660, 2021.

[Chandrasekaran et al., 2022]
S. N. Chandrasekaran, B. A. Cimini, A. Goodale, L. Miller, M. Kost-Alimova,
N. Jamali, J. Doench, B. Fritchman, A. Skepner, M. Melanson, et al.
Three million images and morphological profiles of cells treated with
matched chemical and genetic perturbations.
bioRxiv, pages 2022–01, 2022.

[Chefer et al., 2022]
H. Chefer, I. Schwartz, and L. Wolf.
Optimizing relevance maps of vision transformers improves robustness.
Advances in Neural Information Processing Systems,
35:0 33618–33632, 2022.

[Christie et al., 2018]
G. Christie, N. Fendley, J. Wilson, and R. Mukherjee.
Functional map of the world.
In Proceedings of the IEEE Conference on Computer Vision and
Pattern Recognition, pages 6172–6180, 2018.

[DeVries and Taylor, 2017]
T. DeVries and G. W. Taylor.
Improved regularization of convolutional neural networks with cutout.
arXiv preprint arXiv:1708.04552, 2017.

[12]
A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai,
T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al.
An image is worth 16x16 words: Transformers for image recognition at
In International Conference on Learning Representations.

[13]
I. Gao, S. Sagawa, P. W. Koh, T. Hashimoto, and P. Liang.
Out-of-distribution robustness via targeted augmentations.
In NeurIPS 2022 Workshop on Distribution Shifts: Connecting
Methods and Applications.

[Ghiasi et al., 2022]
A. Ghiasi, H. Kazemi, E. Borgnia, S. Reich, M. Shu, M. Goldblum, A. G. Wilson,
and T. Goldstein.
What do vision transformers learn? a visual exploration.
arXiv preprint arXiv:2212.06727, 2022.

[Goyal et al., 2022]
S. Goyal, A. Kumar, S. Garg, Z. Kolter, and A. Raghunathan.
Finetune like you pretrain: Improved finetuning of zero-shot vision
arXiv preprint arXiv:2212.00638, 2022.

[Haghighi et al., 2022]
M. Haghighi, J. C. Caicedo, B. A. Cimini, A. E. Carpenter, and S. Singh.
High-dimensional gene expression and morphology profiles of cells
across 28,000 genetic and chemical perturbations.
Nature Methods, pages 1–8, 2022.

[Hao et al., 2022]
Y. Hao, Y. Sun, L. Dong, Z. Han, Y. Gu, and F. Wei.
Structured prompting: Scaling in-context learning to 1,000 examples,

[He et al., 2016]
K. He, X. Zhang, S. Ren, and J. Sun.
Deep residual learning for image recognition.
In Proceedings of the IEEE conference on computer vision and
pattern recognition, pages 770–778, 2016.

[19]
D. Hendrycks and T. Dietterich.
Benchmarking neural network robustness to common corruptions and
In International Conference on Learning Representations.

[Hendrycks and Dietterich, 2019]
D. Hendrycks and T. Dietterich.
Benchmarking neural network robustness to common corruptions and
arXiv preprint arXiv:1903.12261, 2019.

[Kaku et al., 2020]
A. Kaku, S. Mohan, A. Parnandi, H. Schambra, and C. Fernandez-Granda.
Be like water: Robustness to extraneous variables via adaptive
feature normalization.
arXiv preprint arXiv:2002.04019, 2020.

[Kingma and Ba, 2014]
D. P. Kingma and J. Ba.
Adam: A method for stochastic optimization.
arXiv preprint arXiv:1412.6980, 2014.

[Koh et al., 2021]
P. W. Koh, S. Sagawa, H. Marklund, S. M. Xie, M. Zhang, A. Balsubramani, W. Hu,
M. Yasunaga, R. L. Phillips, I. Gao, et al.
Wilds: A benchmark of in-the-wild distribution shifts.
In International Conference on Machine Learning, pages
5637–5664. PMLR, 2021.

[Kumar et al., 2022]
A. Kumar, R. Shen, S. Bubeck, and S. Gunasekar.
How to fine-tune vision models with sgd.
arXiv preprint arXiv:2211.09359, 2022.

[Li and Liang, 2021]
X. L. Li and P. Liang.
Prefix-tuning: Optimizing continuous prompts for generation.
arXiv preprint arXiv:2101.00190, 2021.

[Li et al., 2016]
Y. Li, N. Wang, J. Shi, J. Liu, and X. Hou.
Revisiting batch normalization for practical domain adaptation.
arXiv preprint arXiv:1603.04779, 2016.

[Li et al., 2021]
Y. Li, S. Zhang, Z. Wang, S. Yang, W. Yang, S.-T. Xia, and E. Zhou.
Tokenpose: Learning keypoint tokens for human pose estimation.
In Proceedings of the IEEE/CVF International conference on
computer vision, pages 11313–11322, 2021.

[Li et al., 2022]
Z. Li, K. Ren, X. Jiang, Y. Shen, H. Zhang, and D. Li.
Simple: Specialized model-sample matching for domain generalization.
In The Eleventh International Conference on Learning
Representations, 2022.

[Lin and Lu, 2022]
A. Lin and A. Lu.
Incorporating knowledge of plates in batch normalization improves
generalization of deep learning for microscopy images.
In Machine Learning in Computational Biology, pages 74–93.
PMLR, 2022.

[Liu et al., 2021]
Y. Liu, P. Kothari, B. Van Delft, B. Bellot-Gurlet, T. Mordan, and A. Alahi.
Ttt++: When does self-supervised test-time training fail or thrive?
Advances in Neural Information Processing Systems,
34:0 21808–21820, 2021.

[Loshchilov and Hutter, 2019]
I. Loshchilov and F. Hutter.
Decoupled weight decay regularization.
In International Conference on Learning Representations, 2019.
URL <https://openreview.net/forum?id=Bkg6RiCqY7>.

[Mahmood et al., 2021]
K. Mahmood, R. Mahmood, and M. Van Dijk.
On the robustness of vision transformers to adversarial examples.
In Proceedings of the IEEE/CVF International Conference on
Computer Vision, pages 7838–7847, 2021.

[Mao et al., 2022]
X. Mao, Y. Chen, X. Jia, R. Zhang, H. Xue, and Z. Li.
Context-aware robust fine-tuning.
arXiv preprint arXiv:2211.16175, 2022a.

[Mao et al., 2022]
X. Mao, G. Qi, Y. Chen, X. Li, R. Duan, S. Ye, Y. He, and H. Xue.
Towards robust vision transformer.
In Proceedings of the IEEE/CVF conference on Computer Vision
and Pattern Recognition, pages 12042–12051, 2022b.

[Moshkov et al., 2022]
N. Moshkov, M. Bornholdt, S. Benoit, M. Smith, C. McQuin, A. Goodman, R. A.
Senft, Y. Han, M. Babadi, P. Horvath, et al.
Learning representations for image-based profiling of perturbations.
bioRxiv, pages 2022–08, 2022.

[Nado et al., 2020]
Z. Nado, S. Padhy, D. Sculley, A. D'Amour, B. Lakshminarayanan, and J. Snoek.
Evaluating prediction-time batch normalization for robustness under
covariate shift.
arXiv preprint arXiv:2006.10963, 2020.

[Radford et al., 2021]
A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry,
A. Askell, P. Mishkin, J. Clark, et al.
Learning transferable visual models from natural language
In International conference on machine learning, pages
8748–8763. PMLR, 2021.

[38]
S. Sagawa, P. W. Koh, T. Lee, I. Gao, S. M. Xie, K. Shen, A. Kumar, W. Hu,
M. Yasunaga, H. Marklund, et al.
Extending the wilds benchmark for unsupervised adaptation.
In International Conference on Learning Representations.

[Sagawa et al., 2019]
S. Sagawa, P. W. Koh, T. B. Hashimoto, and P. Liang.
Distributionally robust neural networks.
In International Conference on Learning Representations, 2019.

[Sagawa et al., 2021]
S. Sagawa, P. W. Koh, T. Lee, I. Gao, S. M. Xie, K. Shen, A. Kumar, W. Hu,
M. Yasunaga, H. Marklund, et al.
Extending the wilds benchmark for unsupervised adaptation.
arXiv preprint arXiv:2112.05090, 2021.

[Schneider et al., 2020]
S. Schneider, E. Rusak, L. Eck, O. Bringmann, W. Brendel, and M. Bethge.
Improving robustness against common corruptions by covariate shift
Advances in Neural Information Processing Systems,
33:0 11539–11551, 2020.

[Singh et al., 2022]
M. Singh, L. Gustafson, A. Adcock, V. de Freitas Reis, B. Gedik, R. P.
Kosaraju, D. Mahajan, R. Girshick, P. Dollár, and L. Van Der Maaten.
Revisiting weakly supervised pre-training of visual perception
In Proceedings of the IEEE/CVF Conference on Computer Vision
and Pattern Recognition, pages 804–814, 2022.

[43]
S. Sivanandan, M. Salick, B. Leitmann, K. M. Liu, M. Sultan, N. Ranu, C. V.
Hao, O. Chen, J. Bisognano, E. Lubeck, et al.
Machine learning enabled pooled optical screening in human lung
cancer cells.
In NeurIPS 2022 Workshop on Learning Meaningful Representations
of Life.

[Sun et al., 2020]
Y. Sun, X. Wang, Z. Liu, J. Miller, A. Efros, and M. Hardt.
Test-time training with self-supervision for generalization under
distribution shifts.
In International conference on machine learning, pages
9229–9248. PMLR, 2020.

[Tompson et al., 2015]
J. Tompson, R. Goroshin, A. Jain, Y. LeCun, and C. Bregler.
Efficient object localization using convolutional networks.
In Proceedings of the IEEE conference on computer vision and
pattern recognition, pages 648–656, 2015.

[Vaswani et al., 2017]
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez,
Ł. Kaiser, and I. Polosukhin.
Attention is all you need.
Advances in neural information processing systems, 30, 2017.

[White et al., 2022]
J. White, N. Goodman, and R. Hawkins.
Mixed-effects transformers for hierarchical adaptation, 2022.

[Wortsman et al., 2022]
M. Wortsman, G. Ilharco, S. Y. Gadre, R. Roelofs, R. Gontijo-Lopes, A. S.
Morcos, H. Namkoong, A. Farhadi, Y. Carmon, S. Kornblith, et al.
Model soups: averaging weights of multiple fine-tuned models improves
accuracy without increasing inference time.
In International Conference on Machine Learning, pages
23965–23998. PMLR, 2022.

[Xie et al., 2022]
S. M. Xie, A. Raghunathan, P. Liang, and T. Ma.
An explanation of in-context learning as implicit bayesian inference,

[Xu et al., 2022]
L. Xu, W. Ouyang, M. Bennamoun, F. Boussaid, and D. Xu.
Multi-class token transformer for weakly supervised semantic
In Proceedings of the IEEE/CVF Conference on Computer Vision
and Pattern Recognition, pages 4310–4319, 2022.

[Yala et al., 2019]
A. Yala, C. Lehman, T. Schuster, T. Portnoi, and R. Barzilay.
A deep learning mammography-based model for improved breast cancer
risk prediction.
Radiology, 2920 (1):0 60–66, 2019.

[Yala et al., 2021]
A. Yala, P. G. Mikhael, F. Strand, G. Lin, K. Smith, Y.-L. Wan, L. Lamb,
K. Hughes, C. Lehman, and R. Barzilay.
Toward robust mammography-based models for breast cancer risk.
Science Translational Medicine, 130 (578):0
eaba4373, 2021.

[You et al., 2018]
Y. You, Z. Zhang, C.-J. Hsieh, J. Demmel, and K. Keutzer.
Imagenet training in minutes.
In Proceedings of the 47th International Conference on Parallel
Processing, pages 1–10, 2018.

[Zaheer et al., 2017]
M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov, and A. J.
Deep sets.
Advances in neural information processing systems, 30, 2017.

[Zhang et al., 2021]
M. Zhang, H. Marklund, N. Dhawan, A. Gupta, S. Levine, and C. Finn.
Adaptive risk minimization: Learning to adapt to domain shift.
Advances in Neural Information Processing Systems,
34:0 23664–23678, 2021.

[Zhang et al., 2022]
M. Zhang, S. Levine, and C. Finn.
Memo: Test time robustness via adaptation and augmentation.
Advances in Neural Information Processing Systems,
35:0 38629–38642, 2022.