Original: A N I MAGE IS W ORTH 16 X 16 W ORDS : T RANSFORMERS FOR I MAGE R ECOGNITION AT S CALE
Revised: A N I MAGE IS W ORTH 16 X 16 W ORDS : T RANSFORMERS FOR I MAGE R ECOGNITION AT S CALE

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: While the Transformer architecture has become the de-facto standard for natural language processing tasks, its applications to computer vision remain limited.
Revised: While the Transformer architecture has become the de-facto standard for natural language processing tasks, its applications to computer vision remain limited.

Original: In vision, attention is either applied in conjunction with convolutional networks, or used to replace certain components of convolutional networks while keeping their overall structure in place.
Revised: In vision, attention is either applied in conjunction with convolutional networks, or used to replace certain components of convolutional networks while keeping their overall structure in place.

Original: We show that this reliance on CNNs is not necessary and a pure transformer  can perform very well on image classiﬁcation tasks when applied directly to sequences of image patches.
Revised: We show that this reliance on CNNs is not necessary and a pure transformer applied directly to sequences of image patches can perform very well on image classiﬁcation tasks.

Original: When pre-trained on large amounts of data and transferred to multiple  recognition benchmarks (ImageNet, CIFAR-100, VTAB, etc.), Vision Transformer attain excellent results compared to state-of-the-art convolutional networks while requiring substantially fewer computational resources to train.
Revised: When pre-trained on large amounts of data and transferred to multiple mid-sized or small image recognition benchmarks (ImageNet, CIFAR-100, VTAB, etc.), Vision Transformer (ViT) attains excellent results compared to state-of-the-art convolutional networks while requiring substantially fewer computational resources to train.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Self-attention based architectures, in particular Transformers (Vaswani et al., 2017), have become the model of choice in natural language processing (NLP).
Revised: Self-attention-based architectures, in particular Transformers (Vaswani et al., 2017), have become the model of choice in natural language processing (NLP).

Original: The dominant approach is to pre-train on a large text corpus and then ﬁne-tune on a smaller task-speciﬁc dataset (Devlin et al., 2019).
Revised: The dominant approach is to pre-train on a large text corpus and then ﬁne-tune on a smaller task-speciﬁc dataset (Devlin et al., 2019).

Original: Thanks to Transformers’ computational efﬁciency and scalability, it has become possible to train models of unprecedented size, with over 100B parameters.
Revised: Thanks to Transformers’ computational efﬁciency and scalability, it has become possible to train models of unprecedented size, with over 100B parameters (Brown et al., 2020;

Revised: Lepikhin et al., 2020).

Original: With the models and datasets growing, there is still no sign of saturating performance.
Revised: With the models and datasets growing, there is still no sign of saturating performance.

Original: In computer vision, however, convolutional architectures remain dominant (Krizhevsky et al., 2012;
Revised: In computer vision, however, convolutional architectures remain dominant (LeCun et al., 1989; Krizhevsky et al., 2012;

Original: He et al., 2016).
Revised: He et al., 2016).

Original: Inspired by NLP successes, multiple works try combining CNN-like architectures with self-attention (Wang et al., 2018; Carion et al., 2020), some replacing the convolutions entirely (Ramachandran et al., 2019; Wang et al., 2020a).
Revised: Inspired by NLP successes, multiple works try combining CNN-like architectures with self-attention (Wang et al., 2018; Carion et al., 2020), some replacing the convolutions entirely (Ramachandran et al., 2019; Wang et al., 2020a).

Original: The latter models, while theoretically efﬁcient, have not yet been scaled effectively on modern hardware accelerators due to the use of specialized attention patterns.
Revised: The latter models, while theoretically efﬁcient, have not yet been scaled effectively on modern hardware accelerators due to the use of specialized attention patterns.

Original: Therefore, in large-scale image recognition, classic ResNet-like architectures are still state of the art (Mahajan et al., 2018; Xie et al., 2020; Kolesnikov et al., 2020).
Revised: Therefore, in large-scale image recognition, classic ResNetlike architectures are still state of the art (Mahajan et al., 2018; Xie et al., 2020; Kolesnikov et al., 2020).

Original: Inspired by the Transformer scaling successes in NLP, we experiment with applying a standard
Revised: Inspired by the Transformer scaling successes in NLP, we experiment with applying a standard

Original: Transformer directly to images, with the fewest possible modiﬁcations.
Revised: Transformer directly to images, with the fewest possible modiﬁcations.

Original: For this, we split an image into patches and provide the sequence of linear embeddings of these patches as an input to a Transformer.
Revised: To do so, we split an image into patches and provide the sequence of linear embeddings of these patches as an input to a Transformer.

Original: Image patches are treated the same way as tokens (words) in an NLP application.
Revised: Image patches are treated the same way as tokens (words) in an NLP application.

Original: We train the model on image classiﬁcation in supervised fashion.
Revised: We train the model on image classiﬁcation in supervised fashion.

Original: Such models yield modest results when trained on mid-sized datasets such as ImageNet, achieving accuracies of a few percentage points below ResNets of comparable size.
Revised: When trained on mid-sized datasets such as ImageNet, such models yield modest accuracies of a few percentage points below ResNets of comparable size.

Original: This seemingly discouraging outcome may be expected: Transformers lack some  inductive biases inherent to CNNs, such as translation equivariance and locality, and therefore do not generalize well when trained on insufﬁcient amounts of data.
Revised: This seemingly discouraging outcome may be expected: Transformers lack some of the inductive biases inherent to CNNs, such as translation equivariance and locality, and therefore do not generalize well when trained on insufﬁcient amounts of data.

Original: However, the picture changes ifwe train the models  on large datasets (14M-300M images).
Revised: However, the picture changes if the models are trained on larger datasets (14M-300M images).

Original: We ﬁnd that large scale training trumps inductive bias.
Revised: We ﬁnd that large scale training trumps inductive bias.

Original: Transformers attain excellent results when pre-trained at sufﬁcient scale and transferred to tasks with fewer datapoints.
Revised: Our Vision Transformer (ViT) attains excellent results when pre-trained at sufﬁcient scale and transferred to tasks with fewer datapoints.

Original: Our Vision Transformer, pretrained on the  JFT-300M dataset,  approaches or beats state of the art on multiple image recognition benchmarks, reaching accuracy of 88 .
Revised: When pre-trained on the public ImageNet-21k dataset or the in-house JFT-300M dataset, ViT approaches or beats state of the art on multiple image recognition benchmarks. In particular, the best model reaches the accuracy of 88 .

Original: 36% on ImageNet, 90 .
Revised: 55% on ImageNet, 90 .

Original: 77% on ImageNet-ReaL, 94 .
Revised: 72% on ImageNet-ReaL, 94 .

Original: 55% on CIFAR-100, and 77 .
Revised: 55% on CIFAR-100, and 77 .

Original: 16% on the VTAB suite of 19 tasks.
Revised: 63% on the VTAB suite of 19 tasks.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Transformers were proposed by Vaswani et al.
Revised: Transformers were proposed by Vaswani et al.

Original: (2017) for machine translation, and have since become the state of the art method in many NLP tasks.
Revised: (2017) for machine translation, and have since become the state of the art method in many NLP tasks.

Original: Large Transformer-based models are often pre-trained on large corpora and then ﬁne-tuned for the task at hand: BERT (Devlin et al., 2019) uses a denoising self-supervised pre-training task, while the GPT line of work uses language modeling as its pre-training task (Radford et al., 2018; 2019; Brown et al., 2020).
Revised: Large Transformer-based models are often pre-trained on large corpora and then ﬁne-tuned for the task at hand: BERT (Devlin et al., 2019) uses a denoising self-supervised pre-training task, while the GPT line of work uses language modeling as its pre-training task (Radford et al., 2018; 2019; Brown et al., 2020).

Original: Naive application of self-attention to images would require that each pixel attends to every other pixel.
Revised: Naive application of self-attention to images would require that each pixel attends to every other pixel.

Original: With quadratic cost in the number of pixels, this does not scale to realistic input sizes.
Revised: With quadratic cost in the number of pixels, this does not scale to realistic input sizes.

Original: Thus, to apply Transformers in the context of image generation, several approximations have been tried in the past: Parmar et al.
Revised: Thus, to apply Transformers in the context of image processing, several approximations have been tried in the past. Parmar et al.

Original: (2018) applied the self-attention only in local neighborhoods for each query pixel instead of globally.
Revised: (2018) applied the self-attention only in local neighborhoods for each query pixel instead of globally.

Original: Such local multi-head dot-product self attention blocks can completely replace convolutions (Ramachandran et al., 2019; Cordonnier et al., 2020; Zhao et al., 2020).
Revised: Such local multi-head dot-product self attention blocks can completely replace convolutions (Hu et al., 2019; Ramachandran et al., 2019; Zhao et al., 2020).

Original: Alternatively, works such as Sparse Transformers (Child et al., 2019) employ scalable approximations to global self-attention in order to be applicable to images.
Revised: In a different line of work, Sparse Transformers (Child et al., 2019) employ scalable approximations to global selfattention in order to be applicable to images.

Original: An alternative way to scale attention is to apply it in blocks of varying sizes (Weissenborn et al., 2019), in the extreme case only along individual axes (Ho et al., 2019;
Revised: An alternative way to scale attention is to apply it in blocks of varying sizes (Weissenborn et al., 2019), in the extreme case only along individual axes (Ho et al., 2019;

Original: Wang et al., 2020a).
Revised: Wang et al., 2020a).

Original: Many of these specialized attention architectures demonstrate promising results on computer vision tasks, but require complex engineering to be implemented efﬁciently on hardware accelerators.
Revised: Many of these specialized attention architectures demonstrate promising results on computer vision tasks, but require complex engineering to be implemented efﬁciently on hardware accelerators.

Revised: Most related to ours is the model of Cordonnier et al.

Revised: (2020), which extracts patches of size 2 ×from the input image and applies full self-attention on top.

Revised: This model is very similar to ViT, but our work goes further to demonstrate that large scale pre-training makes vanilla transformers competitive with (or even better than) state-of-the-art CNNs.

Revised: Moreover, Cordonnier et al.

Revised: (2020) use a small patch size of 2 × 2 pixels, which makes the model applicable only to small-resolution images, while we handle medium-resolution images as well.

Original: There has also been a lot of interest in combining convolutional neural networks (CNNs) with forms of self-attention, e.g. by augmenting feature maps for image classiﬁcation (Bello et al., 2019) or by further processing the output of a CNN using self-attention, e.g. for object detection (Hu et al., 2018;
Revised: There has also been a lot of interest in combining convolutional neural networks (CNNs) with forms of self-attention, e.g. by augmenting feature maps for image classiﬁcation (Bello et al., 2019) or by further processing the output of a CNN using self-attention, e.g. for object detection (Hu et al., 2018;

Original: Carion et al., 2020), video processing (Wang et al., 2018; Sun et al., 2019), image classiﬁcation (Wu et al., 2020), unsupervised object discovery (Locatello et al., 2020), or uniﬁed text-vision tasks (Chen et al., 2020c; Lu et al., 2019; Li et al., 2019).
Revised: Carion et al., 2020), video processing (Wang et al., 2018; Sun et al., 2019), image classiﬁcation (Wu et al., 2020), unsupervised object discovery (Locatello et al., 2020), or uniﬁed text-vision tasks (Chen et al., 2020c; Lu et al., 2019; Li et al., 2019).

Original: We are not aware of prior application of Transformers with global self-attention to full-sized images.

Original: Closest to our model is iGPT (Chen et al., 2020a), which applies Transformers to image pixels after reducing image resolution and color space.
Revised: Another recent related model is image GPT (iGPT) (Chen et al., 2020a), which applies Transformers to image pixels after reducing image resolution and color space.

Original: The model is trained in an unsupervised fashion as a generative model, and the resulting representation can then be ﬁne-tuned or probed linearly for classiﬁcation performance, achieving a maximal accuracy of 72% on ImageNet.
Revised: The model is trained in an unsupervised fashion as a generative model, and the resulting representation can then be ﬁne-tuned or probed linearly for classiﬁcation performance, achieving a maximal accuracy of 72% on ImageNet.

Original: This work adds to the increasing collection of papers that explore image recognition at larger scales than the standard ImageNet dataset.
Revised: Our work adds to the increasing collection of papers that explore image recognition at larger scales than the standard ImageNet dataset.

Original: To achieve state-of-the-art results, many papers rely on additional data sources  (Mahajan et al., 2018; Touvron et al., 2019; Xie et al., 2020).
Revised: The use of additional data sources allows to achieve state-ofthe-art results on standard benchmarks (Mahajan et al., 2018; Touvron et al., 2019; Xie et al., 2020).

Original:  Sun et al.
Revised: Moreover, Sun et al.

Original: (2017) study how CNN performance scales with dataset size, and Kolesnikov et al.
Revised: (2017) study how CNN performance scales with dataset size, and Kolesnikov et al.

Original: ; Djolonga et al.
Revised: ; Djolonga et al.

Original: (2020) perform an empirical exploration of CNN transfer learning from large scale datasets such as ImageNet-21k and JFT-300M, both of which are also the focus  of this study.
Revised: (2020) perform an empirical exploration of CNN transfer learning from large scale datasets such as ImageNet-21k and JFT-300M. We focus on these two latter datasets as well, but train Transformers instead of ResNet-based models used in prior works.

Original:  M ETHOD
Revised:  M ETHOD

Original: We follow as closely as possible the design of the original Transformer (Vaswani et al., 2017).
Revised: In model design we follow the original Transformer (Vaswani et al., 2017) as closely as possible.

Original: This intentionally simple setup has the advantage that scalable NLP Transformer architectures – and their efﬁcient implementations – can be used almost out of the box.
Revised: An advantage of this intentionally simple setup is that scalable NLP Transformer architectures – and their efﬁcient implementations – can be used almost out of the box.

Original: We aim to show that when scaled appropriately, this approach is sufﬁcient to outperform even the best convolutional neural networks.
Revised: [Figure]

Original: 3. V ISION T RANSFORMER (V I T)
Revised: 3. V ISION T RANSFORMER (V I T)

Original: Our Transformer for images follows the architecture designed for NLP.

Original: Figure 1 depicts the setup.
Revised: An overview of the model is depicted in Figure 1.

Original: The standard Transformer receives as input a 1D sequence of token embeddings.
Revised: The standard Transformer receives as input a 1D sequence of token embeddings.

Original: To handle 2D images, we reshape the image x ∈ R H × W × C into a sequence of ﬂattened 2D patches x p ∈ R N ×
Revised: To handle 2D images, we reshape the image x ∈ R H × W × C into a sequence of ﬂattened 2D patches x p ∈ R N ×

Original: ( P 2 · C ) . ( H, W ) is the resolution of the original image and ( P, P ) is the resolution of each image patch. N = HW/P 2 is then the effective  sequence length for the Transformer.
Revised: ( P 2 · C ) , where ( H, W ) is the resolution of the original image, C is the number of channels, ( P, P ) is the resolution of each image patch, and N = HW/Pis the resulting number of patches, which also serves as the effective input sequence length for the Transformer.

Original: The Transformer uses constant widths through all of its layers, so  a trainable linear projectionmaps each vectorized patch to the model dimension D (Eq. 1), the output of which we refer to as our patch embeddings.
Revised: The Transformer uses constant latent vector size D through all of its layers, so we ﬂatten the patches and map to D dimensions with a trainable linear projection (Eq. 1). We refer to the output of this projection as the patch embeddings.

Original: Similar to BERT’s [class] token, we prepend a learnable embedding to the sequence of embedded patches ( z [-0 
-] = x class ), whose state at the output of the Transformer encoder ( zL 0  ) serves as the[Figure] image representation y
Revised: Similar to BERT’s [class] token, we prepend a learnable embedding to the sequence of embedded patches ( z [-0 
-]00 = x class ), whose state at the output of the Transformer encoder ( z 0 L ) serves as the image representation y

Original: (Eq. 4).
Revised: (Eq. 4).

Original: Both during pre-training and ﬁne-tuning, the classiﬁcation head is attached to z 0 L .
Revised: Both during pre-training and ﬁne-tuning, a classiﬁcation head is attached to z 0 L .

Revised: The classiﬁcation head is implemented by a MLP with one hidden layer at pre-training time and by a single linear layer at ﬁne-tuning time.

Original: Position embeddings are added to the patch embeddings to retain positional information.
Revised: Position embeddings are added to the patch embeddings to retain positional information.

Original: We explore different 2D-aware variants of position embeddings (Appendix C.3) without any signiﬁcant  gains over standard 1D position embeddings.
Revised: We use standard learnable 1D position embeddings, since we have not observed signiﬁcant performance gains from using more advanced 2D-aware position embeddings (Appendix D.3).

Original: The joint embedding  serves as input to the encoder.
Revised: The resulting sequence of embedding vectors serves as input to the encoder.

Original: The Transformer encoder (Vaswani et al., 2017) consists of alternating layers of multiheaded selfattention (MSA, see Appendix A) and MLP blocks (Eq. 2, 3).
Revised: The Transformer encoder (Vaswani et al., 2017) consists of alternating layers of multiheaded selfattention (MSA, see Appendix A) and MLP blocks (Eq. 2, 3).

Original: Layernorm (LN) is applied before every block, and residual connections after every block (Wang et al., 2019; Baevski & Auli, 2019).
Revised: Layernorm (LN) is applied before every block, and residual connections after every block (Wang et al., 2019; Baevski & Auli, 2019).

Original: The MLP contains two layers with a GELU non-linearity.
Revised: The MLP contains two layers with a GELU non-linearity.

Original: [Equation]
Revised: [Equation]

Revised: Inductive bias.

Revised: We note that Vision Transformer has much less image-speciﬁc inductive bias than

Original: Hybrid Architecture.
Revised: CNNs. In CNNs, locality, two-dimensional neighborhood structure, and translation equivariance are baked into each layer throughout the whole model. In ViT, only MLP layers are local and translationally equivariant, while the self-attention layers are global. The two-dimensional neighborhood structure is used very sparingly: in the beginning of the model by cutting the image into patches and at ﬁne-tuning time for adjusting the position embeddings for images of different resolution (as described below).

Revised: Other than that, the position embeddings at initialization time carry no information about the 2D positions of the patches and all spatial relations between the patches have to be learned from scratch.

Original:  As an alternative to dividing the imageinto patches, the input sequence can be formed fromintermediate feature maps of a ResNet (He et al., 2016).
Revised: Hybrid Architecture. As an alternative to raw image patches, the input sequence can be formed from feature maps of a CNN (LeCun et al., 1989).

Original: In this hybrid model, the patch embedding projection E (Eq. 1) is replaced by the early stages of a ResNet.
Revised: In this hybrid model, the patch embedding projection E (Eq. 1) is applied to patches extracted from a CNN feature map.

Original: One of the intermediate 2D feature maps of the ResNet is ﬂattened into a sequence, projected to the Transformer dimension, and then fed as an input sequence to a Transformer.
Revised: As a special case, the patches can have spatial size 1x1, which means that the input sequence is obtained by simply ﬂattening the spatial dimensions of the feature map and projecting to the Transformer dimension.

Original: The classiﬁcation input embedding and position embeddings are added as described above to the input to the Transformer.
Revised: The classiﬁcation input embedding and position embeddings are added as described above.

Original: 3. F INE - TUNING AND H IGHER R ESOLUTION
Revised: 3. F INE - TUNING AND H IGHER R ESOLUTION

Original: Typically, we pre-train ViT on large datasets, and ﬁne-tune to (smaller) downstream tasks.
Revised: Typically, we pre-train ViT on large datasets, and ﬁne-tune to (smaller) downstream tasks.

Original: For this, we remove the pre-trained prediction head and attach a zero-initialized D × K feedforward layer, where K is the number of downstream classes.
Revised: For this, we remove the pre-trained prediction head and attach a zero-initialized D × K feedforward layer, where K is the number of downstream classes.

Original: It is often beneﬁcial to ﬁne-tune at higher resolution than pre-training (Touvron et al., 2019; Kolesnikov et al., 2020).
Revised: It is often beneﬁcial to ﬁne-tune at higher resolution than pre-training (Touvron et al., 2019; Kolesnikov et al., 2020).

Original: When feeding images of higher resolution, we keep the patch size the same, which results in a larger effective sequence length.
Revised: When feeding images of higher resolution, we keep the patch size the same, which results in a larger effective sequence length.

Original: The Vision Transformer can handle arbitrary sequence lengths (up to memory constraints), however, the pre-trained position embeddings may no longer be meaningful.
Revised: The Vision Transformer can handle arbitrary sequence lengths (up to memory constraints), however, the pre-trained position embeddings may no longer be meaningful.

Original: We therefore perform 2D interpolation of the pre-trained position embeddings, according to their location in the original image.
Revised: We therefore perform 2D interpolation of the pre-trained position embeddings, according to their location in the original image.

Original: Note that this resolution adjustment and patch extraction are the only points at which an inductive bias about the 2D structure of the images is manually injected into the Vision Transformer.[Table]
Revised: Note that this resolution adjustment and patch extraction are the only points at which an inductive bias about the 2D structure of the images is manually injected into the Vision Transformer.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We evaluate the representation learning capabilities of ResNet, Vision Transformer (ViT), and the hybrid.
Revised: We evaluate the representation learning capabilities of ResNet, Vision Transformer (ViT), and the hybrid.

Original: To understand the data requirements of each model, we pre-train on datasets of varying size and evaluate many benchmark tasks.
Revised: To understand the data requirements of each model, we pre-train on datasets of varying size and evaluate many benchmark tasks.

Original: When considering the computational cost of pre-training the model, ViT performs very favourably, attaining state-of-the-art on most recognition benchmarks at a lower pre-training cost.
Revised: When considering the computational cost of pre-training the model, ViT performs very favourably, attaining state of the art on most recognition benchmarks at a lower pre-training cost.

Original: Lastly, we perform a small experiment using self-supervision, and show that self-supervised ViT holds promising for the future.
Revised: Lastly, we perform a small experiment using self-supervision, and show that self-supervised ViT holds promise for the future.

Original: 4. S ETUP
Revised: 4. S ETUP

Original: Datasets.
Revised: Datasets.

Original: To explore model scalability, we use  ILSVRC-2012 ImageNet dataset with 1k classes and 1.3M images (we refer to it as ImageNet in what follows), its superset ImageNet-21k with 21k classes and 14M images (Deng et al., 2009), and JFT (Sun et al., 2017) with 18k classes and 303Mmillion high-resolution images.
Revised: To explore model scalability, we use the ILSVRC-2012 ImageNet dataset with 1k classes and 1.3M images (we refer to it as ImageNet in what follows), its superset ImageNet-21k with 21k classes and 14M images (Deng et al., 2009), and JFT (Sun et al., 2017) with 18k classes and 303M high-resolution images.

Original: We de-duplicate the pre-training datasets  following Kolesnikov et al.
Revised: We de-duplicate the pre-training datasets w.r.t. the test sets of the downstream tasks following Kolesnikov et al.

Original: We transfer the models trained on these dataset to many benchmark tasks: ImageNet on the original validation labels, and the cleaned-up ReaL labels (Beyer et al., 2020), CIFAR10/100 (Krizhevsky, 2009), Oxford-IIIT Pets (Parkhi et al., 2012), and Oxford Flowers-102 (Nilsback & Zisserman, 2008).
Revised: We transfer the models trained on these dataset to several benchmark tasks: ImageNet on the original validation labels and the cleaned-up ReaL labels (Beyer et al., 2020), CIFAR-10/100 (Krizhevsky, 2009), Oxford-IIIT Pets (Parkhi et al., 2012), and Oxford Flowers-102 (Nilsback & Zisserman, 2008).

Original: For these datasets, pre-processing follows Kolesnikov et al.
Revised: For these datasets, pre-processing follows Kolesnikov et al.

Original: We also evaluate on the 19-task VTAB classiﬁcation suite (Zhai et al., 2019b).
Revised: We also evaluate on the 19-task VTAB classiﬁcation suite (Zhai et al., 2019b).

Original: VTAB evaluates low-data transfer  using 1 000  examples to diverse tasks.
Revised: VTAB evaluates low-data transfer to diverse tasks, using 1 000 training examples per task.

Original: The tasks are divided into three groups: Natural – tasks like the above, Pets, CIFAR, etc.
Revised: The tasks are divided into three groups: Natural – tasks like the above, Pets, CIFAR, etc.

Original: Specialized – medical and satellite imagery, and Structured – tasks that require geometric understanding like localization.
Revised: Specialized – medical and satellite imagery, and Structured – tasks that require geometric understanding like localization.

Original: Model Variants.
Revised: Model Variants.

Original: We base ViT conﬁgurations on those used for BERT (Devlin et al., 2019).
Revised: We base ViT conﬁgurations on those used for BERT (Devlin et al., 2019), as summarized in Table 1.

Revised: The “Base” and “Large” models are directly adopted from BERT and we add the larger “Huge” model.

Original: We denote our models ViT-L/16 to mean the “Large” variant(in this case), with 16 × 16 input patch size.
Revised: In what follows we use brief notation to indicate the model size and the input patch size: for instance, ViT-L/16 means the “Large” variant with 16 × 16 input patch size.

Original: Note: the Transformer’s sequence length is image res 2 / patch size 2 .
Revised: Note that the Transformer’s sequence length is inversely proportional to the square of the patch size, thus models with smaller patch size are computationally more expensive.

Original: Sufﬁxes include “B” (Base), “L” (Large), and “H” (Huge); see Table 1.
Revised: [Table] [Figure] [Table]

Original: For the baseline CNNs, we use ResNet, but replace the Batch Norm layers  with Group Norm, and used standardized convolutions.
Revised: For the baseline CNNs, we use ResNet (He et al., 2016), but replace the Batch Normalization layers (Ioffe & Szegedy, 2015) with Group Normalization (Wu & He, 2018), and used standardized convolutions (Qiao et al., 2019).

Original: These modiﬁcations improves transfer (Kolesnikov et al., 2020), and we denote the modiﬁed model ResNet (BiT).
Revised: These modiﬁcations improve transfer (Kolesnikov et al., 2020), and we denote the modiﬁed model “ResNet (BiT)”.

Original: For the hybrids, we feed the intermediate feature maps into ViT with patch size of one “pixel”.
Revised: For the hybrids, we feed the intermediate feature maps into ViT with patch size of one “pixel”.

Original: To compare different sequence lengths we try (i) taking the output of stage 4 of a regular R50, and (ii) removing stage 4, and placing the same number of layers in stage 3 (so still 50 in total), and taking the output of  stage 3.
Revised: To experiment with different sequence lengths, we either (i) take the output of stage 4 of a regular ResNet50 or (ii) remove stage 4, place the same number of layers in stage 3 (keeping the total number of layers), and take the output of this extended stage 3.

Original: Option (ii) results in a  longer sequence length, and a more expensive ViT head.
Revised: Option (ii) results in a 4x longer sequence length, and a more expensive ViT model.

Original: Training & Fine-tuning.
Revised: Training & Fine-tuning.

Original: We train all models, including ResNets, using Adam  with β 1 = 0 .
Revised: We train all models, including ResNets, using Adam (Kingma & Ba, 2015) with β 1 = 0 .

Original: 9 , β 2 = 0 .
Revised: 9 , β 2 = 0 .

Original: 999 , a batch size of 4096 and apply a high weight decay of 0 . 1 , which we found to be useful for transfer of all models (Appendix C.1 shows that, in contrast to common practices, Adam works slightly better than SGD for ResNets in our setting).
Revised: 999 , a batch size of 4096 and apply a high weight decay of 0 . 1 , which we found to be useful for transfer of all models (Appendix D.1 shows that, in contrast to common practices, Adam works slightly better than SGD for ResNets in our setting).

Original: We use a linear learning rate warmup and decay, see Appendix B.1 for details.
Revised: We use a linear learning rate warmup and decay, see Appendix B.1 for details.

Original: For ﬁne-tuning we use SGD with momentum, batch size 512, for all models, see Appendix B.1.1.
Revised: For ﬁne-tuning we use SGD with momentum, batch size 512, for all models, see Appendix B.1.1.

Original: For ImageNet results in Table 2, we ﬁne-tuned at higher resolution: 512 for ViT-L/16 and 518 for ViT-H/14, andthe latter also uses Polyak & Juditsky (1992) averaging with a factor of 0 .
Revised: For ImageNet results in Table 2, we ﬁne-tuned at higher resolution: 512 for ViT-L/16 and 518 for ViT-H/14, and also used Polyak & Juditsky (1992) averaging with a factor of 0 .

Original:  (Ramachandran et al., 2019; Wang et al., 2020b).
Revised: 9999 (Ramachandran et al., 2019; Wang et al., 2020b).

Original: Metrics.
Revised: Metrics.

Original: We report results on downstream datasets either through few-shot or ﬁne-tuning accuracy.
Revised: We report results on downstream datasets either through few-shot or ﬁne-tuning accuracy.

Original: Fine-tuning accuracies capture the performance of each model after ﬁne-tuning it on the respective dataset.
Revised: Fine-tuning accuracies capture the performance of each model after ﬁne-tuning it on the respective dataset.

Original: Few-shot accuracies are obtained by solving a regularized linear regression problem that maps the (frozen) representation of a subset of training images to {− 1 , 1 } K target vectors.
Revised: Few-shot accuracies are obtained by solving a regularized least-squares regression problem that maps the (frozen) representation of a subset of training images to {− 1 , 1 } K target vectors.

Revised: This formulation allows us to recover the exact solution in closed form.

Original: Though we mainly focus on ﬁne-tuning performance, we sometimes use linear few-shot accuracies for fast on-the-ﬂy evaluation where ﬁne-tuning would be too costly.[Table] often
Revised: Though we mainly focus on ﬁne-tuning performance, we sometimes use linear few-shot accuracies for fast on-the-ﬂy evaluation where ﬁne-tuning would be too costly.

Original: 4. C OMPARISON TO S TATE OF THE A RT
Revised: 4. C OMPARISON TO S TATE OF THE A RT

Original: We ﬁrst compare our largest models – ViT-H/14 and ViT-L/16pre-trained on JFT-300M – to stateof-the-art CNNs from the literature.
Revised: We ﬁrst compare our largest models – ViT-H/14 and ViT-L/16 – to state-of-the-art CNNs from the literature.

Original: The ﬁrst comparison point is Big Transfer (BiT) (Kolesnikov et al., 2020), which performs supervised transfer learning with large ResNets.
Revised: The ﬁrst comparison point is Big Transfer (BiT) (Kolesnikov et al., 2020), which performs supervised transfer learning with large ResNets.

Original: The second is Noisy Student (Xie et al., 2020), which is a large EfﬁcientNet trained using semi-supervised learning on ImageNet and JFT-300M with the labels removed.
Revised: The second is Noisy Student (Xie et al., 2020), which is a large EfﬁcientNet trained using semi-supervised learning on ImageNet and JFT300M with the labels removed.

Original: Currently, Noisy Student is the state of the art on ImageNet and BiT-L on the other datasets reported here.
Revised: Currently, Noisy Student is the state of the art on ImageNet and BiT-L on the other datasets reported here.

Original: All models we’re trained on TPUv hardware, and we report the number of TPUv3-days taken to pre-train each.
Revised: All models were trained on TPUv3 hardware, and we report the number of TPUv3-core-days taken to pre-train each of them, that is, the number of TPU v3 cores (2 per chip) used for training multiplied by the training time in days.

Original: Table 2 contains the results.
Revised: Table 2 shows the results.

Original: The smaller ViT-L/16 model matches or outperforms BiT-L  on all datasets, while requiring substantially less computational resources to train.
Revised: The smaller ViT-L/16 model pre-trained on JFT-300M outperforms BiT-L (which is pre-trained on the same dataset) on all tasks, while requiring substantially less computational resources to train.

Original: The larger model, ViTH-14, further improves the performance, especially on the more challenging datasets – ImageNet and CIFAR-100, and the VTAB suite. 
Revised: The larger model, ViT-H/14, further improves the performance, especially on the more challenging datasets – ImageNet, CIFAR-100, and the VTAB suite. Interestingly, this

Original: It matches or exceeds state of the art on all of the datasets, in some cases by a substantial margin (e.g. 1% on CIFAR-100).
Revised: [Table] result reported in Touvron et al.

Original: On ImageNet, ViT is roughly 0 .
Revised: [Figure] While large ViT models perform worse than BiT

Original: 1% below Noisy Student with the standard noisy labels, but beats state of the art when evaluated on  the cleaner ReaL labels.
Revised: ResNets (shaded area) when pre-trained on small datasets, they shine when pre-trained on larger datasets. Similarly, larger ViT variants overtake smaller ones as the dataset grows.

Original: Interestingly, our models took substantially less compute to pre-train than prior state of the art, however, we note that pre-training efﬁciency may be affected not only by the architecture choice, but also other parameters, such as training schedule, optimizer, weight decay, etc.
Revised: [Figure] model still took substantially less compute to pre-train than prior state of the art. However, we note that pre-training efﬁciency may be affected not only by the architecture choice, but also other parameters, such as training schedule, optimizer, weight decay, etc.

Original: We provide a controlled study of performance vs. compute for different architectures in Section 4.4.
Revised: We provide a controlled study of performance vs. compute for different architectures in Section 4.4.

Revised: Finally, the ViT-L/16 model pre-trained on the public ImageNet-21k dataset performs well on most datasets too, while taking fewer resources to pre-train: it could be trained using a standard cloud TPUv3 with 8 cores in approximately 30 days.

Original: Figure 2 decomposes the VTAB tasks into their respective groups, and compares to previous SOTA methods on this benchmark: BiT, VIVI – a ResNet co-trained on ImageNet and Youtube (Tschannen et al., 2020), and S4L – supervised plus semi-supervised learning on ImageNet (Zhai et al., 2019a).
Revised: the VTAB tasks into their respective groups, and compares to previous SOTA methods on this benchmark: BiT, VIVI – a ResNet co-trained on ImageNet and Youtube (Tschannen et al., 2020), and S4L – supervised plus semi-supervised learning on ImageNet (Zhai et al., 2019a).

Original: On the Natural tasks, ViT-H/14 is slightly outperformed by BiT-R152x4, although the difference is within repetition noise.
Revised: ViT-H/14 outperforms BiT-R152x4, and other methods, on the Natural and Structured tasks.

Original: OnSpecialized ViT just outperforms BiT (and other methods), but the largest beneﬁt appears to be theStructured task group where ViT is signiﬁcantly superior. [Figure] [Figure]
Revised: On the

Original: While large ViT models perform worse than BiT ResNets (shaded area) when pre-trained on small datasets, they shine when pre-trained on larger datasets. Similarly, larger ViT variants overtake smaller ones as the dataset grows. [Figure]
Revised: Specialized the performance of the top two models is similar.

Original: 4. P RE - TRAINING D ATA R EQUIREMENTS
Revised: 4. P RE - TRAINING D ATA R EQUIREMENTS

Original: The Vision Transformer performs well when pre-trained on a large JFT-300M dataset.
Revised: The Vision Transformer performs well when pre-trained on a large JFT-300M dataset.

Original: With fewer inductive biases for vision than ResNets, how crucial is the dataset size?
Revised: With fewer inductive biases for vision than ResNets, how crucial is the dataset size?

Original: We perform two series of experiments.
Revised: We perform two series of experiments.

Original:  First, we pre-train ViT models on datasets of increasing size: ImageNet, ImageNet-21k, and JFT300M. To get the best possible performance on the smaller datasets, we optimize three regularization parameters – weight decay, dropout, and label smoothing.
Revised: [Figure] First, we pre-train ViT models on datasets of increasing size: ImageNet, ImageNet-21k, and JFT300M. To get the best possible performance on the smaller datasets, we optimize three regularization parameters – weight decay, dropout, and label smoothing.

Original: Figure 3 shows the results after ﬁne-tuning to ImageNet (other datasets  in Table 3).
Revised: Figure 3 shows the results after ﬁnetuning to ImageNet (results on other datasets are shown in Table 5) 2 .

Original: 1 Figure 3 contains the results of ﬁne-tuning on ImageNet.
Revised: Figure 3 contains the results of ﬁne-tuning on ImageNet.

Original: When pre-trained on the smallest dataset, ImageNet, ViT-Large models underperform compared to ViT-Base models, despite heavy regularization.
Revised: When pre-trained on the smallest dataset, ImageNet, ViT-Large models underperform compared to ViT-Base models, despite heavy regularization, and in turn, ViT-Huge do not outperform ViT-Large.

Original: However, with ImageNet-21k pre-training, their performances are similar.
Revised: However, with ImageNet-21k pre-training, their performances are similar.

Original: Only with JFT-300M, do we see the  beneﬁt of larger models.
Revised: Only with JFT-300M, do we see the full beneﬁt of larger models.

Original: Figure 3 also shows the performance region spanned by BiT models of different sizes.
Revised: Figure 3 also shows the performance region spanned by BiT models of different sizes.

Original: The

Original:  BiT CNNs outperform ViT on ImageNet (despite regularization optimization), but with the larger datasets, ViT overtakes.
Revised: The BiT CNNs outperform ViT on ImageNet (despite regularization optimization), but with the larger datasets, ViT overtakes.

Original: Second, we train our models on random subsets of 9M, 30M, and 90M as well as the full JFT300M dataset.
Revised: Second, we train our models on random subsets of 9M, 30M, and 90M as well as the full JFT300M dataset.

Original: We do not perform additional regularization on the smaller subsets and use the same hyper-parameters for all settings.
Revised: We do not perform additional regularization on the smaller subsets and use the same hyper-parameters for all settings.

Original: This way, we assess the intrinsic model properties, and not the effect of regularization.
Revised: This way, we assess the intrinsic model properties, and not the effect of regularization.

Original: We do, however, use early-stopping, and report the best validation accuracy achieved during training.
Revised: We do, however, use early-stopping, and report the best validation accuracy achieved during training.

Original: To save compute, we report few-shot linear accuracy instead of full ﬁnetuning accuracy.
Revised: To save compute, we report few-shot linear accuracy instead of full ﬁnetuning accuracy.

Original: Figure 4 contains the results.
Revised: Figure 4 contains the results.

Original: Vision Transformers overﬁt more than ResNets with comparable computational cost on smaller datasets.
Revised: Vision Transformers overﬁt more than ResNets with comparable computational cost on smaller datasets.

Original: For example, ViT-B/32 is slightly faster than
Revised: For example, ViT-B/32 is slightly faster than

Original: ResNet50; it performs much worse on the 9M subset, but better on 90M+ subsets.
Revised: ResNet50; it performs much worse on the 9M subset, but better on 90M+ subsets.

Original: The same is true for ResNet152x2 and ViT-L/16.
Revised: The same is true for ResNet152x2 and ViT-L/16.

Original: This result reinforces the intuition that the convolutional inductive bias is useful for smaller datasets, but for larger ones, learning the relevant patterns is sufﬁcient, even beneﬁcial.
Revised: This result reinforces the intuition that the convolutional inductive bias is useful for smaller datasets, but for larger ones, learning the relevant patterns is sufﬁcient, even beneﬁcial.

Original: Overall, the few-shot results on ImageNet (Figure 4), as well as the low-data results on VTAB (Table 2) seem promising for very low-data transfer.
Revised: Overall, the few-shot results on ImageNet (Figure 4), as well as the low-data results on VTAB (Table 2) seem promising for very low-data transfer.

Original: Further analysis of few-shot properties of ViT is an exciting direction of future work.
Revised: Further analysis of few-shot properties of ViT is an exciting direction of future work.

Original: 4. S CALING S TUDY
Revised: 4. S CALING S TUDY

Original: We perform a controlled scaling study of different models.
Revised: We perform a controlled scaling study of different models by evaluating transfer performance from

Original: For this, we evaluate transfer performance from JFT-300M. On JFT-300M, data size does not bottleneck the models’ performances, and we assess performance versus pre-training cost of each model.
Revised: JFT-300M. In this setting data size does not bottleneck the models’ performances, and we assess performance versus pre-training cost of each model.

Original: The model set includes: 5 ResNets, 
Revised: The model set includes: 7 ResNets, R50x1,

Original: {+R50x2 R101x1, R152x1, R152x2, pre-trained for 7 epochs, plus R152x
+}
Revised: {+R50x2 R101x1, R152x1, R152x2, pre-trained for 7 epochs, plus R152x
+}

Revised: and R200x3 pre-trained for 14 epochs; 6 Vision Transformers, ViT-B/32, B/16, L/32, L/16, pre-trained for 7 epochs, plus

Revised: L/16 and H/14 pre-trained for 14 epochs; and 5 hybrids, R50+ViT-B/32, B/16, L/32, L/16 pretrained for 7 epochs, plus R50+ViT-L/16 pre-trained for 14 epochs (for hybrids, the number at the

Original: [Figure]
Revised: [Figure]

Original: R50x1, R50x2 R101x1, R152x1, R152x2, pre-trained for 7 epochs; 6 Vision Transformers, ViTB/32, B/16, L/32, L/16, pre-trained for 7 epochs, plus L/16 H/14 pre-trained for 14 epochs; andhybrids, R50 + ViT-B/32, B/16, L/32, L/16 pre-trained for 7 epochs.
Revised: I n p u t p a t c h r o w Position embedding similarity

Original: Figure 5 contains the transfer performance versus total pre-training compute (see Appendix C.4 for details on computational costs computation).
Revised: C o s i n e s i m il a r i t y [Figure] Network depth ( [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: M e a n a tt e n t i o n d

Revised: i s t a n c e ( p i x e l s )

Revised: ViT

Revised: [Figure]

Revised: Figure 5 contains the transfer performance versus total pre-training compute (see Appendix D. for details on computational costs).

Revised: Detailed results per model are provided in Table 6 in the Appendix.

Original: A few patterns can be observed.
Revised: A few patterns can be observed.

Original: First, Vision Transformer dominate ResNets on the performance/compute trade-off.
Revised: First, Vision Transformers dominate ResNets on the performance/compute trade-off.

Original: ViT uses approximately 2  × less compute to attain the same performance (average  5 datasets).
Revised: ViT uses approximately 2 − 4 × less compute to attain the same performance (average over 5 datasets).

Original: Second, hybrids slightly outperform ViT at small computational budgets, but the difference vanishes for larger ones.
Revised: Second, hybrids slightly outperform ViT at small computational budgets but the difference vanishes for larger models.

Original: This result was surprising, since one might expect convolutional local feature processing to assistthe ViT at any size.
Revised: This result is somewhat surprising, since one might expect convolutional local feature processing to assist ViT at any size.

Original: Third, Vision Transformers appear not to saturate within the range tried, motivating future scalability efforts.
Revised: Third, Vision Transformers appear not to saturate within the range tried, motivating future scaling efforts

Original: 4. I NSPECTING V ISION T RANSFORMER
Revised: 4. I NSPECTING V ISION T RANSFORMER

Original: [Figure]
Revised: [Figure]

Original: To begin to understand how the Vision Transformer processes image data, we analyze its internal representations.
Revised: To begin to understand how the Vision Transformer processes image data, we analyze its internal representations.

Original: The ﬁrst layer of the Vision Transformer linearly projects the ﬂattened patches into a lower-dimensional space (Eq. 1).
Revised: The ﬁrst layer of the Vision Transformer linearly projects the ﬂattened patches into a lower-dimensional space (Eq. 1).

Original: Figure 7 (left) shows the top principal components of the the learned embedding ﬁlters.
Revised: Figure 7 (left) shows the top principal components of the the learned embedding ﬁlters.

Original: The components resemble plausible basis functions for a low-dimensional representation of the ﬁne structure within each patch.
Revised: The components resemble plausible basis functions for a low-dimensional representation of the ﬁne structure within each patch.

Original: After the projection, a learned position embedding is added to the patch representations.
Revised: After the projection, a learned position embedding is added to the patch representations.

Original: Figure 7 (center) shows that the model learns to encode distance within the image in the similarity of position embeddings, i.e. closer patches tend to have more similar position embeddings.
Revised: Figure 7 (center) shows that the model learns to encode distance within the image in the similarity of position embeddings, i.e. closer patches tend to have more similar position embeddings.

Original: Further, the row-column structure appears; patches in the same row/column have similar embeddings.
Revised: Further, the row-column structure appears; patches in the same row/column have similar embeddings.

Original: Finally, a sinusoidal structure is sometimes apparent for larger grids (Appendix C).
Revised: Finally, a sinusoidal structure is sometimes apparent for larger grids (Appendix D).

Original: That the position embeddings learn to represent 2D image topology explains why hand-crafted 2D-aware embedding variants do not yield improvements (Appendix C.3).
Revised: That the position embeddings learn to represent 2D image topology explains why hand-crafted 2D-aware embedding variants do not yield improvements (Appendix D.3).

Original: Self-attention allows ViT to integrate information across the entire image even in the lowest layers.
Revised: Self-attention allows ViT to integrate information across the entire image even in the lowest layers.

Original: We investigate to what degree the network makes use of this capability.
Revised: We investigate to what degree the network makes use of this capability.

Original: Speciﬁcally, we compute the average distance in image space across which information is integrated, based on the attention weights (Figure 7, right).
Revised: Speciﬁcally, we compute the average distance in image space across which information is integrated, based on the attention weights (Figure 7, right).

Original: This “attention distance” is analogous to receptive ﬁeld size in CNNs.[Figure] I n p u t p a t c h r o w
Revised: This “attention distance” is analogous to receptive ﬁeld size in CNNs.

Original: Position embedding similarity

Original: C o s i n e s i m il a r i t y

Original: 10 15 20 Network depth (layer)

Original: M e a n a tt e n t i o n d

Original: i s t a n c e ( p i x e l s ) ViT-L/ [Figure] We ﬁnd that some heads attend to most of the image already in the lowest layers, showing that the ability to integrate information globally is indeed used by the model.
Revised: We ﬁnd that some heads attend to most of the image already in the lowest layers, showing that the ability to integrate information globally is indeed used by the model.

Original: Other attention heads have consistently small attention distances in the low layers.
Revised: Other attention heads have consistently small attention distances in the low layers.

Original: This highly localized attention is less pronounced in hybrid models that apply a ResNet before the Transformer (Figure 7, right), suggesting that it may serve a similar function as early convolutional layers in CNNs.
Revised: This highly localized attention is less pronounced in hybrid models that apply a ResNet before the Transformer (Figure 7, right), suggesting that it may serve a similar function as early convolutional layers in CNNs.

Original: Further, the attention distance increases with network depth.
Revised: Further, the attention distance increases with network depth.

Original: Globally, we ﬁnd that the model attends to image regions that are semantically relevant for classiﬁcation (Figure 6).
Revised: Globally, we ﬁnd that the model attends to image regions that are semantically relevant for classiﬁcation (Figure 6).

Original: 4. S ELF - SUPERVISION
Revised: 4. S ELF - SUPERVISION

Original: Transformers show impressive performance on NLP tasks.
Revised: Transformers show impressive performance on NLP tasks.

Original: However, much of their success stems not only from their excellent scalability but also from large scale self-supervised pre-training (Devlin et al., 2019;
Revised: However, much of their success stems not only from their excellent scalability but also from large scale self-supervised pre-training (Devlin et al., 2019;

Original: Radford et al., 2018).
Revised: Radford et al., 2018).

Original: We also perform a preliminary exploration on masked patch prediction for self-supervision, mimicking the masked language modeling task used in BERT.
Revised: We also perform a preliminary exploration on masked patch prediction for self-supervision, mimicking the masked language modeling task used in BERT.

Original: With self-supervised pre-training, our smaller ViT-B/16 model achieves 79.9% accuracy on ImageNet, a signiﬁcant improvement of 2% to training from scratch, but still 4% behind supervised pre-training.
Revised: With self-supervised pre-training, our smaller ViT-B/16 model achieves 79.9% accuracy on ImageNet, a signiﬁcant improvement of 2% to training from scratch, but still 4% behind supervised pre-training.

Original: Appendix B.1.2 contains further details.
Revised: Appendix B.1.2 contains further details.

Original: We leave exploration of contrastive pre-training (Chen et al., 2020b; He et al., 2020; Bachman et al., 2019; H ´ enaff et al., 2020) to future work.
Revised: We leave exploration of contrastive pre-training (Chen et al., 2020b; He et al., 2020; Bachman et al., 2019; H´enaff et al., 2020) to future work.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We have explored the direct application of Transformers to image recognition.
Revised: We have explored the direct application of Transformers to image recognition.

Original: Unlike prior works using self-attention in computer vision, we do not introduceany image-speciﬁc inductive biases into the architecture.
Revised: Unlike prior works using self-attention in computer vision, we do not introduce image-speciﬁc inductive biases into the architecture apart from the initial patch extraction step.

Original: Instead, we interpret an image as a sequence of patches and process it by a standard Transformer encoder as used in NLP.
Revised: Instead, we interpret an image as a sequence of patches and process it by a standard Transformer encoder as used in NLP.

Original: This simple, yet scalable, strategy works surprisingly well when coupled with pre-training on large datasets.
Revised: This simple, yet scalable, strategy works surprisingly well when coupled with pre-training on large datasets.

Original: Thus, Vision Transformer matches or exceeds the state of the art on many image classiﬁcation datasets, whilst being relatively cheap to pre-train.
Revised: Thus, Vision Transformer matches or exceeds the state of the art on many image classiﬁcation datasets, whilst being relatively cheap to pre-train.

Original: While these initial results are encouraging, many challenges remain.
Revised: While these initial results are encouraging, many challenges remain.

Original: One is to apply ViT to other computer vision tasks, such as detection and segmentation.
Revised: One is to apply ViT to other computer vision tasks, such as detection and segmentation.

Original: Our results, coupled with those in Carion et al.
Revised: Our results, coupled with those in Carion et al.

Original: (2020), indicate the promise of this approach.
Revised: (2020), indicate the promise of this approach.

Original: Another – to continue exploring self-supervised pre-training methods.
Revised: Another challenge is to continue exploring selfsupervised pre-training methods.

Original: Our initial experiments show improvement from self-supervised pre-training, but there is still large gap between self-supervised and large-scale supervised pre-training.
Revised: Our initial experiments show improvement from self-supervised pre-training, but there is still large gap between self-supervised and large-scale supervised pretraining.

Original: Finally – to further scale ViT, given that the performance does not seem yet to be saturating with the increased model size.
Revised: Finally, further scaling of ViT would likely lead to improved performance

Original: A PPENDIX A M ULTIHEAD S ELF - ATTENTION
Revised: A PPENDIX A M ULTIHEAD S ELF - ATTENTION

Original: Standard qkv self-attention (SA, Vaswani et al.
Revised: Standard qkv self-attention (SA, Vaswani et al.

Original: (2017)) is a popular building block for neural architectures.
Revised: (2017)) is a popular building block for neural architectures.

Original: For each element in an input sequence z ∈ R N × D , we compute a weighted sum over all values v in the sequence.
Revised: For each element in an input sequence z ∈ R N × D , we compute a weighted sum over all values v in the sequence.

Original: The attention weights A ij are based on the pairwise similarity between two elements of the sequence and their respective query q i and key k j representations.
Revised: The attention weights A ij are based on the pairwise similarity between two elements of the sequence and their respective query q i and key k j representations.

Original: [Equation]
Revised: [Equation]

Original: Multihead self-attention (MSA) is an extension of SA in which we run k self-attention operations, called “heads”, in parallel, and project their concatenated outputs.
Revised: Multihead self-attention (MSA) is an extension of SA in which we run k self-attention operations, called “heads”, in parallel, and project their concatenated outputs.

Original: To keep compute and number of parameters constant when changing k , D h (Eq. 5) is typically set to D/k .
Revised: To keep compute and number of parameters constant when changing k , D h (Eq. 5) is typically set to D/k .

Original: [Equation]
Revised: [Equation]

Original: B E XPERIMENT DETAILS 
Revised: B E XPERIMENT DETAILS B. T RAINING

Original: Pre-training Data

Original: Dataset

Original: ViT-B/

Original: ViT-B/

Original: ViT-L/

Original: ViT-L/

Original: ViT-H/

Original: ImageNet

Original: CIFAR98.

Original: CIFAR87.

Original: ImageNet 77.

Original: ImageNet ReaL 83.

Original: Oxford Flowers89.

Original: Oxford-IIIT-Pets 93.

Original: [Table]

Original: B. T RAINING

Original: Table 4 summarizes our training setups for our different models.
Revised: Table 3 summarizes our training setups for our different models.

Original: We found strong regularization to be key when training models from scratch on ImageNet.
Revised: We found strong regularization to be key when training models from scratch on ImageNet.

Original: Dropout, when used, is applied after every dense layer except for the the qkv-projections and directly after adding positional- to patch embeddings.
Revised: Dropout, when used, is applied after every dense layer except for the the qkv-projections and directly after adding positional- to patch embeddings.

Original: Hybrid models are trained with the exact setup as their ViT counterparts.
Revised: Hybrid models are trained with the exact setup as their ViT counterparts.

Original: Finally, all training is done on resolution 224.
Revised: Finally, all training is done on resolution 224.

Original: B.1. F INE - TUNING
Revised: B.1. F INE - TUNING

Original: For VTAB we follow the protocol in Zhai et al.

Original: (2019b); for each task, we select the best of two learning rates and schedules (Tab. 5) on the 200-example validation set.

Original: We then retrain on the entire 1kexample dataset using the best hyperparameters.

Original: We follow the pre-processing used in Kolesnikov

Original: [Table]

Original: [Table]

Original: et al.

Original: (2020), except that we do not use per-task resolutions.

Original: Instead we ﬁnd that Vision Transformer beneﬁts most from using a high resolution ( 384 × 384 ) for all tasks.

Original: We ﬁne-tune all ViT models using SGD with a momentum of 0.9.
Revised: We ﬁne-tune all ViT models using SGD with a momentum of 0.9.

Original: We run a small grid search over learning rates, see learning rate ranges in Table 5.
Revised: We run a small grid search over learning rates, see learning rate ranges in Table 4.

Original: To do so, we use small sub-splits from the training set (10% for Pets and Flowers, 2% for CIFAR, 1% ImageNet) as development set and train on the remaining data.
Revised: To do so, we use small sub-splits from the training set (10% for Pets and Flowers, 2% for CIFAR, 1% ImageNet) as development set and train on the remaining data.

Original: For ﬁnal results we train on the entire training set and evaluate on the respective test data.
Revised: For ﬁnal results we train on the entire training set and evaluate on the respective test data.

Original: For ﬁne-tuning ResNets and hybrid models we use the exact same setup, with the only exception of ImageNet where we add another value 0 .
Revised: For ﬁne-tuning ResNets and hybrid models we use the exact same setup, with the only exception of ImageNet where we add another value 0 .

Original: 06 to the learning rate sweep.
Revised: 06 to the learning rate sweep.

Original: Additionally,  for ResNets we also run the setup of Kolesnikov et al.
Revised: Additionally, [Table] for ResNets we also run the setup of Kolesnikov et al.

Original: and select the best results across this run and our sweep.
Revised: and select the best results across this run and our sweep.

Original: Finally, if not mentioned otherwise, all ﬁne-tuning experiments run atresolution (running ﬁne-tuning at different resolution than training is common practice (Kolesnikov et al., 2020)).
Revised: Finally, if not mentioned otherwise, all ﬁne-tuning experiments run atresolution (running ﬁne-tuning at different resolution than training is common practice (Kolesnikov et al., 2020)).

Original: When transferring ViT models to another dataset, we remove the whole head (two linear layers) and replace it by a single, zero-initialized linear layer outputting the number of classes required by the target dataset.
Revised: When transferring ViT models to another dataset, we remove the whole head (two linear layers) and replace it by a single, zero-initialized linear layer outputting the number of classes required by the target dataset.

Original: We found this to be a little more robust than simply re-initializing the very last layer.
Revised: We found this to be a little more robust than simply re-initializing the very last layer.

Revised: For VTAB we follow the protocol in Kolesnikov et al.

Revised: (2020), and use the same hyperparameter setting for all tasks.

Original: B.1. S ELF - SUPERVISION
Revised: We use a learning rate of 0 . 01 and train for 2500 steps (Tab. 4). We chose this setting by running a small sweep over two learning rates and two schedules, and selecting the setting with the highest VTAB score on the 200-example validation sets. We follow the pre-processing used in Kolesnikov et al. (2020), except that we do not use task-speciﬁc input resolutions.

Revised: Instead we ﬁnd that Vision Transformer beneﬁts most from a high resolution ( 384 × 384 ) for all tasks.

Revised: B.1. S ELF - SUPERVISION

Original: We employ the masked patch prediction objective for preliminary self-supervision experiments.
Revised: We employ the masked patch prediction objective for preliminary self-supervision experiments.

Original: To do so we corrupt 50% of patch embeddings by either replacing their embeddings with a learnable [mask] embedding (80%), a random other patch embedding (10%) or just keeping them as is (10%).
Revised: To do so we corrupt 50% of patch embeddings by either replacing their embeddings with a learnable [mask] embedding (80%), a random other patch embedding (10%) or just keeping them as is (10%).

Original: This setup is very similar to the one used for language by Devlin et al.
Revised: This setup is very similar to the one used for language by Devlin et al.

Original: Finally, we predict the 3-bit, mean color (i.e., 512 colors in total) of every corrupted patch using their respective patch representations.
Revised: Finally, we predict the 3-bit, mean color (i.e., 512 colors in total) of every corrupted patch using their respective patch representations.

Original: We trained our self-supervised model for 1M steps (ca.
Revised: We trained our self-supervised model for 1M steps (ca.

Original: 14 epochs) with batch size 4096 on JFT.
Revised: 14 epochs) with batch size 4096 on JFT.

Original: We use Adam, with a base learning rate of 2 · 10 − 4 , warmup of 10k steps and cosine learning rate decay.
Revised: We use Adam, with a base learning rate of 2 · 10 − 4 , warmup of 10k steps and cosine learning rate decay.

Original: As prediction targets for pretraining we tried the following settings: 1) predicting only the mean, 3bit color (i.e., 1 prediction of 512 colors), 2) predicting a 4 × 4 downsized version of the 16 × [Table] patch with 3bit colors in parallel (i.e., 16 predictions of 512 colors), 3) regression on the full patch using L2 (i.e., 256 regressions on the 3 RGB channels).
Revised: As prediction targets for pretraining we tried the following settings: 1) predicting only the mean, 3bit color (i.e., 1 prediction of 512 colors), 2) predicting a 4 × 4 downsized version of the 16 ×patch with 3bit colors in parallel (i.e., 16 predictions of 512 colors), 3) regression on the full patch using L2 (i.e., 256 regressions on the 3 RGB channels).

Original: Surprisingly, we found that all worked quite well, though L2 was slightly worse.
Revised: Surprisingly, we found that all worked quite well, though L2 was slightly worse.

Original: We report ﬁnal results only for option 1) because it has shown best few-shot performance.
Revised: We report ﬁnal results only for option 1) because it has shown best few-shot performance.

Original: We also experimented with 15% corruption rate as used by Devlin et al.
Revised: We also experimented with 15% corruption rate as used by Devlin et al.

Original: but results were also slightly worse on our few-shot metrics.
Revised: but results were also slightly worse on our few-shot metrics.

Revised: Lastly, we would like to remark that our instantiation of masked patch prediction doesn’t require such an enormous amount of pretraining nor a large dataset such as JFT in order to lead to similar performance gains on ImageNet classiﬁcation.

Revised: That is, we observed diminishing returns on downstream performance after 100k pretraining steps, and see similar gains when pretraining on

Revised: ImageNet.

Original: C A DDITIONALA NALYSES C. SGD VS . A DAM FOR R ES N ETS
Revised: C A DDITIONAL R ESULTS

Revised: We report detailed results corresponding to the ﬁgures presented in the paper. Table 5 corresponds to Figure 3 from the paper and shows transfer performance of different ViT models pre-trained on datasets of increasing size: ImageNet, ImageNet-21k, and JFT-300M. Table 6 corresponds to ViT-B/

Revised: ViT-B/

Revised: ViT-L/

Revised: ViT-L/

Revised: ViT-H/

Revised: ImageNet

Revised: CIFAR98.

Revised: CIFAR87.

Revised: ImageNet 77.

Revised: ImageNet ReaL 83.

Revised: Oxford Flowers89.

Revised: Oxford-IIIT-Pets 93.

Revised: ImageNet-21k

Revised: CIFAR98.

Revised: CIFAR91.

Revised: ImageNet 83.

Revised: ImageNet ReaL 88.

Revised: Oxford Flowers99.

Revised: Oxford-IIIT-Pets 94.

Revised: JFT-300M

Revised: CIFAR99.

Revised: CIFAR91.

Revised: ImageNet 84.

Revised: ImageNet ReaL 88.

Revised: Oxford Flowers99.

Revised: Oxford-IIIT-Pets 95.

Revised: on various datasets when pre-trained on ImageNet, ImageNet-21k or JFT300M. These values correspond to Figure 3 in the main text.

Revised: Models are ﬁne-tuned at 384 resolution.

Revised: Note that the ImageNet results are computed without additional techniques (Polyak averaging and 512 resolution images) used to achieve results in Table 2.

Revised: [Table]

Revised: Figure 5 from the paper and shows the transfer performance of ViT, ResNet, and hybrid models of varying size, as well as the estimated computational cost of their pre-training.

Revised: D A DDITIONAL A NALYSES D. SGD VS . A DAM FOR R ES N ETS

Original: ResNets are typically trained with SGD and our use of Adam as optimizer is quite unconventional.
Revised: ResNets are typically trained with SGD and our use of Adam as optimizer is quite unconventional.

Original: Here we show the experiments that motivated this choice.
Revised: Here we show the experiments that motivated this choice.

Original: Namely, we compare the ﬁne-tuning performance of two ResNets – 50x1 and 152x2 – pre-trained on JFT with SGD and Adam.
Revised: Namely, we compare the ﬁne-tuning performance of two ResNets – 50x1 and 152x2 – pre-trained on JFT with SGD and Adam.

Original: For
Revised: For

Original: SGD, we use the hyperparameters recommended by Kolesnikov et al.
Revised: SGD, we use the hyperparameters recommended by Kolesnikov et al.

Original: Results are presented in Table 6.
Revised: Results are presented [Table] [Figure] [Table]

Revised: [Figure]

Original:  Adam pre-training outperforms SGD pre-training on most datasets and on average.
Revised: [Table] [Figure] in Table 7. Adam pre-training outperforms SGD pre-training on most datasets and on average.

Original: This justiﬁes the choice of Adam as the optimizer used to pre-train ResNets on JFT.
Revised: This justiﬁes the choice of Adam as the optimizer used to pre-train ResNets on JFT.

Original: Note that the absolute numbers are lower than those reported by Kolesnikov et al.
Revised: Note that the absolute numbers are lower than those reported by Kolesnikov et al.

Original: (2020), since we pre-train only for 7 epochs, not 30 .
Revised: (2020), since we pre-train only for 7 epochs, not 30 .

Original: C. T RANSFORMER SHAPE
Revised: D. T RANSFORMER SHAPE

Original: We ran ablations on scaling different dimensions of the Transformer architecture to ﬁnd out which are best suited for scaling to very large models.
Revised: We ran ablations on scaling different dimensions of the Transformer architecture to ﬁnd out which are best suited for scaling to very large models.

Original: Figure 8 shows 5-shot performance on ImageNet for different conﬁgurations.
Revised: Figure 8 shows 5-shot performance on ImageNet for different conﬁgurations.

Original: All conﬁgurations are based on a ViT model with 8 layers, D = 1024 ,
Revised: All conﬁgurations are based on a ViT model with 8 layers, D = 1024 ,

Original: D MLP = 2048 and a patch size of 32 , the intersection of all lines.
Revised: D MLP = 2048 and a patch size of 32 , the intersection of all lines.

Original: We can see that scaling the depth results in the biggest improvements which are clearly visible up until 64 layers.
Revised: We can see that scaling the depth results in the biggest improvements which are clearly visible up until 64 layers.

Original: However, diminishing returns are already visible after 16 layers.
Revised: However, diminishing returns are already visible after 16 layers.

Original: Interestingly, scaling the width of the network seems to result in the smallest changes.
Revised: Interestingly, scaling the width of the network seems to result in the smallest changes.

Original: Decreasing the patch size and thus increasing the effective sequence length shows surprisingly robust improvements without introducing parameters.
Revised: Decreasing the patch size and thus increasing the effective sequence length shows surprisingly robust improvements without introducing parameters.

Original: These ﬁndings suggest that compute might be a better predictor of performance than the number of parameters, and that scaling should emphasize depth over width if any.
Revised: These ﬁndings suggest that compute might be a better predictor of performance than the number of parameters, and that scaling should emphasize depth over width if any.

Original: Overall, we ﬁnd that scaling all dimensions proportionally results in robust improvements.[Figure]
Revised: Overall, we ﬁnd that scaling all dimensions proportionally results in robust improvements.

Original: C. P OSITIONAL E MBEDDING
Revised: D. P OSITIONAL E MBEDDING

Original: We ran ablations on different ways of encoding spatial information using positional embedding.
Revised: We ran ablations on different ways of encoding spatial information using positional embedding.

Original: We tried the following cases:
Revised: We tried the following cases:

Original: • Providing no positional information: Considering the inputs as a bag of patches .
Revised: • Providing no positional information: Considering the inputs as a bag of patches .

Original: • 1-dimensional positional embedding:
Revised: • 1-dimensional positional embedding:

Original: Considering the inputs as a sequence of patches in the raster order (default across all other experiments in this paper).
Revised: Considering the inputs as a sequence of patches in the raster order (default across all other experiments in this paper).

Original: • 2-dimensional positional embedding:
Revised: • 2-dimensional positional embedding:

Original: Considering the inputs as a grid of patches in two dimensions.
Revised: Considering the inputs as a grid of patches in two dimensions.

Original: In this case, two sets of embeddings are learned, each for one of the axes, X -embedding, and Y -embedding, each with size D/ 2 . 
Revised: In this case, two sets of embeddings are learned, each for one of the axes, X -embedding, and Y -embedding, each with size D/ 2 . Then, based on the coordinate on [Figure] Emb

Original: Then, based on the coordinate on the path in the input, we concatenate the X and Y embedding to get the ﬁnal positional embedding for that patch.
Revised: [Figure] Every Layer [Figure] Pos

Revised: [Figure]

Revised: Emb

Revised: [Figure]

Revised: N/A

Revised: N/A

Revised: [Figure]

Revised: Pos

Revised: [Figure]

Revised: Emb

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: D Pos.

Revised: Emb

Revised: [Figure]

Revised: Rel.

Revised: Pos.

Revised: Emb.

Revised: N/A

Revised: N/A

Revised: [Figure]

Revised: the path in the input, we concatenate the X and Y embedding to get the ﬁnal positional embedding for that patch

Original: • Relative positional embeddings: Considering the relative distance between patches to encode the spatial information as instead of their absolute position.
Revised: • Relative positional embeddings: Considering the relative distance between patches to encode the spatial information as instead of their absolute position.

Original: To do so, we use 1dimensional Relative Attention, in which we deﬁne the relative distance all possible pairs of patches.
Revised: To do so, we use 1dimensional Relative Attention, in which we deﬁne the relative distance all possible pairs of patches.

Original: Thus, for every given pair (one as query, and the other as key/value in the attention mechanism), we have an offset p q − p k , where each offset is associated with an embedding.
Revised: Thus, for every given pair (one as query, and the other as key/value in the attention mechanism), we have an offset p q − p k , where each offset is associated with an embedding.

Original: Then, we simply run extra attention, where we use the original query (the content of query), but use relative positional embeddings as keys.
Revised: Then, we simply run extra attention, where we use the original query (the content of query), but use relative positional embeddings as keys.

Original: We then use the logits from the relative attention as a bias term and add it to the logits of the main attention (content-based attention) before applying the softmax.
Revised: We then use the logits from the relative attention as a bias term and add it to the logits of the main attention (content-based attention) before applying the softmax.

Original: In addition to different ways of encoding spatial information, we also tried different ways of incorporating this information in our model.
Revised: In addition to different ways of encoding spatial information, we also tried different ways of incorporating this information in our model.

Original: For the 1-dimensional and 2-dimensional positional embeddings, we tried three different cases: (1) add positional embeddings to the inputs right after the stem of them model and before feeding the inputs to the Transformer encoder (default across all other experiments in this paper); (2) learn and add positional embeddings to the inputs at the beginning of each layer; (3) add a learned positional embeddings to the inputs at the beginning of each layer (shared between layers).
Revised: For the 1-dimensional and 2-dimensional positional embeddings, we tried three different cases: (1) add positional embeddings to the inputs right after the stem of them model and before feeding the inputs to the Transformer encoder (default across all other experiments in this paper); (2) learn and add positional embeddings to the inputs at the beginning of each layer; (3) add a learned positional embeddings to the inputs at the beginning of each layer (shared between layers).

Original: Table 7 summarizes the results from this ablation study on a ViT-B/16 model.
Revised: Table 8 summarizes the results from this ablation study on a ViT-B/16 model.

Original: As we can see, while there is a large gap between the performances of the model with no positional embedding and models with positional embedding, there is little to no difference between different ways of encoding positional information.
Revised: As we can see, while there is a large gap between the performances of the model with no positional embedding and models with positional embedding, there is little to no difference between different ways of encoding positional information.

Original: We speculate that since our Transformer encoder operates on patch-level inputs, as opposed to pixel-level, the differences in how to encode spatial information is less important.
Revised: We speculate that since our Transformer encoder operates on patch-level inputs, as opposed to pixel-level, the differences in how to encode spatial information is less important.

Original: More precisely, in patch-level inputs, the spatial dimensions are much smaller than the original pixel-level inputs, e.g., 14 × 14 as opposed to 224 × 224 , and learning to represent the spatial relations in this resolution is equally easy for these different positional encoding strategies.
Revised: More precisely, in patch-level inputs, the spatial dimensions are much smaller than the original pixel-level inputs, e.g., 14 × 14 as opposed to 224 × 224 , and learning to represent the spatial relations in this resolution is equally easy for these different positional encoding strategies.

Original: Even so, the speciﬁc pattern of position embedding similarity learned by the network depends on the training hyperparameters (Figure 9).Pos. Emb. Default/Stem
Revised: Even so, the speciﬁc pattern of position embedding similarity learned by the network depends on the training hyperparameters (Figure 9).

Original: Every Layer

Original: Every Layer-Shared

Original: No Pos.

Original: Emb.

Original: N/A

Original: N/A

Original: [Figure]

Original: D Pos.

Original: Emb.

Original: [Figure]
Revised: [Figure]

Original: D Pos.

Original: Emb.

Original: Rel.

Original: Pos.

Original: Emb.

Original: N/A

Original: N/A

Original: [Figure]

Original: [Figure]

Original:  E MPIRICAL C OMPUTATIONAL C OSTS
Revised: D. E MPIRICAL C OMPUTATIONAL C OSTS

Original: We are also interested in real-world speed of the architectures on our hardware, which is not always well predicted by theoretical FLOPs due to details like lane widths and cache sizes.
Revised: We are also interested in real-world speed of the architectures on our hardware, which is not always well predicted by theoretical FLOPs due to details like lane widths and cache sizes.

Original: For this purpose, we perform timing of inference speed for the main models of interest, on a TPUv3 hardware; the difference between inference and backprop speed is a constant model-independent factor.
Revised: For this purpose, we perform timing of inference speed for the main models of interest, on a TPUv3 accelerator; the difference between inference and backprop speed is a constant model-independent factor.

Original: Figure 11 (left) shows how many images one core can handle per second, across various input sizes.
Revised: Figure 11 (left) shows how many images one core can handle per second, across various input sizes.

Original: Every single point refers to the peak performance measured across a wide range of batch-sizes.
Revised: Every single point refers to the peak performance measured across a wide range of batch-sizes.

Original: As can be seen, the theoretical bi-quadratic scaling of ViT with image size only barely starts happening for the largest models at the largest resolutions.
Revised: As can be seen, the theoretical bi-quadratic scaling of ViT with image size only barely starts happening for the largest models at the largest resolutions.

Original: Another quantity of interest is the largest batch-size each model can ﬁt onto a core, larger being better for scaling to large datasets.
Revised: Another quantity of interest is the largest batch-size each model can ﬁt onto a core, larger being better for scaling to large datasets.

Original: Figure 11 (right) shows this quantity for the same set of models.
Revised: Figure 11 (right) shows this quantity for the same set of models.

Original: This shows that large ViT models have a clear advantage in terms of memory-efﬁciency over ResNet models.
Revised: This shows that large ViT models have a clear advantage in terms of memory-efﬁciency over ResNet models.

Original: C. A XIAL A TTENTION
Revised: D. A XIAL A TTENTION

Original: Axial Attention (Huang et al., 2020; Ho et al., 2019) is a simple, yet effective technique to run selfattention on large inputs that are organized as multidimensional tensors.
Revised: Axial Attention (Huang et al., 2020; Ho et al., 2019) is a simple, yet effective technique to run selfattention on large inputs that are organized as multidimensional tensors.

Original: The general idea of axial attention is to perform multiple attention operations, each along a single axis of the input tensor, instead of applying 1-dimensional attention to the ﬂattened version of the input.
Revised: The general idea of axial attention is to perform multiple attention operations, each along a single axis of the input tensor, instead of applying 1-dimensional attention to the ﬂattened version of the input.

Original: In axial attention, each attention mixes information along a particular axis, while keeping information along the other axes independent.
Revised: In axial attention, each attention mixes information along a particular axis, while keeping information along the other axes independent.

Original: Along this line, Wang et al.
Revised: Along this line, Wang et al.

Original: (2020b) proposed the AxialResNet model in which all the convolutions with kernel size 3 × 3 in a ResNet50 are replaced by axial self-attention, i.e.[Figure] a row and column attention, augmented by relative positional encoding.
Revised: (2020b) proposed the AxialResNet model in which all the convolutions with kernel size 3 × 3 in a ResNet50 are replaced by axial self-attention, i.e. a row and column attention, augmented by relative positional encoding.

Original: We have implemented
Revised: We have implemented

Original: AxialResNet as a baseline model.
Revised: AxialResNet as a baseline model.

Original: Moreover

Original: [Figure] we have modiﬁed ViT to process inputs in the 2-dimensional shape, instead of a 1dimensional sequence of patches [Figure] and incorporate Axial Transformer blocks
Revised: Moreover, we have modiﬁed ViT to process inputs in the 2-dimensional shape, instead of a 1dimensional sequence of patches, and incorporate Axial Transformer blocks, in which instead of

Original: [Figure]
Revised: [Figure]

Original: in which instead of a self-attention followed by an MLP, we have a a row-self-attention plus an MLP followed by a column-self-attention plus an MLP.
Revised: a self-attention followed by an MLP, we have a a row-self-attention plus an MLP followed by a column-self-attention plus an MLP.

Original: [Figure]
Revised: [Figure]

Original:  Figure 12, present the performance of Axial ResNet, Axial-ViT-B/32 and Axial-ViT-B/16 on ImageNet 5shot linear, when pretrained on JFT dataset, verses the pretraining compute, both in terms of number of FLOPs and inference time (example per seconds).
Revised: terms of top-1 accuracy on ImageNet 5-shot linear, versus their speed in terms of number of FLOPs ( left ) and inference time ( left ). Figure 12, present the performance of Axial ResNet, Axial-ViT-B/32 and Axial-ViT-B/16 on ImageNet 5shot linear, when pretrained on JFT dataset, verses the pretraining compute, both in terms of number of FLOPs and inference time (example per seconds).

Original: As we can see, both Axial-ViT-B/ and Axial-ViT-B/16 do better than their ViT-B counterpart in terms of performance, but it comes at the cost of more compute.
Revised: As we can see, both Axial-ViT-B/ and Axial-ViT-B/16 do better than their ViT-B counterpart in terms of performance, but it comes at the cost of more compute.

Original: This is because in Axial-ViT models, each Transformer block with global self-attention is replaced by two Axial Transformer blocks, one with row and one with column self- attention and although the sequence length that self-attention operates on is smaller in axial case, there is a extra MLP per Axial-ViT block.
Revised: This is because in Axial-ViT models, each Transformer block with global self-attention is replaced by two Axial Transformer blocks, one with row and one with column selfattention and although the sequence length that self-attention operates on is smaller in axial case, there is a extra MLP per Axial-ViT block.

Original: For the AxialResNet, although it looks reasonable in terms of accuracy/compute trade-off (Figure 12, left), the naive implementation is extremely slow on TPUs (Figure 12, right).
Revised: For the AxialResNet, although it looks reasonable in terms of accuracy/compute trade-off (Figure 12, left), the naive implementation is extremely slow on TPUs (Figure 12, right).

Original: C. A TTENTION D ISTANCE
Revised: D. A TTENTION D ISTANCE

Original: To understand how ViT uses self-attention to integrate information across the image, we analyzed the average distance spanned by attention weights at different layers (Figure 10).
Revised: To understand how ViT uses self-attention to integrate information across the image, we analyzed the average distance spanned by attention weights at different layers (Figure 10).

Original: This “attention distance” is analogous to receptive ﬁeld size in CNNs.
Revised: This “attention distance” is analogous to receptive ﬁeld size in CNNs.

Original: Average attention distance is highly variable across heads in lower layers, with some heads attending to much of the image, while others attend to small regions at or near the query location.
Revised: Average attention distance is highly variable across heads in lower layers, with some heads attending to much of the image, while others attend to small regions at or near the query location.

Original: As depth increases, attention distance increases for all heads.
Revised: As depth increases, attention distance increases for all heads.

Original: In the second half of the network, most heads attend widely across tokens.
Revised: In the second half of the network, most heads attend widely across tokens.

Original: C. A TTENTION M APS
Revised: D. A TTENTION M APS

Original: To compute maps of the attention from the output token to the input space (Figures 6 and 13), we used Attention Rollout (Abnar & Zuidema, 2020).
Revised: To compute maps of the attention from the output token to the input space (Figures 6 and 13), we used Attention Rollout (Abnar & Zuidema, 2020).

Original: Brieﬂy, we averaged attention weights of ViTL/16 across all heads and then recursively multiplied the weight matrices of all layers.
Revised: Brieﬂy, we averaged attention weights of ViTL/16 across all heads and then recursively multiplied the weight matrices of all layers.

Original: This accounts for the mixing of attention across tokens through all layers.
Revised: This accounts for the mixing of attention across tokens through all layers.

Revised: D. VTAB B REAKDOWN

Revised: Table 9 shows the scores attained on each of the VTAB-1k tasks.

Revised: [Table]

Revised: o

Revised: [Table]

Revised: r i

Revised: [Table]

Revised: I

Revised: [Table]

Revised: i s

Revised: [Table]