D ISCRETE R EPRESENTATIONS S TRENGTHEN V ISION T RANSFORMER R OBUSTNESS
Anonymous authors
Paper under double-blind review
A BSTRACT
Vision Transformer ViT) is emerging as the state- of the art architecture for im - age recognition.
While recent studies suggest that ViTs are more robust than their convolutional counterparts , our experiments ﬁnd that ViTs trained on ImageNet are overly reliant on local features e.g.
, nuisances and texture
) and fail to make adequate use of global context
(e.g. , shape and structure) ViTs thus have difﬁculties generalizing to out -of-distribution, real-world data
To address this deﬁciency, we present a simple and effective architecture modiﬁcation to ViT’s input layer by adding discrete tokens produced by a vector-quantized encoder.
Different from the standard continuous pixel tokens, discrete tokens are invariant under small perturbations and contain less information individually, which promote ViTs to learn global information that is invariant.
Experimental results demonstrate that adding discrete representation on four architecture variants strengthens ViT robustness by up to 12% across seven ImageNet robustness benchmarks while maintaining the performance on ImageNet.
I NTRODUCTION
Despite their high performance on in-distribution test sets, deep neural networks fail to generalize under real-world distribution shifts (Barbu et al., 2019).
This gap between training and inference poses many challenges for deploying deep learning models in real-world applications where closedworld assumptions are violated.
This lack of robustness can be ascribed to learned representations that are overly sensitive to minor variations in local texture and insufﬁciently adept at representing more robust scene and object characteristics, such as shape and texture.
Vision Transformer (ViT) (Dosovitskiy et al., 2020) has started to rival Convolutional Neural Networks (CNNs) in many computer vision tasks.
Recent works found that ViTs are more robust than
CNNs (Paul & Chen, 2021; Mao et al., 2021b; Bhojanapalli et al., 2021) and generalize favorably on a variety of visual robustness benchmarks (Hendrycks et al., 2021b; Hendrycks & Dietterich, 2019).
These work suggested that ViTs’ robustness comes from the self-attention architecture that captures a globally-contextualized inductive bias than CNNs.
However, though self-attention is well capable of modeling the global context, recent studies found that ImageNet-trained ViT may not take this advantage.
For example, Chen et al. (2021b)
compared training ViT on ImageNet with and without using position embedding, and showed that spatial structure contributes less than ∼ 3% of ViT’s performance gain.
Without using position embedding, ViT treats the image as orderless bag-of-patches, but achieves similar performance as ViT with position embedding which is designed to model global information. This suggests that ViT unintentionally relies on local detail (e.g., texture) than the shape and structure of the object.
We hypothesize that this deﬁciency on robustness comes from the high-dimensional, individually informative, linear tokenization, that allows ViT to minimize empirical risk without learning too much spatial structure.
In this paper, we propose a simple yet novel input layer for vision transformers, where image patches are represented by discrete tokens .
To be speciﬁc, we discretize an image and represent an image patch as a discrete token or “visual word” in a codebook.
Our key insight is that discrete tokens capture important features in a low-dimensional space (Oord et al., 2017) preserving shape and structure of the object (see Figure 2).
Our approach capitalizes on this discrete representation to promote the robustness of ViT. Using discrete tokens drives ViT towards better modeling of spatial interactions between tokens, given that individual tokens no longer carry enough information to depend on.
We also concatenate a low dimensional pixel token to the discrete token to compensate for the potentially missed local details encoded by discrete tokens, especially for small objects.
Our approach only changes the image patch tokenizer to improve generalization and robustness, which is orthogonal to all existing approaches for robustness, and can be integrated into architectures that extend vision transformer.
We call the ViT model using our Discrete representation the Dr. ViT.
Our experiments and visualizations show that incorporating discrete tokens in ViT signiﬁcantly improves generalization accuracy for all seven out-of-distribution ImageNet benchmarks: ImageNetRendition by up to 10%, Stylized-ImageNet by up to 12%, ImageNet-Sketch by up to 10%, and
ImageNet-C by up to 10%.
Our method establishes the new state-of-the-art by on four benchmarks without any dataset-speciﬁc data augmentation.
Our work is the ﬁrst to connect discrete representations to robustness and demonstrate consistent robustness improvements.
We will release the code and models to support comparisons and future research.
R ELATED W ORK
Vision Transformer (ViT) (Dosovitskiy et al., 2020), inspired by the Transformer (Vaswani et al., 2017) in NLP, is the ﬁrst CNN-free architecture that achieves state-of-the-art image classiﬁcation accuracy.
Since its inception, numerous works have proposed improvements to the ViT architecture (Wang et al., 2021; Chu et al., 2021; Liu et al., 2021; d’Ascoli et al., 2021), objective (Chen et al., 2021a), training strategy (Touvron et al., 2021), etc..
Given the difﬁculty to study all existing ViT models , this paper focuses on the classical ViT model (Dosovitskiy et al., 2020) and its recent published versions.
Speciﬁcally, Steiner et al.
(2021) proposed ViT-AugReg that applies stronger data augmentation and regularization to the ViT model.
Tolstikhin et al.
introduced MLPMixer to replace self-attention in ViT with multi-layer perceptions (MLP).
We select the above ViT model family (Dosovitskiy et al., 2020; Steiner et al., 2021;
Tolstikhin et al., 2021) in our robustness study for three reasons.
First, they represent both the very ﬁrst and one-of-the-best vision transformers in the literature.
Second, these models demonstrated competitive performance when pre-trained on sufﬁciently large datasets such as ImageNet-21K and JFT-300M.
Finally, unlike other Transformer models, they provide architectures consisting of solely Transformer layers as well as a hybrid of CNN and Transformer layers.
These properties improve our understanding of robustness for different types of network layers and datasets.
Robustness.
Recent works established multiple content robustness datasets to evaluate the out-ofdistribution generalization of deep models (Barbu et al., 2019; Hendrycks et al., 2021b;a;
Wang et al., 2019; Geirhos et al., 2019; Hendrycks & Dietterich, 2019; Recht et al., 2019).
In this paper, we consider 7 ImageNet robustness benchmarks of real-world test images (or proxies) where deep models trained on ImageNet are shown to suffer from notable performance drop.
Existing works on robustness are targeted at closing the gap in a subset of these ImageNet robustness benchmarks and were extensively veriﬁed with the CNNs.
Among them, carefully-designed data augmentations (Hendrycks et al., 2021a; Cubuk et al., 2018; Steiner et al., 2021; Mao et al., 2021b;a), model regularization (Wang et al., 2019; Huang et al., 2020b; Hendrycks et al., 2019), and multitask learning (Zamir et al., 2020) are effective to address the issue.
More recently, a few studies (Paul & Chen, 2021; Bhojanapalli et al., 2021; Naseer et al., 2021;
Shao et al., 2021) suggest that ViTs are more robust than CNNs.
Existing works mainly focused on analyzing the cause of superior generalizability in the ViT model and exploring new data augmentations.
As our work focuses on discrete token input, we train the models using the same data augmentation as the ViT-AugReg baseline (Steiner et al., 2021).
While tailoring data augmentation (Mao et al., 2021b) may further improve our results, we leave it out of the scope of this paper.
Discrete Representation was used as a visual representation prior to the deep learning revolution, such as in bag-of-visual-words model (Sivic & Zisserman, 2003; Csurka et al., 2004) and VLAD model (Arandjelovic & Zisserman, 2013).
Recently, (Oord et al., 2017; Vahdat et al., 2018) proposed neural discrete representation to encode an image as integer tokens.
Recent works used discrete representation mainly for image synthesis (Ramesh et al., 2021; Esser et al., 2021).
To the best of our knowledge, our work is the ﬁrst to demonstrate discrete representations strengthening robustness.
The closest work to ours is BEiT (Bao et al., 2021) that pretrains the ViTs to predict the masked tokens.
However, the tokens are discarded after pretraining, where the ViT model can still overﬁt the non-robust nuisances in the pixel tokens at later ﬁnetuning stage, undermining its robustness.
[Figure]
M ETHOD 3. P RELIMINARY ON V ISION T RANSFORMER
Vision Transformer (Dosovitskiy et al., 2020) is a pure transformer architecture that operates on a sequence of image patches.
The 2D image x ∈ R H × W × C is ﬂattened into a sequence of image patches, following the raster scan, denoted by x p ∈ R L × ( P 2 · C ) ,
where L = H × W Pis the effective sequence length and P 2 × C is the dimension of image patch.
A learnable classiﬁcation token x class is prepended to the patch sequence, then the position embedding E pos is added to formulate the ﬁnal input embedding h 0 .
[Equation]
The architecture of ViT follows that of the Transformer (Vaswani et al., 2017), which alternates layers of multi-headed self-attention (MSA) and multi-layer perceptron (MLP) with LayerNorm (LN) and residual connections being applied to every block. We denote the number of blocks as L f .
This paper considers the ViT model family consisting of 4 ViT backbones: the vanilla ViT discussed above, ViT-AugReg (Steiner et al., 2021) which shares the same ViT architecture but applies stronger data augmentation and regularization, MLP-Mixer (Tolstikhin et al., 2021) which replaces self-attention in ViT with MLP, a variant called Hybrid-ViT which replaces the raw image patches in Equation 1 with the CNN features extracted by a ResNet-
(He et al., 2016).
3. A RCHITECTURE
Existing ViTs represent an image patch as a sequence of pixel tokens , which are linear projections of ﬂattened image pixels.
We propose a novel architecture modiﬁcation to the input layer of the vision transformer, where an image patch x p is represented by a combination of two embeddings.
As illustrated in Fig.
1, in addition to the original pixel-wise linear projection, we discretize an image patch into an discrete token in a codebook V ∈ R K × d c , where K is the codebook size and d c is the dimension of the embedding.
The discretization is achieved by a vector quantized (VQ) encoder p θ that produces an integer z for an image patch x as:
[Equation]
where z e ( x ) denotes the output of the encoder network and 1 ( · ) is the indicator function.
The encoder is applied to the patch sequence x p ∈ R L × ( P 2 · C ) to obtain an integer sequence z d ∈ { 1 , 2 , ..., K } L .
Afterward, we use the embeddings of both discrete and pixel tokens to construct the input embedding to the ViT model.
Speciﬁcally, the input embedding in Equation 1 is replaced by:
[Equation]
where f is the function, embodied as a neural network layer, to combine the two embeddings.
We empirically compared four network designs for f and found that the simplest concatenation works best.
Note that our model only modiﬁes the input layer of ViT (Equation 1) and leaves intact the remaining layers depicted in Equation 2-4. [Figure]
Comparison of pixel and discrete embeddings: Pixel and discrete embeddings represent different aspects of the input image.
Discrete embeddings capture important features in a low-dimension space (Oord et al., 2017) that preserves the global structure of an object but lose local details.
Fig.compares the original image (top) and the reconstructed images decoded from the discrete embeddings of our model (bottom).
As shown, the decoded images from discrete embeddings reasonably depict the object shape and global context.
Due to the quantization, the decoder hallucinates the local textures, e.g., in the cat’s eye, or the text in the “vulture” and “ﬂamingo” images.
It is worth noting that the VQ encoder/decoder is only trained on ImageNet 2012 but they can generalize to out-of-distribution images.
Please see more examples in Appendix A.
On the ﬂip side, pixel embeddings capture rich details through the linear projection from raw pixels.
However, given the expressive power of transformers, ViTs can spend capacity on local textures or nuance patterns that are often circumferential to robust recognition.
Since humans recognize images primarily relying on the shape and semantic structure, this discrepancy to human perception undermines ViT’s generalization on out-of-distribution data.
Our proposed model leverages the power of both embeddings to promote the interaction between modeling global and local features.
3. T RAINING P ROCEDURE
The training comprises two stages: pretraining and ﬁnetuning.
First, we pretrain the VQ-VAE encoder and decoder (Oord et al., 2017) on the given training set.
We do not use labels in this step.
In the ﬁnetuning stage, as shown in Fig.
3a, we train the proposed ViT model from scratch and ﬁnetune the discrete embeddings.
Due to the straight-through gradient estimation in VQ-VAE (Oord et al., 2017), we stop the back-propagation after the discrete embedding.
[Figure]
Now we discuss the objective that the proposed ViT model optimizes.
Let q φ ( z | x ) denote the VQ encoder, parameterized by network weights φ , to represent an image x into a sequence of integer tokens z .
The decoder models the distribution p θ ( x | z ) over the RGB image generated from discrete tokens.
p ψ ( y | x, z ) stands for the proposed vision transformer shown in Fig.
We factorize the joint distribution of image x , label y , and the discrete token z by p ( x, y, z ) =
p φ ( x | z ) p
ψ
( y | x, z ) p ( z ) .
Our overall training procedure is to maximize the evidence lower bound (ELBO) on the joint likelihood:
[Equation]
In the ﬁrst stage, we maximize the ELBO with respect to φ and θ , which corresponds to learning the VQ-VAE encoder and decoder.
Following (Oord et al., 2017), we assume a uniform prior for both p ψ ( y | x, z ) and p ( z ) .
Given that q φ ( z | x ) predicts a one-hot output, the regularization term (KL divergence in Equation 7) will be a constant. Note that the DALL-E model (Ramesh et al., 2021) uses a similar assumption to stabilize the training. Our observation is consistent with that of Ramesh et al. (2021), where this strategy works better than jointly training with p ψ ( y | x, z ) .
Our implementation uses VQ-GAN (Esser et al., 2021) which adds a GAN loss for better decoding quality.
We also include results with VQ-VAE (Oord et al., 2017) in Appendix A.8 for reference.
In the ﬁnetuning stage, we optimize ψ while holding θ and φ ﬁxed, which corresponds to learning a
ViT and ﬁnetuning the discrete embeddings.
This can be seen by rearranging the ELBO:
[Equation]
where p ψ ( y | x, z ) is the proposed ViT model.
The ﬁrst term denotes the likelihood for classiﬁcation that is learned by minimizing the multi-class cross-entropy loss.
The second term becomes a constant given the ﬁxed θ and φ .
It is important to note that the discrete embeddings are learned end-to-end in both pretraining and ﬁnetuning. Details are discussed in Appendix A.2.
E XPERIMENTS 4. E XPERIMENTAL S ETUP
Datasets.
In the experiments, we train all the models, including ours, on ImageNet 2012 or
ImageNet-21K under the same training settings, where we use identical training data, batch size, and learning rate schedule, etc.
Afterward, the trained models are tested on the ImageNet robustness benchmarks to assess their robustness and generalization capability.
In total, we evaluate the models on nine benchmarks.
ImageNet and ImageNet-Real are two indistribution datasets.
ImageNet (Deng et al., 2009) is the standard validation set of ILSVRC2012.
ImageNet-
Real (Beyer et al., 2020) revises the ImageNet validation set with a more robust labeling procedure, which measures model’s generalization on different labeling procedures.
Seven out-of-distribution (OOD) datasets are considered.
Fig.
4 shows their example images.
ImageNet-
Rendition (Hendrycks et al., 2021a) is an OOD dataset that contains renditions, such as art, cartoons, of 200 ImageNet classes.
Stylized -ImageNet
(Geirhos et al., 2019) is used to induce the texture vs. the shape bias by stylizing the ImageNet validation set with 79,434 paintings.
The “shape ImageNet labels” are used as the label.
ImageNet-
Sketch (Wang et al., 2019) is a “black and white” dataset constructed by querying the sketches of the 1,000 ImageNet categories.
ObjectNet (Barbu et al., 2019) is an OOD test set that controls the background, context, and viewpoints of the data.
Following the standard practice, we evaluate the performance on the 113 overlapping ImageNet categories.
ImageNet- V2 (Recht et al., 2019) is a new and more challenging test set collected for ImageNet.
The split “matched-frequency” is used.
ImageNet-
A (Hendrycks et al., 2021b) is a natural adversarial dataset that contains real-world, unmodiﬁed, natural examples that cause image classiﬁers to fail signiﬁcantly.
ImageNet-
C (Hendrycks & Dietterich, 2019) evaluates the model’s robustness under common corruptions.
It contains 5 serveries of 15 synthetic corruptions including ‘Defocus,’ ‘Fog,’ and ‘JPEG compression’.
Models.
We verify our method on three ViT backbones: the classical ViT (Dosovitskiy et al., 2020) termed as ViT Vanilla, ViT-AugReg (Steiner et al., 2021), and MLPMixer (Tolstikhin et al., 2021).
By default, we refer ViT to ViT-AugReg considering its superior performance.
For ViT, we study the
Tiny (Ti), the Small (S), and the Base (B) variants, all using the patch size of 16x16.
We also compare with a CNN-based ResNet (He et al., 2016) baseline, and the Hybrid-ViT model (Dosovitskiy et al., 2020) that replaces image patch embedding with a ResNet-50 encoder.
Implementation details On ImageNet , the models are trained with three ViT model variants, i.e.
Ti, S, and B, from small to large.
The codebook size is K = 1024 , and the codebook embeds the discrete token into d c = 256 dimensions.
On ImageNet-21K , the quantizer model is a VQGAN and is trained on ImageNet-21K only with codebook size K = 8 , 192 .
All models including ours use the same augmentation (RandAug and Mixup) as in the ViT baseline (Steiner et al., 2021).
More details can be found in the Appendix A.12.
[Table]
4. M AIN R ESULTS
Table 1 shows the results of the models trained on ImageNet.
All models are trained under the same setting with the same data augmentation from (Steiner et al., 2021) except for the ViT-B Vanilla row, which uses the data augmentation from (Dosovitskiy et al., 2020).
Our improvement is entirely attributed to the additional discrete representations.
By adding discrete representations, all ViT variants, including Ti, S, and B, improve robustness across all eight benchmarks.
When only using discrete representations as denoted by “(discrete only)”, we observe a larger margin of improvement ( 10%-12% ) on the datasets depicting object shape: Rendition, Sketch, and Stylized-ImageNet.
It is worth noting that it is a challenging robustness test to obtain gains across all the benchmarks as different datasets may capture distinct types of data distributions [Table] [Table]
Our approach scales as more training data is available.
Table 2 shows the performance for the models pretrained on ImageNet-21K and ﬁnetuned on ImageNet.
Training on sufﬁciently large datasets inherently improves robustness (Bhojanapalli et al., 2021) but also renders further improvement even more challenging.
Nevertheless, our model consistently improves the baseline ViT-B model across all robustness benchmarks (by up to 10% on ImageNet-C).
The results in Table 1 and Tablevalidate our model as a generic approach that is highly effective for the models pretrained on the sufﬁciently large ImageNet-21k dataset
[Table]
Comparison to the State-of-the-Art We compare our model with the state-of-the-art results on four datasets in Table 3 and Table 4.
It is noteworthy that different from the compared approaches that are tailored to speciﬁc datasets, our method is generic and uses the same data augmentation as our ViT baseline (Steiner et al., 2021), i.e. RandAug and Mixup, for all datasets. On ImageNet-Rendition , we compare with the leaderboard numbers from (Hendrycks et al., 2021a).
Trained on ImageNet, our approach beats the state-of-the-art DeepAugment+AugMix approach by I m age N e t R end i t i on S t y li z ed car original
S k e t c h O b j e c t N e t I m age N e t - A I m age N e t - C I m age N e t - Vdecoded original decoded original decoded original decoded
[Table]
StylizedTop
[Table]
Notice that the augmentation we used—RandAug+MixUp—is 17% worse than the DeepAugment+AugMix on ResNets.
Our performance can be further improved by another 6% when pretrained on ImageNet-21K.
On ImageNet-Sketch , we surpass the state-of-the-art (Hermann et al., 2020) by 8%.
On Stylized-ImageNet , our approach improves 18% top-5 accuracy by switching to discrete representation.
On ImageNet-C , our approach slashes the prior mCE by up to 10%.
Note that most baselines are trained with CNN architectures that have smaller capacity than ViT
-B.
The comparison is fair for the bottom entries that are all built on the same ViT-B backbone. 4. I N - DEPTH A NALYSIS In this subsection, we analyze both quantitatively and qualitatively that discrete representations strengthen ViT robustness by better capturing object shape and global contexts. Quantitative results.
Result in Table 3c studies the OOD generalization setting “IN-SIN” used in (Geirhos et al., 2019), where the model is trained only on the ImageNet (IN) and tested on the
Stylized ImageNet (SIN) images with conﬂicting shape and texture information.
The StylizedImageNet is designed to measure the model’s ability to recognize shapes rather than textures, e.g. images of cat with elephant texture, and higher performance in Table 3c is a direct evidence of our model’s superiority in recognizing object by shape.
While ViT outperforms CNN on the task, our discrete representations achieves another 10+% gain.
However, if we replace the discrete token with the global continuous CNN features , then ViT does not have robustness gain (see Table 8).
This substantiates the beneﬁt of our discrete representation in helping ViT model shape.
Additionally, we conduct the “shape vs. texture biases” analysis following (Geirhos et al., 2019) under the setting “IN-SIN” (cf.
Appendix A.5).
Figure.
5a compares shape bias between humans and
[Table]
human mean=0.
Dr. ViT-B/16 (Ours) mean = 0.
ViT-B/16 mean = 0.
ResNet-50 mean = 0.
(a) Shape vs. Texture Biases.
(i) attention on test images (ii) average attention
Vi T
Our s
Vi T
Our s I mage
I mage
Vi T
Our s (b) Attention Visualization.
three models: ResNet-50, ViT-B, and Ours.
As shown, humans are highly biased towards shape with an average fraction of 0.96 to correctly recognize an image by shape.
ViT is more biased towards shape than the CNN-based ResNet-50, which is consistent with the prior work studies (Naseer et al., 2021).
Our discrete representation (0.62) further shrinks the gap between the ViT (0.42) and human baseline (0.96), signiﬁcantly improving the shape-bias of the baseline ViT model.
Finally, we validate discrete representation’s ability in modeling global context via position embedding.
Following (Chen et al., 2021b), we compare training ViT with and without using position embedding.
As position embedding is the only vehicle to equip ViT with global input information, its contribution suggests to what degree the model makes use of global contexts for recognition.
As shown in Table 5, removing position embedding from the ViT model only leads to a marginal performance drop (2.8%) on ImageNet, which is consistent with (Chen et al., 2021b).
Removing position embeddding from ViT with global continuous CNN feature input also does not change results (See
Table 14 Hybrid). However, without position embedding, our model accuracy drops by 29%, and degrades by a signiﬁcant 36%-94% on the robustness benchmarks. This result shows that spatial information is more crucial when discrete representation is used, which also suggests our model relies on more global shape rather than texture for recognition. This result quantitatively substantiates that discrete representations improve ViT to better capture global contexts.
[Table]
Qualitative results. First, following (Dosovitskiy et al., 2020) [Table] we visualize the learned pixel embeddings, called ﬁlters, of the ViT-B model and compare them with ours in Fig.
3b with respect to (1) randomly selected ﬁlters (the top row in Fig.
3b) and (2) the ﬁrst principle components (the bottom).
More visualization is included in Appendix A.1.3.
The evident visual differences suggest that our learned ﬁlters capture structural patterns that are more global.
Second, we compare the attention from the classiﬁcation tokens in Fig.
5b, where (i) visualizes the individual images and (ii) averages attention values of all the image in the same mini-batch.
Our model attends to image regions that are semantically relevant for classiﬁcation and captures more global contexts relative to the central object.
Finally, we investigate what information is preserved in discrete representations.
Speciﬁcally, we reconstruct images by the VQ-GAN decoder from the discrete tokens and ﬁnetuned discrete embed- dings.
Fig.
4 visualizes representative examples from ImageNet and the seven robustness benchmarks.
By comparing the original and decoded images in Fig.
4, we ﬁnd the discrete representation preserves object shape but can perturb unimportant local signals like textures.
For example, the decoded image changes the background but keeps the dog shape in the last image of the ﬁrst row.
Similarly, it hallucinates the text but keeps the bag shape in the ﬁrst image of the third row.
In the ﬁrst image of ImageNet-C, the decoding deblurs the bird image by making the shape sharper.
Note that the discrete representation is learned only on ImageNet, but it can generalize to the other out-of-distribution test datasets.
We present more results with failure cases in Appendix A.1.1.
4. A BLATION S TUDIES
Model capacity vs. robustness.
Does our robustness come from using larger models?
Fig.
6 shows the robustness vs. the number of parameters for the ViT baselines (blue) and our models (orange).
We use ViT variants Ti, S, B, and two hybrid variants Hybrid-S and Hybrid-B, as baselines.
We use “+Our” to denote our method and “+Small” to indicate that we use the small version of quantization encoder (cf.
Appendix A.11.1).
With a similar amount of model parameters, our model outperforms the ViT and Hybrid-ViT in robustness.
Codebook size vs. robustness.
In Table 6, we conduct an ablation study on the size of the codebook, from 1024 to 8192, using the small variant of our quantized encoder trained on ImageNet.
A larger codebook size gives the model more capacity, making the model closer to a continuous one.
Overall,
Stylized-ImageNet beneﬁts the most when the feature is highly quantized.
With a medium codebook size, the model strikes a good balance between quantization and capacity, achieving the best overall performance.
A large codebook size learned on ImageNet can hurt performance. [Figure]
[Table]
Discrete vs. Continuous global representation.
In Appendix A.4, instead of using our discrete token, we study whether concatenate continuous CNN representations with global information can also improve robustness.
Results show that simply using continuous global CNN representation fails to improve robustness, which demonstrate the importance of using our discrete representation.
Network designs for combining the embeddings.
In Appendix A.9, we experiment with 4 different designs, including addition, concatenation, gate ﬁltering, and cross-attention, to combine the discrete and pixel embeddings.
Among them, the simple concatenation yields the best overall result.
C ONCLUSION
This paper introduces a simple yet highly effective input representations for vision transformers, in which an image patch is represented as the combined embeddings of pixel and discrete tokens.
The results show the proposed method is generic and works with several vision transformer architectures, improving robustness across seven out-of-distribution ImageNet benchmarks.
Our new ﬁndings connect the robustness of vision transformer to discrete representation, which hints towards a new direction for understanding and improving robustness.
Ethics statement: The authors attest that they have reviewed the ICLR Code of Ethics for the 2022 conference and acknowledge that this code applies to our submission.
Our explicit intent with this research paper is to improve the state-of-the-art in terms of accuracy and robustness for Vision
Transformers.
Our work uses well-established datasets and benchmarks and undertakes detailed evaluation of these with our novel and improved approaches to showcase state-of-the-art improvements.
We did not undertake any subject studies or conduct any data-collection for this project.
We are committed in our work to abide by the eight General Ethical Principles listed at ICLR Code of
Ethics ( https://iclr.cc/public/CodeOfEthics ).
Reproducibility statement Our approach is simple in terms of implementation, as we only change the input embedding layer for the standard ViT while keeping everything else, such as the data augmentation, the same as the established ViT work.
We use the standard, public available training and testing dataset for all our experiments.
We include ﬂow graph for our architecture in Fig.
1, pseudo JAX code in Fig.
3a and implementation details in Section 4.1 and Appendix A.11.
We will release our model and code to assist in comparisons and to support other researchers in reproducing our experiments and results.
A A PPENDIX
The Appendix is organized as follows.
Appendix A.2 extends the discussion on learning objective.
Appendix A.1 presents more ﬁgures to visualize the attention, the decoded images, and the decoding failure cases.
Appendix A.9 compares designs for combing the pixel and discrete embeddings.
In the end, Appendix A.11 discusses the implementation details.
[Figure]
A. V ISUALIZATION A.1. R ECONSTRUCTED I MAGES FROM DISCRETE REPRESENTATION
Fig.
7 shows more examples of reconstructed images decoded from the ﬁnetuned discrete embedding.
Generally, the discrete representation reasonably preserves object shape but can perturb local signals.
Besides, Fig 7 also shows the distribution diversity in the experimented robustness benchmarks, e.g., the objects in ImageNet-A are much smaller.
Note that while the discrete representation is learned only on the ImageNet 2012 training dataset (the ﬁrst row in Fig. 7), it can generalize to the other out-of-distribution test datasets.
[Figure]
A.1. L IMITATIONS OF DISCRETE REPRESENTATION
We ﬁnd four failure cases that the discrete representation are unable to capture: text, human faces, parallel lines, and small objects or parts.
The example images are illustrated in Fig.
Without prior knowledge, the discrete representation has difﬁculty to reconstruct text in image.
On the other hand, it is interesting to ﬁnd the decoder can model animal faces but not human faces.
This may be a result of lacking facial images in the training data of ImageNet and ImageNet-21K.
The serious problem we found for recognition is its failure to capture small objects or parts.
Though this can be beneﬁcial sometimes, e.g. force the model to recognize the tire without using the “BMW” logo (see the ﬁrst image in the last row of Fig.
8), it can cause problems in many cases, e.g. recognizing the small whales in the last image.
As a result, our proposed model leverages the power of both embeddings to promote the interaction between modeling global and local features.
A.1. F ILTERS
In Fig.
3b of the main paper, we illustrate the visual comparison of the learned pixel embeddings between the standard ViT and our best performing model.
In this subsection, we extend the visualization to our models with varying number of pixel dimension when concatenating the discrete and pixel embedding, where their classiﬁcation performances are compared in Table 13.
As the dimension of pixel embedding grows, the ﬁlters started to pick up more high-frequency signals.
However, the default setting in the paper using only 32 pixel dimensions (pixel dim=32), which induces an inductive bias by limiting the budget of pixel representation, turns out to be the best performing model, suggesting the importance to capturing low-frequency patterns for OOD generalization.
A.1. A TTENTIONS
We visualize the attention following the steps in (Dosovitskiy et al., 2020).
To be speciﬁc, to compute maps of the attention from the output token to the input space, we use Attention Rollout (Abnar & Zuidema, 2020).
Brieﬂy, we average attention weights of ViT-B and ViT-B+Ours across all heads and then recursively multiply the weight matrices of all layers.
Both ViT and ViT+Ours are trained on ImageNet.
In Fig.
11, Fig. 12 and Fig.
13, we visualize the attention of individual images in the ﬁrst mini-batch of each test set.
As shown, our model attends to image regions that are semantically relevant for [Figure] classiﬁcation and captures more global contexts relative to the central object.
This can be better seen from Fig.
10, where the heat map averages the attention values of all images in the ﬁrst minibatch of each test set.
As found in prior works (Dosovitskiy et al., 2020; Naseer et al., 2021), ViTs put much attentions on the corners of the image, our attentions are more global and do not exhibit the same defect. [Figure]
The heat map averages the attention values of all images in the ﬁrst mini-batch of each test set.
The results show our attention capture more global context relative to the central object.
A. D ETAILS ON LEARNING OBJECTIVE
Let x denote an image with the class label y , and z represent the discrete tokens z ∈ Z .
For notational convenience, we use a single random variable to represent the discrete latent variables whereas our implementation actually represents an image as a ﬂattened 1-D sequence following the raster scan.
• q φ ( z | x ) denotes the VQ -VAE encoder, parameterized by φ , to represent an image x into the discrete token z .
• p θ ( x | z ) is the decoder that models the distribution over the RGB image generated from discrete tokens.
• p ψ ( y | x, z ) stands for the vision transformer model shown in Fig.
of the main paper.
We model the joint distribution of image x , label y , and the discrete token z using the factorization p ( x, y, z )
=
p θ ( x | z ) p
ψ
( y | x, z ) p ( z ) .
Our overall procedure can be viewed as maximizing the evidence lower bound (ELBO) on the joint likelihood, which yields:
[Equation]
Since the f -divergence is non-negative, we have:
[Equation]
Using the factorization, we have:
[Equation]
Equation 13 yields the ELBO:
[Equation]
In the ﬁrst stage, we maximize the ELBO with respect to φ and θ , which corresponds to learning the VQ-VAE encoder and decoder.
We assume a uniform prior for both p ψ ( y | x, z ) and p ( z ) .
Given that q φ ( z | x ) predicts a one-hot output, the regularization term (KL divergence) will be a constant log K , where K is the size of the codebook V .As in the VQ-VAE model, the distribution p θ is deterministic. We note a similar assumption as in the DALL-E model (Ramesh et al., 2021) is used to stabilize the training, which means the transformer is not learned at this stage. We have the same observation as in DALL-E that this strategy works better than jointly training with p ψ ( y | x, z ) .
In addition to the reconstruction loss, the VQ-VAE’s training objective adds a dictionary learning loss and a commitment loss, calculated from:
[Equation]
(16) where z e ( x ) and z q ( x ) represent the encoder output and the decoder input, respectively.
v denotes the discrete embedding for the image x . sg ( · ) is the stop gradient function.
The training used straight-through estimation which just copies gradients from z q ( x ) to z e ( x ) .
Notice that our implementation uses VQ-GAN (Esser et al., 2021) which adds an additional GAN loss:
[Equation]
where β
= 0 .
25 and λ = 0 .
1 are used to balance different loss terms.
In the second stage of ﬁnetuning, we optimize ψ while holding φ and θ ﬁxed, which corresponds to learning a vision transformer for classiﬁcation and ﬁnetuning the discrete embedding V .
By rearranging Equation 14, we have:
[Equation]
where p ψ ( y | x, z ) is the proposed ViT model as illustrated in Fig.
1 of the main paper.
The ﬁrst term denotes the likelihood for classiﬁcation that is learned by minimizing the multi-class cross-entropy loss.
The second term becomes a constant given the ﬁxed θ and φ .
In the ﬁnetuning stage, two sets of parameters are updated: the vision transformer ψ is learned from scratch and the discrete embedding V is ﬁnetuned.
We ﬁx the encoder and stop the gradient propagating back to the encoder considering the instability of straight-through estimation and the stop gradient operation in the dictionary learning loss, i.e. Equation 16.
Empirically, we also compared the proposed pretraining-ﬁnetuning strategy with the joint learning strategy in which all parameters are optimized simultaneously.
We ﬁnd the joint learning signiﬁcantly underperforms the proposed learning strategy.
A. C OMPARING TO D ATA A UGMENTATIONS
One pivotal advantage of our approach is that it only changes the token representations in the model architecture, which enables it to be compatible to different kind of data augmentation strategies.
As shown in Table 7, Our Dr. ViT method achieves signiﬁcant improve with and without data augmentations.
In addition, we experiment with 7 different combinations of the state-of-the-art data augmentation strategy, include CutOut, CutMix, RandAug, and Mixup.
Our Dr. ViT method can further improve up to 10% robust accuracy on the best augmentation strategy — ‘RandAug +
Mixup’.
Stronger data augmentations in the further can further improve the performance of our approach.
The box highlights the best accuracy.
While the state-of-the-art data augmentation (CutMix+RandAug) can improve Imagenet in-distribution accuracy, they lose up to 10% on OOD test accuracy.
[Table]
A. T HE I MPORTANCE OF D ISCRETE I NFORMATION
[Table]
From our results, we show that using discrete representations derived from VQGAN can improve robustness.
Note that to calculate the discrete representations, the VQGAN encoder uses deep CNN architecture.
To study whether it is the CNN encoder or the discrete representation that contribute to our robustness gain, we directly concatenate the global feature extracted by CNN to the input token of ViT.
While deep CNN can capture global features, our results show that concatenating the continuous feature from deep CNN decrease the overall robustness.
In Table 8, we run extensive experiments that concatenates global CNN feature with dimension 64, 128, 256, 384, 512, and 640 to the original token in ViT. With continuous representation, even if they are global features from CNN, none of the models can improve robustness.
This is because the continuous representations from CNN can still carry enough local information that the ViT can depend on.
Just by making the representations discrete, our method improves robustness by up to 14%, which demonstrate the importance of discrete representation for robustness.
A. T EXTURE VS . S HAPE STUDY
In Fig.
5a of the main paper, we performs the “shape vs. texture biases” analysis following (Geirhos et al., 2019)
(see Figure 4 in their paper).
It uses the score “fraction of shape decision” to quantify a model’s shape-bias between 0 and 1.
To compute the score, we follow the steps on their Github 1 as follows: 1) evaluate the model on all 1280 images in Stylized ImageNet; 2) map the model decision to 16 classes; 3) exclude images without a cue conﬂict; 4) take the subset of “correctly” classiﬁed images (either shape or texture category correctly predicted); (5) compute “shape bias” as the following fraction: (correct shape decisions) /
(correct shape decisions + correct texture decisions).
Fig.
5a in the main paper presents the score on the 16 categories along their average denoted by the colored horizontal lines.
We compute the scores for three models trained on ImageNet, i.e. ResNet50, ViT-B/16, and ours, and quote human scores from (Geirhos et al., 2019).
Note that the OOD generalization setting “IN-SIN” (Geirhos et al., 2019) is used, where the model is trained only on the ImageNet (IN) and tested on the Stylized ImageNet (SIN) images with conﬂicting shape and texture cue.
A. T HE E FFECT OF C ODEBOOK E NCODER C APACITY ON R OBUSNTESS
[Table]
We analyze how the capacity of the discrete encoder affects the model robustness.
We train a small encoder and a big encoder and compare their performance.
We describe the VQ encoder’s architecture conﬁguration in Table 17.
We run robust experiment in Table 9.
our results show that using a larger encoder for calculating discrete representation improves more robustness.
A. T HE E FFECT OF P RETRAINING C ODEBOOK WITH M ORE D ATA ON R OBUSNTESS
[Table]
[Table]
We analyze if pretraining codebook with more data improves robustness.
We use the same model architecture, but pretrain two codebook encoders on the small ImageNet and large ImageNet21K, respectively.
We then apply the pretrained discrete token to ViT and train Dr. ViT on ImageNet21K.
We show the robustness of the resulting Dr. ViT models in Table 10, where pretraining on more data improves Dr. ViT’s robustness.
A. A DDITIONAL M ODELS FOR L EARNING D ISCRETE R EPRESENTATION
VQVAE (Oord et al., 2017) is another major model to learn neural discrete representation, which uses a different loss function and different encoder architecture from VQ-GAN.
In Table 11, we also evaluate ViT’s robustness using discrete tokens from a standard VQVAE model.
The model also achieved consistent robustness improvement with using VQ-GAN discrete tokenizer.
Speciﬁcally,
VQ-VAE can further improve the robustness on Stylized ImageNet by 1% than VQ-GAN, and the robustness on ImageNet-C by 3% than VQ-GAN.
In addition, VQ-VAE improves Rendition by 8% than ViT baseline, and Sketch by 7% than ViT baseline.
OUr results show that our apporach is general, different architectures than VQ-GAN can also produce discrete representation that helps
ViT’s robusstness.
[Table]
A. C OMPARING D ESIGNS FOR C OMBINING E MBEDDINGS
In this subsection [Table] we compare different designs to combine the pixel and discrete embeddings.
See the combination operation in Fig.
1 and in Equation 6 of the main paper.
The results are used verify our design of using a simple concatenation presented in the paper.
Four designs are considered and discussed below.
Their performances are compared in Table 12.
[Table]
Addition.
We ﬁrst use linear projection matrix E to obtain the pixel embedding that shares the same dimension as the discrete embedding which is 256.
We then add the two embedding, and project the resulting embedding to the dimension required by the ViT if needed.
Formally, we compute
[Equation]
) where V z d and x p E represent the discrete and pixel embeddings, respectively, and E 2 is a new MLP projection layer that is needed when the resulting dimension is different to the ViT input.
Concatenation.
We concatenate the discrete embedding with the pixel embedding, and then feed the resulting embedding into the transformer encoder.
By default, we use a dimension of 32 for the pixel embedding, a dimension of 256 for the discrete embedding, thus we pad 0 to the vector if the input dimension of the ViT is higher than 288.
[Equation]
where “;” indicates the concatenation operation.
Residual Gating.
As pixel embeddings may contain nuisances, inspired by (Vo et al., 2019), we learn a gate from both embeddings, and then multiply it with the pixel embeddings to ﬁlter out details that are unimportant to classiﬁcation.
Speciﬁcally, we calculate the gate by a 2-layer MLP, and apply a softmax to the output from the last layer.
[Equation]
Then the pixel embeddings are gated by:
[Equation]
Then we concatenate the gated embedding with the discrete embedding:
[Equation]
Cross-Attention.
We use the discrete embedding as the query, and the pixel token as the key and value in a standard multi-head cross-attention module (MCA).
[Equation]
We then concatenate the attended feature output with the discrete embedding.
We also pad 0 if the
ViT requires larger input dimensionality.
[Equation]
A.9. T HE R OBUSTNESS E FFECT OF U SING C ONTINUOUS R EPRESENTATION WITH D ISCRETE R EPRESENTTION .
The above study shows that concatenation is the most effective way to combine the discrete and continuous representations.
As we ﬁxed the dimensionality for the discrete token to be 256, we can study the optimal ratio between continuous and discrete representations by changing the dimensionality on the pixel token.
The robust performance under different pixel token dimension in concatenation combining method.
By increasing the dimension of the pixel embeddings, the model uses more continuous representations than discrete representations.
There is an inherent trade off between the out of distribution robustness performance.
The more continuous representations the model use, the lower robustness on the out-of-distribution set that depicting object shape, but also aciheves higher accuracy on the in-distribution ImageNet and out of distribution variants ImageNet-A that requires detailed information.
However, using only continuous embeddings also hurt robustness.
Thus we choice to use dimension 32, which balance the trade-off and achieves the best overall robustness. [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table] [Table] [Table] A. T HE IMPORTANCE OF SPATIAL INFORMATION FOR DIFFERNT ARCHITECTURES
In addition to the analysis in the main paper
[Table]
we also investigate whether it is the convolutional operation in the vector quantized encoder that makes the model use the spatial information better.
We remove the position embedding from the Hybrid-B model, whose input is also produced by a convolution encoder.
Our results show that removing the positional information from the Hybrid model with a convolutional encoder does not decrease the performance.
However, for both our models, the discrete only and the combined one, removing positional embedding causes a large drop in performance.
This shows that our robustness gain via the spatial structure is from our discrete design, not convolution
[Table] [Table]
A. I MPLEMENTATION DETAILS A.11. A RCHITECTURE
We use three variants of ViT transformer encoder backbone in our experiments.
We show the conﬁguration details and the number of parameters in Table 15.
We also use hybrid model with ResNet as the tokenization backbone.
We show the conﬁguration in Table
[Table]
We use B-16 for the MLPMixer Tolstikhin et al.
We use two kinds of VQ Encoder in our experiment: Small VQ Encoder and VQ Encoder.
The VQ encoder uses the same encoder architecture as the VQGAN encoder (Esser et al., 2021).
For Small
VQ Encoder, we decrease both the number of the resisual blocks and the number for ﬁlter channel by half, which is a lightweight model that decrease the model size by around 8 times.
We show the conﬁguration detail in Table 17.
A. T RAINING
Implementation details We implement our model in Jax and optimize all models with Adam (Kingma & Ba, 2014).
Unless speciﬁed otherwise, the input images are resized to 224x224, trained with a batch size of 4,096, with a weight decay of 0.1.
We use a linear learning rate warmup and cosine decay.
On ImageNet , the models are trained for 300 epochs using three ViT model variants, i.e. Ti, S, and B, from small to large.
The VQ-GAN model is also trained on ImageNet for 100K steps using a batch size of 256.
The codebook size is K = 1024 , and the codebook embeds the discrete token into d c = 256 dimensions.
For discrete ViT-B, we use average pooling on the ﬁnal features.
On ImageNet-21K , we train 90 epochs.
We ﬁnetune on a higher resolution on ImageNet, with Adam and 1e-4 learning rate, batch size 512, for 20K iterations.
The quantizer model is a VQGAN and is trained on ImageNet-21K only with codebook size K = 8 , 192 .
All models including ours use the same augmentation (RandAug and Mixup) as in the ViT baseline (Steiner et al., 2021).
More details can be found in the Appendix.
[Table]
[Table]
ImageNet Training We train our model with the Adam (Kingma & Ba, 2014) optimizer.
Our batchsize is 4096 and trained on 64 TPU cores.
We start from a learning rate of 0.001 and train 300 epoch.
We use linear warm up for 10k iterations and then a cosine annealing for the learning rate schedule.
We use 224 × 224 resolution for the input image.
We conduct experiment on three variants of the
ViT model, Ti, S, and B, from small to large.
In addition, we also experiment on the newly released
MLPMixer, which also uses image patch as the token embeddings.
We denote our approach as using both representations, and we also run a discrete only variant where we only uses the discrete embeddings without pixel embeddings.
Due to the small input dimension of Ti, we only use the code representation without concatenate the pixel representation, and we project the 256 dimension to 192 via a linear projection layer.
For the other variants, we concatenate both representations, and pad zeros if additional dimension required for the transformer.
We use the same augmentation and model regularization as (Steiner et al., 2021), where we use RandAug with hyper-parameter (2,15) and mixup with α = 0 . 5 , and we apply a dropout rate of 0.1 and stochastic block dropout of 0.1.
We download pretrained VQGAN model, which is trained on only ImageNet.
We use VQ-GAN’s encoder only with a codebook size of 1,024.
The codebook integer is embedded into a 256 dimension vector.
ImageNet-21K Training Our models use learning rate of 0.001, weight decay of 0.03, and trainepoch.
We use linear warm up for 10k iterations and then a cosine annealing for the learning rate schedule.
We use 224 × 224 resolution for the input image.
We use the same augmentation as (Steiner et al., 2021).
For the quantizer model, we use VAGAN that is trained on unlabeled ImageNet-21K only.
We use a codebook size of 8,192, and the codebook integer is embedded into a 256 dimension vector.
To evaluate on the standard test set, we further ﬁnetune the pretrained ImageNet-21K model on ImageNet.
We use Adam and optimize for 20k with 500 steps of warm up.
We use a learning rate of 0.0001.
We use a larger resolution 384 × 384 and 512 × 512 .
We only experiment the ViT-B variant.
For training VQ Encoder and Decoder, we train with batchsize 256 until it converges.
We use perceptual loss with weight of 0.1, adversarial loss with weight 0.1.
We use L1 gradient penalty of 10 during the optimization.
The model is trained with resolution of 256 × 256 .
The model can scale to different image resolution without ﬁnetuning.
[Figure]
[Table]
[Figure]
[Table] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure] [Table]
[Figure]
[Table]
23.7M
[Figure] [Figure] [Figure]