SPARO: Selective Attention for Robust and Compositional Transformer Encodings for Vision
Ankit Vani Bac Nguyen Samuel Lavoie Ranjay Krishna Aaron Courville





Sparo: Selective Attention in Transformer Encodings for Vision



A. Vani et al.

Mila, Université de Montréal Sony AI University of Washington, Allen Institute for Artificial Intelligence CIFAR AI Chair

SPARO: Selective Attention for Robust and Compositional Transformer Encodings for Vision
    Ankit Vani1Correspondence at: ankit.vani@umontreal.ca.0009-0007-4781-9995 Bac Nguyen20000-0001-9193-1908 Samuel Lavoie10009-0005-8370-2415 Ranjay Krishna30000-0001-8784-2531 Aaron Courville1,40000-0001-6223-0301
    April 29, 2024
=======================================================================================================================================================================================================================



Selective attention helps us focus on task-relevant aspects in the constant flood of our sensory input. This constraint in our perception allows us to robustly generalize under distractions and to new compositions of perceivable concepts. Transformers employ a similar notion of attention in their architecture, but representation learning models with transformer backbones like CLIP and DINO often fail to demonstrate robustness and compositionality. We highlight a missing architectural prior: unlike human perception, transformer encodings do not separately attend over individual concepts. In response, we propose , a read-out mechanism that partitions encodings into separately-attended slots, each produced by a single attention head. Using with CLIP imparts an inductive bias that the vision and text modalities are different views of a shared compositional world with the same corresponding concepts. Using , we demonstrate improvements on downstream recognition, robustness, retrieval, and compositionality benchmarks with CLIP (up to +14% for ImageNet, +4% for SugarCrepe), and on nearest neighbors and linear probe for ImageNet with DINO (+3% each). We also showcase a powerful ability to intervene and select individual concepts to further improve downstream task performance (up from +4% to +9% for SugarCrepe) and use this ability to study the robustness of 's representation structure. Finally, we provide insights through ablation experiments and visualization of learned concepts.





§ INTRODUCTION



Selective attention is an intrinsic property of human perception  <cit.>, enabling people to focus on task-relevant aspects of their surroundings while tuning out the rest <cit.>. For instance, it allows a driver to focus on traffic signals, road signs, and other vehicles, or an individual searching for their friend in a crowd to attend to details like people's heights and clothing.
More broadly, it empowers us to comprehend the vast complexity of our world with limited mental resources <cit.> by constraining our perception to one concept at a time.
When a task requires information for multiple concepts (, looking at traffic signals and pedestrians), this constraint calls for sequentially attending to separate salient aspects of our stimuli <cit.>.
The cognitive bottleneck of selective attention, with the ability to separately represent concepts from our environment, enables human perception to remain insensitive to irrelevant distractions and generalize to new compositions of perceivable concepts.

In machine learning, the notion of selective attention is party emulated by the attention mechanism <cit.>, which is a core component of transformers <cit.> such as the vision transformer (ViT) <cit.>.
However, transformer encodings learnt using approaches like CLIP <cit.> and DINO <cit.> still struggle with robustness <cit.> and compositional generalization <cit.>. For instance, despite CLIP emerging as the de facto backbone for many vision tasks, it can perform close to chance when evaluated for compositionality <cit.>.
While the attention mechanism equips these models with the ability to emphasize relevant aspects of their intermediate representations, it does not endow them with the means to produce encodings for separately-attended concepts. Consequently, transformer-based representation learning models lack an important prior from human cognition: that perception of a complex input can be broken down into perception of its salient concepts.

This limitation deprives downstream tasks easy access to task-relevant aspects of the data's underlying structure. In a task like ImageNet classification, where most images contain one primary object of interest (, “dog”), information like relationships between objects and attributes of the background are irrelevant and can lead to overfitting if considered. However, these details can become crucial to attain high performance in other tasks like image retrieval (, “dog chasing a frisbee in a park”). Similar to the generalization benefits of selective attention in human perception, encodings that allow easy separation of attended concepts in the data can generalize to a diverse set of downstream tasks. Without additional constraints, however, the emergence of such compositional structure in encodings is unlikely, especially in the common setting of training on noisy internet data.

We propose  (Separate-head attention read-out)[Source code: <https://github.com/ankitkv/sparo-clip>.] as an improved read-out mechanism for transformer encoders in vision inspired by selective attention.
replaces the last transformer block to provide a mechanism for partitioning encodings into slots of separately-attended concepts.
We design each slot encoding to be a low-dimensional result of a single-head attention operation, producing a bottleneck that encourages each slot to “selectively attend.”
In this mechanism, producing multiple slots can be interpreted as selective attention over different concepts occurring in parallel.
When training CLIP, imparts an inductive bias that both modalities are views onto a shared compositional world with the corresponding slots of both encoders representing its concepts. In contrast, standard CLIP merely embeds the encodings of both modalities in a shared vector space without imposing any additional structure.

Training using CLIP and DINO exhibits improved generalization, robustness, and compositionality while retaining the same model size.
Our experiments show that using CLIP with improves generalization for zero-shot recognition, robustness, retrieval, and compositionality (, +14% for ImageNet, +4% for SugarCrepe, when trained with 15M examples; +3% each when trained with 400M examples). We also demonstrate improved linear probe performance of CLIP (+10% trained with 15M examples; +2% with 400M examples), as well as both linear probe and nearest neighbors performance of DINO (+3% each). We showcase how 's representational structure enables the ability to manually intervene and select relevant concepts for downstream tasks, leading to improved compositionality and generalization (up from +4% to +9% for SugarCrepe trained with 15M examples; +3% to +6% with 400M examples). We then study the robustness of this structure, validate the choice of separate-head attention through ablations, and provide insights into the effect of the number and dimensionality of concepts. Finally, we provide visualizations for some of the concepts learns to attend to.



§ RELATED WORK



We situate our work amongst other transformer read-out mechanisms and slot representation learning methods.



  
Transformer read-outs.
For transformer <cit.> decoders with causal masking, the last hidden state, corresponding to the end of sequence () input token, is used as the output representation. BERT <cit.> proposed inserting a special  token before input sequences for transformer encoders where information from all positions is gathered to produce the encoding. Using the  token output is also popular for Vision transformer (ViT) <cit.> backbones, along with alternatives such as global average pooling <cit.>. In the context of language, Multi-CLS BERT <cit.> argues that an input text sequence can have multiple facets, and that samples can be similar along one facet and dissimilar along another. The authors aim to represent these facets by adding multiple  tokens to the input sequence. Using separate attention heads for each slot can also encourage attending to different facets of the inputs, but without requiring expensive accumulations through each layer of the backbone. By keeping independent of the backbone, we also allow for its easy potential placement on top of frozen pre-trained models without needing to train for added  tokens.

is most related to the attentional pooler <cit.> architecture, which uses multi-head attention with embedded learned queries to produce output encodings. However, employs a critical structural constraint of not mixing the information from separate attention heads, a design choice we justify empirically in <ref>. Slot attention <cit.> also produces slot encodings through the attention mechanism, but rather than having the input positions compete for each slot, the slots compete for each input position. Slot attention also proposes an expensive iterative read-out mechanism where the role of each slot is determined by sampling from a per-sample Gaussian distribution, which can be limiting for large real-world datasets.



  
Slot representations. Object-centric slot representations have been studied extensively in the literature <cit.>. Often, these representations are expected to represent predefined objects and attributes, and evaluated in small-scale or synthetic settings. Recently, the use of slot attention on features produced by a learned encoder <cit.>, and utilizing unsupervised saliency masks for object discovery <cit.> has enabled unsupervised learning of object-centric representations in larger scale settings. Related to these ideas, attends to features produced by a transformer using a single attention head per slot, and each head is capable of learning saliency maps without supervision <cit.>. However, we do not constrain our slots to correspond to predefined types of objects in the data, nor do we impose any independence conditions between the slots. Self-supervised simplicial embeddings (SEM) <cit.>, with only a soft-discrete slot structure constraint, enable improved generalization on downstream tasks. replaces the separation of  application between SEM slots with a separation of underlying attention operations.



§ METHOD





Having motivated our approach, we now detail the design of . We start by introducing the relevant notation, then discuss the module. Finally, we discuss the addition of in CLIP and DINO, and the applicability of beyond these settings.




 §.§ Notation



We envision as a special attention-based layer that modifies the top block of a transformer encoder <cit.> backbone (including the ViT <cit.>). Transformers take as input a sequence = {_1, …, _n} of length n, and produce an output state per position {_1, …, _n}, where _i ∈^input dim, _i ∈^d, for any i ∈ [n], with d the model width. Finally, a pooling operation reduces the sequence of output states to output encodings with a fixed size ∈^M. Examples of pooling operations include attentional pooling <cit.>, global average pooling, and extracting the  or  token representations for images or text respectively.



 §.§ Separate-head attention read-out ()



We replace the pooling operation of transformer encoders with a concatenation of outputs of L single-head attention mechanisms. Our module, , acts on the transformer outputs = {_1,…,_n}∈^n× d to produce the encoder output = (_1,…,_L) ∈^LV, where slot _l ∈^V, l ∈ [L]. We illustrate the architecture of in <ref>. Each slot is produced as follows:

    _l    = _l ^(_l^_l/√(D)).

Here, _l ∈^D× d is a learned key projection weight, and _l ∈^D is a learned query embedding for attention. We parameterize the value projection as _l where ∈^V× D is learned parameter shared between all L slots. Decomposing the value projection in this manner helps reduce the number of model parameters, allowing a larger choice of L with fixed resources.

Generally, we pick V and D to be significantly smaller than the transformer width. In practice, we use the typical value of D=64 used in standard multi-head attention modules, and set V to the same value in our experiments. Therefore, each slot has limited expressivity and representational capacity. The bottleneck underpinning each slot is a result of these constraints and the mechanistic bottleneck of a single attention mechanism instantiating competition between its input features.







can be thought of as multi-head attention <cit.> where the queries are learned embeddings, key and value projection weights are shared, and the output projection weight is block-diagonal with each block containing the same parameters. Assuming L=V=D=√(d), requires a total of d^2 + 2d parameters, compared to multi-head attention's 4d^2.




 §.§ CLIP with 



We add to both the image and text encoders of CLIP <cit.>, setting the same values of L and V in both encoders. Each slot is ℓ_2-normalized separately, and we divide their concatenation by √(L) to produce the global ℓ_2-normalized encoding. With this scheme, the cosine similarity (^i, ^t) between the image encoding ^i and the text encoding ^t becomes the expected cosine similarity between corresponding slot encodings:

    (^i, ^t)
        = ^i^⊤^t
        = ∑_l=1^L ^i_l^t_l^i_l^⊤^t_l/^i_l^t_l
        = ∑_l=1^L 1/L(^i_l, ^t_l)
        = _l∼ [L][(^i_l, ^t_l)].


We use the standard CLIP loss to train the model, which maximizes (^i, ^t) of encodings ^i and ^t for aligned image-text pairs and minimizes it across all other pairings. With 's separate-head attention bottleneck, maximizing _l∼ [L][(^i_l, ^t_l)] for aligned pairs across the training dataset encourages both modalities to learn similar attention semantics per slot l ∈ [L]. Consequently, imparts a prior that both modalities encode a shared world with L selectively attendable concepts, and similarities between inputs can be expressed as their average similarity for these concepts.








 §.§ DINO with 



The online and momentum encoders of DINO <cit.> are each comprised of a ViT backbone which produces the encoding, and a DINO head that transforms it into a categorical distribution for distillation. Unlike CLIP, these encoders have an inherent bias for learning similar encoding functions due to one being an exponentially moving average of the other. However, in standard DINO, there are no constraints for selective attention in the encoding structure. We add to DINO to impart the prior for representing separately-attendable concepts in its encodings. Concretely, we replace the final transformer block of the ViT backbone with the , but do not modify the DINO head.



§ RESULTS



Over <ref>, we evaluate our method on a variety of downstream tasks, analyze 's representation structure with the ability to intervene on selected slots, perform ablations to support our design, illustrate the effect of changing 's hyperparameters, and visualize example learned concepts.

In this section, we validate that partitioning the representation structure as a collection of separately-attended concepts leads to improved generalization for downstream recognition, robustness, compositionality, and retrieval, using with CLIP <cit.> in <ref> and DINO <cit.> in <ref>.



  
Datasets. We train our CLIP models on one of Conceptual Captions 3M (CC3M) <cit.>, Conceptual Captions 12M (CC12M) <cit.>, a combination of CC3M and CC12M (CC15M), or LAION-400M (L400M) <cit.>. <ref> provides the statistics for our training datasets.



  
Models. For all experiments in this section, we train CLIP models using the open-source OpenCLIP <cit.> project. We consider two model sizes: those with a ViT-B/16 visual backbone (CLIP16) and those with a ViT-B/32 backbone (CLIP32). For models trained with , we replace the last transformer block of the base transformer backbones. For our chosen settings, this ensures that the resulting model size is comparable to the original model size. Further, since only the  token output from the image transformer and the  token output from the text transformer are used for the standard CLIP encodings, removing one transformer block ensures that we do not gain an unfair advantage by attending to positions from a layer that standard CLIP discards. Note that standard CLIP also retains the MLP as a part of its final transformer block, which is absent in the module. We also compare our encodings with those produced by CLIP with global average pooling (CLIP+GAP), which enjoys the advantage of being able to use all of the final layer outputs. When using , we use L=V=64 when training on Conceptual Captions, and L=128,V=64 when training on LAION-400M.




 §.§ Zero-shot recognition, robustness, and compositionality





We evaluate the zero-shot classification accuracy of trained CLIP models on ImageNet <cit.> and a set of robustness benchmarks including ImageNet-V2 <cit.>, ImageNet-Sketch <cit.>, ImageNet-A <cit.>, ImageNet-R <cit.>, and ObjectNet <cit.>. Additionally, to evaluate the compositionality of our learned encodings in terms of objects, attributes, and relations, we test our trained models on SugarCrepe <cit.>. We report our results in <ref>, and provide more fine-grained SugarCrepe numbers in <ref>. We find that using encodings outperforms CLIP and CLIP+GAP in all settings considered. The improved results on the SugarCrepe benchmark suggests that encodings exhibit compositionality at a level that enable better binding of object, attribute, and relation properties than baseline encodings on average.



Finally, we also use the VTAB benchmark <cit.> to evaluate on a diverse set of datasets <cit.> that vary more significantly from the training data. We present the relative zero-shot classification accuracy improvements of CLIP+over CLIP+GAP in <ref>. We see that outperforms standard CLIP encodings for a majority of VTAB tasks in all four settings.





 §.§ Zero-shot image and text retrieval





To separately evaluate the quality of the learned image and text encodings, we consider zero-shot retrieval based on image and text on MS COCO <cit.>, Flickr8k <cit.>, and Flickr30k <cit.>. Our results are presented in <ref>. To compute Recall@5 for `Text' we find the nearest 5 texts for an image, calculate the fraction of those texts that match the image, and average this metric over all the images (similar for Recall@5 for `Image'). We find that encodings improve retrieval across all three datasets for both modalities.




 §.§ Linear probe and DINO nearest neighbors classification





We evaluate the classification accuracy of linear probes trained on the encodings of our trained models.
For CLIP models, we follow <cit.> where we take a subset of the training data to use as a validation set, and sweep for the best weight decay hyperparameter on it for each model setting. Additionally, we evaluate for self-supervised learning in vision using DINO <cit.>. We train DINO with a ViT-S/16 backbone on ImageNet <cit.> without labels with 8 GPUs for 100 epochs without changing any other hyperparameters. Here, we use with L=V=64, replacing the last transformer block of the ViT encoder. The trained DINO models are evaluated on k-nearest neighbours with k=20 and linear classification on ImageNet with labels. We report our results for CLIP and DINO in <ref>. Our results for CLIP align with our previous results, with outperforming CLIP and CLIP+GAP in all linear probe settings. Furthermore, we see that a straightforward incorporation of in the DINO encoder outperforms standard DINO for both nearest neighbors and linear classification. We provide additional experimental details in <ref>.




§ ANALYSIS



In this section, we provide insights into encodings learned by . We demonstrate the ability of using 's slot structure to perform post-hoc selection of slots for improving downstream performance. We then show that without this structure, such selection is more prone to overfitting. Next, we validate the design choice of using single-head attention to model the notion of selective attention through controlled ablations. We analyze the effect of 's hyperparameters on performance and model size. Finally, we visualize example concepts that CLIP models trained with learn to attend to.




 §.§ Post-hoc concept selection





enables downstream tasks access to a collection of concepts. However, not all concepts may be important for a specific downstream task. For example, attending to watermarks or caption style can be important when trying to determine the source of a sample, but is unnecessary and often harmful for most downstream tasks that focus on the elements in the scenes. The separation of concepts in the structure of a encoding makes it possible to manually intervene to select desirable slots.

We demonstrate this ability with a simple heuristic: we perform concept selection by computing the zero-shot ImageNet validation accuracy of each slot separately, and pick only the top-performing 32 slots. In <ref>, we show that such an intervention exposes improved compositionality when evaluating CLIP+models trained on CC15M and LAION-400M on SugarCrepe.




 §.§ Robustness of 's slot structure





We further utilize the notion of post-hoc concept selection to study the robustness of 's representational structure, and compare it with that of standard encodings. We set up the experiment as an inverse of <ref> — we perform concept selection using the smaller SugarCrepe benchmark and evaluate the resulting encodings on real-world benchmarks. SugarCrepe contains 7,512 example pairs of positive and negative captions, differing only in compositional interpretation in 7 ways, over a set of 1,561 images. Compositionality is a shared high-level notion present in most real-world data, yet there is much information that exists in data beyond what is necessary to solve the SugarCrepe tasks. The small size of the SugarCrepe as a training dataset, together with the ubiquitous nature of compositionality in machine learning tasks, makes it an ideal candidate to overfit concept selection to for stress-testing the representational structure of by evaluating on other tasks.

We train a global mask ∈ [0,1]^M of the image encodings ∈^M to maximize SugarCrepe performance, and evaluate the masked encodings ⊙ on VTAB tasks and ImageNet for zero-shot classification. We train masks for encodings of frozen CLIP32+GAP and CLIP32+models trained on CC15M with M=512 and M=4096 each. For the models, we use L=64 and V=M/L. For each of these settings, we consider two types of masking: per-dimension and per-slot. To compare the CLIP32+GAP encodings on per-slot masking, we consider their L contiguous equal-sized partitions as the slot encodings. We provide full implementation details of our mask training setup in <ref>.

We expect the compositionality of the data to be encoded in concepts that are useful for real-world downstream benchmarks. Therefore, the aggressive selection of only the concepts that perform well on a compositionality benchmark should not significantly impact downstream performance. However, since SugarCrepe does not represent all possible compositional manipulations, we expect to see a drop in performance on some tasks due to overfitting. We present our results in <ref>. We find that when the mask structure is aligned with the slot structure of , overfitting concept selection improves the average performance on VTAB tasks, and incurs only a minimal drop in performance on ImageNet. Furthermore, even in settings where the masks are not trained to align with 's slot structure, we see that remains more resilient to overfitting than standard encodings in a majority of settings.





 §.§ Ablating the bottleneck and cross-modal alignment of heads





We evaluate the generalization benefits of the separate-head attention bottleneck and of aligning the attention heads between the modalities in CLIP. We set up on one end of the ablation spectrum and the cross-attention read-out of attentional pooler (AttPool) <cit.> to the other. AttPool uses multi-head attention to attend to the backbone encoder's outputs using embedded learned queries whereas uses separate single-head attention mechanisms. Additionally, AttnPool performs layer normalization <cit.> followed by a linear projection on the output of attention to produce the output encoding, whereas directly uses the output of separate-head attention as the encoding.

We train CLIP32, CLIP32+GAP, and variants of CLIP32+on CC15M with L=128,V=64 and evaluate for ImageNet zero-shot accuracy and model size. We start by considering the choice of either separate-head or multi-head cross-attention over the backbone outputs. For each cross-attention setting, we apply layer normalization and linear projection operations to the attention module's outputs, similar to AttnPool. However, we additionally also consider slot-wise variants of these operations. Here, a slot-wise operation f_ϕ is one that acts separately on each input slot _l using the same parameters ϕ, and produces the slot-structured output (f_ϕ(_1), …, f_ϕ(_L)). When applied to an input that is not slot-structured, we assume a slot structure by partitioning the input into L contiguous equal-sized partitions. With , using both slot-wise layer normalization and slot-wise projection enables each read-out attention head from one modality to align with the corresponding head from the other without being influenced by the other heads.

We present our results in <ref>, showing a clear advantage of over AttnPool. Additionally, both the choices of separate-head attention over multi-head and slot-wise operations over the alternatives result in improved generalization and a reduced model size. While the former supports the effectiveness of the separate-head attention bottleneck, the latter validates the benefits of the prior for shared separately-attendable concepts between the modalities enabled by .





 §.§ Hyperparameters L and V





encodings are concatenations of L separately-attended slots of V dimensions each. We study the effect of L and V with CLIP32+trained on CC15M, comparing its ImageNet zero-shot accuracy and model size with that of CLIP32 and CLIP32+GAP in <ref>. Generally, performance is poor for very low values of L and V, and improves when they are increased. Improvements remain monotonic with L in our experimental range, but the returns diminish for larger values of L. For each increment of L, we add additional _l and _l parameters to the model, and considering the trade-off between model size and performance gains becomes vital. Unlike L, increasing V beyond optimal values leads to a performance drop. With very large values of V, we produce low-rank slots that can potentially encode arbitrary harmful biases for downstream tasks.




 §.§ Visualizations





To conclude our analysis, we qualitatively visualize the concepts represented by using the CLIP16+model trained on CC15M from <ref>. A limitation of visualizing attended positions is that it only provides partial insight into the represented concept — we can tell where the information is taken from, but not how the information is used. Different slots can have similar attention maps for the same sample, but reveal different semantics when comparing patterns across other images. For instance, a slot that attends to animals (dogs, cats, elephants) and another that attends to transportation (trains, boats, cars) can both have the same attention mask over patches and words for a horse being used for transportation. Furthermore, not all concepts align between the two modalities for every sample due to information being present in one that is not possible to infer from the other.

To ease our interpretation of attended concepts, we choose concepts to visualize that have sharp attention in the text modality, attend to non-overlapping text tokens, and have cosine similarity  more than 0.75 with the corresponding image slots. We present examples of these filtered concepts in <ref>. In this example, we see that is capable of attending to different concepts such as subject of the scene, the activity represented, and the surrounding location, for both vision and text. We provide additional visualizations in <ref>.



§ DISCUSSION



Humans utilize a prior for compositionality to coherently represent different subsets of salient aspects of the world. We can enrich our comprehension by selectively attending to new aspects or specialize to a task by filtering out distracting ones. We introduce with a goal of imparting a similar prior to transformers in representation learning frameworks by partitioning encodings into separately-attended concepts. Although we see evidence of disentanglement through positive results of post-hoc concept selection and visualization of attention maps, we do not impose any explicit independence or disentanglement constraints. In our representation learning frameworks, the pressure for learning distinct concepts arises from the need to explain the factors of variation through the training objective while remaining bottlenecked to represent features that can be accumulated by a single head of attention with embedded queries and limited dimensionality. However, auxiliary training objectives to impose stronger distributional conditions on the learned concepts can be explored as promising future directions. Beyond pretraining, we highlight that more sophisticated post-hoc concept selection approaches than explored in our work, like using human interaction or set cover algorithms, can further improve the downstream utility of .



§ ACKNOWLEDGMENTS
  

This research was funded by Sony and enabled in part by compute resources provided by the Digital Research Alliance of Canada, Mila, and Sony.



    splncs04
    


figuresection
tablesection
equationsection



CHAPTER: APPENDIX



§ PYTORCH IMPLEMENTATION OF 



We provide an example implementation of in PyTorch <cit.> as follows:
[basicstyle=]
class Sparo(nn.Module):
    def __init__(self, d, L, V, D, grp_size=1):
        """See Sec. 3 for descriptions of d, L, V, D.
        Each key/value projection weight is shared between grp_size slots."""
        super().__init__()
        self.L, self.V, self.D = L, V, D
        self.grp_size, self.nkeys = grp_size, L // grp_size
        self.scale = D ** -0.5

        self.q = nn.Parameter(torch.randn(1, self.nkeys, grp_size, 1, D))
        self.KVproj = nn.Linear(d, self.nkeys * D)
        self.Wproj = nn.Linear(D, V)

        nn.init.xavier_uniform_(self.KVproj.weight)
        nn.init.zeros_(self.KVproj.bias)
        nn.init.zeros_(self.Wproj.bias)

    def forward(self, x, eos_idx=None):
        """Provide eos_idx for text to prevent attending to positions after it."""
        batch_size, n, d = x.shape
        kv = self.KVproj(x).view(batch_size, n, self.nkeys, 1, self.D)
        kv = kv.expand(-1, -1, -1, self.grp_size, -1).permute(0, 2, 3, 1, 4)

        attn = (self.q * self.scale) @ kv.transpose(-2, -1)
        if eos_idx is not None:
            attn_mask = (
                torch.arange(n, device=eos_idx.device)[None, None, None, None, :]
                > eos_idx[:, None, None, None, None]
            )
            attn_mask = torch.where(
                attn_mask, -torch.inf, torch.zeros_like(attn_mask, dtype=attn.dtype)
            )
            attn += attn_mask
        attn = attn.softmax(dim=-1)
        x = (attn @ kv).view(batch_size, self.L, self.D)
        return self.Wproj(x)


Here, we optionally utilize multi-query attention <cit.> to reduce the model size further by setting grp_size>1, allowing larger values of L with the same resources. For example, when grp_size=2 and L=64, only 32 _l parameters are created.
We encourage practitioners to explore optimal grp_size settings for their applications.





§ ADDITIONAL RESULTS







 §.§ Fine-grained SugarCrepe performance





The SugarCrepe <cit.> benchmark comprises of seven compositional manipulations of captions for a subset of MS COCO <cit.>. The benchmark evaluates the zero-shot accuracy of vision-language models for picking the ground-truth captions over the generated hard negatives. We provide the fine-grained results of our models for each of the seven SugarCrepe manipulations in <ref>. Individual categories often demonstrate trade-offs amongst each other, but the averaged performance reveals consistent trends of outperforming the baselines. Additionally, we show that the compositionality exposed in learned encodings across different categories can be enhanced by performing post-hoc concept selection, demonstrated in <ref>.




 §.§ Comparison with other bottlenecks





imposes a representational bottleneck by constraining each slot to the output of a single head of attention with fixed queries and limited dimensionality. In this section, we compare our bottleneck with two other representational bottlenecks: a linear (affine) information bottleneck and a discrete bottleneck using Finite Discrete Tokens (FDT) <cit.>.



  
Setup. We add the considered bottlenecks and to CLIP32 and CLIP16, and train on different sizes of Conceptual Captions. We empirically verified that removing a block of the transformer backbone reduces the performance for all models. Still, we follow the typical setup of replacing the last transformer block of the backbone when using . However, we do not handicap other bottleneck baselines, adding the corresponding bottlenecks after the last transformer block.



  
Linear information bottleneck details. We implement a linear information bottleneck by setting the embedding size of standard CLIP to m and adding an affine transformation to an output size of M on top of it, with m < M. We evaluate two settings for the m → M bottleneck: 64 → 512 and 512 → 4096.



  
Discrete bottleneck details. We use Finite Discrete Tokens (FDT) <cit.> to implement a discrete bottleneck. FDT encodes images and texts as discrete codes belonging to a shared learned codebook. During training, a convex combination of the discrete codes is produced using the Sparsemax operation <cit.>, which the authors show performs better than Softmax. We follow the hyperparameters used by the authors: the codebook contains 16,384 codes of 512 dimensions each.



  
Results. We present our results in <ref>. First, we see that a simple linear bottleneck performs worse than standard CLIP in a majority of settings, indicating that the mechanism behind the representational bottleneck is important. As expected, the discrete bottleneck widely outperforms or stays comparable to standard CLIP. However, note that global average pooling (GAP) continues to shine as a strong baseline, outperforming the discrete bottleneck in most settings. Finally, we verify that outperforms or remains comparable to the best baselines across all settings.









 §.§ Additional visualizations





We present additional examples of learned concepts in <ref> following the filtering described in <ref>. Each row illustrates the attended positions of the single-head attention operation for the same slot across examples from MS COCO <cit.>. We reuse the same last input for the first two visualized slots to highlight cases where visually similar attention maps can correspond to distinct concepts. Looking at attended positions allows us to interpret where each slot takes information from, but not how that information is used.




§ REPRODUCIBILITY







We present the details of the training datasets used in our experiments in <ref>. Note that the datasets we use for training our CLIP models are potentially smaller than those used by prior work, as images become unavailable for download over time for the datasets we consider. <ref> lists the open-source repositories we utilized and modified for training and for the evaluations presented in our work. Finally, note that <ref> provides example PyTorch code for implementing the module.




 §.§ CLIP







  
CLIP training. We use the open-source OpenCLIP <cit.> project for training our CLIP models. The important hyperparameters that we change from the OpenCLIP defaults for training our models are listed in <ref>.




  
Linear probe training. We use the open-source CLIP benchmark <cit.> project for training and evaluating linear probes for our trained OpenCLIP checkpoints. We train our linear probes on ℓ_2-normalized CLIP encodings using the AdamW <cit.> optimizer with a learning rate of 0.1 for 10 epochs. For each setting, we first sweep for the optimal weight decay hyperparameter by training on a subset of the training split and evaluating on the remaining samples. Finally, we train on the full training set with the chosen weight decay value. The design of the hyperparameter sweep follows that of <cit.>.




  
Training masks on SugarCrepe. 
Our training set is prepared by combining the splits of the SugarCrepe benchmark without any extra balancing. For each training sample, we obtain three encodings: the image encoding ^i, positive text encoding ^p, and the negative text encoding ^n. We produce the masked image encoding ^i = ⊙, where  is generated as = σ(1/4max(100, exp(α)) ). Here, σ is the element-wise sigmoid function, α∈ and ∈^M' are learned parameters, and max(100, exp(α))/4 is the (inverse) temperature for the sigmoid that we found to help in faster convergence and stability of training. We have M'=L for per-slot masking and M'=M for per-dimension masking. We use the cosine similarities ^i·^p / ^i^p and ^i·^n / ^i^n, scaled by max(100, exp(α)) similar to CLIP <cit.>, as the logits for the cross-entropy loss that we minimize. For each setting, we train for 100 epochs using SGD, with learning rate 0.02 and momentum 0.9. All settings converge before 100 epochs, and we pick the mask which achieves the best fit on standard SugarCrepe evaluation during training.




 §.§ DINO



We use the official DINO repository <cit.> to train and evaluate our DINO models. For training, we follow the instructions for “Vanilla DINO training,” using a ViT-Small backbone trained on a single node with 8 GPUs for 100 epochs. Similarly, we use the default hyperparameters for k-nearest neighbors and linear probe evaluations. The linear probe is trained on the layer normalized  token outputs from the last 4 ViT blocks. When training the linear probe on DINO+, we use the unnormalized encoding instead of the last block's normalized  output.


