<figure><img src="image_1.png" /></figure>

HAL
open science

# Contributions to the Design and Training of Transformers in Computer Vision

Alaaeldin Mohamed Elnouby Abdallah Ali

<figure><img src="image_2.png" /></figure>

## To cite this version:

Alaaeldin Mohamed Elnouby Abdallah Ali. Contributions to the Design and Training of Transformers in Computer Vision. Computer Science [cs]. INRIA Paris; ENS Paris - Ecole Normale Supérieure de Paris; PSL University, 2023. English. NNT: tel-04477587

HAL Id: tel-04477587

https://inria.hal.science/tel-04477587

Submitted on 26 Feb 2024

HAL is a multi-disciplinary open access archive for the deposit and dissemination of scientific research documents, whether they are published or not. The documents may come from teaching and research institutions in France or abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est destinée au dépôt et à la diffusion de documents scientifiques de niveau recherche, publiés ou non, émanant des établissements d'enseignement et de recherche français ou étrangers, des laboratoires publics ou privés.

<figure><img src="image_3.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# Contributions to the Design and Training of Transformers in Computer Vision

Soutenue par

**Alaaeldin Ali**

Le 11 Juillet 2023

École doctorale n°386

**Sciences Mathématiques de Paris-Centre**

Spécialité

**Informatique**

Préparée au

Willow - Meta AI

## Composition du jury :

Tinne TUYTELAARS
KU Leuven

Président du jury
Examinatrice

Paolo FAVARO
University of Bern

Rapporteur

Cees SNOEK
University of Amsterdam

Rapporteur

Alexei EFROS
University of California Berkeley

Examinateur

Natalia NEVEROVA
Meta AI

Membre Invitée

Ivan LAPTEV
Inria and DI/ENS

Directeur de thèse

Hervé JEGOU
Meta AI

Co-directeur de thèse
<figure><img src="image_4.png" /></figure>
Abstract

Transformers have revolutionized representation learning across modalities, achieving state-of-the-art results in natural language processing, computer vision, speech, and beyond. This thesis explores the potential of Transformer models for computer vision. We propose architectural innovations to overcome their limitations, developing sample-efficient self-supervised pre-training methods, and advancing multimodal learning with Transformers. First, we propose Cross-Covariance Attention to reduce the quadratic complexity of self-attention achieving similar performance as vision transformers with lower memory footprint and computational cost, enabling the application of vision transformers to higher-resolution images. We then investigate self-supervised pre-training for vision transformers. We propose SplitMask, a denoising autoencoding method based on masked image modeling. Unlike joint embedding methods, SplitMask does not require large-scale pre-training datasets and can be applied to diverse visual data. SplitMask matches the performance of joint embedding methods when pre-trained on datasets two orders of magnitude smaller, highlighting its improved sample efficiency. Moreover, we apply masked image modeling to neural image compression in the form of an improved entropy model yielding a strong rate-distortion performance and enabling the compression of images to the size of a short SMS or tweet. Finally, we propose ImageBind, a method for learning a shared embedding space across six modalities. ImageBind leverages the abundance of images and text on the web to enable transfer to modalities with scarce annotations like depth, thermal, audio, and IMU. In summary, this thesis demonstrates the potential of Transformers for computer vision through architectural innovations, new self-supervised objectives, and multimodal knowledge transfer. The methods proposed in this thesis push the boundaries of transformers in vision by enhancing their scalability and generality, enabling more sample-efficient representation learning, and facilitating transfer across modalities.

**Keywords :** vision transformers, self-supervised learning, weakly supervised learning, multimodal learning, image compression
# Résumé

Les transformateurs ont révolutionné l'apprentissage de la représentation dans de nombreuses modalités, obtenant des résultats de pointe dans le traitement du langage naturel, la vision par ordinateur, la parole et bien d'autres domaines. Cette thèse explore le potentiel des modèles de transformateurs pour la vision par ordinateur. Nous proposons des innovations architecturales pour surmonter certaines de leurs limites. Nous développons des méthodes de pré-entraînement auto-supervisé efficaces en termes d'échantillons, et considérons l'utilisation de ces transformateurs dans un contexte d'apprentissage multimodal.

Dans un premier temps, nous proposons l'attention à covariance croisée pour réduire la complexité quadratique de l'attention d'origine et obtenir des performances similaires avec une empreinte mémoire et un coût de calcul moindres, ce qui permet d'appliquer les transformateurs de vision à des images à plus haute résolution. Nous étudions ensuite le pré-entraînement auto-supervisé pour les transformateurs de vision. Nous proposons SplitMask, une méthode de débruitage automatique basée sur la modélisation d'images masquées. Contrairement aux méthodes de plongements conjointes, SplitMask ne nécessite pas d'ensembles de données de pré-entraînement à grande échelle et peut être appliqué à diverses données visuelles. SplitMask est aussi performant que les méthodes de plongements conjoints lorsqu'il est entraîné sur des ensembles de données deux fois plus petits, ce qui met en évidence l'amélioration de l'efficacité d'apprentissage avec peu de données. En outre, nous appliquons la modélisation d'images masquée à la compression d'images neuronales sous la forme d'un modèle entropique amélioré. Cela permet d'obtenir de bonnes performances en matière de débit-distorsion dans les régimes où la compression d'image est extrême, tels la taille d'un SMS ou d'un tweet. Enfin, nous proposons ImageBind, une méthode d'apprentissage d'un espace de plongement partagé entre six modalités. En résumé, cette thèse démontre le potentiel des transformateurs pour la vision par ordinateur grâce à des innovations architecturales, de nouveaux objectifs auto-supervisés et un transfert de connaissances multimodal. Les méthodes proposées dans cette thèse repoussent les limites des transformateurs en vision en améliorant leur passage à l'échelle et leur généralité, en permettant un apprentissage de la représentation plus efficace en termes d'échantillons, et en facilitant le transfert entre les modalités.

**Mots clés** : transformateurs de vision, apprentissage auto-supervisé, apprentissage faiblement supervisé, apprentissage multimodal, compression d'image.
# Acknowledgements

I would like to express my deepest gratitude to the many people who have made this thesis possible. First and foremost, I would like to thank my wife Sara for her unwavering support throughout my academic journey. Her encouragement, patience, and love have been my pillars of strength, and I could not have done this without her. I am deeply grateful to my parents for their constant support and encouragement. Their sacrifices and belief in me have been instrumental in shaping my academic and personal life. I owe them a debt of gratitude that can never be repaid. I would also like to thank my lovely sisters for their support and encouragement, especially during the tough times. Their love, guidance, and motivation have been invaluable to me.

I would like to express my sincere appreciation to my PhD advisors Hervé, Ivan, and Natalia for their guidance and support throughout my PhD journey. I feel truly fortunate to have had such exceptional mentors who dedicated their time, expertise, and energy to help me achieve my goals and have advocated for me countless times. In particular, I cannot overstate my gratitude to Hervé for his exceptional mentorship. He taught me to be ambitious, to never settle for easy wins, and most importantly, how to be a kind and compassionate leader. His unwavering support and guidance were instrumental in my success during the PhD program. I will forever be grateful for his mentorship and the impact he has had on my professional as well as personal development. I also want to thank Ivan for his support and guidance during the PhD program. His expertise and insights were invaluable in shaping the direction of my research and helping me navigate the challenges of the program. He consistently challenged me to ask the hard questions, pushing me to grow and excel in my research. Finally, I want to express my deep appreciation to Natalia, who played a crucial role in my PhD journey by presenting me with the opportunity and believing in my potential. Her guidance and mentorship during the first half of the program were invaluable, and she continued to be a great mentor throughout the program and hopefully beyond. I am also very grateful to the members of the jury. It is a great pleasure to present and defend my PhD work in front of researchers I admire and have been looking up to ever since I started my research journey. I would like to thank the rapporteurs, Cees Snoek and Paolo Favaro, and the examiners, Alexei Efros and Tinne Tuytelaars, for dedicating the time to review and provide feedback for the manuscript and my PhD work.

I can't express enough gratitude to my collaborators, who've played a crucial part in my academic journey. Working with Gautier was one of the highlights of my PhD program. He was an amazing collaborator, and we had some of the most fun and engaging projects together. Our discussions were always insightful, and I learned a great deal from working with him. I am grateful to Ishan, his ability to approach complex problems with clarity and creativity has been truly remarkable. From my interactions with Ishan, I have come to see him as a gold standard
for post-PhD research career success. I feel fortunate to have had the opportunity to work and learn from him during my PhD program. I also want to thank Rohit who was a very dedicated and inspiring collaborator. His contributions were invaluable and contributed significantly to the quality of our research. I want to thank Jakob for his mentorship and for providing valuable feedback and insights that helped me to shape my research. Our weekly discussions were always constructive and thought-provoking, and I learned a great deal from his expertise and experience. I would like to thank Matthew for being a great collaborator and supporter, he has a great work ethic and strong drive that was especially helpful when projects did not go as planned. I wish to thank Edouard, Armand, Piotr, Matthijs and Gabriel whom presence was very influential and inspiring for me throughout my PhD journey. I want to thank Federico and Joao, the amazing interns that I had the pleasure of working with. Finally, I would like to thank my fellow PhD students, Hugo, Mathilde, Adrien, Mohamed, Baptiste, Charlotte, Roberto, Pierre, Lina, Badr, Guillaume and Alex.

Once again, I would like to express my deepest gratitude to all these individuals who have contributed to the completion of this PhD thesis. I could not have done this without their support, guidance, and encouragement.
# Table of Contents

Abstract i

Résumé ii

Acknowledgements iii

Table of Contents v

1 Introduction 1

1.1 Challenges 2

1.1.1 Computational Complexity 2

1.1.2 Scalable and Sample Efficient Representation Learning 3

1.1.3 Multimodal Learning 3

1.2 Outline and Contributions 3

1.2.1 Taming the quadratic complexity of vision transformers 3

1.2.2 Sample efficient self-supervised pre-training with Transformers 4

1.2.3 Masked Image modeling for improved image compression 4

1.2.4 ImageBind for learning a shared embedding space for six modalities 5

1.3 Publications 5

2 Background 7

2.1 Modern Architectures in Computer Vision 7

2.1.1 Convolutional Neural Networks 7

2.1.2 Self-Attention and Transformers 8

2.1.3 Efficient Self-Attention 9

2.1.4 CNNs Augmented with Self-Attention 10

2.1.5 Vision Transformers 11

2.2 Self-Supervised Learning 12

2.2.1 Autoencoders 12

2.2.2 Pretext tasks 12

2.2.3 Joint embeddings methods 13

2.3 Vision-Language pre-training 15

2.4 Lossless Compression 15

3 Cross Covariance Image Transformers 16

3.1 Introduction 17

3.2 Related work 19

3.3 Method 20
Table of Contents

<table><tr><td>3.3.1 Background</td><td>20</td></tr><tr><td>3.3.2 Cross-covariance attention</td><td>22</td></tr><tr><td>3.3.3 Cross-covariance image transformers</td><td>24</td></tr><tr><td>3.4 Preliminary study on Vision Transformers (ViT)</td><td>26</td></tr><tr><td>3.4.1 Impact of resolution versus patch size</td><td>26</td></tr><tr><td>3.4.2 Approximate attention models in ViT with DeiT training</td><td>27</td></tr><tr><td>3.4.3 Training and testing with varying resolution</td><td>27</td></tr><tr><td>3.5 Experimental evaluation</td><td>28</td></tr><tr><td>3.5.1 Image classification</td><td>28</td></tr><tr><td>3.5.2 Object detection and instance segmentation</td><td>32</td></tr><tr><td>3.5.3 Semantic segmentation</td><td>34</td></tr><tr><td>3.5.4 Implementation details</td><td>35</td></tr><tr><td>3.6 Conclusion and Future Work</td><td>35</td></tr><tr><td>4 Sample Efficient Self-Supervised Pre-training with Vision Transformers</td><td>36</td></tr><tr><td>4.1 Introduction</td><td>37</td></tr><tr><td>4.2 Related Work</td><td>39</td></tr><tr><td>4.3 Analysis</td><td>40</td></tr><tr><td>4.3.1 Sample Efficiency</td><td>41</td></tr><tr><td>4.3.2 Learning using non object-centric images</td><td>42</td></tr><tr><td>4.3.3 Tokenizers</td><td>42</td></tr><tr><td>4.4 Methodology</td><td>44</td></tr><tr><td>4.4.1 SplitMask</td><td>44</td></tr><tr><td>4.4.2 Encoder-Decoder Architecture</td><td>44</td></tr><tr><td>4.4.3 Global Contrastive Loss</td><td>45</td></tr><tr><td>4.5 Experiments</td><td>45</td></tr><tr><td>4.5.1 Datasets</td><td>45</td></tr><tr><td>4.5.2 Dense Prediction</td><td>47</td></tr><tr><td>4.5.2.1 Object detection and Instance Segmentation</td><td>47</td></tr><tr><td>4.5.2.2 Semantic Segmentation</td><td>47</td></tr><tr><td>4.5.3 Image Classification</td><td>47</td></tr><tr><td>4.5.4 Pre-training using ImageNet</td><td>49</td></tr><tr><td>4.5.5 Ablations</td><td>49</td></tr><tr><td>4.5.5.1 SplitMask vs BEiT</td><td>49</td></tr><tr><td>4.5.5.2 Encoder-Decoder vs BEiT</td><td>50</td></tr><tr><td>4.5.5.3 Overfitting during pre-training</td><td>51</td></tr><tr><td>4.5.6 Implementation Details</td><td>51</td></tr><tr><td>4.6 Conclusion</td><td>52</td></tr><tr><td>5 Image Compression with Product Quantized Masked Image Modeling</td><td>53</td></tr><tr><td>5.1 Introduction</td><td>54</td></tr><tr><td>5.2 Related work</td><td>56</td></tr><tr><td>5.3 Product Quantized Masked Image Modeling</td><td>57</td></tr><tr><td>5.3.1 High-level architecture: PQ-VAE</td><td>58</td></tr><tr><td>5.3.2 Image entropy model</td><td>59</td></tr><tr><td>5.3.3 Training the PQ-MIM</td><td>61</td></tr><tr><td>5.4 Experiments</td><td>62</td></tr><tr><td>5.4.1 Experimental setup</td><td>62</td></tr><tr><td>5.4.2 Main experimental results</td><td>64</td></tr><tr><td>5.4.3 Analysis and Ablations</td><td>65</td></tr></table>
<table><tbody><tr><td>5.4.4 Limitations</td><td>67</td></tr><tr><td>5.5 Conclusion</td><td>68</td></tr><tr><td><b>6</b> Learning a shared embedding space for six modalities with ImageBind</td><td><b>70</b></td></tr><tr><td>6.1 Introduction</td><td>71</td></tr><tr><td>6.2 Related Work</td><td>72</td></tr><tr><td>6.3 Method</td><td>73</td></tr><tr><td>6.3.1 Preliminaries</td><td>73</td></tr><tr><td>6.3.2 Binding modalities with images</td><td>74</td></tr><tr><td>6.3.3 Implementation Details</td><td>75</td></tr><tr><td>6.3.3.1 Inference implementation details</td><td>75</td></tr><tr><td>6.4 Experiments</td><td>76</td></tr><tr><td>6.4.1 Emergent zero-shot classification</td><td>77</td></tr><tr><td>6.4.1.1 Zero-shot evaluation details</td><td>77</td></tr><tr><td>6.4.2 Comparison to prior work</td><td>78</td></tr><tr><td>6.4.3 Few-shot classification</td><td>79</td></tr><tr><td>6.4.3.1 Few-shot evaluation details</td><td>80</td></tr><tr><td>6.4.4 Analysis and Applications</td><td>80</td></tr><tr><td>6.4.4.1 Qualitative evaluation details</td><td>81</td></tr><tr><td>6.4.5 Pretraining details</td><td>82</td></tr><tr><td>6.4.6 Datasets and Metrics</td><td>82</td></tr><tr><td>6.5 Ablation Study</td><td>84</td></tr><tr><td>6.5.1 Scaling the Image Encoder</td><td>85</td></tr><tr><td>6.5.2 Training Loss and Architecture</td><td>85</td></tr><tr><td>6.6 Discussion and Limitations</td><td>88</td></tr><tr><td><b>7</b> Conclusion</td><td><b>89</b></td></tr><tr><td>7.1 Summary of Contributions</td><td>89</td></tr><tr><td>7.1.1 Reducing the quadratic complexity of Vision Transformers</td><td>89</td></tr><tr><td>7.1.2 Masked Image Modeling as a powerful tool</td><td>90</td></tr><tr><td>7.1.3 Learning representations for data-scarce modalities</td><td>91</td></tr><tr><td>7.2 Open Problems and Future Directions</td><td>92</td></tr><tr><td>7.2.1 Efficient Attention for Video Data</td><td>92</td></tr><tr><td>7.2.2 Improving the scalability of Self-supervised Learning</td><td>92</td></tr><tr><td><b>8</b> Résumé Substantiel</td><td><b>94</b></td></tr><tr><td>8.1 Défis</td><td>95</td></tr><tr><td>8.1.1 Complexité computationnelle</td><td>96</td></tr><tr><td>8.1.2 Apprentissage de la représentation à grande échelle et efficace en termes d'échantillons</td><td>96</td></tr><tr><td>8.1.3 Apprentissage multimodal</td><td>96</td></tr><tr><td>8.2 Présentation et contributions</td><td>97</td></tr><tr><td>8.2.1 Maîtriser la complexité quadratique des transformateurs de vision</td><td>97</td></tr><tr><td>8.2.2 Pré-entraînement auto-supervisé efficace avec Transformer</td><td>97</td></tr><tr><td>8.2.3 Modélisation d'images masquées pour une meilleure compression d'images</td><td>98</td></tr><tr><td>8.2.4 ImageBind pour l'apprentissage d'un espace d'intégration partagé pour six modalités</td><td>98</td></tr></tbody></table>
Table of Contents

<table><tr><td>A</td><td>Additional results for XCiT</td><td>100</td></tr><tr><td></td><td>A.1 More XCiT models</td><td>100</td></tr><tr><td></td><td>A.2 Image Retrieval</td><td>100</td></tr><tr><td></td><td><strong>References</strong></td><td><strong>104</strong></td></tr></table>
Table of Contents
# Chapter 1

## Introduction

Deep learning has had a profound impact on the field of computer vision. This impact has revolutionized the way we perceive, understand, and interact with the digital world. The rapid progress has driven the development of new algorithms and architectures that have continuously advanced the state of the art in computer vision tasks. The advancements in deep learning have outperformed classical computer vision methods, and have enabled a wide range of applications such as autonomous driving [Chen, 2015], image and video generation [Ramesh, 2022a; Singer, 2022], and medical image analysis [Shin, 2016], to name but a few. One of the key factors contributing to the success of deep learning in computer vision has been the innovation in network architectures [He, 2016; Szegedy, 2015; Hu, 2018; Wang, 2018]. These innovations have resulted in models that can better learn and represent complex patterns from large-scale data. Moreover, they have enabled the development of models that can learn hierarchical representations, which are crucial for understanding the structure in visual data.

Convolutional Neural Networks (CNNs) [Fukushima, 1980; LeCun, 1989] have long been considered the go-to architecture for computer vision tasks. CNNs are characterized by their inductive biases and sample efficiency, which enable them to learn and generalize effectively from a limited number of training examples. Their ability to learn local features and robustly represent spatial hierarchies has made them the dominant architecture for vision tasks. The success of CNNs can be attributed to their ability to exploit the spatial structure inherent in images, allowing them to efficiently learn meaningful representations. Despite their effectiveness, the same inductive biases and design restrictions that make CNNs efficient can potentially limit their generality as the amount of training data increases. For example, CNNs are known to have a strong bias towards texture [Geirhos, 2018], which can limit their ability to generalize to new and diverse image datasets.

The landscape of deep learning began to change with the introduction of the Transformer architecture [Vaswani, 2017b]. Originally designed for machine translation and natural language processing (NLP) tasks, Transformers have proven to be highly effective in capturing long-range dependencies and complex patterns in sequential data. The key innovation in the Transformer architecture is the self-attention mechanism, which allows the model to dynamically weigh the importance of different input elements in the context of the entire sequence. The success of Transformers in NLP tasks led researchers to explore their potential in the realm of computer
vision. Initially, self-attention mechanisms were integrated into existing CNN architectures, re-
sulting in hybrid models [Wang, 2018] that combined the strengths of both approaches. The
next logical step in this evolution was the development of Vision Transformers (ViT) [Doso-
vitskiy, 2021a], which marked a significant departure from traditional CNN-based approaches.
ViT models are composed entirely of self-attention and Multi-Layer Perceptron (MLP) mod-
ules, completely discarding convolutional layers. Despite this radical shift in architecture, ViTs
have achieved state-of-the-art results on a wide range of computer vision tasks, surpassing the
performance of their CNN-based counterparts. Moreover, Vision Transformers exhibit different
properties compared to CNNs, including higher robustness towards occlusions and perturbations
as well as lower bias towards texture [Naseer, 2021]. These properties make ViTs a promising
candidate for pushing the generalization performance for a wide range of vision tasks.

Transformers have opened up a vast landscape of possibilities for computer vision research. The expressiveness and generality of Transformer models make them highly adaptable to a wide range of tasks and domains. Moreover, their ability to model long-range dependencies enables capturing global context and complex relationships within visual data more effectively than CNNs. In this thesis, we delve deeper into the world of Transformers, exploring their potential, and applications in various computer vision tasks. We will investigate the factors that have contributed to their success and how to adapt them better for vision applications. Furthermore, we explore how the unification of architectures across modalities leads to innovative approaches for providing supervision to modalities with scarce data by leveraging strong vision models.

1.1 Challenges

Transformers pave the way to exciting new possibilities for learning stronger visual repre-
sentations via supervised and unsupervised learning as well as opening the door for more ho-
mogenous approaches for multimodal learning. In this manuscript, we address several challenges
related to these topics.

1.1.1 Computational Complexity

One significant challenge posed by the Transformers, as introduced by Vaswani et al. [Vaswani, 2017b], is the quadratic complexity with respect to the sequence length. The self-attention mechanism employed by transformers involves computing pairwise interactions between every element in the input sequence, resulting in a complexity of $O(N^2)$, where N is the sequence length. In the context of computer vision, this complexity poses a challenge for the direct application of transformers to high-resolution images, which is often required for tasks such as object detection, semantic segmentation, and image compression. To overcome this limitation, there is a need for architectural innovations that can reduce the computational cost of transformers for large-resolution image tasks.
1.1.2 Scalable and Sample Efficient Representation Learning

Transformers excel at capturing long-range dependencies and can learn hierarchical repre-
sentations effectively. Similar to their transformative impact on representation learning in Nat-
ural Language Processing with approaches like BERT [Devlin, 2018] and GPT [Radford, 2018],
Transformers have the potential to revolutionize self-supervised pre-training for vision tasks.
However, to fully realize this potential, it is critical to develop more sample-efficient methods
that can scale more easily in terms of data and compute compared to existing joint embedding
approaches. The expressive power of the Transformer architecture, when used in conjunction
with more generic and less biased denoising objectives, has the potential to deliver the efficiency
and scalability required to substantially improve self-supervised pre-training of vision models.

1.1.3 Multimodal Learning

As a universal architecture, the Transformer model has seen successful application across
various modalities. This versatility opens the door for shared design principles and components
that can facilitate multimodal learning. By using Transformers as a common framework for
multimodal learning, we can develop methods that effectively integrate information from diverse
modalities and enable seamless transfer of knowledge between different tasks and domains.
However, while there may be an abundance of data for some modalities, others may suffer
from severe data scarcity. In such cases, the knowledge transfer enabled by using a common
framework like Transformers can have a significant impact. By leveraging the strengths of
different modalities and transferring knowledge learned from one domain to another, we can
improve the performance of models in domains where data scarcity is a major issue.

1.2 Outline and Contributions

The manuscript begins with a discussion of the background and relevant literature in Chap-
ter 2. We then delve into the details of each of our four main contributions in the subsequent
chapters as outlined below.

1.2.1 Taming the quadratic complexity of vision transformers

In Chapter 3, we delve into the computational complexity of vision transformers [Dosovitskiy, 2021a]. Specifically, we identify that the self-attention operation, which lies at the heart of the Transformer architecture [Vaswani, 2017b], displays a quadratic growth rate with respect to image resolution. This can render vision transformers computationally prohibitive, particularly when tasked with processing higher-resolution images, as is frequently encountered in essential downstream computer vision applications such as semantic segmentation and object detection. To address this challenge, we propose a novel alternative formulation of the self-attention operation, which we refer to as *Cross-Covariance Attention* (XCA) [El-Nouby, 2021c]. Cross-covariance attention exhibits a linear complexity with respect to input size and can be seamlessly integrated as a replacement for self-attention in vision transformers. We introduce
XCiT, a novel architecture for computer vision with XCA at its core. We demonstrate that XCiT offers significant enhancements in terms of memory consumption and throughput while retaining the strong performance of vision transformers.

1.2.2 Sample efficient self-supervised pre-training with Transformers

The Transformer architecture has sparked innovation in self-supervised pre-training in Natu-
ral Language Processing, leading to the development of influential models such as BERT [Devlin,
2018] and GPT [Radford, 2018]. In computer vision, joint embedding methods [Chen, 2020c;
He, 2020; Caron, 2021] have emerged as dominant due to their strong off-the-shelf performance
and competitive performance with supervised methods for finetuning. However, joint embedding
methods suffer from some limitations, such as their dependence on hand-crafted data augmenta-
tion techniques tailored for ImageNet [Deng, 2009b]. Moreover, these methods can be challenging
to scale in terms of model size, as pointed out by [Chen, 2021b]. Motivated by the success of
BERT in NLP, in Chapter 4 we investigate the efficacy of denoising autoencoding methods when
used in conjunction with vision transformers for self-supervised pre-training. Firstly, we pro-
pose SplitMask [El-Nouby, 2021a], a novel self-supervised pre-training method based on masked
image modeling. Secondly, we find that SplitMask offers improved sample efficiency and can be
employed to learn robust representations using datasets orders of magnitude smaller than those
required by joint embedding methods. Additionally, we observe that SplitMask is well-suited for
training using a wider range of visual data as it is not biased towards a specific distribution of
images, unlike joint embedding methods that are typically biased towards object-centric images
of ImageNet.

1.2.3 Masked Image modeling for improved image compression

Neural image compression has emerged as a promising alternative to traditional codecs due to its superior perceptive and psychovisual image quality. Typically, neural image compression systems consist of three main components: a deep neural encoder and decoder, which learn a mapping and its inverse between pixel and latent representations of the image; a quantizer, which maps the continuous latent representations to a set of discrete symbols; and an entropy model, which exploits redundancies in the set of discrete symbols to reduce the final length of the bitstream. In Chapter 5, we propose (1) Adopting XCiT for designing the encoder and decoder; (2) Employing product quantization in place of vector quantization for better scaling in bit-rates; (3) A novel entropy model for neural image compression based on masked image modeling. Our method, PQ-MIM [El-Nouby, 2023], achieves strong rate reduction compared to simple frequency-based baselines while providing a significant speedup compared to autoregressive methods that can be prohibitively slow for high-resolution images. As a result, PQ-MIM enables extreme compression of images to a size of a short tweet or SMS.
1.2.4 ImageBind for learning a shared embedding space for six modalities

Transformers have emerged as a versatile architecture that achieves exceptional performance across various modalities, such as text [Devlin, 2018; Radford, 2018], images [Dosovitskiy, 2021a; Touvron, 2020], video [Gedas Bertasius, 2021; Tong, 2022], audio [Xu, 2022], and graphs [Yun, 2019], among others. This presents exciting possibilities for developing multimodal systems with shared components and similar designs. However, the predominant approach for training these modalities is supervised learning, which relies on learning a mapping from the sensory inputs to a set of categorical labels. Unfortunately, supervised learning is constrained by the challenge of efficiently and scalably collecting annotations. To overcome this challenge, weakly supervised learning has emerged as an alternative paradigm that relies on collecting large amounts of data with noisy labels that can be easier to acquire. This paradigm has shown immense success for learning visual representations [Joulin, 2016; Radford, 2021; Zhai, 2022] powered by extremely large-scale collections of image and text pairs scraped from the open web. While there is no shortage in images accompanied by textual descriptions on the open web, generating analogous collections for other modalities of interest, such as audio, thermal images and depth images, is considerably more challenging. In Chapter 6 we address this issue introducing ImageBind, a novel method for training encoding of six different modalities in a shared latent space leveraging the strong performance of existing vision and language models.

The manuscript concludes with Chapter 7 summarizing our key findings and insights, discussing the limitations of our approaches and future work.

1.3 Publications

The following publications are included in whole or in part within this manuscript:

- **Alaaeldin El-Nouby**, Hugo Touvron, Mathilde Caron, Piotr Bojanowski, Matthijs Douze, Armand Joulin, Ivan Laptev, Natalia Neverova, Gabriel Synnaeve, Jakob Verbeek and Hervé Jegou. "XCiT: Cross-Covariance Image Transformers". In *Conference on Neural Information Processing Systems (NeurIPS), 2021*. The code is available at https://github.com/facebookresearch/xcit. (see Chapter 3)

- **Alaaeldin El-Nouby***, Gautier Izacard*, Hugo Touvron, Ivan Laptev, Hervé Jegou, Edouard Grave. "Are Large-scale Datasets Necessary for Self-Supervised Pre-training?". *Under Review*. (see Chapter 4)

- **Alaaeldin El-Nouby**, Matthew J. Muckley, Karen Ullrich, Ivan Laptev, Jakob Verbeek and Hervé Jégou. "Image Compression with Product Quantized Masked Image Modeling", In *Transactions of Machine Learning Research (TMLR), 2023* (see Chapter 5)

- Rohit Girdhar*, **Alaaeldin El-Nouby***, Zhuang Liu, Mannat Singh, Kalyan Vasudev Alwala, Armand Joulin, Ishan Misra*. "ImageBind: One Embedding Space To Bind Them All". To presented at *Conference on Computer Vision and Pattern Recognition (CVPR), 2023* as a **highlighted paper**. The code is available at https://github.com/facebookresearch/imagebind. (see Chapter 6)
Chapter 1. Introduction

Additionally, the publications listed below were a result of research conducted during the
PhD program, but are not directly included within the thesis:

- Rohit Girdhar*, **Alaaeldin El-Nouby***, Mannat Singh*, Kalyan Vasudev Alwala*, Armand Joulin and Ishan Misra*. “OmniMAE: Single Model Masked Pretraining on Images and Videos”. To be presented at *Conference on Computer Vision and Pattern Recognition (CVPR)*, 2023. The code is available at https://github.com/facebookresearch/omnivore/tree/main/omnimae.

- **Alaaeldin El-Nouby**, Natalia Neverova, Ivan Laptev, Hervé Jégou. “Training vision transformers for image retrieval”. arXiv 2021.

- Ben Graham, **Alaaeldin El-Nouby**, Hugo Touvron, Pierre Stock, Armand Joulin, Hervé Jégou, Matthijs Douze. “LeViT: a Vision Transformer in ConvNet’s Clothing for Faster Inference”. *nternational Conference on Computer Vision, 2021*. The code is available at https://github.com/facebookresearch/LeViT

- Hugo Touvron, Matthieu Cord, **Alaaeldin El-Nouby**, Jakob Verbeek, Hervé Jégou. “Three things everyone should know about Vision Transformers”. *European Conference on Computer Vision, 2022*.
# Chapter 2

## Background

### 2.1 Modern Architectures in Computer Vision

In this section, we provide an overview of modern architectures in computer vision starting from Convolutional Neural Networks to Transformers passing by hybrid models that combines components from both. Additionally, we will summarize multiple efforts that aimed to reduce the computation complexity of Transformers. Finally, we delve into the most recent family of vision models that were developed after and motivated by Vision Transformers.

#### 2.1.1 Convolutional Neural Networks

Convolutional Neural Networks, in their current form, were first proposed by LeCun et al. [LeCun, 1989]. Convolutional layers operate by applying a set of learnable kernels to an input. Each kernel has dimensions that are relatively small, but slides across the full width and height of the input computing the dot product between the weights of the kernel and the input at every spatial position, allowing it to effectively capture local patterns. The work of LeCun et al. [LeCun, 1989] demonstrated that the kernel weights can be effectively trained using backpropagation. The LeNet-5 network [LeCun, 1998] was subsequently designed for digit recognition and marked the inception of CNNs. LeNet-5 was primarily composed of alternating convolutional and pooling layers, followed by a fully connected layer for output classification. The fundamental architecture of CNNs remained unchanged for years until the introduction of AlexNet [Krizhevsky, 2012]. AlexNet significantly increased the depth and width of CNNs, adopted Rectified Linear Units (ReLU) as activation functions and popularized data augmentation as a regularization technique. The success of AlexNet marked the start of the deep learning revolution with many innovations in the CNN design to follow. VGGNet [Simonyan, 2014] and GoogLeNet [Szegedy, 2015] further built upon the foundation laid by AlexNet. VGGNet demonstrated that the depth of a CNN is a crucial factor for its performance, while GoogLeNet introduced the inception module, which allowed for increased network depth and width without an explosion in computational cost.

The introduction of ResNet [He, 2016] was another important milestone in the evolution of CNNs in particular and deep learning in general. ResNets introduced residual connections, which helped solve the vanishing gradient problem, allowing for the training networks with hun-
dreds of layers. With CNNs demonstrating a dominant performance for various vision tasks, the efficiency and scalability of network architectures then became the focus of research. MobileNet [Howard, 2017] and EfficientNet [Tan, 2019] were designed to be lightweight and efficient while maintaining high performance. MobileNet introduced depthwise separable convolutions, while EfficientNet used a systematic approach to scale up CNNs, considering depth, width, and resolution. RegNets [Radosavovic, 2020] introduced the idea utilizes a design space exploration approach to identify network architectures that offer excellent trade-offs between computational cost and accuracy. Most recently, NfNets Brock et al. [Brock, 2021] have been introduced as another step forward in efficient network design. NfNets employ a normalization-free design with model scaling, and they have achieved state-of-the-art performance on ImageNet while being more efficient than existing models. To sum up, the architecture of CNNs has retrospectively come a long way from the early days of LeNet, with each new development leading to improved performance and efficiency.

## 2.1.2 Self-Attention and Transformers

The attention mechanism was initially proposed by Bahdanau et al. [Bahdanau, 2014] for sequence-to-sequence models, especially in neural machine translation tasks. By dynamically distributing the model’s “attention” over different parts of the input sequence based on their relevance to the current task. Attention improved upon fixed-length context vectors and enabled more effective learning of long-range dependencies. A significant evolution in the self-attention mechanism was achieved with the advent of the Transformer model by Vaswani et al. [Vaswani, 2017b]. The Transformer discarded the traditional recurrent and convolutional layers and relied solely on attention mechanisms for processing the input. This approach, which allows parallel computation and captures long-range dependencies effectively, has resulted in remarkable improvements in performance on several NLP tasks.

The self-attention mechanism allows each token in the input sequence to compute a weighted sum of all tokens, including itself. In the simplest form, it’s defined by:

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V \qquad (2.1) $$

where:

- $Q$ is the matrix of queries;

- $K$ is the matrix of keys;

- $V$ is the matrix of values;

- $d_k$ is the dimension of the key.

The self-attention mechanism is applied to the input sequence, where $Q$, $K$, and $V$ are linear projections of the same input

$$ Q = XW_Q, \quad K = XW_K, \quad V = XW_V. \qquad (2.2) $$
In this setting, each token in the sequence attends to all other tokens to compute its representation. However, in practice the self-attention mechanism uses multiple sets of learned linear transformations, resulting in the *Multi-Head* attention as proposed by Vaswani et al. [Vaswani, 2017b]. This allows the model to focus on different types of information. For *h* heads, multi-head attention is defined as follows:

$$ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W_O \qquad (2.3) $$

where each head *i* is computed as:

$$ \text{head}_i = \text{Attention}(QW_{Qi}, KW_{Ki}, VW_{Vi}) \qquad (2.4) $$

Here:

- $W_{Qi}$, $W_{Ki}$, and $W_{Vi}$ are the parameter matrices for each head

- $W_O$ is the output projection.

The Transformer model since its inception has been generalized to other modalities and tasks. One of its most notorious applications has been language understanding and generation tasks, such as BERT [Devlin, 2018] and GPT [Radford, 2018]. BERT uses the Transformer's bidirectional self-attention to encode a rich understanding of contextual relationships between words, demonstrating state-of-the-art performance on several NLP benchmarks. GPT, on the other hand, uses an autoregressive causal Transformer to generate text, illustrating the potential for large-scale language models. Transformers have also been extended to non-textual modalities. In the realm of speech recognition and synthesis, models like Conformer [Gulati, 2020] have integrated self-attention into a hybrid architecture, combining the best of convolutional and self-attention mechanisms to deal with time-series data effectively. Self-attention and Transformers have evolved from being solutions to specific problems in machine translation to cornerstone architectures across different modalities. Their potential for capturing complex patterns and long-range dependencies, coupled with their suitability for parallel computation, position them as a driving force in the evolution of machine learning models.

### 2.1.3 Efficient Self-Attention

While Transformers have achieved state-of-the-art performance in various domains, they often come with costly computational and memory requirements due to the quadratic complexity of self-attention with respect to the input sequence length. Many variants of the self-attention operation has been proposed to address this particular issue. An overview of such methods is summarized in Figure 2.1. One family of efficient Transformer models are Sparse Transformers [Child, 2019], which aim to reduce the computational complexity by limiting the number of attended positions in self-attention. Models like the Longformer [Beltagy, 2020] and Big-Bird [Zaheer, 2020] introduced novel sparse attention patterns that significantly reduce complexity while maintaining high performance. Another line of research, referred to as Kernelized Transformers, leverages kernel methods to approximate the attention mechanism. These meth-
<figure><img src="image_5.png" /><figcaption>Figure 2.1 – A Venn diagram of different approaches for reducing the computational and memory com-
plexity of attention as illustrated and detailed by [Tay, 2020].</figcaption></figure>

ods, such as the Performer [Choromanski, 2020] and the Nyströmformer [Xiong, 2021b]. The
Linformer [Wang, 2020c] reduce the quadratic complexity by projecting the attention matrix
into a lower-dimensional space or by using random feature maps.

Moreover, Factorized Transformers such as the Transformer-XL [Dai, 2019] introduce re-
currence into the self-attention mechanism, enabling the model to handle longer context than
standard Transformers. Methods like Reversible Transformers, including Reformer [Kitaev,
2020], use reversible layers to reduce the memory requirements. This way, they can accom-
modate longer sequences during training. Finally, work on hybrid models, like the Conformer
[Gulati, 2020], combines elements of Transformers with other neural architectures like convo-
lutional neural networks, capturing local and global dependencies more efficiently. Overall,
efficient Transformer architectures continue to be an active area of research. The improvements
in computational and memory efficiency brought about by these methods open up possibilities
for applying Transformers to larger and more complex datasets and tasks. For a more in-depth
study of efficient attention variants the reader can refer to the survey by Tay et al. [Tay, 2020].

2.1.4 CNNs Augmented with Self-Attention

The success of the self-attention mechanism in various settings encourages researchers in
computer vision to augment the popular CNN-based architectures with self-attention. This
augmentation aims to capitalize on the strengths of both mechanisms: the strong locality bias of
CNNs and the global context awareness of self-attention. Such an approach potentially enables
more robust handling of long-range dependencies. An early influential work in this domain,
Non-local Neural Networks Wang et al. [Wang, 2018] where the authors incorporated a non-
local operation into existing CNN architectures. This operation, inspired by the self-attention
mechanism of Transformers, computes a weighted sum of all the input feature maps' features at
a given position. The non-local block is capable of capturing dependencies regardless of their
relative distance in the input data, significantly enhancing the model’s ability to handle complex
spatial relationships. In the same spirit, Squeeze-and-Excitation Networks [Hu, 2018], proposed
a channel-wise self-attention mechanism. This approach focuses on interdependencies between
the channels of the convolutional layers, allowing for an adaptive recalibration of channel-wise
feature responses. The resulting mechanism provides the model with the ability to emphasize
informative features selectively.

The work of Ramachandran et al. [Ramachandran, 2019] explored using self-attention as a stand-alone layer, replacing all convolutions with a form of self-attention in a ResNet model. BoTNets by Srinivas et al. [Srinivas, 2021] replaces the spatial convolutions of a ResNet by a novel bottleneck Transformer blocks which helps control the computational cost of self-attention. Meanwhile, Bello [Bello, 2021] introduced LambdaNetworks, a model that effectively replaces self-attention with lambda layers which provide a method of capturing long-range interactions without the associated computational costs of self-attention, making it more suitable for larger inputs. Carion et al. [Carion, 2020] proposed a novel approach to object detection, treating it as a direct set prediction problem. By leveraging self-attention and Transformers, this model coined DETR bypassed the need for several components traditionally used in object detection, such as non-maximum suppression and anchor boxes. Another interesting application of attention for semantic segmentation is Axial-Attention [Wang, 2020b] which applies self-attention sequentially in the height and width axes of the image. This mechanism effectively handles high-resolution images, demonstrating the advantages of self-attention in semantic segmentation tasks.

2.1.5 Vision Transformers

Vision Transformers (ViT) have become a paradigm-shifting development in computer vision, signifying a move away from traditional convolutional neural networks. Introduced by Dosovitskiy et al. [Dosovitskiy, 2021a], ViTs treat images as sequences of patches and apply Transformer-style self-attention to the task of image recognition. Several impactful variants of ViT soon followed. The DeiT work of [Touvron, 2020] focused on improving the data efficiency of ViT, managing to match its performance with fewer training images and setting the standard for optimization and regularization methods for training Vision Transformers (see also [Touvron, 2022b; Touvron, 2022a]). The Swin Transformer [Liu, 2021b], introduced a hierarchical structure with shifted windows to handle images of varying resolution, further bridging the gap between Transformers and CNNs. Pyramid Vision Transformer [Wang, 2021a] incorporated a pyramid structure to extract features at different scales, akin to CNNs, while maintaining the global self-attention mechanism of Transformers.

Subsequent iterations like CaiT [Touvron, 2021c] tackled the problem of making Vision
Transformers deeper by proposing LayerScale for stabilizing training for models 48 layers deep.
Additionally, Touvron et al. proposed a class attention final layer enabling the model to focus on
class-related features. Graham et al. [Graham, 2021] introduces LeViT which offers an important
step towards reducing the computational burden and model size associated with ViTs. LeViT
provides favorable trade-offs in terms of throughput even when compared to popular backbones
like EfficientNet. MViT [Fan, 2021] introduces an efficient method of spatial reduction which
helps control the computational complexity and enabling MViT to be applied to successfully to
video recognition tasks.

The success of Vision Transformers has motivated innovative designs that further blend the
advantages of Transformers and CNNs. For example, ConvNext [Liu, 2022b] revisits the design
of CNNs by employing key design choices from ViTs like the columnar structure, replacing batch
normalization by layer normalization and larger kernel sizes. ConvNext shows that CNNs can be
competitive to vision transformers after integrating some of the modern design choices developed
for ViTs. In summary, Vision Transformers and their successors have had a profound influence
on the field of computer vision, inspiring a wide range of new model architectures that merge the
advantages of CNNs and Transformers. This evolving landscape promises a wealth of potential
for further exploration and innovation.

## 2.2 Self-Supervised Learning

Self-supervised learning has recently seen a lot of interest due to its ability to exploit the inherent structure of data to learn useful representations without the dependency on costly and domain-specific annotations. Self-supervised methods come in various shapes and forms, each with unique strengths and applications. In this section, we will explore some of the popular families of self-supervised learning methods including autoencoders, contrastive methods, non-contrastive methods, and clustering methods.

### 2.2.1 Autoencoders

Autoencoders are neural networks that are trained with the objective of reconstructing their input, typically with some form of constraint in order to learn useful features. They have been applied in a wide range of tasks, including image reconstruction, noise reduction, and representation learning. Denoising Autoencoders [Vincent, 2008] are designed to recover a clean input from a corrupted version, thereby learning robust representations of the data. Contractive Autoencoders [Rifai, 2011] enforce a form of robustness by adding a penalty term to the standard reconstruction error that constrains the derivatives of the encoder functions. The Context Encoder by Pathak et al. [Pathak, 2016] is a modification of the traditional autoencoder paradigm to perform inpainting of missing image regions. The autoencoding family of methods have found even more success in the context of Natural Language Processing. BERT [Devlin, 2018] marked a significant advancement by learning contextual relations between words in a text by training a bidirectional Transformer to predict missing (masked) words in a sentence. Subsequently, the highly popular GPT [Radford, 2018] model can be thought of as a Denoising Autoencoder that uses causal masking as the form of noise.

### 2.2.2 Pretext tasks

Early work in self-supervised learning relied on the idea of exploiting inherent spatial and temporal cues in images and videos. One such approach, Noroozi et al. [Noroozi, 2016] proposed
<figure><img src="image_6.png" /><figcaption>Figure 2.2 – An overview of popular joint embedding methods including contrastive, non-contrastive, and clustering based methods as illustrated by Chen et al. [Chen, 2021a]. Figure reproduced by permission of the authors.</figcaption></figure>

the Jigsaw Puzzle, where permuted image patches are rearranged to their correct configuration, serving as a powerful source of supervisory signal for learning rich visual representations. The Context Prediction task by Doersch et al. [Doersch, 2015] leverages spatial context as a free and abundant supervisory signal where a CNN is trained to predict the relative positions of random pairs of patches extracted from each image. Meanwhile, Gidaris et al. [Gidaris, 2018] developed RotNet, which involves predicting the rotation of an image to learn visual representations by inferring the correct orientation of an object. Larsson et al. [Larsson, 2016] introduced a task of automatic colorization, which uses the task of filling in color in grayscale images to learn representations that capture both semantics and texture of the input images. Some other works made use of the inherent temporal cues in video such as [Wang, 2015] which proposed a self-supervised learning method that learns visual representations using videos. This approach exploits the temporal coherence and context in videos to learn robust visual representations.

### 2.2.3 Joint embeddings methods

A highly popular family of self-supervised approaches are joint embedding methods where an encoder is trained to be invariant to certain types of distortions and transformations. This can be achieved explicitly via contrastive or clustering methods or implicitly through non-contrastive methods that rely on applying regularization to the representations to avoid trivial solutions and collapse. An overview of notable examples of different joint embeddings methods is illustrated in Figure 2.2.

**Contrastive Methods.** Dosovitskiy et al. [Dosovitskiy, 2014] proposed instance discrimination where a model is tasked to classify an image correctly as itself under different transformations via an *N*-way classification objective where *N* corresponds to the number of unique images in the dataset. A more efficient the formulation is proposed in CPC [Oord, 2018] in the form of the InfoNCE objective. In the context of self-supervised learning, given a query *q* and a set of *K* keys consisting of one positive key *k*<sup>+</sup> (from the same sample as the query) and *K* − 1 negative keys *k*<sup>-</sup> (from different samples), the InfoNCE loss can be defined as follows:

$$ \mathcal{L}_{\text{NCE}} = -\mathbb{E} \left[ \log \frac{\exp(f(q, k^+))}{\sum_{k \in k^+, k^-} \exp(f(q, k))} \right] \quad (2.5) $$
The InfoNCE loss encourages the model to assign higher similarity to positive pairs than to negative pairs. As a result, during the learning process, the model learns to map positive pairs closer together and negative pairs further apart in the embedding space. This property makes the InfoNCE loss particularly well-suited for contrastive learning tasks. CPC is primarily proposed for learning representations from audio, but InfoNCE was widely adopted by the majority of the subsequent contrastive approaches for vision tasks. Most notably, Chen et al. [Chen, 2020c] presented a simple yet effective method termed SimCLR, which significantly improved the state-of-the-art self-supervised learning performance. The SimCLR framework employs data augmentation as a key ingredient to generate positive pairs as well as identifying that using a large batch size, longer training and MLP head can significantly improve the performance of such methods. He et al. [He, 2020] introduced MoCo which constructs the dictionary of negative samples dynamically with a queue and a momentum-updated encoder. This approach allowed training with a large number of negative examples, which improved the quality of the learned representations.

**Clustering Methods.** Another family of effective self-supervised learning approaches are clustering-based methods. effective. DeepCluster [Caron, 2018] leverages clustering to enforce a pseudo-label based learning objective, thereby learning rich visual features in an unsupervised manner. Subsequently, Asano et al. [Asano, 2019b] introduced a method called SeLa to combine the process of clustering and representation learning in a single framework. The most notable of the clustering-based methods is SwAV [Caron, 2020a] which enhances the self-supervised learning process by promoting consistency between cluster assignments produced for different transformations of the same image.

**Non-contrastive Methods.** Non-contrastive methods have become increasingly popular because they sidestep the need for sampling a high number of negative pairs for every update. BYOL [Grill, 2020] was the first method to show a competitive performance can be achieved without relying on negative pairs. This was achieved by using an online and offline encoders where the weights of the offline encoder are updated via an exponential moving average. Combining this with a new formulation for the model head by adding a predictor, BYOL achieved a surprising result by outperforming many contrastive methods with no negative pairs and without collapse. SimSiam [Chen, 2021a] extended this idea, utilizing two identical networks that predict each other's outputs, but with an added stop-gradient operation to prevent the collapse of the representations. Barlow Twins [Zbontar, 2021] proposed a method where the outputs of two identical networks are decorrelated, enforcing diversity in the learned representations and preventing collapse. VicReg [Bardes, 2021] introduced an objective that regularizes the representation's variance along each dimension independently. Finally, Caron et al. [Caron, 2021] proposed DINO which is based on a self-distillation framework. The student is encouraged to match the output of a teacher that is updated with an exponential moving average. DINO was also the first method to show that self-supervised learning when used in conjunction with Vision Transformers can lead to outstanding results.
## 2.3 Vision-Language pre-training

In addition to learning representations through self-supervised learning, as discussed in the previous section, weakly supervised pre-training has been a highly successful paradigm in recent years. It leverages noisy labels, which are easier to obtain, as a substitute for expensive curated annotations. Early work by [Joulin, 2016] made use of image collections with paired textual captions to learn rich semantic representations. the full potential of this idea was realized through aggressive scaling of the image and text collection to millions [Radford, 2021] or even billions [Ilharco, 2021] of pairs. CLIP [Radford, 2021] and ALIGN [Jia, 2021a] use a simple contrastive objective to align matching visual and textual pairs. CLIP models can be very flexible as they learn about a wide range of visual concepts directly from natural language. CoCa [Yu, 2022] incorporates a text decoder and an additional captioning loss, which contributes to its strong performance in Visual Question Answering (VQA) and captioning tasks. Zhai et al. [Zhai, 2022] demonstrated that a strong pre-trained vision encoder can be easily augmented with open-set recognition capabilities by training a text encoder to align with the visual features resulting in performance often superior to training both encoders jointly from scratch. Alayrac et al. [Alayrac, 2022] explores interweaving images and text utterances by adapting a vision encoder to work jointly with a pre-trained Large Language Model yielding outstanding results in VQA and serving as a general-purpose visual chatbot.

## 2.4 Lossless Compression

The goal of lossless compression is to encode samples from a discrete probability distribution $x \sim p_d(x)$ $x \in \mathcal{X}$ as bit strings $m = \text{enc}(x) \in \{0,1\}^*$ of shortest possible length $\ell(m)$ such that $x$ can be decoded without loss of information $x = \text{dec}(m)$. Block codes assign a unique binary code of equal length to each event in $\mathcal{X}$. The expected length per symbol of such codes is $l(m) = \lceil \log_2 |\mathcal{X}| \rceil$ bits. We can improve this result if we consider to give shorter code words to more frequently occurring symbols and longer codes to less frequently occurring ones. In fact, the smallest expected achievable average code length per point using an approximation to the true data distribution $p(x)$, is given by the cross-entropy:

$$H(p_d, p) := \mathbb{E}_{x \sim p_d} [-\log p(x)]. \qquad (2.6)$$

This quantity is bounded from below (seen via Jensen's inequality) by the entropy of the "true" probability distribution $p_d$, i.e. when $p = p_d$ [MacKay, 2003; Cover, 1991]. Entropy coders reach this bound up to a small constant $\epsilon$. An entropy codec is a tuple of an inverse pair of functions, $\text{enc}_{p_X}$ and $\text{dec}_{p_X}$ that achieve near optimal compression rates on sequences of symbols.

$$ \begin{aligned} \text{enc}_p &: m, x \mapsto \hat{m} \\ \text{dec}_p &: \hat{m} \mapsto (m, x), \end{aligned} \qquad (2.7) $$

such that $\ell(\hat{m}) = \ell(m) + \log_2 p(x) + \epsilon$.
# Chapter 3

## Cross Covariance Image Transformers

### Objectives

Following tremendous success in natural language processing, transformers have recently shown much promise for computer vision. The self-attention operation underlying transformers yields global interactions between all tokens, i.e. words or image patches, and enables flexible modelling of image data beyond the local interactions of convolutions. This flexibility, however, comes with a quadratic complexity in time and memory, hindering application to long sequences and high-resolution images. We propose a “transposed” version of self-attention that operates across feature channels rather than tokens, where the interactions are based on the cross-covariance matrix between keys and queries. The resulting cross-covariance attention (XCA) has linear complexity in the number of tokens, and allows efficient processing of high-resolution images. Our cross-covariance image transformer (XCiT) – built upon XCA – combines the accuracy of conventional transformers with the scalability of convolutional architectures. We validate the effectiveness and generality of XCiT by reporting excellent results on multiple vision benchmarks, including (self-supervised) image classification on ImageNet-1k, object detection and instance segmentation on COCO, and semantic segmentation on ADE20k.

### Contents

<table><tr><td>3.1</td><td>Introduction</td><td>17</td></tr><tr><td>3.2</td><td>Related work</td><td>19</td></tr><tr><td>3.3</td><td>Method</td><td>20</td></tr><tr><td>3.3.1</td><td>Background</td><td>20</td></tr><tr><td>3.3.2</td><td>Cross-covariance attention</td><td>22</td></tr><tr><td>3.3.3</td><td>Cross-covariance image transformers</td><td>24</td></tr><tr><td>3.4</td><td>Preliminary study on Vision Transformers (ViT)</td><td>26</td></tr><tr><td>3.4.1</td><td>Impact of resolution versus patch size</td><td>26</td></tr><tr><td>3.4.2</td><td>Approximate attention models in ViT with DeiT training</td><td>27</td></tr><tr><td>3.4.3</td><td>Training and testing with varying resolution</td><td>27</td></tr><tr><td>3.5</td><td>Experimental evaluation</td><td>28</td></tr><tr><td>3.5.1</td><td>Image classification</td><td>28</td></tr><tr><td>3.5.2</td><td>Object detection and instance segmentation</td><td>32</td></tr><tr><td>3.5.3</td><td>Semantic segmentation</td><td>34</td></tr><tr><td>3.5.4</td><td>Implementation details</td><td>35</td></tr><tr><td>3.6</td><td>Conclusion and Future Work</td><td>35</td></tr></table>
3.1 Introduction

Transformers architectures [Vaswani, 2017b] have provided quantitative and qualitative break-
throughs in speech and natural language processing. Recently, Dosovitskiy et al. [Dosovitskiy,
2021a] established transformers as a viable architecture for learning visual representations, re-
porting competitive results for image classification while relying on large-scale pre-training.
Touvron et al. [Touvron, 2020] have shown on par or better accuracy/throughput compared
to strong convolutional baselines such as EfficientNets [Tan, 2019] when training transformers
on ImageNet-1k using extensive data augmentation and improved training schemes. Promising
results have been obtained for other vision tasks, including image retrieval [El-Nouby, 2021b],
object detection and semantic segmentation [Liu, 2021b; Wang, 2021a; Zhang, 2021; Zheng,
2020], as well as video understanding [Arnab, 2021; Bertasius, 2021; Fan, 2021].

One major drawback of transformers is the time and memory complexity of the core self-
attention operation, that increases quadratically with the number of input tokens, or similarly
number of patches in computer vision. For w×h images, this translates to a complexity of
O(w²h²), which is prohibitive for most tasks involving high-resolution images, such as object
detection and segmentation. Various strategies have been proposed to alleviate this complexity,
for instance using approximate forms of self-attention [Liu, 2021b; Zhang, 2021], or pyramidal
architectures which progressively downsample the feature maps [Wang, 2021a]. However, none
of the existing solutions are fully satisfactory, as they either trade complexity for accuracy, or
their complexity remains excessive for processing very large images.

We replace the self-attention, as originally introduced by Vaswani et al. [Vaswani, 2017b],
with a “transposed” attention that we denote as “cross-covariance attention” (XCA). Cross-
covariance attention substitutes the explicit full pairwise interaction between tokens by self-
attention among features, where the attention map is derived from the cross-covariance matrix
computed over the key and query projections of the token features. Importantly, XCA has a
linear complexity in the number of patches. To construct our Cross-Covariance Image Trans-
formers (XCiT), we combine XCA with local patch interaction modules that rely on efficient
depth-wise convolutions and point-wise feedforward networks commonly used in transformers,
see Figure 3.1. XCA can be regarded as a form of a dynamic 1 × 1 convolution, which multi-
plies all tokens with the same data-dependent weight matrix. We find that the performance
of our XCA layer can be further improved by applying it on blocks of channels, rather than
directly mixing all channels together. This “block-diagonal” shape of XCA further reduces the
computational complexity with a factor linear in the number of blocks.

Given its linear complexity in the number of tokens, XCiT can efficiently process images
with more than thousand pixels in each dimension. Notably, our experiments show that XCiT
does not compromise the accuracy and achieves similar results to DeiT [Touvron, 2020] and
CaiT [Touvron, 2021c] in comparable settings. Moreover, for dense prediction tasks such as
object detection and image segmentation, our models outperform popular ResNet [He, 2016]
backbones as well as the recent transformer-based models [Liu, 2021b; Wang, 2021a; Zhang,
2021]. Finally, we also successfully apply XCiT to the self-supervised feature learning using
Chapter 3. Cross Covariance Image Transformers

<figure><img src="image_7.png" /><figcaption>Figure 3.1 – Our XCiT layer consists of three main blocks, each preceded by LayerNorm and followed by a residual connection: (i) the core cross-covariance attention (XCA) operation, (ii) the local patch inter- action (LPI) module, and (iii) a feed-forward network (FFN). By transposing the query-key interaction, the computational complexity of XCA is linear in the number of data elements N , rather than quadratic as in conventional self-attention.</figcaption></figure>

DINO [Caron, 2021], and demonstrate improved performance compared to a DeiT-based back-
bone [Touvron, 2020].

Overall, we summarize our contributions as follows:

— We introduce cross-covariance attention (XCA), which provides a “transposed” alternative to conventional self-attention, attending over channels instead of tokens. Its complexity is linear in the number of tokens, allowing for efficient processing of high-resolution images, see Figure 3.2.

— XCA attends to a fixed number of channels, irrespective of the number of tokens. As a result, our models are significantly more robust to changes in image resolution at test time, and are therefore more amenable to process variable-size images.

— For image classification, we demonstrate that our models are on par with state-of-the-art vision transformers for multiple model sizes using a simple columnar architecture, *i.e.*, in which we keep the resolution constant across layers. In particular, our XCiT-L24 model achieves 86.0% top-1 accuracy on ImageNet, outperforming its CaiT-M24 [Touvron, 2021c] and NFNet-F2 [Brock, 2021] counterparts with comparable numbers of parameters.

— For dense prediction tasks with high-resolution images, our models outperform ResNet and multiple transformer-based backbones. On the COCO benchmark, we achieve a strong performance of 48.5% and 43.7% mAP for object detection and instance segmentation respectively. Moreover, we report 48.4% mIoU for semantic segmentation on the ADE20k benchmark, outperforming the state-of-the-art Swin Transformer [Liu, 2021b] backbones across all comparable model sizes.

— Finally, our XCiT model is highly effective in self-supervised learning setups, achieving 80.9% top-1 accuracy on ImageNet-1k using DINO [Caron, 2021].
## 3.2 Related work

**Deep vision transformers.** Training deep vision transformers can be challenging due to instabilities and optimization issues. Touvron et al. [Touvron, 2021c] successfully train models with up to 48 layers using LayerScale, which weighs contributions of residual blocks across layers and improves optimization. Additionally, the authors introduce class attention layers which decouple the learning of patch features and the feature aggregation stage for classification.

**Spatial structure in vision transformers.** Yuan et al. [Yuan, 2021b] propose applying a soft split for patch projection with overlapping patches which is applied repeatedly across model layers, reducing the number of patches progressively. Han et al. [Han, 2021] introduce a transformer module for intra-patch structure, exploiting pixel-level information and integrating with an inter-patch transformer to attain higher representation power. d'Ascoli et al. [dAscoli, 2021] consider the initialization of self-attention blocks as a convolutional operator, and demonstrate that such initialization improves the performance of vision transformers in low-data regimes. Graham et al. [Graham, 2021] introduce LeViT, which adopts a multi-stage architecture with progressively reduced feature resolution similar to popular convolutional architectures, allowing for models with high inference speed while retaining a strong performance. Moreover, the authors adopt a convolution-based module for extracting patch descriptors. Yuan et al. [Yuan, 2021a] improve both the performance and the convergence speed of vision transformers by replacing the linear patch projection with convolutional layers and max-pooling, as well as modifying the feed-forward networks in each transformer layer to incorporate depth-wise convolutions.

**Efficient attention.** Numerous methods for efficient self-attention have been proposed in the literature to address the quadratic complexity of self-attention in the number of input tokens. These include restricting the span of the self-attention to local windows [Parmar, 2018; Qiu, 2019], strided patterns [Child, 2019], axial patterns [Ho, 2019], or an adaptive computation across layers [Sukhbaatar, 2019]. Other methods provide an approximation of the self-attention matrix which can be achieved by a projection across the token dimension [Wang, 2020c], or through a factorization of the softmax-attention kernel [Choromanski, 2020; Katharopoulos, 2020; Shen, 2021; Xiong, 2021a], which avoids explicit computation of the attention matrix. While conceptually different, our XCA performs similar computations without being sensitive to the choice of the kernel. Similarly, Lee-Thorp et al. [Lee-Thorp, 2021] achieve faster training by substituting self-attention with unparametrized Fourier Transform. Other efficient attention methods rely on local attention and adding a small number of global tokens, thus allowing interaction among all tokens only by hopping through the global tokens [Ainslie, 2020; Beltagy, 2020; Jaegle, 2021; Zaheer, 2020].

**Transformers for high-resolution images.** Several works adopt visual transformers to high-resolution image tasks beyond image classification, such as object detection and image segmentation. Wang et al. [Wang, 2021a] design a model with a pyramidal architecture and address complexity by gradually reducing the spatial resolution of keys and values. Similarly, for
video recognition Fan et al. [Fan, 2021] utilize pooling to reduce the resolution across the spatial
and temporal dimensions to allow for an efficient computation of the attention matrix. Zhang et
al. [Zhang, 2021] adopt global tokens and local attention to reduce the model complexity, while
Liu et al. [Liu, 2021b] provide an efficient method for local attention with shifted windows. In
addition, Zheng et al. [Zheng, 2020] and Ranftl et al. [Ranftl, 2021] study problems like semantic
segmentation and monocular depth estimation with the quadratic self-attention operation.

**Data-dependent layers.** Our XCiT layer can be regarded as a “dynamic” 1×1 convolution, which multiplies all token features with the same data-dependent weight matrix, derived from the key and query cross-covariance matrix. In the context of convolutional networks, Dynamic Filter Networks [Brabandere, 2016] explore a related idea, using a filter generating subnetwork to produce convolutional filters based on features in previous layers. Squeeze-and-Excitation networks [Hu, 2018] use data dependent 1×1 convolutions in convolutional architectures. Spatially average-pooled features are fed to a 2-layer MLP which produces per channel scaling parameters. Closer in spirit to our work, Lambda layers propose a way to ensure global interaction in ResNet models [Bello, 2021]. Their “content-based lambda function” is computing a similar term as our cross-covariance attention, but differing in how the softmax and $\ell_2$ normalizations are applied. Moreover, Lambda layers also include specific position-based lambda functions, and LambdaNetworks are based on ResNets while XCiT follows the ViT architecture. Recently data-independent analogues of self-attention have also been found to be an effective alternative to convolutional and self-attention layers for vision tasks [Ding, 2021b; Melas-Kyriazi, 2021; Tolstikhin, 2021; Touvron, 2021a]. These methods treat entries in the attention map as learnable parameters, rather than deriving the attention map dynamically from queries and keys, but their complexity remains quadratic in the number of tokens. Zhao et al. [Zhao, 2020] consider alternative attention forms in computer vision.

## 3.3 Method

In this section, we first recall the self-attention mechanism, and the connection between
the Gram and covariance matrices, which motivated our work. We then propose our cross-
covariance attention operation (XCA) – which operates along the feature dimension instead of
token dimension in conventional transformers – and combine it with local patch interaction and
feedforward layers to construct our Cross-Covariance Image Transformer (XCiT). See Figure 3.1
for an overview.

### 3.3.1 Background

**Token self-attention.** Self-attention, as introduced by Vaswani et al. [Vaswani, 2017b], operates on an input matrix $X \in \mathbb{R}^{N \times d}$, where $N$ is the number of tokens, each of dimensionality $d$. The input $X$ is linearly projected to queries, keys and values, using the weight matrices $W_q \in \mathbb{R}^{d \times d_q}$, $W_k \in \mathbb{R}^{d \times d_k}$ and $W_v \in \mathbb{R}^{d \times d_v}$, such that $Q=XW_q$, $K=XW_k$ and $V=XW_v$, where $d_q=d_k$. Keys and values are used to compute an attention map $A(K,Q) = \mathrm{Softmax}(QK^{\top}/\sqrt{d_k})$,
<figure><img src="image_8.png" /><figcaption>Figure 3.2 – Inference memory usage of vision transformer variants. Our XCiT models scale linearly in the number of tokens, which makes it possible to scale to much larger image sizes, even in comparison to approaches employing approximate self-attention or a pyramidal design. All measurements are performed with a batch size of 64 on a single V100-32GB GPU.</figcaption></figure>

and the output of the self-attention operation is defined as the weighted sum of *N* token features in *V* with the weights corresponding to the attention map: $Attention(Q, K, V) = \mathcal{A}(K, Q)V$. The computational complexity of self-attention scales quadratically in *N*, due to pairwise interactions between all *N* elements.

**Relationship between Gram and covariance matrices.** To motivate our cross-covariance attention operation, we recall the relation between Gram and covariance matrices. The unnormalised $d \times d$ covariance matrix is obtained as $C=X^\top X$. The $N \times N$ Gram matrix contains all pairwise inner products: $G=XX^\top$. The non-zero part of the eigenspectrum of the Gram and covariance matrix are equivalent, and the eigenvectors of *C* and *G* can be computed in terms of each other. If *V* are the eigenvectors of *G*, then the eigenvectors of *C* are given by $U=XV$. To minimise the computational cost, the eigendecomposition of either the Gram or covariance matrix can be obtained in terms of the decomposition of the other, depending on which of the two matrices is the smallest.<fnref n="1" />

We draw upon this strong connection between the Gram and covariance matrices to consider if it is possible to avoid the quadratic cost to compute the $N \times N$ attention matrix, which is computed from the analogue of the $N \times N$ Gram matrix $QK^\top=ZW_qW_k^\top X^\top$. Below we consider how we can use the $d_k \times d_q$ cross-covariance matrix, $K^\top Q = W_k^\top X^\top Z W_q$, which can be computed in linear time in the number of elements *N*, to define an attention mechanism.

<footnote n="1">1. For *C* to represent the covariance, *X* should be centered, i.e. $X1=0$. For the relation between *C* and *G*, however, centering is not required.</footnote>
<figure><img src="image_9.png" /><figcaption>Figure 3.3 – Performance when changing the resolution at test-time for models with a similar number of parameters. All networks were trained at resolution 224, w/o distillation. XCIT is more tolerant to changes of resolution than the Gram-based DeiT and benefit more from the “FixRes” effect Touvron et al. [Touvron, 2019] when inference is performed at a larger resolution than at train-time.</figcaption></figure>

### 3.3.2 Cross-covariance attention

We propose a cross-covariance based self-attention function that operates along the feature dimension, rather than along the token dimension as in token self-attention. Using the definitions of queries, keys and values from above, the cross-covariance attention function is defined as:

$$ \mathrm{XC-Attention}(Q, K, V) = V \mathcal{A}_{\mathrm{XC}}(K, Q), \quad \mathcal{A}_{\mathrm{XC}}(K, Q) = \mathrm{Softmax} \left( \hat{K}^{\top} \hat{Q} / \tau \right), \quad (3.1) $$

where each output token embedding is a convex combination of the $d_v$ features of its corresponding token embedding in $V$. The attention weights $\mathcal{A}$ are computed based on the cross-covariance matrix.

**$\ell_2$-Normalization and temperature scaling.** In addition to building our attention operation on the cross-covariance matrix, we make a second modification compared to token self-attention. We restrict the magnitude of the query and key matrices by $\ell_2$-normalising them, such that each column of length $N$ of the normalised matrices $\hat{Q}$ and $\hat{K}$ has unit norm, and every element in $d \times d$ cross-covariance matrix $\hat{K}^{\top}\hat{Q}$ is in the range $[-1, 1]$. We observed that controlling the norm strongly enhances the stability of training, especially when trained with a variable numbers of tokens. However, restricting the norm reduces the representational power of the operation by removing a degree of freedom. Therefore, we introduce a learnable temperature parameter $\tau$ which scales the inner products before the Softmax, allowing for sharper or more uniform distribution of attention weights.
**Block-diagonal cross-covariance attention.** Instead of allowing all features to interact among each other, we divide them into a *h* groups, or “heads”, in a similar fashion as multi-head token self-attention. We apply the cross-covariance attention separately per head where for each head, we learn separate weight matrices to project *X* to queries, keys and values, and collect the corresponding weight matrices in the tensors *W*<sub>q</sub> ∈ ℝ<sup>h×d×d<sub>q</sub></sup>, *W*<sub>k</sub> ∈ ℝ<sup>h×d×d<sub>k</sub></sup> and *W*<sub>v</sub> ∈ ℝ<sup>h×d×d<sub>v</sub></sup>, where we set *d*<sub>k</sub>=*d*<sub>q</sub>=*d*<sub>v</sub>=*d*/*h*. Restricting the attention within heads has two advantages: (i) the complexity of aggregating the values with the attention weights is reduced by a factor *h*; (ii) more importantly, we empirically observe that the block-diagonal version is easier to optimize, and typically leads to improved results. This observation is in line with observations made for Group Normalization [Wu, 2018a], which normalizes groups of channels separately based on their statistics, and achieves favorable results for computer vision tasks compared to Layer Normalization [Ba, 2016], which combines all channels in a single group. Figure 3.6 shows that each head learns to focus on semantically coherent parts of the image, while being flexible to change what type of features it attends to based on the image content.

**Complexity analysis.** The usual token self-attention with *h* heads has a time complexity of *O*(N²d) and memory complexity of *O*(hN²+Nd). Due to the quadratic complexity, it is problematic to scale token self-attention to images with a large number of tokens. Our cross-covariance attention overcomes this drawback as its computational cost of *O*(Nd²/h) scales linearly with the number of tokens, as does the memory complexity of *O*(d²/h+Nd). Therefore, our model scales much better to cases where the number of tokens *N* is large, and the feature dimension *d* is relatively small, as is typically the case, in particular when splitting the features into *h* heads.

**Runtime and Memory Usage** We present the peak memory usage as well as the throughput of multiple models including full-attention and efficient vision transformers in Table 3.1. Additionally, in Figure 3.4 we plot the processing speed represented as millisecond per image as a function of image resolution for various models. We can observe that XCiT provides a strong trade-off, possessing the best scalability in terms of peak memory, even when compared to ResNet-50. Additionally, the processing time scales linearly with respect to resolution, with only ResNet-50 providing a better trade-off on that front.

<table><caption>Table 3.1 – **Inference throughput and peak GPU memory usage** for our XCiT small model compared to other models of comparable size that include token self-attention. All models tested using batch size of 64 on a V100 GPU with 32GB memory.</caption><thead><tr><th rowspan="3">Model</th><th rowspan="3">#params<br>(×10<sup>6</sup>)</th><th rowspan="3">ImNet<br>Top-1<br>@224</th><th colspan="8">Image Resolution</th></tr><tr><th colspan="2">224<sup>2</sup></th><th colspan="2">384<sup>2</sup></th><th colspan="2">512<sup>2</sup></th><th colspan="2">1024<sup>2</sup></th></tr><tr><th>im/sec</th><th>mem (MB)</th><th>im/sec</th><th>mem (MB)</th><th>im/sec</th><th>mem (MB)</th><th>im/sec</th><th>mem (MB)</th></tr></thead><tbody><tr><td>ResNet-50</td><td>25</td><td>79.0</td><td>1171</td><td>772</td><td>434</td><td>2078</td><td>245</td><td>3618</td><td>61</td><td>14178</td></tr><tr><td>DeiT-S</td><td>22</td><td>79.9</td><td>974</td><td>433</td><td>263</td><td>1580</td><td>116</td><td>4020</td><td>N/A</td><td>OOM</td></tr><tr><td>CaiT-S12</td><td>26</td><td>80.8</td><td>671</td><td>577</td><td>108</td><td>2581</td><td>38</td><td>7117</td><td>N/A</td><td>OOM</td></tr><tr><td>PVT-Small</td><td>25</td><td>79.8</td><td>777</td><td>1266</td><td>256</td><td>3142</td><td>134</td><td>5354</td><td>N/A</td><td>OOM</td></tr><tr><td>Swin-T</td><td>29</td><td>81.3</td><td>704</td><td>1386</td><td>220</td><td>3890</td><td>120</td><td>6873</td><td>29</td><td>26915</td></tr><tr><td>XCiT-S12/16</td><td>26</td><td>82.0</td><td>781</td><td>731</td><td>266</td><td>1372</td><td>151</td><td>2128</td><td>37</td><td>7312</td></tr></tbody></table>
<figure><img src="image_10.png" /><figcaption>Figure 3.4 – We present the millisecond per image during inference of multiple models. Our XCiT-S12/16 model provides a speed up for images with higher resolution compared to existing vision transformers, especially the ones with quadratic complexity like DeiT and CaiT.</figcaption></figure>

3.3.3 Cross-covariance image transformers

To construct our cross-covariance image transformers (XCiT), we adopt a columnar architec-
ture which maintains the same spatial resolution across layers, similarly to [Dosovitskiy, 2021a;
Touvron, 2020; Touvron, 2021c]. We combine our cross-covariance attention (XCA) block with
the following additional modules, each one being preceded by a LayerNorm [Ba, 2016]. See
Figure 3.1 for an overview. Since in this section we specifically design the model for computer
vision tasks, tokens correspond to image patches in this context.

**Local patch interaction.** In the XCA block communication between patches is only implicit through the shared statistics. To enable explicit communication across patches we add a simple Local Patch Interaction (LPI) block after each XCA block. LPI consists of two depth-wise 3×3 convolutional layers with Batch Normalization and GELU non-linearity in between. Due to its depth-wise structure, the LPI block has a negligible overhead in terms of parameters, as well as a very limited overhead in terms of throughput and memory usage during inference.

**Feed-forward network.** As is common in transformer models, we add a point-wise feedforward network (FFN), which has a single hidden layer with 4d hidden units. While interaction between features is confined within groups in the XCA block, and no feature interaction takes place in the LPI block, the FFN allows for interaction across all features.

**Global aggregation with class attention.** When training our models for image classifica-
tion, we utilize the class attention layers as proposed by Touvron *et al.* [Touvron, 2021c]. These
layers aggregate the patch embeddings of the last XCiT layer through writing to a CLS token
<table><caption>Table 3.2 – **XCiT models.** Design choices include model depth, patch embeddings dimensionality *d*, and the number of heads *h* used in XCA. By default our models are trained and tested at resolution 224 with patch sizes of 16×16. We also train with distillation using a convolutional teacher (denoted Υ) as proposed by Touvron *et al.* [Touvron, 2020]. Finally, we report performance of our strongest models obtained with 8×8 patch size, fine-tuned (↑) and tested at resolution 384×384 (column @384/8), using distillation with a teacher that was also fine-tuned @384.</caption><thead><tr><th rowspan="2">Model</th><th rowspan="2">Depth</th><th rowspan="2">d</th><th rowspan="2">#heads</th><th rowspan="2">#params</th><th colspan="2">GFLOPs</th><th colspan="3">ImageNet-1k-val top-1 acc. (%)</th></tr><tr><th>@224/16</th><th>@384/8</th><th>@224/16</th><th>@224/16Υ</th><th>@384/8Υ ↑</th></tr></thead><tbody><tr><td>XCiT-N12</td><td>12</td><td>128</td><td>4</td><td>3M</td><td>0.5</td><td>6.4</td><td>69.9</td><td>72.2</td><td>77.8</td></tr><tr><td>XCiT-T12</td><td>12</td><td>192</td><td>4</td><td>7M</td><td>1.2</td><td>14.3</td><td>77.1</td><td>78.6</td><td>82.4</td></tr><tr><td>XCiT-T24</td><td>24</td><td>192</td><td>4</td><td>12M</td><td>2.3</td><td>27.3</td><td>79.4</td><td>80.4</td><td>83.7</td></tr><tr><td>XCiT-S12</td><td>12</td><td>384</td><td>8</td><td>26M</td><td>4.8</td><td>55.6</td><td>82.0</td><td>83.3</td><td>85.1</td></tr><tr><td>XCiT-S24</td><td>24</td><td>384</td><td>8</td><td>48M</td><td>9.1</td><td>106.0</td><td>82.6</td><td>83.9</td><td>85.6</td></tr><tr><td>XCiT-M24</td><td>24</td><td>512</td><td>8</td><td>84M</td><td>16.2</td><td>188.0</td><td>82.7</td><td>84.3</td><td>85.8</td></tr><tr><td>XCiT-L24</td><td>24</td><td>768</td><td>16</td><td>189M</td><td>36.1</td><td>417.9</td><td>82.9</td><td>84.9</td><td>86.0</td></tr></tbody></table>

by one-way attention between the CLS tokens and the patch embeddings. The class attention is also applied per head, i.e. feature group.

**Handling images of varying resolution.** In contrast to the attention map involved in token self-attention, in our case the covariance blocks are of fixed size independent of the input image resolution. The softmax always operates over the same number of elements, which may explain why our models behave better when dealing with images of varying resolutions (see Figure 3.3). In XCiT we include additive sinusoidal positional encoding [Vaswani, 2017b] with the input tokens. We generate them in 64 dimensions from the 2d patch coordinates and then linearly project to the transformer working dimension *d*. This choice is orthogonal to the use of learned positional encoding, as in ViT [Dosovitskiy, 2021a]. However, it is more flexible since there is no need to interpolate or fine-tune the network when changing the image size.

**Model configurations.** In Table 3.2 we list different variants of our model which we use in our experiments, with different choices for model width and depth. For the patch encoding layer, unless mentioned otherwise, we adopt the alternative used by Graham *et al.* [Graham, 2021] with convolutional patch projection layers. We also experimented with a linear patch projection as described in [Dosovitskiy, 2021a], see our ablation in Table 4.8. Our default patch size is 16×16, as in other vision transformer models including ViT [Dosovitskiy, 2021a], DeiT [Touvron, 2020] and CaiT [Touvron, 2021c]. We also experiment with smaller 8×8 patches, which has been observed to improve performance [Caron, 2021]. Note that this is efficient with XCiT as its complexity scales linearly which the number of patches, while ViT, DeiT and CaiT scale quadratically.

**Pseudo-code** We provide a PyTorch-style pseudo code of the Cross-covariance attention operation. The pseudo code resembles the Timm library [Wightman, 2019] implementation of token self-attention. We show that XCA only requires few modifications, namely the $l_2$ normalization, setting the learnable temperature parameters and a transpose operation of the keys, queries and
values.

<caption>**Algorithm 3.1:** Pseudocode of XCA in a PyTorch-like style.</caption>

```python
# self.qkv: nn.Linear(dim, dim * 3, bias=qkv Bias) # self.temp: nn.Parameter(torch.ones(num_headss, 1, 1))
def forward(self, x): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_headss, C // self.
num_headss) qkv = qkv.permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] # split into query, key and value
q = q transpose(-2, -1) k = k transpose(-2, -1) # Transpose to shape (B, h, C, N) v = v transpose(-2,
-1)
q = F normalization across the token dimension k = F normalization(k, dim
=-1, p=2)
attn = (k @ q transpose(-2, -1)) # Computing the block diagonal cross-covariance matrix attn = attn *
self.temp # Adjusting the activations scale with temperature parameter attn = attn softmax(dim=-1) #
d x d attention map
x = attn @ v # Apply attention to mix channels per token x = x permute(0, 3, 1, 2).reshape(B, N, C)
x = self.proj(x) return x
```

## 3.4 Preliminary study on Vision Transformers (ViT)

In this section, we report the results associated with our preliminary study on high-resolution transformers. Most of the experiments were carried out on the ViT architecture [Dosovitskiy, 2021a] with DeiT training [Touvron, 2020], and intended to analyze different aspects of transformers when considering images with varying resolution or high-resolution images specifically.

### 3.4.1 Impact of resolution versus patch size

<figure><img src="image_11.png" /></figure>

<table><thead><tr><th></th><th colspan="6">Variable patch size</th></tr></thead><tbody><tr><th scope="row">Image Size</th><td>80</td><td>112</td><td>160</td><td>256</td><td>320</td><td>384</td></tr><tr><th scope="row">Patch Size</th><td>5</td><td>7</td><td>10</td><td>16</td><td>20</td><td>24</td></tr><tr><th scope="row">Top-1</th><td>78.2</td><td>79.7</td><td>80.5</td><td>80.7</td><td>80.9</td><td>80.7</td></tr></tbody></table><table><thead><tr><th></th><th colspan="6">Variable number of tokens size</th></tr></thead><tbody><tr><th scope="row">Image Size</th><td>160</td><td>224</td><td>256</td><td>288</td><td>320</td><td>384</td></tr><tr><th scope="row"># of tokens</th><td>100</td><td>196</td><td>256</td><td>324</td><td>400</td><td>576</td></tr><tr><th scope="row">Top-1</th><td>77.4</td><td>79.9</td><td>80.7</td><td>81.2</td><td>81.5</td><td>82.3</td></tr></tbody></table>

Figure 3.5 – **Impact of input resolution on accuracy for DeiT-S.** We consider different image resolutions, and either (1) **increase the patch size while keeping the number of tokens fixed**; or (2) **keep the patch size fixed and use more tokens**. Larger input images are beneficial if the number of tokens increases. The impact of a change of a resolution for a constant number of patches (of varying size) is almost neutral. As one can observe, the main driver of performance is the number of patches. The patch size has a limited impact on the accuracy, except when considering very small ones. We have observed and confirmed similar trends with XCiT models.
### 3.4.2 Approximate attention models in ViT with DeiT training

In Table 3.3, we report the results that we obtain by replacing the Multi-headed Self-attention operation with efficient variants [Ho, 2019; Shen, 2021; Wang, 2020c; Wang, 2021a] in the DeiT-S backbone. First, we can notice that for all efficient self-attention choices there is a clear drop in performance compared to the DeiT-S baseline. The spatial reduction attention (SRA) proposed in PVT [Wang, 2021a] has a significantly weaker performance compared to the full-attention with a quadratic complexity that is more efficient than full-attention by only a constant factor $R^2$. Linformer [Wang, 2020c] provides a better accuracy compared to SRA, however, it is also clearly weaker than full-attention. Moreover, Linformer does not have the flexibility of processing variable length sequences which limits its application in many computer vision tasks. Efficient attention [Shen, 2021] provides a better trade-off than the aforementioned methods, with improved accuracy and linear complexity. However, it has a 3.6% drop in performance compared to full-attention. Finally, axial attention [Ho, 2019] provides the strongest performance among the efficient attention variants we studied with a 1.5% drop in accuracy compared to the baseline. We observe a saving in memory usage, but a drop in speed due to the separate row and column attention operations. Our observations are consistent with [Dosovitskiy, 2021a].

<table><caption>Table 3.3 – ImageNet Top-1 accuracy of efficient self-attention variants (after 300 epochs of training).</caption><thead><tr><th>Model</th><th>Complexity</th><th>Top-1</th></tr></thead><tbody><tr><td>DeiT-S [Touvron, 2020]</td><td>O(N<sup>2</sup>)</td><td>79.9</td></tr><tr><td>SRA (Average Pool) [Wang, 2021a]</td><td>O(N<sup>2</sup>/R<sup>2</sup>)</td><td>73.5</td></tr><tr><td>SRA (Convolutional) [Wang, 2021a]</td><td>O(N<sup>2</sup>/R<sup>2</sup>)</td><td>74.0</td></tr><tr><td>Linformer (k=√n) [Wang, 2020c]</td><td>O(kN)</td><td>75.7</td></tr><tr><td>Efficient Transformer [Shen, 2021]</td><td>O(N)</td><td>76.3</td></tr><tr><td>Axial [Ho, 2019]</td><td>O(N√N)</td><td>78.4</td></tr></tbody></table>

### 3.4.3 Training and testing with varying resolution

For several tasks, it is important that the network is able to handle images of varying resolutions. This is the case, for instance, for image segmentation, image detection, or image retrieval where the object of interest may have very different sizes. We present an analysis of train/test resolution trade-off in Table 3.4.

<table><caption>Table 3.4 – Trade-off between train and test resolutions for DeiT. MS refers to multi-scale training, where the models have seen images from different resolutions at training time.</caption><thead><tr><th>Test / Train</th><th>160</th><th>224</th><th>256</th><th>288</th><th>320</th><th>MS</th></tr></thead><tbody><tr><td>160</td><td><strong>77.2</strong></td><td>75.9</td><td>73.3</td><td>68.2</td><td>59.6</td><td>76.3</td></tr><tr><td>224</td><td>78.0</td><td><strong>79.9</strong></td><td>79.9</td><td>79.0</td><td>77.9</td><td>79.6</td></tr><tr><td>256</td><td>77.3</td><td>80.4</td><td><strong>80.7</strong></td><td>80.2</td><td>79.9</td><td>80.6</td></tr><tr><td>288</td><td>76.3</td><td>80.4</td><td>81.0</td><td><strong>81.2</strong></td><td>80.8</td><td>81.0</td></tr><tr><td>320</td><td>75.0</td><td>80.1</td><td>80.9</td><td>81.3</td><td><strong>81.5</strong></td><td>81.3</td></tr></tbody></table>
## 3.5 Experimental evaluation

In this section we demonstrate the effectiveness and versatility of XCiT on multiple computer vision benchmarks, and present ablations providing insight on the importance of its different components.

<table><caption>Table 3.5 – ImageNet classification. Number of parameters, FLOPs, image resolution, and top-1 accuracy on ImageNet-1k and ImageNet-V2. Training strategies vary across models, transformer-based models and the reported RegNet mostly follow recipes from DeiT [Touvron, 2020].</caption><thead><tr><th>Model</th><th>#params</th><th>FLOPs</th><th>Res.</th><th>ImNet</th><th>V2</th></tr></thead><tbody><tr><td>EfficientNet-B5 RA [Cubuk, 2020]</td><td>30M</td><td>9.9B</td><td>456</td><td>83.7</td><td>—</td></tr><tr><td>RegNet-Y-4GF [Radosavovic, 2020]</td><td>21M</td><td>4.0B</td><td>224</td><td>80.0</td><td>72.4</td></tr><tr><td>DeiT-SΥ [Touvron, 2020]</td><td>22M</td><td>4.6B</td><td>224</td><td>81.2</td><td>68.5</td></tr><tr><td>Swin-T [Liu, 2021b]</td><td>29M</td><td>4.5B</td><td>224</td><td>81.3</td><td>—</td></tr><tr><td>CaiT-XS24Υ ↑ [Touvron, 2021c]</td><td>26M</td><td>19.3B</td><td>384</td><td>84.1</td><td>74.1</td></tr><tr><td>XCiT-S12/16Υ</td><td>26M</td><td>4.8B</td><td>224</td><td>83.3</td><td>72.5</td></tr><tr><td>XCiT-S12/16Υ ↑</td><td>26M</td><td>14.3B</td><td>384</td><td>84.7</td><td>74.1</td></tr><tr><td>XCiT-S12/8Υ ↑</td><td>26M</td><td>55.6B</td><td>384</td><td><b>85.1</b></td><td><b>74.8</b></td></tr><tr><td>EfficientNet-B7 RA [Cubuk, 2020]</td><td>66M</td><td>37.0B</td><td>600</td><td>84.7</td><td>—</td></tr><tr><td>NFNet-F0 [Brock, 2021]</td><td>72M</td><td>12.4B</td><td>256</td><td>83.6</td><td>72.6</td></tr><tr><td>RegNet-Y-8GF [Radosavovic, 2020]</td><td>39M</td><td>8.0B</td><td>224</td><td>81.7</td><td>72.4</td></tr><tr><td>Swin-S [Liu, 2021b]</td><td>50M</td><td>8.7B</td><td>224</td><td>83.0</td><td>—</td></tr><tr><td>CaiT-S24Υ ↑ [Touvron, 2021c]</td><td>47M</td><td>32.2B</td><td>384</td><td>85.1</td><td>75.4</td></tr><tr><td>XCiT-S24/16Υ</td><td>48M</td><td>9.1B</td><td>224</td><td>83.9</td><td>73.3</td></tr><tr><td>XCiT-S24/16Υ ↑</td><td>48M</td><td>26.9B</td><td>384</td><td>85.1</td><td>74.6</td></tr><tr><td>XCiT-S24/8Υ ↑</td><td>48M</td><td>105.9B</td><td>384</td><td><b>85.6</b></td><td><b>75.7</b></td></tr><tr><td>RegNet-Y-16GF [Radosavovic, 2020]</td><td>84M</td><td>16.0B</td><td>224</td><td>82.9</td><td>72.4</td></tr><tr><td>Swin-B↑ [Liu, 2021b]</td><td>88M</td><td>47.0B</td><td>384</td><td>84.2</td><td>—</td></tr><tr><td>DeiT-BΥ ↑ [Touvron, 2020]</td><td>87M</td><td>55.5B</td><td>384</td><td>85.2</td><td>75.2</td></tr><tr><td>CaiT-S48Υ ↑ [Touvron, 2021c]</td><td>89M</td><td>63.8B</td><td>384</td><td>85.3</td><td><b>76.2</b></td></tr><tr><td>XCiT-M24/16Υ</td><td>84M</td><td>16.2B</td><td>224</td><td>84.3</td><td>73.6</td></tr><tr><td>XCiT-M24/16Υ ↑</td><td>84M</td><td>47.7B</td><td>384</td><td>85.4</td><td>75.1</td></tr><tr><td>XCiT-M24/8Υ ↑</td><td>84M</td><td>187.9B</td><td>384</td><td><b>85.8</b></td><td>76.1</td></tr><tr><td>NFNet-F3 [Brock, 2021]</td><td>255M</td><td>114.8B</td><td>416</td><td>85.7</td><td>75.2</td></tr><tr><td>CaiT-M24Υ ↑ [Touvron, 2021c]</td><td>186M</td><td>116.1B</td><td>384</td><td>85.8</td><td>76.1</td></tr><tr><td>XCiT-L24/16Υ</td><td>189M</td><td>36.1B</td><td>224</td><td>84.9</td><td>74.6</td></tr><tr><td>XCiT-L24/16Υ ↑</td><td>189M</td><td>106.0B</td><td>384</td><td>85.8</td><td>75.8</td></tr><tr><td>XCiT-L24/8Υ ↑</td><td>189M</td><td>417.8B</td><td>384</td><td><b>86.0</b></td><td><b>76.6</b></td></tr></tbody></table>

### 3.5.1 Image classification

We use ImageNet-1k [Deng, 2009b] to train and evaluate our models for image classification. It consists of 1.28M training images and 50k validation images, labeled across 1,000 semantic categories. Our training setup follows the DeiT recipe [Touvron, 2020]. We train our model for 400 epochs with the AdamW optimizer [Loshchilov, 2017] using a cosine learning rate decay. In order to enhance the training of larger models, we utilize LayerScale [Touvron, 2021c] and adjust the stochastic depth [Huang, 2016] for each of our models accordingly. Following [Touvron, 2021c], images are cropped with crop ratio of 1.0 for evaluation. In addition to the ImageNet-1k validation set, we report results for ImageNet-V2 [Recht, 2019] which has a distinct test set. Our implementation is based on the Timm library [Wightman, 2019].
<table><caption>Table 3.6 – Evaluation on transfer learning.</caption><thead><tr><th>Architecture</th><th>CIFAR<sub>10</sub></th><th>CIFAR<sub>100</sub></th><th>Flowers<sub>102</sub></th><th>Cars</th><th>iNat<sub>18</sub></th><th>iNat<sub>19</sub></th></tr></thead><tbody><tr><td>EfficientNet-B7 [Tan, 2019]</td><td><u>98.9</u></td><td><b>91.7</b></td><td><b>98.8</b></td><td><b>94.7</b></td><td>—</td><td>—</td></tr><tr><td>ViT-B/16 [Dosovitskiy, 2021a]</td><td>98.1</td><td>87.1</td><td>89.5</td><td>—</td><td>—</td><td>—</td></tr><tr><td>ViT-L/16 [Dosovitskiy, 2021a]</td><td>97.9</td><td>86.4</td><td>89.7</td><td>—</td><td>—</td><td>—</td></tr><tr><td>DeiT-B/16 [Touvron, 2020] Υ</td><td><b>99.1</b></td><td>91.3</td><td><b>98.8</b></td><td>92.9</td><td><u>73.7</u></td><td><u>78.4</u></td></tr><tr><td>XCiT-S24/16 Υ</td><td><b>99.1</b></td><td>91.2</td><td>97.4</td><td>92.8</td><td>68.8</td><td>76.1</td></tr><tr><td>XCiT-M24/16 Υ</td><td><b>99.1</b></td><td><u>91.4</u></td><td>98.2</td><td>93.4</td><td>72.6</td><td>78.1</td></tr><tr><td>XCiT-L24/16 Υ</td><td><b>99.1</b></td><td>91.3</td><td><u>98.3</u></td><td><u>93.7</u></td><td><b>75.6</b></td><td><b>79.3</b></td></tr></tbody></table>

**Results on ImageNet.** We present a family of seven models in Table 3.2 with different operating points in terms of parameters and FLOPs. We observe that the performance of the XCiT models benefits from increased capacity both in depth and width. Additionally, consistent with [Touvron, 2020; Touvron, 2021c] we find that using hard distillation with a convolutional teacher improves the performance. Because of its linear complexity in the number of tokens, it is feasible to train XCiT at 384×384 resolution with small 8×8 patches, i.e. 2304 tokens, which provides a strong boost in performance across all configurations.

We compare to the state-of-the-art convolutional and transformer-based architectures [Brock, 2021; Liu, 2021b; Radosavovic, 2020; Tan, 2019; Touvron, 2021c] in Table 3.5. By varying the input image resolution and/or patch size, our models provide competitive or superior performance across model sizes and FLOP budgets. First, the models operating on 224×224 and 16×16 (e.g. XCiT-S12/16) enjoy high accuracy at relatively few FLOPs compared to their counterparts with comparable parameter count and FLOPs. Second, our models with 16×16 and 384×384 resolution images (e.g. XCiT-S12/16↑) yield an improved accuracy at the expense of higher FLOPs, and provide superior or on-par performance compared to state-of-the-art models with comparable computational requirements. Finally, XCiT linear complexity allows us to scale to process 384×384 images with 8×8 patch sizes (e.g. XCiT-S12/8↑), achieving the highest accuracy across the board, albeit at a relatively high FLOPs count.

**Transfer learning** In order to further demonstrate the flexibility and generality of our models, we report transfer learning experiments in Table 3.6 for models that have been pre-trained using ImageNet-1k and finetuned for other datasets including CIFAR-10, CIFAR-100 [Krizhevsky, 2009], Flowers-102 [Nilsback, 2008], Stanford Cars [Krause, 2013] and iNaturalist [Horn, 2017]. We observe that the XCiT models provide competitive performance when compared to strong baselines like ViT-B, ViT-L, DeiT-B and EfficientNet-B7.

**Class attention visualization.** In Figure 3.6 we show the class attention map obtained in the feature aggregation stage. Each head focuses on different semantically coherent regions in the image (e.g. faces or umbrellas). Furthermore, heads tend to focus on similar patterns across images (e.g. bird head or human face), but adapts by focusing on other salient regions when such patterns are absent.
<figure><img src="image_12.png" /><figcaption>Figure 3.6 – Visualization of the attention map between the CLS token and individual patches in the class-attention stage. For each column, each row represents the attention map w.r.t. one head, corresponding to the image in the first raw. Each head seems salient to semantically coherent regions. Heads are sensitive to similar features within the same or across images (e.g. people or bird faces). They trigger on different concepts when such features are missing (e.g., cockpit for race cars).</figcaption></figure>

**Queries and Keys magnitude visualizations** Our XCA operation relies on the cross-covariance matrix of the queries $\hat{Q}$ and keys $\hat{K}$ which are $\ell_2$ normalized across the patch dimension. Therefore, each element in the $d \times d$ matrix represents a cosine similarity whose value is strongly influenced by the magnitude of each patch. In Figure 3.7 we visualize the magnitude of patch embeddings in the queries and keys matrices. We observe that patch embeddings with higher magnitude corresponds to more salient regions in the image, providing a very cheap visualization and interpretation of which regions in the image contribute more in the cross-covariance attention.

**Robustness to resolution changes.** In Figure 3.3 we report the accuracy of XCiT-S12, DeiT-S and ResNet-50 trained on 224×224 images and evaluated at different image resolutions. While DeiT outperforms ResNet-50 when train and test resolutions are similar, it suffers from a larger drop in performance as the image resolution deviates farther from the training resolution. XCiT displays a substantially increased accuracy when train and test resolutions are similar, while also being robust to resolution changes, in particular for the model with 8×8 patches.
<figure><img src="image_13.png" /><figcaption>Figure 3.7 – Visualization of the queries $\hat{Q}$ and keys $\hat{K}$ norm across the feature dimension. We empirically observe that magnitude of patch embeddings in the queries and keys correlates with the saliency of their corresponding region in the image.</figcaption></figure>

<table><caption>Table 3.7 – **Self-supervised learning.** Top-1 acc. on ImageNet-1k. We report with a crop-ratio 0.875 for consistency with DINO. For the last row it is set to 1.0 (improves from 80.7% to 80.9%). All models are trained for 300 epochs.</caption><thead><tr><th>SSL Method</th><th>Model</th><th>#params</th><th>FLOPs</th><th>Linear</th><th>k-NN</th></tr></thead><tbody><tr><td>MoBY [Xie, 2021]</td><td>Swin-T [Liu, 2021b]</td><td>29M</td><td>4.5B</td><td>75.0</td><td>-</td></tr><tr><td>DINO [Caron, 2021]</td><td>ResNet-50 [He, 2016]</td><td>23M</td><td>4.1B</td><td>74.5</td><td>65.6</td></tr><tr><td>DINO [Caron, 2021]</td><td>ViT-S/16 [Dosovitskiy, 2021a]</td><td>22M</td><td>4.6B</td><td>76.1</td><td>72.8</td></tr><tr><td>DINO [Caron, 2021]</td><td>ViT-S/8 [Dosovitskiy, 2021a]</td><td>22M</td><td>22.4B</td><td><b>79.2</b></td><td><b>77.2</b></td></tr><tr><td>DINO [Caron, 2021]</td><td>XCiT-S12/16</td><td>26M</td><td>4.9B</td><td>77.8</td><td>76.0</td></tr><tr><td>DINO [Caron, 2021]</td><td>XCiT-S12/8</td><td>26M</td><td>18.9B</td><td><b>79.2</b></td><td>77.1</td></tr><tr><td>DINO [Caron, 2021]</td><td>ViT-B/16 [Dosovitskiy, 2021a]</td><td>87M</td><td>17.5B</td><td>78.2</td><td>76.1</td></tr><tr><td>DINO [Caron, 2021]</td><td>ViT-B/8 [Dosovitskiy, 2021a]</td><td>87M</td><td>78.2B</td><td>80.1</td><td>77.4</td></tr><tr><td>DINO [Caron, 2021]</td><td>XCiT-M24/16</td><td>84M</td><td>16.2B</td><td>78.8</td><td>76.4</td></tr><tr><td>DINO [Caron, 2021]</td><td>XCiT-M24/8</td><td>84M</td><td>64.0B</td><td><b>80.3</b></td><td><b>77.9</b></td></tr><tr><td>DINO [Caron, 2021]</td><td>XCiT-M24/8†384</td><td>84M</td><td>188.0B</td><td><b>80.9</b></td><td>-</td></tr></tbody></table>

**Self-supervised learning.** We train XCiT in a self-supervised manner using DINO [Caron, 2021] on ImageNet-1k. In Table 3.7 we report performance using the linear and k-NN protocols as in [Caron, 2021]. Across model sizes XCiT obtains excellent accuracy with both protocols, substantially improving DINO with ResNet-50 or ViT architectures, as well as over those reported for Swin-Transformer trained with MoBY [Xie, 2021]. Comparing the larger models to ViT, we also observed improved performance for XCiT achieving a strong 80.3% accuracy. For fair comparison, all reported models have been trained for 300 epochs. Further improved performance of small models is reported by Caron et al. [Caron, 2021] when training for 800 epochs,
<table><caption>Table 3.8 – **Ablations** of various architectural design choices on the task of ImageNet-1k classification using the XCiT-S12 model. Our baseline model uses the convolutional projection adopted from LeVit.</caption><thead><tr><th>Model</th><th>Ablation</th><th>ImNet top-1 acc.</th></tr></thead><tbody><tr><td>XCiT-S12/16</td><td rowspan="2">Baseline</td><td>82.0</td></tr><tr><td>XCiT-S12/8</td><td>83.4</td></tr><tr><td>XCiT-S12/16</td><td rowspan="2">Linear patch proj.</td><td>81.1</td></tr><tr><td>XCiT-S12/8</td><td>83.1</td></tr><tr><td rowspan="2">XCiT-S12/16</td><td>w/o LPI layer</td><td>80.8</td></tr><tr><td>w/o XCA layer</td><td>75.9</td></tr><tr><td rowspan="2">XCiT-S12/16</td><td>w/o l<sub>2</sub>-normal.</td><td>failed</td></tr><tr><td>w/o learned temp. τ</td><td>81.8</td></tr></tbody></table>

which we expect to carryover to XCiT based on the results presented here.

**Analysis and ablations.** In Table 3.8 we provide ablation experiments to analyse the impact of different design choices for our XCiT-S12 model. First, we observe the positive effect of using the convolutional patch projection as compared to using linear patch projection, for both 8×8 and 16×16 patches. Second, while removing the LPI layer reduces the accuracy by only 1.2% (from 82.0 to 80.8), removing the XCA layer results in a large drop of 6.1%, underlining the effectiveness of XCA. We noticed that the inclusion of two convolutional components – convolutional patch projection and LPI – not only brings improvements in accuracy, but also accelerates training. Third, although we were able to ensure proper convergence without $l_2$-normalization of queries and keys by tweaking the hyper-parameters, we found that it provides stability across model size (depth and width) and other hyper-parameters. Finally, while the learnable softmax temperature parameter is not critical, removing it drops accuracy by 0.2%.

### 3.5.2 Object detection and instance segmentation

Our XCiT models can efficiently process high-resolution images (see Figure 3.2). Additionally, XCiT has a better adaptability to varying image resolutions compared to ViT models (see Figure 3.3). These two properties make XCiT a good fit for dense prediction tasks including detection and segmentation.

We evaluate XCiT for object detection and instance segmentation using the COCO benchmark [Lin, 2014] which consists of 118k training and 5k validation images including bounding boxes and mask labels for 80 categories. We integrate XCiT as backbone in the Mask R-CNN [He, 2017] detector with FPN [Lin, 2017]. Since the XCiT architecture is inherently columnar, we make it FPN-compatible by extracting features from different layers (e.g., [4, 6, 8, 12] for XCiT-S12). All features have a constant stride of 8 or 16 based on the patch size, and the feature resolutions are adjusted to have strides of [4, 8, 16, 32], similar to ResNet-FPN backbones, where the downsampling is achieved by max pooling and the upsampling is obtained using a single transposed convolution layer. The model is trained for 36 epochs (3x schedule) using the AdamW optimizer with learning rate of $10^{-4}$, 0.05 weight decay and 16 batch size. We adopt the multiscale training and augmentation strategy of DETR [Carion, 2020]. Our implementation is
3.5. Experimental evaluation

<table><caption>Table 3.9 – **COCO object detection and instance segmentation** performance on the mini-val set. All backbones are pre-trained on ImageNet-1k, use Mask R-CNN model [He, 2017] and are trained with the same 3x schedule.</caption><thead><tr><th>Backbone</th><th>#params</th><th>AP<sup>b</sup></th><th>AP<sup>b</sup><sub>50</sub></th><th>AP<sup>b</sup><sub>75</sub></th><th>AP<sup>m</sup></th><th>AP<sup>m</sup><sub>50</sub></th><th>AP<sup>m</sup><sub>75</sub></th></tr></thead><tbody><tr><td>ResNet18 [He, 2016]</td><td>31.2M</td><td>36.9</td><td>57.1</td><td>40.0</td><td>33.6</td><td>53.9</td><td>35.7</td></tr><tr><td>PVT-Tiny [Wang, 2021a]</td><td>32.9M</td><td>39.8</td><td>62.2</td><td>43.0</td><td>37.4</td><td>59.3</td><td>39.9</td></tr><tr><td>ViL-Tiny [Zhang, 2021]</td><td>26.9M</td><td>41.2</td><td>64.0</td><td>44.7</td><td>37.9</td><td>59.8</td><td>40.6</td></tr><tr><td>XCiT-T12/16</td><td>26.1M</td><td>42.7</td><td>64.3</td><td>46.4</td><td>38.5</td><td>61.2</td><td>41.1</td></tr><tr><td>XCiT-T12/8</td><td>25.8M</td><td><strong>44.5</strong></td><td><strong>66.4</strong></td><td><strong>48.8</strong></td><td><strong>40.3</strong></td><td><strong>63.5</strong></td><td><strong>43.2</strong></td></tr><tr><td>ResNet50 [He, 2016]</td><td>44.2M</td><td>41.0</td><td>61.7</td><td>44.9</td><td>37.1</td><td>58.4</td><td>40.1</td></tr><tr><td>PVT-Small [Wang, 2021a]</td><td>44.1M</td><td>43.0</td><td>65.3</td><td>46.9</td><td>39.9</td><td>62.5</td><td>42.8</td></tr><tr><td>ViL-Small [Zhang, 2021]</td><td>45.0M</td><td>43.4</td><td>64.9</td><td>47.0</td><td>39.6</td><td>62.1</td><td>42.4</td></tr><tr><td>Swin-T [Liu, 2021b]</td><td>47.8M</td><td>46.0</td><td>68.1</td><td>50.3</td><td>41.6</td><td>65.1</td><td>44.9</td></tr><tr><td>XCiT-S12/16</td><td>44.3M</td><td>45.3</td><td>67.0</td><td>49.5</td><td>40.8</td><td>64.0</td><td>43.8</td></tr><tr><td>XCiT-S12/8</td><td>43.1M</td><td><strong>47.0</strong></td><td><strong>68.9</strong></td><td><strong>51.7</strong></td><td><strong>42.3</strong></td><td><strong>66.0</strong></td><td><strong>45.4</strong></td></tr><tr><td>ResNet101 [He, 2016]</td><td>63.2M</td><td>42.8</td><td>63.2</td><td>47.1</td><td>38.5</td><td>60.1</td><td>41.3</td></tr><tr><td>ResNeXt101-32</td><td>62.8M</td><td>44.0</td><td>64.4</td><td>48.0</td><td>39.2</td><td>61.4</td><td>41.9</td></tr><tr><td>PVT-Medium [Wang, 2021a]</td><td>63.9M</td><td>44.2</td><td>66.0</td><td>48.2</td><td>40.5</td><td>63.1</td><td>43.5</td></tr><tr><td>ViL-Medium [Zhang, 2021]</td><td>60.1M</td><td>44.6</td><td>66.3</td><td>48.5</td><td>40.7</td><td>63.8</td><td>43.7</td></tr><tr><td>Swin-S [Liu, 2021b]</td><td>69.1M</td><td><strong>48.5</strong></td><td><strong>70.2</strong></td><td><strong>53.5</strong></td><td><strong>43.3</strong></td><td><strong>67.3</strong></td><td><strong>46.6</strong></td></tr><tr><td>XCiT-S24/16</td><td>65.8M</td><td>46.5</td><td>68.0</td><td>50.9</td><td>41.8</td><td>65.2</td><td>45.0</td></tr><tr><td>XCiT-S24/8</td><td>64.5M</td><td>48.1</td><td>69.5</td><td>53.0</td><td>43.0</td><td>66.5</td><td>46.1</td></tr><tr><td>ResNeXt101-64 [Xie, 2017]</td><td>101.9M</td><td>44.4</td><td>64.9</td><td>48.8</td><td>39.7</td><td>61.9</td><td>42.6</td></tr><tr><td>PVT-Large [Wang, 2021a]</td><td>81.0M</td><td>44.5</td><td>66.0</td><td>48.3</td><td>40.7</td><td>63.4</td><td>43.7</td></tr><tr><td>ViL-Large [Zhang, 2021]</td><td>76.1M</td><td>45.7</td><td>67.2</td><td>49.9</td><td>41.3</td><td>64.4</td><td>44.5</td></tr><tr><td>XCiT-M24/16</td><td>101.1M</td><td>46.7</td><td>68.2</td><td>51.1</td><td>42.0</td><td>65.6</td><td>44.9</td></tr><tr><td>XCiT-M24/8</td><td>98.9M</td><td><strong>48.5</strong></td><td><strong>70.3</strong></td><td><strong>53.4</strong></td><td><strong>43.7</strong></td><td><strong>67.5</strong></td><td><strong>46.9</strong></td></tr></tbody></table>

based on the mm detection library [Chen, 2019].

**Results on COCO.** In Table 3.9 we report object detection and instance segmentation results of four variants of XCiT using 16 × 16 and 8 × 8 patches. We compare to ResNets [He, 2016] and concurrent efficient vision transformers [Liu, 2021b; Wang, 2021a; Zhang, 2021]. All models are trained using the 3x schedule after ImageNet-1k pre-training. Note that other results with higher absolute numbers have been achieved when pre-training on larger datasets [Liu, 2021b] or with longer schedules [Bello, 2021], and are therefore not directly comparable to the reported results. First, across all model sizes XCiT outperforms the convolutional ResNet [He, 2016] and ResNeXt [Xie, 2017] by a large margin with either patch size. Second, we observe a similar increase in accuracy compared to PVT [Wang, 2021a] and ViL [Zhang, 2021] backbones. Finally, XCiT provides a competitive performance with Swin [Liu, 2021b]<fnref n="2" />. For relatively small models, XCiT-S12/8 outperforms its Swin-T counterpart with a decent margin. On the other hand, Swin-S provides slightly stronger results compared to XCiT-S24/8. Utilizing smaller 8 × 8 patches leads to a consistent gain across all models.

<footnote n="2">2. We use report the results provided by the authors in their open-sourced code https://github.com/SwinTransformer/Swin-Transformer-Object-Detection</footnote>
Chapter 3. Cross Covariance Image Transformers

<table><caption>Table 3.10 – ADE20k semantic segmentation performance using Semantic FPN [Kirillov, 2019] and UperNet [Xiao, 2018] (in comparable settings). We do not include comparisons with other state-of-the-art models that are pre-trained on larger datasets [Liu, 2021b; Ranftl, 2021; Zheng, 2020].</caption><thead><tr><th rowspan="2">Backbone</th><th colspan="2">Semantic FPN</th><th colspan="2">UperNet</th></tr><tr><th>#params</th><th>mIoU</th><th>#params</th><th>mIoU</th></tr></thead><tbody><tr><td>ResNet18 [He, 2016]</td><td>15.5M</td><td>32.9</td><td>-</td><td>-</td></tr><tr><td>PVT-Tiny [Wang, 2021a]</td><td>17.0M</td><td>35.7M</td><td>-</td><td>-</td></tr><tr><td>XCiT-T12/16</td><td>8.4M</td><td>38.1</td><td>33.7M</td><td>41.5</td></tr><tr><td>XCiT-T12/8</td><td>8.4M</td><td><b>39.9</b></td><td>33.7</td><td><b>43.5</b></td></tr><tr><td>ResNet50 [He, 2016]</td><td>28.5M</td><td>36.7</td><td>66.5M</td><td>42.0</td></tr><tr><td>PVT-Small [Wang, 2021a]</td><td>28.2M</td><td>39.8</td><td>-</td><td>-</td></tr><tr><td>Swin-T [Liu, 2021b]</td><td>-</td><td>-</td><td>59.9M</td><td>44.5</td></tr><tr><td>XCiT-S12/16</td><td>30.4M</td><td>43.9</td><td>52.4M</td><td>45.9</td></tr><tr><td>XCiT-S12/8</td><td>30.4M</td><td><b>44.2</b></td><td>52.3M</td><td><b>46.6</b></td></tr><tr><td>ResNet101 [He, 2016]</td><td>47.5M</td><td>38.8</td><td>85.5M</td><td>43.8</td></tr><tr><td>ResNeXt101-32 [Xie, 2017]</td><td>47.1M</td><td>39.7</td><td>-</td><td>-</td></tr><tr><td>PVT-Medium [Wang, 2021a]</td><td>48.0M</td><td>41.6</td><td>-</td><td>-</td></tr><tr><td>Swin-S [Liu, 2021b]</td><td>-</td><td>-</td><td>81.0M</td><td>47.6</td></tr><tr><td>XCiT-S24/16</td><td>51.8M</td><td>44.6</td><td>73.8M</td><td>46.9</td></tr><tr><td>XCiT-S24/8</td><td>51.8M</td><td><b>47.1</b></td><td>73.8M</td><td><b>48.1</b></td></tr><tr><td>ResNeXt101-64 [Xie, 2017]</td><td>86.4M</td><td>40.2</td><td>-</td><td>-</td></tr><tr><td>PVT-Large [Wang, 2021a]</td><td>65.1M</td><td>42.1</td><td>-</td><td>-</td></tr><tr><td>Swin-B [Liu, 2021b]</td><td>-</td><td>-</td><td>121.0M</td><td>48.1</td></tr><tr><td>XCiT-M24/16</td><td>90.8M</td><td>45.9</td><td>109.0M</td><td>47.6</td></tr><tr><td>XCiT-M24/8</td><td>90.8M</td><td><b>46.9</b></td><td>108.9M</td><td><b>48.4</b></td></tr></tbody></table>

### 3.5.3 Semantic segmentation

We further show transferability of our models with semantic segmentation experiments on the ADE20k dataset [Zhou, 2017], which consists of 20k training and 5k validation images with labels over 150 semantic categories. We integrate our backbones in two segmentation methods: Semantic FPN [Kirillov, 2019] and UperNet [Xiao, 2018]. We train for 80k and 160k iterations for Semantic FPN and UperNet respectively. Following [Liu, 2021b], the models are trained using batch size 16 and an AdamW optimizer with learning rate of $6 \times 10^{-5}$ and 0.01 weight decay. We apply the same method of extracting FPN features as explained in Section 3.5.2. We report the performance using the standard single scale protocol (without multi-scale and flipping). Our implementation is based on the mmsegmentation library [Contributors, 2020].

**Results on ADE20k.** We present the semantic segmentation performance using XCI T backbones in Table 3.10. First, for Semantic FPN [Kirillov, 2019], XCI T provides a superior performance compared to ResNet, ResNeXt and PVT backbones using either option of patch size. Second, compared to Swin Transformers using the same UperNet decoder [Xiao, 2018], XCI T with 8×8 patches consistently achieves a higher mIoU for different models. XCI T with 16×16 patches provides a strong performance especially for smaller models where XCI T-S12/16 outperforms Swin-T.
3.5.4 Implementation details

**Sinusoidal Positional Encoding** We adopt a sinusoidal positional encoding as proposed by Vaswani *et al.* [Vaswani, 2017b] and adapted to the 2D case by Carion *et al.* [Carion, 2020]. However we depart from this method in that we first produce this encoding in an intermediate 64-d space before projecting it to the working space of the transformers. More precisely, in our implementation each of the *x* and *y* coordinates is encoded using 32 dimensions corresponding to cosine and sine functions with different frequencies (16 frequency for each function). The encoding of both coordinates are eventually concatenated to obtain a 64 dimension 2D positional encoding. Finally, the 64 dimension positional encoding is linearly projected to the working dimension of the model *d*.

**Obtaining Feature Pyramid for Dense Prediction** For state-of-the-art detection and segmentation models, FPN is an important component which provides features of multiple scales. We adapt XCiT to be compatible with FPN detection and segmentation methods through a simple re-scaling of the features extracted from different layers. In particular, for models with 12 layers, we extract features from the 4^th^, 6^th^, 8^th^ and 12^th^ layers respectively. As for models with 24 layers, we extract features from 8^th^, 12^th^, 16^th^ and 24^th^ layers. Concerning the re-scaling of the features, the 4 feature levels are downsized by a ratio of 4, 8, 16 and 32 compared to the input image size. Feature downsizing is performed with max pooling and upsampling is achieved using a single layer of transposed convolutions with kernel size *k* = 2 and stride *s* = 2.

**3.6 Conclusion and Future Work**

**Contributions.** We present an alternative to token self-attention which operates on the feature dimension, eliminating the need for expensive computation of quadratic attention maps. We build our XCiT models with the cross-covariance attention as its core component and demonstrate the effectiveness and generality of our models on various computer vision tasks. In particular, it exhibits a strong image classification performance on par with state-of-the-art transformer models while similarly robust to changing image resolutions as convnets. XCiT is effective as a backbone for dense prediction tasks, providing excellent performance on object detection, instance and semantic segmentation. Finally, we showed that XCiT can be a strong backbone for self-supervised learning, matching the state-of-the-art results with less compute. XCiT is a generic architecture that can readily be deployed in other research domains where self-attention has shown success.

**Limitations.** Our models enable training with smaller patches and on higher-resolution images, which leads to clear performance gains. However, for tasks like image classification this gain comes at a cost of relatively high number of FLOPs. In order to address this issue, other components, like FFN, could also be re-examined. Another point is that XCiT models seem to overfit more than their CaiT counterparts, see Table 3.5. They are more similar to some convnets in that respect.
# Chapter 4

## Sample Efficient Self-Supervised Pre-training with Vision Transformers

### Objectives

Pre-training models on large-scale datasets, like ImageNet, is a standard practice in computer vision. This paradigm is especially effective for tasks with small training sets, for which high-capacity models tend to overfit. In this work, we consider a self-supervised pre-training scenario that only leverages the target task data. We consider datasets, like Stanford Cars, Sketch or COCO, which are order(s) of magnitude smaller than Imagenet. Our study shows that denoising autoencoders, such as BEiT or a variant that we introduce in this chapter, are more robust to the type and size of the pre-training data than popular joint embedding self-supervised methods. We obtain competitive performance compared to ImageNet pre-training on a variety of classification datasets, from different domains. On COCO, when pre-training solely using COCO images, the detection and instance segmentation performance surpasses the supervised ImageNet pre-training in a comparable setting.

### Contents

<table><tr><td>4.1</td><td>Introduction</td><td>37</td></tr><tr><td>4.2</td><td>Related Work</td><td>39</td></tr><tr><td>4.3</td><td>Analysis</td><td>40</td></tr><tr><td>4.3.1</td><td>Sample Efficiency</td><td>41</td></tr><tr><td>4.3.2</td><td>Learning using non object-centric images</td><td>42</td></tr><tr><td>4.3.3</td><td>Tokenizers</td><td>42</td></tr><tr><td>4.4</td><td>Methodology</td><td>44</td></tr><tr><td>4.4.1</td><td>SplitMask</td><td>44</td></tr><tr><td>4.4.2</td><td>Encoder-Decoder Architecture</td><td>44</td></tr><tr><td>4.4.3</td><td>Global Contrastive Loss</td><td>45</td></tr><tr><td>4.5</td><td>Experiments</td><td>45</td></tr><tr><td>4.5.1</td><td>Datasets</td><td>45</td></tr><tr><td>4.5.2</td><td>Dense Prediction</td><td>47</td></tr><tr><td>4.5.3</td><td>Image Classification</td><td>47</td></tr><tr><td>4.5.4</td><td>Pre-training using ImageNet</td><td>49</td></tr><tr><td>4.5.5</td><td>Ablations</td><td>49</td></tr><tr><td>4.5.6</td><td>Implementation Details</td><td>51</td></tr><tr><td>4.6</td><td>Conclusion</td><td>52</td></tr></table>
4.1 Introduction

Modern computer vision neural networks are heavily parametrized: they routinely have tens or hundreds of millions of parameters [He, 2016; Dosovitskiy, 2021a; Liu, 2021b; Radosavovic, 2020]. This has been the key to their success for leveraging large-scale image collections such as ImageNet. However these high-capacity models tend to overfit on small, or even medium-sized datasets consisting of hundreds of thousands of images. This problem was pointed out by Oquab et al. [Oquab, 2014] in 2014:

"Learning CNNs [...] amounts to estimating millions of parameters and requires a very large number of annotated image samples. This property currently prevents the application of CNNs to problems with limited training data."

The authors describe a learning setting [Oquab, 2014; Yosinski, 2014] that is nowadays the dominant learning paradigm for data-starving problems:

(1) pre-train a model on a large dataset like Imagenet [Deng, 2009b], and in turn (2) finetune the weights of the models on the target task for which we have a limited amount of data. The second training stage typically adopts a shorter optimization procedure than the one employed when training from scratch (i.e., from randomly generated weights).

This simple approach has led to impressive results, which are state-of-the-art in many tasks such as detection [He, 2017; Carion, 2020], segmentation [Chen, 2014] and action recognition [Carreira, 2017]. Despite this success, we point out that it is difficult to disentangle the benefits offered by such a large-scale curated label dataset from the limitations of this pre-training paradigm. Putting aside the discussion on the collection effort (cost, requiring in-domain expertise, etc), we point out that pre-training a model on a dataset and fine-tuning it on another can introduce two sort of discrepancies.

First, this setting introduces a domain shift between the images used to pre-train the model and those targeted by the fine-tuning stage. Imagenet images may be sufficiently representative of natural images (despite the collecting bias). To date, most researchers consider that the benefit of having a large amount of images vastly compensates the domain discrepancy on benchmarks involving natural images, such as the fine-grained iNaturalist datasets [Van Horn, 2018; Horn, 2017] or even out-of-domain distributions such as sketches, paintings or clipart.

The second question, discussed by Doersch et al. [Doersch, 2020], is the so-called *supervision collapse*. This phenomenon is inherent to pre-training with a fixed set of labels: the network learns to focus on the mapping between images and the labels of the pre-training stage, but can discard information that is relevant to other downstream tasks. In other terms, pre-training on large-scale classification datasets does not necessarily align with the goal of learning general-purpose features, as it uses only a subset of the available information controlled by the given dataset categorization bias [Rosch, 1973].

These limitations have motivated the development of self-supervised pre-training methods which learn directly from data, without relying on annotations. Most notably, the contrastive and joint embedding approaches [He, 2020; Caron, 2020a; Caron, 2021; Chen, 2020c; Grill, 2020] can serve as effective pre-training strategies. While obtaining a strong performance on numerous tasks, such methods have a strong bias towards ImageNet data since the transformations have
<figure><img src="image_14.png" /><figcaption>Figure 4.1 – We demonstrate that self-supervised pre-training using denoising autoencoders like BEiT and our variant SplitMask are more robust to the type and/or size of pre-training data used. For example, the object detection performance of such models, when pre-trained only using COCO images and a Mask R-CNN pipeline, outperforms both supervised and BEiT self-supervised baselines pre-trained on ImageNet, as well as a randomly initialized baseline trained for a long schedule.</figcaption></figure>

been hand-designed to perform well on the ImageNet benchmark. Some of the most effective transformations, like cropping, rely on the images being object-centric [Purushwalkam, 2020]. When applied on uncurated data, these methods degrade significantly and require larger datasets to obtain similar performance [Goyal, 2021].

This is in contrast with natural language processing, where nowadays, most applications use large models which were pre-trained on uncurated data. In particular, the (masked) language modeling loss has been applied to transformer networks, leading to the BERT model [Devlin, 2018], which is now the foundation of most NLP models. Inspired by this success, Bao *et al.* [Bao, 2021] have shown the potential of the Masked Image Modeling (MIM) task to pre-train vision transformers. Such a model can be thought of as a denoising autoencoder [Vincent, 2008] where the noise corresponds to the patch masking operation. This technique has been successfully applied to ImageNet, but research questions remain:

(1) How much does this pre-training technique rely on the number of pre-training samples, and in particular, does it require millions of images to be useful?

(2) Is this technique robust to different distributions of training images? In particular, is it an effective paradigm to learn with non object-centric or uncurated images?

If the answer to both questions is positive, it will enable pre-training using a larger variety of datasets, including the training sets of many tasks that are smaller or belong to a different domain than ImageNet.

Overall, we make the following contributions:

— First, we demonstrate that denoising autoencoders are more sample efficient than joint embedding techniques, enabling pre-training without relying on large-scale datasets (e.g.
<figure><img src="image_15.png" /><figcaption>Figure 4.2 – Pre-training using different ImageNet subsets. Transfer performance does not improve beyond using a subset as small as 5% when trained for the same number of iterations.</figcaption></figure>

<figure><img src="image_16.png" /><figcaption>Figure 4.3 – Varying the number of pre-training epochs for the 10% subset. The performance first increases with longer training. Then we observe a plateau and a slight overfitting.</figcaption></figure>

ImageNet);

- Second, as a consequence of the better sample efficiency, we show on multiple datasets that it is possible to pre-train directly on the target task data and obtain a competitive performance, even with datasets that are orders of magnitude smaller than ImageNet;

- Third, we demonstrate that denoising autoencoders can be successfully applied to non object-centric images such as COCO, achieving performance similar to the one obtained when pre-training with ImageNet, unlike joint embedding techniques which seem to suffer a drop in performance.

## 4.2 Related Work

In this section, we briefly review some previous work on self-supervised learning, including autoencoders and instance discrimination methods.

**Pre-training with autoencoders** has a long history in deep learning, where it was initially used as a greedy layer-wise method to improve optimization [Hinton, 2006; Bengio, 2007; Ranzato, 2007; Vincent, 2010; Vincent, 2008]. In the context of unsupervised feature learning for image classification, different tasks related to denoising autoencoders have been considered, such as in-painting [Pathak, 2016], colorization [Zhang, 2016] or de-shuffling of image patches [Noroozi, 2016]. In natural language processing, denoising autoencoders have been applied by masking or randomly replacing some tokens of the input, and reconstructing the original sequence, leading to the BERT model [Devlin, 2018]. Similar methods have been proposed to pre-train sequence-to-sequence models, by considering additional kind of noises such as word shuffling or deleting [Raffel, 2019; Lewis, 2019].

There has been efforts to adopt such successful ideas in NLP to computer vision, but with limited success. Chen et al. [Chen, 2020b] proposed iGPT, a transformer-based autoregressive model that operates over image pixels, while Atito et al. [Atito, 2021] trained a ViT model on denoising of images where the noise is applied at pixel level. More recently, Bao et al. [Bao,
[2021] introduced the Masked Image Modeling loss in computer vision, where image patches are masked, and the goal is to predict the discretized label of the missing patches corresponding to their visual words as defined by a pre-trained discrete VAE [Ramesh, 2021b].

**Instance discrimination** is a set of self-supervised techniques which consider that each image corresponds to its own class [Dosovitskiy, 2014; Dosovitskiy, 2015]. A set of data augmentations (or transformations) is then applied to each image to generate multiple examples for each class. The global image representations are trained in a contrastive framework, typically using the InfoNCE loss [Wu, 2018b], to have high similarity for instances transformed from the same source image and low similarity with all other images. As the performance of these methods depends on the number of negatives, it either requires large batches or memory banks to work well [Wu, 2018b; He, 2020; Chen, 2020c]. It was later shown that when using a momentum encoder [He, 2020], simpler loss functions that did not directly discriminate against other images could be used [Grill, 2020; Caron, 2021; Zbontar, 2021; Bardes, 2021; Chen, 2021a]. Finally, a related line of work is to use clustering techniques to pre-train deep neural networks [Xie, 2016; Yang, 2016; Caron, 2018; Asano, 2019b; Caron, 2020b].

**Transformer networks** were originally introduced in the context of machine translation, replacing recurrent neural networks by an attention-based mechanism [Vaswani, 2017b]. Transformers were later applied to image recognition, by splitting images into patches, embedding these independently, and then processing the obtained representations as a sequence [Dosovitskiy, 2021a]. Initially, only vision transformers pre-trained on very large collections obtained good performance, but smaller models trained on ImageNet with heavy augmentation can also yield competitive tradeoffs [Touvron, 2020].

**Pre-training data** is an important ingredient of self-supervised learning, and multiple works have studied its impact on the transfer performance of models. While it is possible to learn high quality features from non-curated (eg. YFCC or IG) data using instance discrimination, this usually requires an order of magnitude more data than ImageNet [Caron, 2019; Goyal, 2021]. Similarly, one can perform supervised pre-training using weakly supervised data, such as using hashtags as labels, but this strategy also requires a large amount of data to work well [Joulin, 2016; Mahajan, 2018; Dosovitskiy, 2021a]. On the other hand, it was shown that for many natural language processing tasks, increasing the size of the pre-training dataset did not lead to strong improvement when using denoising autoencoders [Raffel, 2019]. Finally, some work studied how much could be learned from a single pre-training image [Asano, 2019a] or from synthetic data [Kataoka, 2020; Krishna, 2021].

## 4.3 Analysis

In this section, we study the impact of the pre-training data on the performance of denoising autoencoder, and how they compare to those of joint embedding methods. More precisely, we
<table><caption>Table 4.1 – Analysis of different self-supervision methods transfer performance to the iNaturalist-2019 dataset when varying the size of the ImageNet subset used in the pre-training stage, in addition to using non object-centric dataset like COCO for pre-training. We observe that denoising autoencoders have a more robust behaviour w.r.t. pre-training data size or nature compared to joint embedding methods like DINO as well as supervised pre-training.</caption><thead><tr><th rowspan="2">Method</th><th>IMNet 1%</th><th>IMNet 10%</th><th>IMNet Full</th><th>COCO</th></tr><tr><th>epochs: 30k</th><th>epochs: 3k</th><th>epochs: 300</th><th>epochs: 3k</th></tr></thead><tbody><tr><td>Supervised</td><td>71.6</td><td>75.0</td><td>75.8</td><td>—</td></tr><tr><td>DINO [Caron, 2021]</td><td>70.1</td><td>73.1</td><td><strong>78.4</strong></td><td>71.9</td></tr><tr><td>BEiT [Bao, 2021]</td><td>74.1</td><td>74.5</td><td>75.2</td><td>74.4</td></tr><tr><td>SplitMask</td><td><strong>74.8</strong></td><td><strong>75.4</strong></td><td>75.4</td><td><strong>76.3</strong></td></tr></tbody></table>

investigate how the number of images, and their nature, influence the quality of self-supervised models. In this preliminary analysis, we consider the recent method BEiT and SplitMask, our variant as detailed in Section 4.4, as representatives of denoising autoencoders, and DINO [Caron, 2021] of a joint embedding method, respectively.

### 4.3.1 Sample Efficiency

**Denoising autoencoders vs Supervised/DINO** First, we start by studying the impact of the pre-training dataset size, by varying the number of ImageNet examples we use to train models. We consider subsets of ImageNet containing 10% and 1% of the total number of examples, and use the balanced (in terms of classes) subsets from [Assran, 2021]. To decouple the effect of using smaller datasets and the effect of doing less training updates, we adapt the number of epochs to keep the number of iterations constant. This means that we perform 3k and 30k epochs on ImageNet 10% and 1% respectively. We report results in Table 4.1. Observe how pre-training with an autoencoder loss such as masked image modeling is robust to the reduction in dataset size. In contrast, like for supervised pre-training, the performance of models pre-trained with DINO self-supervision degrades when training with smaller datasets.

**Pre-training number of samples** We plot the iNaturalist-2019 transfer performance as a function of ImageNet subset size used during pre-training using SplitMask in Figure 4.2. We observe that the peak performance is achieved using only 5% of the ImageNet samples and adding more samples does not provide additional boost, given the number of updates are kept constant. We also observe that using only a single image per class, which corresponds to the 0.1% subset containing 1000 samples, leads to a non-trivial boost (+4 points) over training from scratch. This is a strong indication that denoising autoencoders are highly sample efficient unsupervised learning methods.

**Pre-training schedule length** Furthermore, we plot the transfer performance as a function of number of pre-training epochs in Figure 4.3 using the 10% ImageNet subset. It can be observed that training for long schedules of nearly 3k epochs, matching the total number of updates for that of full ImageNet with 300 epochs, is crucial to achieve such strong performance
Chapter 4. Sample Efficient Self-Supervised Pre-training with Vision Transformers

<table><caption>Table 4.2 – Ablation study on the effect of different tokenization methods. We compare the DALL-E tokenizer originally used in BEiT with patch level techniques: random projection, random patches and k-means clustering. We observe that the DALL-E tokenizer can be effectively replaced by simpler methods that do not require training on a large dataset.</caption><thead><tr><th></th><th>DALL-E</th><th>Rand. Proj.</th><th>Rand. Patches</th><th>K-Means</th></tr></thead><tbody><tr><td>iNat19</td><td>75.2</td><td>75.2</td><td>75.3</td><td>75.0</td></tr></tbody></table>

for smaller subsets. However, we observe slight overfitting for very long schedules. This problem is more predominant for pre-training using very small datasets like Stanford-Cars as illustrated in Figure 4.6.

### 4.3.2 Learning using non object-centric images

We now study the impact of changing the nature of the pre-training data. In particular we use images that are not object-centric, like in Imagenet. To this end, instead of pre-training using ImagetNet, we pre-train with images from the COCO dataset only. As COCO contains roughly 118k images, this dataset is approximately equivalent in terms of size to the ImageNet 10% subset. Again, to disentangle the effect of training with a different number of iterations, we adapt the number of epochs: we use 3k epochs on COCO.

We report the results of this experiments in Table 4.1. When pre-trained on COCO, DINO drops significantly compared to full ImageNet pre-training (-8.3). Interestingly, the drop is higher than using 10% ImageNet even though the numbers of samples is roughly the same. We hypothesis this is because COCO images are not biased to be object-centric, while this joint embedding method was designed and developed using ImageNet as benchmark. In contrast, BEiT’s performance only decreases slightly while SplitMask attains +0.7 improvement over full ImageNet pre-training. This is an interesting property which makes such models prime candidates for learning effectively from uncurated images in the wild.

### 4.3.3 Tokenizers

The BEiT method, as proposed by Bao et al. [Bao, 2021], relies on the discrete VAE tokenizer from DALL-E, which has been pretrained on a large weakly supervised dataset. Since we want to study whether it is possible to pre-train models solely on small datasets, or non object-centric ones, we replace the DALL-E tokenizer by a simple alternative. To this end, we consider different simple alternatives to discretize images at the patch level without any pre-training. Each of these techniques is applied on each patch independently, making them relatively lightweight and more efficient than the original tokenizer considered in BEiT.

Given a vocabulary of size $V$, each element of the vocabulary is represented by a unit vector $\mathbf{e}_i \in \mathbb{R}^d$, where $i \in \{1, ..., V\}$ and $d$ is the dimension of patches (in the case of 8x8 patches, $d = 192$). Then, to tokenize an image, we associate each patch to the element of the vocabulary which has the highest cosine similarity with the patch in the pixel space. Hence, for a patch $\mathbf{x}$,
<figure><img src="image_17.png" /><figcaption>Figure 4.4 – SplitMask consists of three steps. First, the input image patches are split into two disjoint subsets. Second, a shared deep ViT encoder processes each subset separately. The encoder outputs on each branch are augmented with a set of special mask tokens, representing the positions of the missing patches, and fed to a shallow ViT decoder. The decoder output corresponding to the mask tokens is used to solve a MIM task similar to BEiT. Finally, a global image descriptor is extracted from the decoder outputs of each branch by means of average pooling. The descriptors are trained to have high similarity using a contrastive loss (InfoNCE).</figcaption></figure>

its corresponding token $t$ is obtained as

$$t = \operatorname{argmax}_{i \in \{1, \dots, V\}} \mathbf{x}^\top \mathbf{e}_i. \qquad (4.1)$$

We now discuss three simple ways to obtain the elements of the vocabulary $\mathbf{e}_i$. First, we can sample random vectors with uniform element-wise distribution, and call the corresponding tokenizer *random projection*. Second, we can sample $V$ random patches uniformly in the set of all patches of images from the training set, and refer to the tokenizer as *random patches*. Finally, we can perform $k$-means clustering on the patches of images from the training set, and use the centroids as elements of the vocabulary. We refer to this last tokenizer, which was once widely employed in computer vision for bag-of-words representations, as *k-means*.

We train a ViT-base model on the ImageNet dataset, using these three tokenizers, as well as the DALL-E tokenizer originally considered by BEiT. We report results in Table 4.2. We observe that replacing the DALL-E tokenizer by simpler choices does not lead to any significant degradation in accuracy. This also provides a 26% relative runtime improvement for base models over its counterpart using the DALL-E tokenizer on 16 GPUs with a batch size of 1024.
## 4.4 Methodology

In this section, we introduce SplitMask, a variant of denoising autoencoders based on vision transformers. An overview of our method is illustrated in Figure 5.2.

### 4.4.1 SplitMask

Our approach is based on three steps, which we refer to as *split*, *inpaint* and *match*. As in standard vision transformers, an image is first broken down into patches of 16×16 pixels. Then, we *split* the patches into two disjoint subsets *A* and *B*, which are processed independently by our deep ViT encoder. Next, using the patch representations of the subset *A* and a shallow decoder (e.g. 2 layers), we *inpaint*<fnref n="1" /> the patches of the subset *B*, by solving a *MIM* task, and vice versa. Finally, we obtain a global image descriptor by average pooling of the patch representations from the decoder output corresponding to each branch.

The feature aggregation is over both observed and hallucinated patches. We try to *match* the global descriptors of the image obtained from subset *A* to that obtained from subset *B*. In other words, we use the masking operation of the mask image modeling loss as a data augmentation for a contrastive learning loss similar to NPID or SimCLR. Note, SplitMask does not add any significant computational cost over *MIM* methods like BEiT to produce this global contrastive training signal.

### 4.4.2 Encoder-Decoder Architecture

We now discuss in more details the architecture of the model that we use to implement the SplitMask pipeline described in the previous subsection. Our method relies on an encoder-decoder architecture. The encoder of our model is a standard vision transformer, with absolute positional embeddings. In contrast to BEiT method, our encoder does not process representations of the *masked* tokens, but only of the observed ones². Hence, an image is divided into patches, which are linearly embedded, and positional embeddings are added to these representations. These representations are split into two subsets *A* and *B*, which are processed independently by standard transformer layers. Before feeding the output representations to the decoder, we insert mask embeddings that includes the position information of the missing patches in the sequences *A* and *B*. Finally, using the decoded representations of the *masked* patches, we predict their corresponding visual words using a cross entropy loss function.

Thus, if an image contains *n* patches, the encoder processes two sequences of size *n*/2, while the decoder processes two sequences of size *n*. Since in practice we use decoder which is much more lightweight than standard vision transformers, the computational complexity of our models is similar to a standard ViT. One advantage of our approach compared to BEiT is that at each iteration, the encoder processes all the patches of the image. The loss function is also computed over all the patches of the image, instead of only on a subset.

<footnote n="1">1. Inpainting in this context is implemented by solving a Masked Image Modeling task rather than the typical inpainting by reconstruction of pixels.
2. Concurrent to this work, He et al. [He, 2021b] proposed MAE. This is an encoder-decoder architecture where the encoder processing the observed patches only, similar to what we do in our SplitMask variant.</footnote>
### 4.4.3 Global Contrastive Loss

In addition to the MIM loss, which is computed at the patch level, our approach also uses a contrastive loss at the image level. To this end, we apply an average pooling operation over all the output representations of the decoder (including representations of the masked patches). For each image, we obtain two representations $\mathbf{x}_a$ and $\mathbf{x}_b$, corresponding to the subsets $\mathcal{A}$ and $\mathcal{B}$ of observed patches. We then apply the InfoNCE loss [Oord, 2018] over these representations:

$$ \ell(\mathbf{x}_a) = \frac{\exp(\mathbf{x}_a^\top \mathbf{x}_b / \tau)}{\sum_{\mathbf{y} \in \{\mathbf{x}_b\} \cup \mathcal{N}} \exp(\mathbf{x}_a^\top \mathbf{y} / \tau)}, \quad (4.2) $$

where $\tau$ is a temperature hyper-parameter and $\mathcal{N}$ is a set of negatives, corresponding to the representations of the other images in the batch. Following previous work [Chen, 2020c], we symmetrize the contrastive loss, and apply it similarly on the representation $\mathbf{x}_b$ from the subset $\mathcal{B}$. The motivation for adding this contrastive loss is to encourage the model to produce globally coherent features that are consistent across different choices of observed subsets without relying on any hand-designed transformations. Using our design of SplitMask, we attain such signal with almost no overhead.

## 4.5 Experiments

In this section, we perform empirical evaluations of denoising autoencoders, and the impact of the pre-training data on downstream task performance. In particular, we study how well pre-training performs when only the target task data is used instead of relying on a large-scale dataset such as ImageNet. We perform experiments on different tasks, such as classification, detection and instance segmentation. We consider datasets of varying size, including some significantly smaller than ImageNet. We also compare our variant SplitMask method to BEiT, either pre-trained on target task data or ImageNet, in addition to the supervised pre-training baselines. Finally, we perform an ablation study on our method to investigate the impact of its different components on finetuning and linear evaluation.

### 4.5.1 Datasets

We study the pre-training and finetuning of computer vision models on a variety of datasets, see Table 4.3 for details. For image classification, we consider the iNaturalist 2018 and 2019 [Van Horn, 2018], Stanford Cars [Krause, 2013] and Food101 [Bossard, 2014] datasets, which all contain fine-grained categories. We also consider three subsets from the DomainNet dataset [Peng, 2019], *clipart*, *painting* and *sketch*, which are not natural images and hence from different domains than ImageNet. For object detection and instance segmentation, we use the COCO dataset [Lin, 2014]. Finally, we also use the ADE20k dataset [Zhou, 2017] for semantic segmentation. The training set sizes of these different datasets vary from 8k to 437k images, thus all being significantly smaller than ImageNet, some more than two order of magnitude smaller. This allows to investigate under different data regimes how feasible it is to pre-train directly on
Chapter 4. Sample Efficient Self-Supervised Pre-training with Vision Transformers

<table><caption>Table 4.3 – Data size, number of classes and number of pre-training epochs details for all datasets used for pre-training.</caption><thead><tr><th>Dataset</th><th>#Train</th><th>#Test</th><th>#Classes</th><th>Epochs</th></tr></thead><tbody><tr><td>ImageNet [Deng, 2009b]</td><td>1,281,167</td><td>50,000</td><td>1000</td><td>300</td></tr><tr><td>iNaturalist 2018 [Van Horn, 2018]</td><td>437,513</td><td>24,426</td><td>8,142</td><td>800</td></tr><tr><td>iNaturalist 2019 [Horn, 2017]</td><td>265,240</td><td>3,003</td><td>1,010</td><td>1,400</td></tr><tr><td>Food 101 [Bossard, 2014]</td><td>75,750</td><td>25,250</td><td>101</td><td>5,000</td></tr><tr><td>Stanford Cars [Krause, 2013]</td><td>8,144</td><td>8,041</td><td>196</td><td>5,000</td></tr><tr><td>Clipart [Peng, 2019]</td><td>34,019</td><td>14,818</td><td>345</td><td>5,000</td></tr><tr><td>Painting [Peng, 2019]</td><td>52,867</td><td>22,892</td><td>345</td><td>5,000</td></tr><tr><td>Sketch [Peng, 2019]</td><td>49,115</td><td>21,271</td><td>345</td><td>5,000</td></tr><tr><td>ADE20k [Zhou, 2017]</td><td>20,210</td><td>2,000</td><td>150</td><td>21,000</td></tr><tr><td>COCO [Lin, 2014]</td><td>118,287</td><td>5,000</td><td>80</td><td>3,000</td></tr></tbody></table>

<table><caption>Table 4.4 – COCO detection and instance segmentation performance, using a Mask R-CNN pipeline, for models with different pre-training recipes. We see that BEiT and SplitMask pre-training using COCO outperform supervised ImageNet pre-training of DeiT as well as self-supervised ImageNet pre-training using BEiT. †: Method uses a longer 6x schedule instead of the default 3x following He et al. [He, 2018].</caption><thead><tr><th rowspan="2">Method</th><th rowspan="2">Backbone</th><th colspan="3">Pre-training</th><th rowspan="2">AP<sup>b</sup></th><th rowspan="2">AP<sub>50</sub><sup>b</sup></th><th rowspan="2">AP<sub>75</sub><sup>b</sup></th><th rowspan="2">AP<sup>m</sup></th><th rowspan="2">AP<sub>50</sub><sup>m</sup></th><th rowspan="2">AP<sub>75</sub><sup>m</sup></th></tr><tr><th>Supervised</th><th>IMNet</th><th>COCO</th></tr></thead><tbody><tr><td>Random Initialization</td><td rowspan="6">ViT-S</td><td>✗</td><td>✗</td><td>✗</td><td>38.3</td><td>60.1</td><td>41.4</td><td>35.6</td><td>57.1</td><td>37.7</td></tr><tr><td>Random Initialization†</td><td>✗</td><td>✗</td><td>✗</td><td>42.8</td><td>64.5</td><td>45.6</td><td>39.1</td><td>61.5</td><td>41.7</td></tr><tr><td>DeiT [Touvron, 2020]</td><td>✓</td><td>✓</td><td>✗</td><td>44.2</td><td>66.6</td><td>47.9</td><td>40.1</td><td>63.2</td><td>42.7</td></tr><tr><td>BEiT [Bao, 2021]</td><td>✗</td><td>✓</td><td>✗</td><td>44.5</td><td>66.2</td><td>48.8</td><td>40.3</td><td>63.2</td><td>43.1</td></tr><tr><td>DINO [Caron, 2021]</td><td>✗</td><td>✗</td><td>✓</td><td>43.7</td><td>65.5</td><td>47.7</td><td>39.6</td><td>62.3</td><td>42.3</td></tr><tr><td>BEiT</td><td>✗</td><td>✗</td><td>✓</td><td>44.7</td><td>66.3</td><td>48.8</td><td>40.2</td><td>63.1</td><td>43.2</td></tr><tr><td>SplitMask</td><td></td><td>✗</td><td>✗</td><td>✓</td><td><strong>45.3</strong></td><td><strong>66.9</strong></td><td><strong>49.4</strong></td><td><strong>40.6</strong></td><td><strong>63.6</strong></td><td><strong>43.5</strong></td></tr><tr><td>Random Initialization</td><td rowspan="6">ViT-B</td><td>✗</td><td>✗</td><td>✗</td><td>40.7</td><td>62.7</td><td>44.2</td><td>37.1</td><td>59.1</td><td>39.4</td></tr><tr><td>Random Initialization†</td><td>✗</td><td>✗</td><td>✗</td><td>43.0</td><td>64.2</td><td>46.9</td><td>38.8</td><td>61.3</td><td>41.6</td></tr><tr><td>DeiT [Touvron, 2020]</td><td>✓</td><td>✓</td><td>✗</td><td>45.5</td><td><strong>67.9</strong></td><td>49.2</td><td>41.0</td><td>64.6</td><td>43.8</td></tr><tr><td>BEiT [Bao, 2021]</td><td>✗</td><td>✓</td><td>✗</td><td>46.3</td><td>67.6</td><td>50.6</td><td>41.6</td><td>64.5</td><td>44.9</td></tr><tr><td>DINO [Caron, 2021]</td><td>✗</td><td>✗</td><td>✓</td><td>43.1</td><td>64.4</td><td>46.9</td><td>38.9</td><td>61.4</td><td>41.4</td></tr><tr><td>BEiT</td><td>✗</td><td>✗</td><td>✓</td><td>46.7</td><td>67.7</td><td>51.2</td><td>41.8</td><td>65.0</td><td>44.6</td></tr><tr><td>SplitMask</td><td></td><td>✗</td><td>✗</td><td>✓</td><td><strong>46.8</strong></td><td><strong>67.9</strong></td><td><strong>51.5</strong></td><td><strong>42.1</strong></td><td><strong>65.3</strong></td><td><strong>45.1</strong></td></tr></tbody></table>

the target task data, alleviating the need for a large scale curated dataset as ImageNet.

As previously mentioned, we want to perform a constant number of updates during pre-training, and we thus adapt the number of epochs when training on target task data to match the number of updates corresponding to 300 epochs on ImageNet. For smaller classification datasets, we limit the number of pre-training epochs to 5000 since we observed pre-training for longer generally does not result in further improvement in terms of downstream performance. For very small datasets, like Stanford Cars, we observed an overfitting behaviour with training for very long schedules (e.g. more than 5k epochs, see Figure 4.6). Note that the adjusted number of pre-training epochs is provided in Table 4.3.
4.5.2 Dense Prediction

4.5.2.1 Object detection and Instance Segmentation

First, we evaluate our approach on the COCO object detection and instance segmentation dataset using the Mask R-CNN pipeline [He, 2017] and report our results in Table 4.4. We compare models pre-trained on the COCO dataset alone with their equivalent counterparts that were pre-trained on ImageNet, either in a supervised or self-supervised fashion. First, we observe that BEiT models which were pre-trained on the COCO dataset alone obtain better downstream task performance than the same models pre-trained on ImageNet. For example, when using a ViT-base backbone, pre-training on COCO instead of ImageNet leads to a boost of +0.4 in box AP.

Additionally, we observe that a similar pre-training of DINO using COCO images provides
a relatively weak performance, only outperforming random initialization. This indicates that
strong pre-training on COCO is a unique property of denoising autoencoders and it does not
extend to other self-supervised learning methods.

Finally, we observe that SplitMask leads to a consistent improvement compared to the BEiT
baseline, such as +0.6 box AP when using a ViT-small and +0.3 mask AP for ViT-base back-
bones. All put together, in a comparable setting, we obtain a +1.1 box AP increase while not
using ImageNet. Since COCO contains one order of magnitude less images than ImageNet, this
suggests that large scale datasets are not necessary for pre-training.

4.5.2.2 Semantic Segmentation

For semantic segmentation, we compare our denoising autoencoder models, pre-trained solely using ADE20k images, to their counterparts pre-trained on ImageNet. The results are reported in Table 4.5. All models use an UperNet pipeline [Xiao, 2018]. We observe that denoising autoencoders can provide a very competitive performance on such a challenging task even when pre-trained using a relatively small sample size of 20k images. The performance matches that of BEiT self-supervised pre-training using ImageNet and only marginally lower than supervised ImageNet pre-training.

We have found that adapting the random cropping strategy is a crucial implementation
detail that helps improve the denoising autoencoders pre-training performance on such dataset.
In particular, we reduce the maximal size of the crop from 100% to 25% of the raw image size.

4.5.3 Image Classification

We perform empirical evaluation on a number classification datasets and report our results in Table 4.6. Overall, we find that BEiT or SplitMask pre-training, using solely the target datasets images, consistently obtains either the strongest or, at worst, the second strongest performance when compared to different options of self-supervised and supervised pre-training using ImageNet as well as training from scratch [Liu, 2021a].
Chapter 4. Sample Efficient Self-Supervised Pre-training with Vision Transformers

<table><caption>Table 4.5 – Semantic segmentation performance for different pre-trained models on ADE20k using an UperNet pipeline [Xiao, 2018]. All models reported use a ViT-B architecture. In spite of the small size of the ADE20k dataset, performance of our models provides a performance competitive to those pre-trained using ImageNet.</caption><thead><tr><th rowspan="2">Method</th><th colspan="3">Pre-training</th><th rowspan="2">mIoU</th></tr><tr><th>Supervised</th><th>IMNet</th><th>ADE20k</th></tr></thead><tbody><tr><td>Random Init.</td><td>✗</td><td>✗</td><td>✗</td><td>25.4</td></tr><tr><td>DeiT [Touvron, 2020]</td><td>✓</td><td>✗</td><td>✗</td><td>46.1</td></tr><tr><td>BEiT [Bao, 2021]</td><td>✗</td><td>✓</td><td>✗</td><td>45.6</td></tr><tr><td>BEiT</td><td>✗</td><td>✗</td><td>✓</td><td>45.6</td></tr><tr><td>SplitMask</td><td>✗</td><td>✗</td><td>✓</td><td>45.7</td></tr></tbody></table>

<table><caption>Table 4.6 – Comparison between finetuning performance on the target datasets of different sizes and domains when pre-trained using the target datasets themselves, ImageNet pre-training (both supervised and self-supervised), and training from scratch. Both denoising autoencoders (BEiT and SplitMask) obtain competitive performance when solely using the target data. $\ddagger$: Liu et al. [Liu, 2021a] use a different pre-training setup and backbones.</caption><thead><tr><th rowspan="2">Method</th><th rowspan="2">Backbone</th><th rowspan="2">Supervised pre-training</th><th colspan="2">Data Used</th><th>iNat-18</th><th>iNat-19</th><th>Food 101</th><th>Cars</th><th>Clipart</th><th>Painting</th><th>Sketch</th></tr><tr><th>IMNet</th><th>Target</th><th>437k</th><th>265k</th><th>75k</th><th>8k</th><th>34k</th><th>52k</th><th>49k</th></tr></thead><tbody><tr><td rowspan="2">Liu et al. [Liu, 2021a]<sup>‡</sup></td><td>CVT-13</td><td>✗</td><td>✗</td><td>✓</td><td>—</td><td>—</td><td>—</td><td>—</td><td>60.6</td><td>55.2</td><td>57.6</td></tr><tr><td>ResNet-50</td><td>✗</td><td>✗</td><td>✓</td><td>—</td><td>—</td><td>—</td><td>—</td><td>63.9</td><td>53.5</td><td>59.6</td></tr><tr><td>Random Init.</td><td></td><td>✗</td><td>✗</td><td>✓</td><td>59.6</td><td>67.5</td><td>84.7</td><td>35.3</td><td>41.0</td><td>38.4</td><td>37.2</td></tr><tr><td>DeiT [Touvron, 2020]</td><td></td><td>✓</td><td>✓</td><td>✓</td><td>69.9</td><td>75.8</td><td><strong>91.5</strong></td><td>92.2</td><td><strong>79.6</strong></td><td><strong>74.2</strong></td><td><strong>72.5</strong></td></tr><tr><td>BEiT [Bao, 2021]</td><td>ViT-S</td><td>✗</td><td>✓</td><td>✓</td><td>68.1</td><td>75.2</td><td>90.5</td><td>92.4</td><td>75.3</td><td>68.7</td><td>68.5</td></tr><tr><td>BEiT</td><td></td><td>✗</td><td>✗</td><td>✓</td><td>68.8</td><td>76.1</td><td>90.7</td><td>92.7</td><td>—</td><td>69.0</td><td>—</td></tr><tr><td>SplitMask</td><td></td><td>✗</td><td>✗</td><td>✓</td><td><strong>70.1</strong></td><td><strong>76.3</strong></td><td><strong>91.5</strong></td><td><strong>92.8</strong></td><td>78.3</td><td>69.2</td><td><strong>70.7</strong></td></tr><tr><td>Random Init.</td><td></td><td>✗</td><td>✗</td><td>✓</td><td>59.6</td><td>68.1</td><td>83.3</td><td>36.9</td><td>41.9</td><td>37.6</td><td>34.9</td></tr><tr><td>DeiT [Touvron, 2020]</td><td></td><td>✓</td><td>✓</td><td>✓</td><td>73.2</td><td>77.7</td><td><strong>91.9</strong></td><td>92.1</td><td><strong>80.0</strong></td><td><strong>73.8</strong></td><td><strong>72.6</strong></td></tr><tr><td>BEiT [Bao, 2021]</td><td>ViT-B</td><td>✗</td><td>✓</td><td>✓</td><td>71.6</td><td>78.6</td><td>91.0</td><td><strong>93.9</strong></td><td>78.0</td><td>71.5</td><td>71.4</td></tr><tr><td>BEiT</td><td></td><td>✗</td><td>✗</td><td>✓</td><td>72.4</td><td>79.3</td><td><strong>91.7</strong></td><td>92.7</td><td>—</td><td>70.7</td><td>—</td></tr><tr><td>SplitMask</td><td></td><td>✗</td><td>✗</td><td>✓</td><td><strong>74.6</strong></td><td><strong>80.4</strong></td><td>91.2</td><td><strong>93.1</strong></td><td>79.3</td><td>72.0</td><td><strong>72.1</strong></td></tr></tbody></table>

<table><caption>Table 4.7 – Finetuning performance on ImageNet. Here, epochs refer to the number of pre-training epochs on ImageNet.</caption><thead><tr><th>Method</th><th>Backbone</th><th>Epochs</th><th>Top-1</th></tr></thead><tbody><tr><td>MocoV3 [Chen, 2021b]</td><td rowspan="4">ViT-S</td><td>300</td><td>81.4</td></tr><tr><td>DINO [Caron, 2021]</td><td>300</td><td><strong>81.5</strong></td></tr><tr><td>BEiT [Bao, 2021]</td><td>300</td><td>81.3</td></tr><tr><td>SplitMask</td><td>300</td><td><strong>81.5</strong></td></tr><tr><td>MocoV3 [Chen, 2021b]</td><td rowspan="5">ViT-B</td><td>300</td><td>83.2</td></tr><tr><td>DINO [Caron, 2021]</td><td>400</td><td><strong>83.6</strong></td></tr><tr><td>BEiT [Bao, 2021]</td><td>300</td><td>82.8</td></tr><tr><td>BEiT [Bao, 2021]</td><td>800</td><td>83.2</td></tr><tr><td>SplitMask</td><td>300</td><td><strong>83.6</strong></td></tr></tbody></table>

**BEiT pre-training: ImageNet vs Target** First, we compare ImageNet pre-training to the target data pre-training with BEiT and observe that for many cases, pre-training on the target data alone leads to better results. This is true for the ViT-small backbone across all the datasets including Stanford cars (+1.1% acc), which consists of only 8k images. When using a ViT-base backbone, pre-training on the target task data outperforms BEiT self-supervised
ImageNet pre-training for datasets as small as Food101 (+0.7 acc), which is more than 10x smaller than ImageNet. Second, we observe that SplitMask leads to further improvement in performances for multiple datasets: for example, on the iNaturalist 2018 dataset, we see +3.0 in accuracy with a ViT-base model.

**Supervised ImageNet pre-training** As it was already observed in previous work [Chen, 2020c; Chen, 2021b; Caron, 2021], we also see in many cases that self-supervised training outperforms supervised pre-training on ImageNet. For example, on the iNaturalist datasets, training with the target task data alone (including a pre-training step) gives better results than pre-training on ImageNet with labels: with a ViT-base model and the SplitMask method, we see an improvement of +2.7% in top-1 accuracy. As for the *clipart*, *painting* and *sketch* datasets, we see that SplitMask provides a competitive performance, outperforming an ImageNet pre-trained BEiT across all datasets for ViT-S. However, for the aforementioned datasets, supervised pre-training achieves the best performance for both ViT-S and ViT-B.

We note that when pre-training using the *clipart* and *sketch* datasets with the BEiT method, we experienced numerical instability that prevented the model from converging with long schedules (e.g. 5000 epochs). However, the instability problem was not observed for SplitMask models. Nevertheless, more investigation might be needed to fully understand how to optimize pre-training of such models.

### 4.5.4 Pre-training using ImageNet

In addition to our main study concerning the robustness of denoising autoencoders w.r.t the size and type of pre-training data, we study SplitMask in the more commonly used setting of pre-training and finetuning using ImageNet.

In Table 4.7 we show the performance of our SplitMask method using the ViT-S and ViT-B backbones and 300 epochs pre-training compared to other recent transformer-based self-supervised learning methods. It can be observed that SplitMask provides a strong performance, outperforming both BEiT and MocoV3 for all backbones. Additionally, SplitMask achieves a performance on par with DINO while being significantly cheaper and simpler to train. Note that while SplitMask and BEiT attain a strong finetuning performance, denoising autoencoding methods typically fall behind in terms of linear probing compared to instance discrimination methods like DINO.

### 4.5.5 Ablations

#### 4.5.5.1 SplitMask vs BEiT

We ablate our proposed components in SplitMask compared to a BEiT baseline in Table 4.8. All models use a ViT-B backbone and pre-trained for 300 epochs. First, we observe that the ImageNet finetuning performance improves with a margin (+0.5) by simply adopting the encoder-decoder architecture and processing two disjoint subsets per iteration. Second, the global contrastive loss on its own, without the MIM objective, provides a very weak performance.
Chapter 4. Sample Efficient Self-Supervised Pre-training with Vision Transformers

This is expected since there is no training signal for the local patch representations, and a global matching objective with 50% masking of patches may be too hard, providing a noisy training signal and hindering the model's ability to learn informative features.

Our full SplitMask model that uses both the MIM and contrastive objectives obtains the best performance and outperforms BEiT by a large margin of +0.8. The Linear probing performance of SplitMask is stronger than BEiT. However, both models provide a relatively weak performance on this benchmark compared to instance discrimination methods, whose final layers are more aligned to the classification task. Note, SplitMask adds a negligible computing overhead compared to the BEiT baseline: its wall-clock training time is marginally higher as detailed in Table 4.8. All models are trained using 16 GPUs and batch size of 2048.

<table><caption>Table 4.8 – Ablations of different components in our SplitMask model in comparison with a BEiT baseline. All models including the baseline have been trained for 300 epochs using a ViT-B backbone.</caption><thead><tr><th>Method</th><th>Split</th><th>Inpaint</th><th>Match</th><th>Finetune</th><th>Lin.</th><th>Hours</th></tr></thead><tbody><tr><td>BEiT [Bao, 2021]</td><td>✗</td><td>✓</td><td>✗</td><td>82.8</td><td>41.0</td><td>32.5</td></tr><tr><td rowspan="3">SplitMask</td><td>✓</td><td>✓</td><td>✗</td><td>83.3</td><td>46.4</td><td><strong>31.0</strong></td></tr><tr><td>✓</td><td>✗</td><td>✓</td><td>79.3</td><td>4.0</td><td>32.5</td></tr><tr><td>✓</td><td>✓</td><td>✓</td><td><strong>83.6</strong></td><td><strong>46.5</strong></td><td>34.0</td></tr></tbody></table>

### 4.5.5.2 Encoder-Decoder vs BEiT

<figure><img src="image_18.png" /><figcaption>Figure 4.5 – Linear probing accuracy on ImageNet for SplitMask and BEiT using features extracted from different layers.</figcaption></figure>

An advantage of the encoder-decoder design we propose in 4.4.2 is that it encourages decoupling of general-purpose encoding of image features, which is required for the downstream tasks, and features specific to solving the pretext task of MIM. In particular, compared to BEiT the encoder is not capable of solving the pretext task on its own since it does not have access to the mask token. Therefore, it can only help solve the task by providing informative representations.
tation to the decoder which is the component responsible of solving the pretext task. We can
see in Figure 4.5 that this property improves the transferability of later layers representation to
downstream tasks compared to BEiT which has a stronger drop in linear probing performance
in later layers.

4.5.5.3 Overfitting during pre-training

We observed that for pre-training of very small datasets (e.g. Stanford-Cars), longer pre-
training schedules can be counterproductive. For example, if we follow the assumption we need
to pre-training for the same number of updates of ImageNet pre-training for 300 epochs, the
Stanford-Cars equivilant schedule would be 45k epochs. However, as we see in Figure 4.6,
pre-training longer than 5k epochs leads to a severe drop in finetuning performance.

<figure><img src="image_19.png" /><figcaption>Figure 4.6 – Finetuning performance for the Stanford Cars datasets as a function of number of pre-training epochs using the same datasets images.</figcaption></figure>

4.5.6 Implementation Details

**Tokenizers.** Similarly to the tokenizer used in [Bao, 2021], all tokenizers presented in Table 4.2 have a vocabulary of size 8192. For the random tokenizer, we sample 8192 vectors with uniform component-wise distribution. For the random patches tokenizer we sample 8192 patches from different images. For the K-means tokenizer, the 8192 elements of the vocabulary are obtained by applying the K-means algorithm to 3 millions patches sampled from the dataset.

**Pre-training.** We use the original ViT formulation as proposed by Dosovitskiy et al. [Dosovitskiy, 2021a] and we follow the pre-training hyperparameters of Bao et al. [Bao, 2021]. All baselines reported use the same backbone implementation and trained in similar settings. For SplitMask, by default, we use random block masking [Bao, 2021] of 50% masking ratio to obtain a mask and its complement to extract the two subsets. The maximum and minimum number of patches per block is 75 and 16 respectively. We use the standard random cropping and horizontal flipping as data augmentations. We use 2 transformer layers for the decoder with embedding dimensions matching that of the encoder.
However, for the smallest datasets (i.e. Stanford-Cars, ClipArt, Sketch and Paintings), we
found that stronger data augmentation and more aggressive masking prevents early overfitting.
In particular, we use a uniform masking of 75% (like in the work by He et al. [He, 2021b]),
as well as using random greyscale, solarization, Gaussian blur and color jittering as additional
forms of data augmentation.

The BEiT baselines pre-trained on ImageNet and reported in Table 4.4 and 4.6 use the DALL-E tokenizer. Other BEiT and SplitMask models have been pre-trained using our random projection tokenizer. For the InfoNCE loss we use τ = 0.2 following Chen et al. [Chen, 2021b].

**Object detection and Instance segmentation.** We use the Mask R-CNN detection method [He, 2017] with ViT backbone as our detection method. In order to obtain features compatible with the Feature Pyramid Network (FPN) design [Lin, 2017], we use max pooling and transposed convolution operations similar to El-Nouby et al. [El-Nouby, 2021d]. To accommodate for the variable resolution we replace the absolute positional encoding for our models and the baselines with sinusoidal positional encoding [Vaswani, 2017b]. All models are trained using the 3x schedule (36 epochs) unless mentioned otherwise. We use the training hyper-parameters used by Liu et al. [Liu, 2021b].

4.6 Conclusion

In this chapter, we have raised the question of how to pre-train models with self-supervised
learning, wondering in particular on whether large scale datasets such as Imagenet are necessary
for pre-training. Our study on ImageNet shows that taking a smaller pre-training dataset
does not lead to a big performance drop for denoising autoencoders, as opposed to instance
discrimination self-supervised techniques or supervised pre-training. Similarly, training on non
object-centric images does not impact the downstream task performance significantly.

Building upon these observations, we have pre-trained models directly on the target task
data, instead of ImageNet, and performed evaluations on datasets of various sizes. We have
shown that it is possible to pre-train on datasets 10x smaller than ImageNet, for example
obtaining +0.5 box AP gains by solely using COCO images. We believe that this is strong
evidence that large-scale datasets, such as ImageNet, are not necessary for self-supervised pre-
training when using denoising autoencoders.
# Chapter 5

## Image Compression with Product Quantized Masked Image Modeling

### Objectives

Recent neural compression methods have been based on the popular hyperprior framework. It relies on Scalar Quantization and offers a very strong compression performance. This contrasts from recent advances in image generation and representation learning, where Vector Quantization is more commonly employed.

In this chapter, we attempt to bring these lines of research closer by revisiting vector quantization for image compression. We build upon the VQ-VAE framework and introduce several modifications. First, we replace the vanilla vector quantizer by a product quantizer which implicitly defines high-quality quantizers that would otherwise require intractably large codebooks. Second, inspired by the success of Masked Image Modeling in the context of self-supervised learning, we propose a novel conditional entropy model which improves entropy coding by modeling the co-dependencies of the quantized latent codes. The resulting PQ-MIM model is surprisingly effective. It outperforms HiFiC in terms of FID and KID when optimized with perceptual losses. Finally, since PQ-MIM is compatible with image generation frameworks, we show qualitatively that it can operate under a hybrid mode between compression and generation. As a result, we explore the extreme compression regime where an image is compressed into 200 bytes, i.e., less than a tweet.

### Contents

<table><tr><td>5.1</td><td>Introduction</td><td>54</td></tr><tr><td>5.2</td><td>Related work</td><td>56</td></tr><tr><td>5.3</td><td>Product Quantized Masked Image Modeling</td><td>57</td></tr><tr><td>5.3.1</td><td>High-level architecture: PQ-VAE</td><td>58</td></tr><tr><td>5.3.2</td><td>Image entropy model</td><td>59</td></tr><tr><td>5.3.3</td><td>Training the PQ-MIM</td><td>61</td></tr><tr><td>5.4</td><td>Experiments</td><td>62</td></tr><tr><td>5.4.1</td><td>Experimental setup</td><td>62</td></tr><tr><td>5.4.2</td><td>Main experimental results</td><td>64</td></tr><tr><td>5.4.3</td><td>Analysis and Ablations</td><td>65</td></tr><tr><td>5.4.4</td><td>Limitations</td><td>67</td></tr><tr><td>5.5</td><td>Conclusion</td><td>68</td></tr></table>
## 5.1 Introduction

Image compression played a crucial role in accelerating multiple events in history: it was a major component of the Voyager mission [Ludwig, 2016]. Efficient image codecs have accelerated the rapid growth of the internet by enabling the transmission of images in a few dozen of kilobytes, thanks to the emergence of effective lossy methods. This democratization was accompanied by standardization efforts to facilitate interoperability, which led to the emergence of standards such as the Joint Photographic Experts Groups (JPEG). Subsequent formats have leveraged scientific advances on all components of source coding, ranging from transforms [Antonini, 1992], and quantization [Gray, 1998], to entropy coding [Witten, 1987; Taubman, 2000], eventually leading to modern video compression codecs enabling streaming and video-conferencing applications.

Neural methods have recently become increasingly popular for image compression as well as other image processing tasks, such as denoising [Tian, 2020], super-resolution [Bruna, 2016; Dong, 2015; Ledig, 2017; Wang, 2021b] or image reconstruction [Wang, 2020a; Knoll, 2020]. In typical scenarios, neural image compression is not necessarily mature enough to take over standard techniques like the BPG format inherited from the High-Efficiency Video Coding standard [Sullivan, 2012]. This is because they do not offer a significant quantitative advantage over prior works that would justify the higher complexity, which depends on the context and operational constraints. A key advantage of neural compression methods is their enhanced qualitative reconstruction when incorporating an adversarial loss or likewise psycho-visual objectives favoring visually appealing reconstructions [Agustsson, 2019; Mentzer, 2020]. From this perspective, neural compression is related to image generation. The two subfields, however, are currently dominated by different approaches, noticeably they employ different discretization procedures. Indeed, while earlier neural compression methods utilized vector quantization [Agustsson, 2017, VQ], recent methods mostly employ scalar quantization (SQ). In contrast, the recent literature on image generation [Chang, 2022; Yu, 2021; Esser, 2021; Rombach, 2021] relies on Vector Quantization jointly with a distortion criterion akin to those used in compression.

In this work we aim to reduce the methodological gap and to make a step towards the unification of neural image compression and image generation, and allowing image compression to more directly benefit from the rapid advances in image generation methods. Patch-based masking methods for self-supervised learning [Bao, 2022; He, 2021a; El-Nouby, 2021a] have recently demonstrated their potential for image generation [Chang, 2022]. Inspired by this work we propose a compression approach built upon Vector Quantized Variational Auto-Encoders [Oord, 2017; Razavi, 2019]. In this context, we focus on two intertwined questions: (1) How to define a vector quantizer offering a range of rate-distortion operating points? (2) How to define an entropy model minimizing the cost of storing the quantization indexes, while avoiding the prohibitive complexity of an auto-regressive model?

To address the challenges above, we revisit vector quantization in image compression, and investigate product quantization [Jégou, 2010] (PQ) in a compression system derived from VQ-VAE [Oord, 2016b]. We show that PQ offers a strong and scalable rate-distortion trade-off. We
<figure><img src="image_20.png" /><figcaption>Figure 5.1 – **Qualitative example of PQ-MIM compression.** PQ-MIM provides a strong compression performance. We retain many of the details present original image with minimal blurring effect even with compression rate as low as 0.196 bpp. Moreover, compared to HiFiC we achieve a lower rate for the same image. PQ-MIM provides colors that are more faithful to the original image, while HiFiC has a darkening effect and some high-frequency artifacts. On the other hand, PQ-MIM can have some smoothing effect that can cause loss of detail for particular regions (e.g. some of the Arabic letters in the example above). More qualitative examples are provided in Figure 5.11.</figcaption></figure>

then we focus on the spatial entropy modeling and coding of the quantization indexes in the VQ
or PQ latent layer, hence we name our method as (Vector/Product)-Quantized Masked Image
Modeling (VQ-MIM and PQ-MIM). To this end, we introduce a multi-stage vector-quantized
image model: we gradually reduce the conditional entropy of the patch latent codes by increasing
the number of observed patches we condition on for each stage. The conditional distribution over
patches is estimated by a transformer model and provided to an entropy coder, symmetrically
on the emitter and receiver sides.

In summary, we make the following contributions:

- We introduce a novel Masked Image Modeling conditional entropy model that significantly reduces the rates by leveraging the spatial inter-dependencies between latent codes.

- We introduce product quantization for VQ-VAE. This simple PQ-VAE variant offers a strong and scalable rate-distortion trade-off.

- When trained with adversarial and perpetual losses, PQ-MIM exhibits a strong performance in terms of perceptual metrics like FID and KID, outperforming HiFiC [Mentzer, 2020].

- We qualitatively show that PQ-MIM is capable of operating in a hybrid mode, between generative and compression, without requiring further training and finetuning. This allows for higher resilience to corrupted or missing signal where our model can fill-in the missing information.
## 5.2 Related work

**Neural image compression** Early approaches to neural image compression reach back to the late 1980s [Sonehara, 1989; Sicuranza, 1990; Bottou, 1998]. Recent rapid advances in explicit and implicit density modelling [Goodfellow, 2014; Kingma, 2014; Larochelle, 2011; Oord, 2016a; Rezende, 2015; Salimans, 2017] have renewed interest in posing image compression as a learning problem. Due to the connection to variational learning [Gregor, 2016; Frey, 1998; Alemi, 2018], variational auto-encoders have been the primary choice for lossy image compression. In contrast to standard variational models, the evaluation of neural compression models focuses on achievable bitrates, multi-scale applicability and computational complexity. The initial works in the field used fully convolutional architectures for encoding/decoding [Ballé, 2017; Theis, 2017; Mentzer, 2018]. The resulting encoded latent image representations are quantized and compressed via an entropy coder with learned explicit density model or “entropy bottleneck”. Initial variational approaches directly modeled a single level of code densities. Ballé extended these models by introducing a second “hyperprior” that yielded improved performance [Ballé, 2018]. Hyperprior models have been the basis for several subsequent advances with further improvements for density modeling, such as joint autoregressive models [Minnen, 2018], Gaussian mixture/attention [Cheng, 2020], and channel-wise auto-regressive models [Minnen, 2020]. Another line of work has proposed to use vector quantization with histogram-based probabilities for image compression [Agustsson, 2017; Lu, 2019]. Contrary to VQ-VAE models, these models typically optimize the rate (or a surrogate of the rate) directly and may include a spatial component for the quantized vectors. Yang et al. [Yang, 2020] showed multiple ways to improve the encoding process, including the fine-tuning of the discretization process and employing bits-back coding in the entropy bottleneck. Finally, [Santurkar, 2018; Mentzer, 2020; Rippel, 2017; Agustsson, 2019] showed that altering the distortion metric to include an additional adversarial loss can make a large difference for compression rate. Another interesting line of work considers image compression by training image-specific networks, or network adapters, that map image coordinates to RGB values, and compressing the image-specific parameters [Dupont, 2021; Dupont, 2022; Strüpler, 2022].

**VQ-models for image generation.** There has been significant interest in generative models based in discrete image representations, as introduced by VQ-VAE [Oord, 2017; Razavi, 2019]. A discrete representation of reduced spatial resolution is learned by means of an autoencoder which quantizes the latent representations. This discrete representation is coupled with a strong prior, for example implemented as an autoregressive pixel-CNN model. VQ-GAN [Esser, 2021] replaces the prior architecture with a transformer model [Vaswani, 2017a], and introduces an adversarial loss term to learn an autoencoder with more visually pleasing reconstructions and improved sample quality. This approach has been extended to text-based generative image models by extending the prior to model a longer sequence that combines the discrete image representation with a prefix that encoding the conditioning text. This has yielded impressive results by scaling the model capacity and training data to tens or hundreds of million text-image pairs [Ding, 2021a; Gafni, 2022; Ramesh, 2021a]. A fundamental limitation of autoregressive generative
<figure><img src="image_21.png" /><figcaption>Figure 5.2 – PQ-MIM overview. Our model consists of (i) a transformer based encoder and decoder, (ii) a masked image model (MIM) for conditional entropy modeling, and (iii) an entropy coder, e.g. an arithmetic coder (AE/AD). The input image $\mathbf{x}$ is projected to a set of latent features, followed by product quantization to yield quantization indices $\mathbf{q}$. The arithmetic coder encodes (and decodes) $\mathbf{q}$ into a bitstream $\mathbf{b}$ in a lossless manner. The elements in $\mathbf{q}$ are spatially split into groups, as detailed in Figure 5.4. Conditional Entropy Modeling. Our model estimates the conditional probabilities of the discrete indices in $S$ steps. Every step, a subset of the tokens $q_i$ is selected using the quincunx pattern. Our MIM transformer estimates $p(q_i|q_{s<i})$ and passes it to the Arithmetic encoder as a CDF, effectively reducing the lossless compression cost.</figcaption></figure>

models is that they sample data sequentially, requiring separate non-parallel evaluation of the predictive (transformer) model to sample each token. To alleviate this, several data items can be sampled independently in parallel, conditioning on all previously sampled tokens. This has been leveraged to speed-up pixel-CNNs for small images and video by two to three orders of magnitude [Reed, 2017]. More recently, MaskGIT [Chang, 2022] and follow-up work [Lezama, 2022] explored this for generative models of VQ-VAE representations, and find that similar or better sample quality is obtained by parallel sampling of image patch subsets in few steps, reducing the generation time significantly. Despite their success for image synthesis we are not aware of the use of such models for image compression in earlier work.

## 5.3 Product Quantized Masked Image Modeling

This work takes a step towards to closing the gap between neural image compression and image generation methodologies. We revisit vector quantization for image compression and propose an entropy model inspired by masked image modelling. Our compression pipeline, depicted in Figure 5.2, relies on three neural networks:

1. The Encoder network $\hat{E}: \mathcal{X} \to \mathcal{Z}$ maps input images $\mathbf{x} \in \mathcal{X}$ to a quantized representation $\mathbf{z} \in \mathcal{Z}$.

2. The Mask Image Model (MIM) compresses the quantized representations without loss of information. This network is involved both on the compression and decompression side.
3. The Decoder network $G: \mathcal{Z} \to \mathcal{X}$ produces an estimate $\hat{\mathbf{x}} = G(\hat{E}(\mathbf{x}))$ of the original image $\mathbf{x}$.

We now detail the architecture, in particular our PQ proposal, of the image model that we employ in the statistical lossless coding, as well as the training scheme.

### 5.3.1 High-level architecture: PQ-VAE

**High-level architecture.** We follow recent work on discrete generative image models for the design of image encoder and decoder [Chang, 2022; Esser, 2021; Oord, 2017; Razavi, 2019; Yu, 2021]. The encoder $E: \mathcal{X} \to \mathbb{R}^{w \times h \times d}$ takes an RGB image $\mathbf{x}$ of resolution $W \times H$ as input and maps it to a latent representation $E(\mathbf{x})$ with $d$ feature channels and a reduced spatial resolution $w \times h$, downsampling the input resolution by a factor $f = W/w = H/h$. The $T = w \times h$ elements of the latent representation $E(\mathbf{x})$ are quantized with a vector quantizer $Q(\cdot)$ to produce the quantized latent representation $\mathbf{z} = Q(E(\mathbf{x})) = \hat{E}(\mathbf{x})$, where each element in $E(\mathbf{x})$ is replaced with its nearest cluster center. The decoder $G$ uses the quantized latents $\mathbf{z}$ to reconstruct the image.

**Product Quantization.** In VQ-VAE, the quantizer $Q$ is simply an online k-means quantizer that produces quantization indices from real-valued vectors. We denote by $\mathbf{q} \in \{1, \dots, V\}^T$ the map of quantization indices indicating for each element of $\mathbf{z}$ which of the $V$ centroids is selected. The higher $V$ the more precise is the approximation $\mathbf{z}$, leading to higher bit rates. For instance, assuming that indices are coded with a naive coding scheme (see next section), the bit rate is doubled when moving to $V = 256$ to $V = 65536$ centroids.

However, scaling the number of centroids $K$ is possible up to thousands of centroids, but beyond that it is computationally prohibitive. Additionally, it is challenging to train large codebooks where each centroid has a very low probability of being updated. To address this problem, we replace the online k-means quantizer by a product quantizer [Jégou, 2010] (PQ): the latent vector $\mathbf{z}$ is split into $M$ subvectors as $\mathbf{z} = [\mathbf{z}^1, \dots, \mathbf{z}^j, \dots, \mathbf{z}^M]$ of dimension $M/d$. Each subvector is quantized by a distinct quantizer having $V_s$ quantization value. The set of quantizers implicitly defines a vector quantizer in the latent space with $V = V_s^M$ distinct centroids. Hence, we can easily define very large codebooks without the computational and optimization problems mentioned above, because both the assignment and learning are marginalized

over the different subspaces. Empirically, we observe in Figure 5.3 that PQ provides a better scaling behaviour for higher rates compared to VQ whose codebook size needs to grow exponentially to achieve the same rates.

<figure><img src="image_22.png" /><figcaption>Figure 5.3 – PQ and VQ comparison. While VQ provides a comparable performance to PQ for extremely low rates where the codebook size is small, PQ exhibits better scaling behaviour for higher rates.</figcaption></figure>
<figure><img src="image_23.png" /><figcaption>Figure 5.4 – Illustration of PQ-MIM with a quincunx pattern. We employ the quincunx pattern both on the encoder and decoder side. Each panel represents one of five stages in which we en/de-code a set of tokens in parallel using a probability model $p_s$ implemented by a transformer with parameters $\theta_s$. The transformer predicts the tokens in $q_s$, displayed in grayscale and marked by “?”, and takes as input the preceding groups of tokens $q_1, \dots, q_{s-1}$ that are displayed in color. The distribution provided by this neural network is fed to an arithmetic en/de-coder.</figcaption></figure>

**Neural network.** Without loss of generality, we choose all neural network models to be identical. This is not a requirement but this offers the property that the encoder and decoder have identical complexities, and that the memory and compute peaks are identical. More specifically we choose a cross-variance transformer [El-Nouby, 2021c] (XCiT), whose complexity is linear with respect to image resolution. In contrast, standard vision transformers [Dosovitskiy, 2021b] (ViT) are quadratic in the image surface, which is prohibitive for high resolution images that can typically require strong compression. We point out that recent work has shown that Swin-Transformers [Zhu, 2021] could be a compelling choice as well in the context of image compression. Formally they also have a quadratic complexity, but this is amortized by the hierarchical structure of this architecture.

### 5.3.2 Image entropy model

In this section we present PQ-MIM. The objective is to compress the discrete representations $q$ without loss of information, producing a bitstream of the compressed representation that can be transmitted or stored. During the decoding stage, we invert the aforementioned lossless compression. This model can be regarded as the VQ/PQ-VAE counterpart of adaptive contextual arithmetic coders, like EBCOT [Taubman, 2000] or CABAC [Richardson, 2004], proposed in early compression standards, in that it couples a conditional probabilistic model with an arithmetic coder.

**Lossless compression.** A naive manner for lossless compression of the discrete image codes $q = \{q_t\}_{t=1}^T$ is to use fixed-length codes. In that case, each code word is assigned to a unique binary representation of equal length, resulting in $\lceil \log_2 V \rceil$ bits per element $q_t$. This approach is computationally very efficient as fixed-length codes are not model-based, and as such do not require likelihood estimation, and because all codes are of equal length by construction, the computation is perfectly parallelizable. However, theoretically this coding scheme could be Shannon optimal only if codewords are uniformly and independently distributed. Those assumptions are not met in practice due to the architecture choices we have made previously: