Discrete Representations Strengthen Vision Transformer Robustness
Authors: Chengzhi Mao
Lu Jiang
Mostafa Dehghani
Carl Vondrick
Rahul Sukthankar
Irfan Essa
Venue: ICLR 2022 Poster (2022)
Keywords: vision transformer
robustness
image recognition

Abstract: Vision Transformer (ViT) is emerging as the state-of-the-art architecture for image recognition. While recent studies suggest that ViTs are more robust than their convolutional counterparts, our experiments find that ViTs are overly reliant on local features (\eg, nuisances and texture) and fail to make adequate use of global context (\eg, shape and structure). As a result, ViTs fail to generalize to out-of-distribution, real-world data. To address this deficiency, we present a simple and effective architecture modification 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.

Published as a conference paper at ICLR 2022
DISCRETE REPRESENTATIONS STRENGTHEN VISION TRANSFORMER ROBUSTNESS
Chengzhi Mao?; Lu Jiang', Mostafa Dehghani', Carl Vondrick?, Rahul Sukthankar', IIffn Essa
1 Google Research
2 Computer Science, Columbia University
3 School of Interactive Computing, Georgia Insitute of Technology
{mcz, vondri snndricckee ck co lumbi a. edu (lujiang, dehghani, sukthankar, [rffanessa$&googlee com
ABSTRACT
Vision Transformer (ViT) emerging as the state-of-the-ar architecture for im- age recognition. While recent studies suggest that VITs are more robust than their convolutional counterparts, our experiments find that ViTs trained on ImageNet are overly reliant on local textures and fail to make adequate of shape infor- mation. ViTs thus have difficulties generalizing to outoo--distrbuutin, real-world data. To address this deficiency, we present simple and effective architecture modification 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 indi- vidually, which promote ViTs to learn global information that is invariant. Experi- mental results demonstrate that adding discrete representation on four architecture variants strengthens ViT robustness by up to 12% across seven ImageNet robust- ness benchmarks while maintaining the performance on ImageNet.
INTRODUCTION
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 closed- world assumptions that are overly sensitive to minor variations in local texture and insufficiently adept at representing more robust scene and object characteristics, such the shape.
violated. This lack of robustness can be ascribed to learned representations
Vision Transformer (ViT) (Dosovitskiy et al. 2020) has started to rival Convolutional Neural Net- works (CNNs) in many computer vision tasks, Recent works found that ViTs are more robust than CNNs Paal & Chen 2021; Mao et al. 202 Bhojanapalli et al. [021]) and generalize favorably on a variety of visual robustness benchmarks (Hendrycks al. [021 | Hendrycks & Dietterich 2019). These work suggested that ViTs' robustness comes from the self aetf-attentoon architecture that captures globally contttllee inductive bias than CNNs.
However, though self-attention can model the shape information, we found that ImageNet-trained VIT is still biased to textures than shape (Geirhos al., 2019). We hypothesize that this deficiency in robustness comes from the high-dimensional, individually informative, linear tokenization which biases ViT to minimize empirical risk via local signals without learning much shape information.
In this paper, we propose simple yet novel layer for vision transformers, where image patches are represented by discrete tokens. To be specific, we discretize an image and represent an image patch as a discrete token "visual word" in codebook. Our key insight is that discrete tokens capture important features in a low-dimensional space (Oord et al.) (017)) 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
Work done while interning Google Research.

---

Published as a conference paper at ICLR 2022
Figure 1: Overview of the proposed VIT us- ing discrete representa- tions. addition to the pixel embeddings (or- ange), we introduce dis- crete tokens and em- heddings (pink) as the input to the standard Transformer Encoder of the VIT model (Doso- vitskiy al. 2020).
Fire-turee ad Codelloox
PECHLER
Decode
Recorimmcte: mage
Closelicetion Prodicto
Impul Image
Vector Quantree Modell
Closs Tokon
1
Tranformer Broader
Token
Diecrene
Eni Eribadding
Linear Proj Projecton
Fosition
Plaw Token
Powl Embeeding
Embediing
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 all existing approaches for robustness, and can be integrated into architectures that extend vision transformer. We call the ViT model using our Discrete representation or Dr. VIT.
Our experiments and visualizations show that incorporating discrete tokens in ViT significantly im- proves generalization accuracy for all seven out-of-distributiin ImageNet benchmarks: Stylized- ImageNet by up 12%, ImageNet- Sketch by up to 10%, and ImageNet-C by up to 10%. Our method establishes the new soute-of-the-are on four benchmarks without any dataset-specific data augmentation. The success of discrete representation has been limiting to image generation (Ramesh) et al. 2021; [Esse et al. (02I). Our work is the first to connect discrete representations to robustness and demonstrate consistent robustness improvements. We hope our work helps pave the road to a joint vision transformer for both image classification and generation.
2 RELATED WORK
Vision Transformer (ViT) (Dosovitskiy et al. 2020), inspired by the Transformer (Vaswani et al., 2017 in NLP, is the first CNN- architecture that achieves ofate-of-the-art image classification accuracy. Since its inception, numerous works have proposed improvements to the VIT architec- ture (Wang et al.) 202T; Chu et al., 202T; Liu al.) [021]] d d'Ascoli et al.] 202I). objective (Chen
eet al. 2021a). training strategy (Touvron et al. 202 1) etc.. Given the difficulty to study all existing
VIT models, this paper focuses on the classical ViT model (Dosovitskiy et al.) 2020) and its recent published versions. Specifically, (Steiner al. (2021) proposed Vig-AugReg that applies stronger data augmentation and regularization to the ViT model. Tolstikhin et al. (2021) introduced MLP- Mixer 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 eet al., (021) in our robustness study for three reasons. First, they represent both the very first and one-of-the-best vision transformers in the literature. Second, these models demonstrated competitive performance when pre-trained on sufficiently large datasets such as ImageNet-21 K and JFT-300M. Finally, unlike other Transformer models, they provide architectures consisting of solely Trans- former 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-of- distribution generalization of deep models (Barbu et al. 2019 Hendryeks et al. 202 a Wang al.) 2019 Geirhos et al., 2019 Hendrycks Dietterich] 2019 Recht et al. 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 verified with the CNNs. Among them, carefully-designed data augmenta- tions (Hendrycks al. 202 a] Cubuk al., 2018 Steiner al. 2021 Mao al., 202Tbla), model regularization (Wang et al. 2019] Huang et al. 2020b) Hendrycks et al. 2019), and multitask learn- effective to address the issue.
ing (Zamir et al. 2020)
More recently, a few studies (Paul & Chen 202T; Bhojanapalli et al. 2021 Naseer et al. 2021 Shao et al. 2021) suggest that ViT's are more robust than CNNs. Existing works mainly focused
2

---

Published as a conference paper at ICLR 2022
on analyzing the cause of superior generalizability in the ViT model. 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. 202T). 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 visual representation prior to the deep learning revolution, such as in bag-of-visual-veer model (Arandjelovic & Zisserman 2013). Recently, (Oord et al. 2 et al. al. 2018 proposed neural discrete representation to encode an image as integer tokens. Recent works used discrete representation mainly for image synthesis (Ramesh al. 2021: Esseretal (02I). To the best of our knowledge, our work is the first 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 overfit the non-robust nuisances in the pixel tokens at later finetuning stage, undermining its robustness.
model (Sivic & Zisserman [2003] Csurka et al.] (004) and VLAD
METHOD
3.1 PRELIMINARY ON VISION TRANSFORMER
Vision Transformer (Dosovitskiy al. 2020 is a pure transformer architecture that operates on sequence of image patches. The 2D image x E R xC is flattened into a sequence of image patches, following the raster scan, denoted by Xp E L22((22.)) where I, #xW is the effective sequence length and p2 x C is the dimension of image patch. A learnable classification token Xelass is prepended to the patch sequence, then the position embedding Epos is added to formulate the final input embedding ho-
ho = [Xetass: E; E; E.... x
E ppos:
E E R R( Epos E R((+1)xD
(1)
h-- MSA(LN((b- hg-1,
é = ,lr é = L....dl
(2)
h. MLP(LN((()) + h/e
(3)
y L.())),
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 every block. We denote the number of blocks Ly.
This paper considers the ViT model family consisting of 4 ViT backbones: the vanilla ViT dis- cussed above, ViT-AugReg (Steiner al., 2021) which shares the same VIT architecture but applies stronger data augmentation and regularization, MLP-Mixer (Tolstikhin et al., (02[] 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 ResNet-50 (He et al. 2016).
3.2 ARCHITECTURE
Existing ViT's represent an image patch as sequence of pixel tokens, which are linear projections of flattened image pixels. We propose a novel architecture modification to the input layer of the vision transformer, where an image patch 2p is represented by combination of two embeddings. As illustrated Fig. D in addition to the original pixel-wise linear projection, we discretize an image patch into discrete token in codebook V E kxd, where K is the codebook size and de is the dimension of the embedding. The discretization is achieved by a vector quantized (VQ) encoder pa that produces an integer for image patch as:
(=(z kl(z))
1(k arg min ||ze V5ll2),
(5)
j=1:K
where (x) denotes the output of the encoder network and 1() is the indicator function.
The encoder is applied to the patch sequence xp (RL((P*CC obtain an integer sequence žạ 21, ... K)². Afterward, we use the embeddings of both discrete and pixel tokens to construct the input embedding to the ViT model. Specifically, the input embedding in Equation is replaced by:
ho Xelass: S(V.3. x E); J(V 23 E) S(V 2 E) + Epos,
(6)
where of 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
3

---

Published as a conference paper at ICLR 2022
Figure 2: Comparison of pixel tokens (top) and the reconstructed image decoded from the discrete tokens (bottom). Discrete to- kens capture important shapes and structures but may lose local texture.
THEM READES
16
Dlanee palnt)
FLAMINGO
CakRy LUC
Viliere
1imi
Finstago
UIGANS
best. Note that our model only modifies the input layer of ViT (Equation
I and leaves intact the
remaining layers depicted in Equation 2
Comparison of pixel and discrete embeddings: Pixel and discrete embeddings represent differ- ent aspects of the input image. Discrete embeddings capture important features in low-dimension space (Oord et al.) (017)) 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 embed- dings 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 "flamingo" images. It is worth noting that the VQ encoder/decoder is only trained ImageNet 2012 but they can generalize to out-o--distribution images. Please see more examples in Appendix
On the flip 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 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 out-of-distribution data. Our proposed model leverages the power of both embeddings to promote the interaction between modeling global and local features.
3.3 TRAINING PROCEDURE
The training comprises two stages: pretraining and finetuning. First, we pretrain the VQ-VAR encoder and decoder (Oord et al.) (017) on the given training set. We do not use labels in this step. In the finetuning stage, as shown in Fig. 3al we train the proposed ViT model from scratch and finetune the discrete embeddings. Due the straight-through gradient estimation in VQ-VAE (Ord et al. 2017). we stop the back- propppropagttion after the discrete embedding.
x: input image mini-batch, pixei_embed:: pixei snbeddings x. vegan, encoder and codebook initialized form the pretraining.
4 import jax.numpy as
5 discrote_token - jaz, 1ax, 1ax. stopp_grddeent teggan, tvggan.encoder encoder (x)
discrete_enbed - not (dl serete_rooen, codebook)
takens cp..conatenata:
jsiserate__mbed, pixel_enbed), dim-2)
9 predictions Transsooreernnodderrre Itokenal
ViT
ours
(a) Pseudo JAX code
(b) Pixel/RGB embedding filters
Figure 3: (a) Pseudo code for training the proposed VIT model. (b) Comparing visualized pixel embeddings of the ViT and our model. Top row shows the randomly selected filters and Bottom shows the first 28 principal components. Our filters capture more structural and shape patterns. Now we discuss the objective that the proposed ViT model optimizes. Let (z(z)) denote the VQ encoder, parameterized by network weights ő, represent image into sequence of integer tokens z. The decoder models the distribution (z]2) over the RGB image generated from discrete tokens. Pa (x 2) stands for the proposed vision transformer shown in Fig.
We factorize the joint distribution of image x, label y, and the discrete token 2 by p(x, y, 2) (>(x))>) (v/ (00l,,22ppp =))( Our overall training procedure is maximize the evidence lower bound
4

---

Published as a conference paper at ICLR 2022
(ELBO) the joint likelihood:
pog y) 2 g=( [lo po [x]=)]]
po (o 2)p(2]]
(7)
In the first stage, we maximize the ELBO with respect to e and 0, which corresponds to learning the VQ- VAE encoder and decoder. Following (Oord et al., 2017),
assume uniform prior for both (2/22,2 2) and p(z). Given that (=(2/z) predicts a one-hot output, the regularization term (KL divergence in Equation / will be constant. Note that the DALL-E model (Ramesh et al., 2021) uses a similar assumption to stabilize the training. Our implementation uses VQ-GAN (Esser et al, 2021) which adds GAN loss and a perceptual loss. We also include results with VQ-VAE (Oord et al. 2017) in Appendix A. 8 for reference.
In the finetuning stage, we optimize while holding 0 and fixed, which corresponds learning ViT and finetuning the discrete embeddings. This can be seen by rearranging the ELBO:
Po (0(x22))(2)
logp(a, v) 2 (xg [log (w/ar, =)] (=(()) [log
(2(2)))
where p& (y]x, 2) is the proposed ViT model. The first term denotes the likelihood for classifica- tion that is learned by minimizing the multi-class cross-entropy loss. The second term becomes constant given the fixed and ø. is important to note that the discrete embeddings are learned end-to-end in both pretraining and finetuning. Details are discussed in Appendix A.2
EXPERIMENTS
4.1 EXPERIMENTAL SETUP
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-RRes are two in- distribution datasets. ImageNet (Deng al.) 2009) is the standard validation set of ILSVRC2012. ImageNet-Real (Beyer et al., 2020) corrects the label errors in the ImageNet validation set (Northcut) ct al. 202T). which measures model's generalization on different labeling procedures.
whichhy
Seven out-of diitr---dittiiutionn (OOD) datasets are considered. Fig. 4 shows their example images. BrugeeNe--Renditiinn (Hendrycks et al. 2021a) an OOD dataset that contains renditions, such as art, cartoons, of 200 ImageNet classes. the texture vs. the shape bias by stylizing the ImageNet validation set with 79,434 paintings. The "shape ImageNet labels" are used the label. ImageNer-SKeted Sketch (Wang et al., 2019) is a "black and white" dataset constructed by querying the sketches of the 1, 1,000) ImageNet categories. Object- Net (Barbu et al. (019) is OOD test that controls the background, context, and viewpoints of the data. Following the standard practice, we evaluate the performance on the 113 overlapping Ima- geNet categories. ImageNees-V V2 (Recht al. 2019) is a new and more challenging test set collected for ImageNet. The split "matched-frequenccy is used. ImageNet-A (Hendrycks et al.) (2021)) is natural adversarial dataset that contains real-world, unmodified, natural examples that cause image classifiers to fail significantly. ImageNet- C (Hendrycks & Dietterich, 2019) evaluates the model's robustness under common corruptions. It contains 5 serveries of 15 synthetic corruptions including DDfocuus" 'Fog,' and 'JPEG compression'.
(Geirhos et al. 2019) is used to induce
Models. We verify our method on three ViT backbones: the classical ViT (Dosovitskiy et al. 2020) termed VIT Vanilla, ViT-AugReg (Steiner et al.) [021]. and MLPMixer (Tolstikhin et al. 202 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 com- pare with a CONN-based ResNet (He et al. 2016) baseline, and the Hybrid- ViT model (Dosovitskiy ct al. 2020 that replaces image patch embedding with ResNet-50 encoder.
Implementation details On ImageNet, the models are trained with three ViT model variants, i.c. Ti, S, and B, from small to large. The codebook size is K 1024, and the codebook embeds the discrete token into de = 256 dimensions. On ImageNet-22 22 K, the quantizer model is a VQGAN and is trained on ImageNet-21K only with codebook size K 8. 192. All models including ours 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
5

---

Published as a conference paper at ICLR 2022
Table 1: Model performance trained on ImageNet. All columns indicate the top-1 classification accuracy except the last column ImageNet- which indicates the mean Corruption Error (the lower the better). The bold number indicates higher accuracy than the corresponding bascline. The box highlights the best accuracy.
Out of Distribation Robustness Test Sketch
Model
ImageNet
Real
Rendition
Stylized
ObjectNet
V2
C4
ResNet-50
76.61
83.01
36.35
6.56
23.06
26.52
64.70
4,83
75.07
WT-BB Vanilla
72.47
78.69
24.56
5.94
14.34
13.44
59.32
5.23
88,72
Ours (discrete only)
73.73
79.63
34.61
8.94
23.66
20.84
60.30
6,45
74.82
VIT-Ti
58.75
66.30
21.37
6.17
10:76
12.63
46.89
3.73
86,62
+Ours (discrete only)
61.74
32.36
13.44
19.94
15.35
50.60
3.81
83.62
MM.MMixer
21.54
13.47
53.52
5.20
Ours (discrete only)
68.00
74.36
33,85
8.98
25.36
15.44
54,12
4,88
80.75
Ours
69.23
76.04
36.27
13.05
26.34
16.45
55.68
4,84
72.06
Hybrid VIT-S
75.10
81.45
34.01
7.42
26.69
24.50
62,53
6.17
69.26
VIT-S
24.50
63.00
Ours (discrete only)
72.42
80.14
42.58
18.44
31.16
23.95
60.89
7.52
66.82
Ours
77.03
83.59
39.02
14.22
28.78
26.49
64.49
11.85
$6.89
Hybrid VIT-B
74.94
80.54
33.03
7,50
25. 25.33
23.08
61.30
7,44
69
VIT-B
84.85
38.15
28.
28.71
67.34
16.92
53.51
Ours (discrete only)
78.67
84.28
48.82
22.19
39.10
30.27
66.52
14.77
$5.21
Ours
79.48
84.86
44.77
34.59
30.55
68.05
17.20
46.22
VITB (384x384)
81,63
85.06
38.23
7.58
28.07
32.36
68.57
24.01
59.01
35.72
36.01
70.33
27.19
4.2
MAIN RESULTS
Table[]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. [202]) except for the VIT-B Vanilla row, which uses the data augmentation from (Dosovitskiy et al. 2020). Our improvement is en- tirely attributed to the proposed 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 denoted by "(discrete only)". observe larger margin of improvement -10%) ) on the datasets depicting object shape: Rendition, Sketch, and worth noting that it is a challenging for a single model obtain robustness gains across all the benchmarks as different datasets may capture distinct types of data distributions.
Table 2: Model performance when pretrained on ImageNe12 IK and finetuned on ImageNet with 384x384 resolution. See the caption of Table T for description.
Out of Distribution Robustness Test Sketch
Model
ImageNet
Rendition
Stylized
ObjectNet
V2
WT-B
84.20
88.61
51.23
13.59
37.83
44.30
74.54
46.59
49.62
+ Ours (discrete only)
83.40
88.44
55.26
19.69
44.72
43.92
72.68
36.64
45.86
+ Ours
84.43
88.88
54.64
18.05
41.91
44.61
75.17
44.97
38.74
Ours (512x512)
85.07
54.27
16.02
41.92
46.62
75.55
52.64
38.97
Our approach scales more training data is available. Table pretrained on ImageNet-2IIK and finetuned on ImageNet. Training on sufficiently large datasets inherently improves robustness (Bhojanapalli 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 and Table validate our model as generic approach that is highly effective for the models pretrained on the sufficiently large ImageNet-211 dataset.
shows the performance for the models
al. 2021) but also renders further improvement
Comparison to the State-of-the-Art We compare model with the tftte-of-thh-arr results, in- cluding ViT with CutOut (DeVries & Taylor 2017) and CutMix (Yun et al.) 2019), on four datasets in Table 3 and Table 4 It is noteworthy that different from the compared approaches that are tai- lored to specific 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 AuepAugmenttAAugMe approach by

---

Published as a conference paper at ICLR 2022
Figure 4: Visualization of eight evaluation benchmarks. Each image consists of the original test image (Left) and the decoded image from the finctuned discrete embeddingg Note that the encoder and decoder trained only ImageNet 2012 data but generalize on out- of-distribution datasets. See more examples in Appendix(A.I.I
(Right).
Table 3: Comparison the Mate-offthc-are classification accuracy three ImageNet robustness datascts.
Model
Rendition
Model
Sketch
Model
Stylized Top5
ResNet 50
36.1
Huang
2020
16.1
ResNet #21K
37.2
Kand/Au Mixup
17.7
BugNet-9
1.4
DeepAngment
42.2
Xu et 2020
18.1
BagNet-17 17
2.5
DeepAngment AugMix
46.8
Misara m 2020
24.5
BagNet-33
4.2
RandAug Mixup
29.6
30.9
ResNet-50
16.4
38.2
VIT-B
22.2
VIT-B CurOut
38.1
WT-B CutOut
26.9
VIT-B CutOut
24.7
VIE-B + CuiMix
38.4
WTFB CutMix
27.5
WITB + CutMix
22.7
Our VT-B (discrete only)
48.8
VIT-B (discrete only)
39.1
ViT-B *21K
31.3
Our VIFB *21K
55.3
Our VII-B (discrete only) *21K
44.7
Our ViT-B (discrete only)
40.3
(a) ImageNet Rendition
(b) ImageNet. Sketch
Sketch
(c) Stylized ImageNet
2%. Notice that the augmentation we used-RandAAug- RandAAug-MixUp-ii 17% worse than the DeepAug- munt+AugMix on ResNets. Our performance can be further improved by another 6% when pre- trained on ImageNet-2IK. On ImageNet-Sketch, we surpass the state-of-the-art (Hermann et al.) 2020 by 8%. On Stylized- Myylized-maggeNe, 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 a smaller capacity than ViT-B. The comparison is fair for the bottom entries that are all built the same ViT-B backbone.
4.3 IN-DEPTH ANALYSIS
In this subsection, we demonstrate, both quantitatively and qualitatively, that discrete representations facilitate VIT better capture object shape and global contexts.
Quantitative results. Result 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 conflicting shape and texture information, e.g. the images of cat with elephant texture. The Stylized-ImageNett dataset is designed to measure the model's ability to recognize shapes rather than textures, and higher performance model's 's efficacy in recognizing objects by shape. While ViT outperforms CNN on the task, discrete representations yield another 10+% gain. However, the discrete token replaced with the global, continuous CNN features, such robustness gain is gone (cf. Table R). This substantiates the benefit of discrete representations in recognizing object shapes.
Table 3d is a direct proof of
Additionally, conduct the "shape texture biases" analysis following (Geirhos al.) 2019) under the OOD setting "IN-SIN" (cf. Appendix A.5) Figure. 5a compares shape bias between
Table 4: Statc-of-the-art mean Corruption Error (mCF) rate on ImageNet-C. (the smaller the hetter).
mCE
Gauss
Impol Defoc Glass Motion Zoom Snow
Fog Bright| Elastic Pixel JPEG
R5snet152 He al. 2016
69.27
72.5
73.4
76.3
66.9
81.4
65.7
74.5
70.7
67.8
62.1
51.0
67.1
75.6
68.9
65.
+Syylized aa.
66.1
77.0
62.4
73.2
GGennnt Mao el
20
61.70
59.2
60.2
62.4
60.7
70.8
59.5
69.9
64.4
63.8
58.3
48.7
61.5
70.9
55.2
60.0
(Hendrycks
202
53.60
VIT.B
$2.71
43.0
47.2
44.4
63.2
73.4
55.1
70.8
51.6
45.6
35,2
44.2
41.3
61.6
54.0
59.4
VIT.B Ours
46.22
36.9
38.6
36.0
54.6
57.4
53.4
63.2
45.
38.7
34.1
40.9
39.6
56.6
45.0
VIT-B "21K
49.62
47.0
48.3
46.0
55.7
65.6
46.7
54.0
40.0
40.6
32.2
42.3
42.0
57.1
63.1
63.6
30.5
34.5
52.9
39.5

---

Published as a conference paper at ICLR 2022
VIT
Our'
Irrage
VIT
ours
Inage
VIT
Proctin shape dockaora
(1 attentin on images
average attenton
(a) Shape vs. Texture Biases.
(b) Attention Visualization.
Figure 5: We show the fraction of shape decisions on Stylized- ImageNet in Figure (a), and attention on OOD images in Figure (b), where (i) is the attention map, and (ii) is the heat map of averaged attention from images mini-batch. See Appendix (. A..14 for details.
Stylized-ImageNet
humans and three models: ResNet-50, ViT-B, and Ours. It shows the scores on 16 categories along their average denoted by the colored horizontal line. Humans highly biased towards shape with an average fraction of 0.96 to correctly recognize an image by shape. VIT (0.42) is more shape- biased than CNN ResNet-50 (0.20), which consistent with the prior studies (Naseer al., [02I]. Adding discrete representation (0.62) greatly shrinks the gap between the ViT (0.42) and human baseline (0.96). Such behavior not observed when adding global CNN features whose average fraction of shape decisions is 0.3, lower than the baseline VIT, hence is not displayed in the figure.
Finally, we validate discrete representation's ability in modeling shape information via position em- bedding. Following (Chen al. 2021b), we compare training ViT with and without using position embedding. As position embedding is the only vehicle to equip VIT with shape information, its contribution suggests to what degree the model makes use of shape 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 al., 2021b). However, without po- sition embedding, our model accuracy drops by 29%, and degrades by a significant 36%-94% on the robustness benchmarks. This result shows that spatial information becomes crucial only when discrete representation is used, which also suggests our model relies on more shape information for recognition.
Table 5: Contribution of position embedding for robust recognition as measured by the relative performance drop when the position embedding removed. Position embedding much more erucial in model.
Out of Distribution Robusiness Test
Model
ImagoNet
Real
Rendition
Stylized
Sketch
ObjectNer
78.73
28.71
67.34
16.92
-wo, PosEmb
76.51
77.01
28.25
5,86
15.17
24.02
63,22
13.13
69.99
Relative drop (18)
2.8%
9.2%
26.0%
43.6%
47.0%
16.3%
6.1%
22.4%
30.8%
Ours VIT-B
79.48
84.86
44.77
19.38
34.59
30.55
68.05
17.20
46.22
Postamb
56.27
11.36
3.51
Relative drup (96
29.2%
30.4%
61.5%
78.6%
81.0%
62.8%
36.8%
79.6/%
94.25
Qualitative results. First, following (Dosovitskiy et al. 2020), we visualize the learned pixel em- beddings, called filters, of the ViT-B model and compare them with ours Fig. with respect to (1) randomly selected filters (the top row in Fig. 3b) and (2) the first principle components (the bottom). We visualize the filters of our default model here and include more visualization in Appendix A.1.3 The evident visual differences suggest that our learned filters capture structural patterns.
Second, we compare the attention from the classification 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 classification and captures more global contexts relative the central object.
Finally, we investigate what information preserved in discrete representations. Specifically, we reconstruct images by the VQ-GAN decoder from the discrete tokens and finetuned discrete embed- dings. Fig. 4 visualizes representative examples from ImageNet and the seven robustness bench- marks. By comparing the original and decoded images in Fig. we find the discrete representation
8

---

Published as a conference paper at ICLR 2022
preserves object shape but 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 first row. Similarly, it hallucinates the text but keeps the bag shape in the first image of the third row. In the first image of ImageNet-C, the decoding deblurs the bird image by making the shape sharper. Note that the discrete representation is learned only ImageNet, but out-of-distribution test datasets. We present more results with failure cases in Appendix A.III
generalize to the other
4.4 ABLATION STUDIES
Model capacity vs. robustness. Does 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, 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 a smaller version of quantization encoder (cf. Appendix II. With similar amount of model parameters, our model outperforms the VIT and Hybrid- ViT in robustness.
Codebook size vs. robustness. In Table from 1024 to 8 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 continuous one. Overall, Stylized-ImageNett benefits the most when the feature is highly quantized. With medium codebook size, the model strikes good balance between quantization and capacity, achieving the best overall performance, A large codebook size learned on ImageNet can hurt performance.
we conduct an ablation study on the size of the codebook,
BrOu
Redicon Acc/No
SrOur
StOur
E
VC Dcc-Errur
Jyy Sybriss
Tibric-t
Hykrist
Our
yor Styyrrd
--Our Small deDur
params (Millon))
params (Milion)
para (Millinn
params (Milloor)
Figure 6: The robustness vs. #model-paramettrs on 4 robust test set. Our models (orange) achieve better robustness with similar model capacity.
Table 6: The impact of codehook size for robustness.
CodeBoook Size
Distribution Robusiness
ImageNet
Real
Rendition
Stylized
Sketch
ObjectNet
V2
A
C.
1024
76.63
77.69
45.92
21.95
35.02
26.03
64.18
10.68
58.64
4096
77.31
78.17
47.04
21.33
35.71
27.72
64.79
11.37
57.26
35.72
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 improve robustness. Results show that concatenating global CNN representation performs no better than the baseline, which demonstrates the necessity of discrete representations.
Network designs for combining the embeddings. In Appendix we experiment with 4 different designs, including addition, concatenation, residual gating (Vo al. 2019), and cross-attention, to combine the discrete and pixel embeddings. Among them, the simple concatenation yields the best overall result.
CONCLUSION
This paper introduces a simple yet highly effective input representations for vision transformers, in which an image patch 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-tf-distributiinn ImageNet benchmarks. Our new findings connect the robustness of vision transformer discrete representation, which hints towards direction for understanding and improving robustness as well as a joint vision transformer for both image classification and generation.

---

Published as a conference paper at ICLR 2022
6 ACKNOWLEDGEMENTS
We thank the discussions and feedback from Han Zhang, Hao Wang, and Ben Poole.
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 to improve the state-of-the-art 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 improve- ments. 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 cnttps://icl:.
terms of accuracy and robustness for Vision
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 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 flow graph for our architecture in Fig. [ pseudo JAX code in Fig. [3a] and implementation details in Section 4. and Appendix A1.1 We will release our model and code assist in comparisons and support other researchers reproducing our experiments and results.
REFERENCES
Samira Abnar and Willem Zuidema. Quantifying attention flow in transformers. arXiv preprint arXiv:2005.00928. 2020.
Relja Arandjelovic and Andrew Zisserman. All about VLAD. In CVPR, 2013.
Hangbo Bao, Li Dong, and Furu Wei. Beit: Bert pre-training of image transformers. arXiv preprint 0:X15:2106.08254, 2021.
Andrei Barbu, David Mayo, Julian Alverio, William Luo, Christopher Wang, Dan Gutfreund, Josh Tenenbaum, and Boris Katz. ObjectNet: large-scale bias-controlled dataset for pushing the limits of object recognition models. In NeurlPS, 2019.
Lucas Beyer, Olivier Hénaff, Alexander Kolesnikov, Xiaohua Zhai, and Aäron van den Oord. Are we done with imagenet? arXiv preprint arXiv:2006.00 000.660779959 55, 2020.
Srinadh Bhojanapalli, Ayan Chakrabarti, Daniel Glasner, Daliang Li, Thomas Unterthiner, and An- dreas Veit. Understanding robustness of transformers for image classification. arXiv preprint arXiv:2103. 14586, 2021.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pretraining strong data augmentations. arXiv preprint arXiv:2106.0/548 2021a.
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. arXiv preprint arXiv:2104 02057, 2021b.
Xiangxiang Chu, Zhi Tian, Yuqing Wang, Bo Zhang, Haibing Ren, Xiaolin Wei, Huaxia Xia, and Chunhua Shen. Twins: Revisiting the design of spatial attention in vision transformers. arXiv preprint a0Xiv:2104. 13840, 1(2):3 2021.
Gabriella Csurka, Christopher R. Dance, Lixin Fan, Jutta Willamowski, and Cédric Bray. Visual categorization with bags of keypoints. In In Workshop Statistical Learning in Computer Vision, ECCV, pp. 1-22, 2004.
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc Le. Autoaugment: Learning augmentation policies from data. arXiv preprint / 2018.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with reduced search space. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pp. 702-703, 2020.
10

---

Published as a conference paper at ICLR 2022
Stéphane d' Assooli Hugo Touvron, Matthew Leavitt, Ari Morcos, Giulio Biroli, and Levent Sagun. Convit: Improving vision transformers with soft convolutional inductive biases. arXiv preprint arXiv:2103.10697, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei, Imagenet: hierarchical image database. In CVPR, 2009.
large-scale
Terrance Vries and Graham W Taylor. Improved regularization convolutional neural networks with cutout. arXiv preprint 1708:::7770005522 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXhv:2010.1192) 2020.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In CVPR, 2021.
Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness. In ICLR, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In CVPR, 2016.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common cor- ruptions and perturbations. In ICLR, 2019.
Dan Hendrycks, Mantas Mazeika, Saurav Kadavath, and Dawn Song. Using self-supervised learning can improve model robustness and uncertainty. arXiv preprint arXiv: 12340, 2019.
Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, Dawn Song. Jacob Steinhardt, and Justin Gilmer. The many faces of robustness: A critical analysis of -il-t--diittibution generalization. In ICCV, 2021a.
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. CVPR, 2021b.
Katherine Hermann, Ting Chen, and Simon Korbllith. The origins and prevalence of texture bias in convolutional neural networks. In NeurIPS, 2020.
Zeyi Huang, Haohan Wang, Eric P Xing, and Dong Huang. Self-challenging improves cross-domain generalization. In ECCV, 2020a.
Zeyi Huang, Haohan Wang, Eric P. Xing, and Dong Huang. Self-challenging improves cross-domain generalization. In ECCV, 2020b.
Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In ECCV, 2016.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Ze Liu, Yutong Lin. Yue Cao. Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2/03.1 /4030, 2021.
Chengzhi Mao, Augustine Cha, Amogh Gupta, Hao Wang. Junfeng Yang, and Carl Vondrick. Gen- erative interventions for causal learning. In CVPR, 2021a.
Xiaofeng Mao, Gege Qi, Yuefeng Chen, Xiaodan Li, Ranjie Duan, Shaokai Ye, Yuan He, and Hui Xue. Towards robust vision transformer. arXiv preprint arXiv: 100v:2105.07990 07966,, 2021b.
11

---

Published as a conference paper at ICLR 2022
Shlok Mishra, Anshul Shah, Ankan Bansal, Jonghyun Choi, Abhinav Shrivastava, Abhishek Sharma, and David Jacobs. Learning visual representations for transfer learning by suppress- ing texture. arXiv preprint arXiv:2011.01901, 2020.
Muzammal Naseer, Kanchana Ranasinghe, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, and Ming-Hsuan Yang. Intriguing properties of vision transformers. arXiv preprint arXiv:2105.10497, 2021.
Curtis Northcutt, Lu Jiang, and Isaac Chuang. Confident learning: Estimating uncertainty in dataset labels. Journal of Artificial Intelligence Research, 10:1373-1411, 2021.
Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learn- ing. arXiv preprint arXiv:1711.00931 2017.
Sayak Paul and Pin-Yu Chen. Vision transformers are robust learners. arXiv:2105.07580 075811 2021.
arXiv preprint
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alee Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. arXiv preprint arXiv:2102.12092, 2021.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In ICML, 2019.
Rulin Shao, Zhouxing Shi, Jinfeng Yi, Pin-Yu Chen, and Cho-Jui Hsieh. On the adversarial robust- ness of visual transformers. arXiv preprint 1507:2 2021.
J. Sivic and A. Zisserman. Video Google: a text retrieval approach object matching in videos. In CVPR, 2003.
Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv preprint ar 106. 10270, 2021.
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung. Daniel Keysers, Jakob Uszkoreit, Mario Lucic, et al. Mlp-mixer: An all-mlp architecture for vision. arXiv preprint arXiv:2105.01601, 2021.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In ICML, 2021.
Arash Vahdat, Evgeny Andriyash, and William G Macready. Dvac#: Discrete variational autoen- coders with relaxed boltzmann priors. arXiv preprint arXiv:1805.07445. 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz Kaiser, and Illia Polosukhin. Attention all you need. In NewrlPS, 2017.
Nam Vo, Lu Jiang, Chen Sun, Kevin Murphy, Li-Jia Li, Li Fei-Fei, and James Hays. Composing text and image for image retrieval-an empirical odyssey. CVPR, 2019.
Haohan Wang, songwei Ge, Zachary Lipton, and Eric P Xing. Learning robust global representa- tions by penalizing local predictive power. NeurIPS, 2019.
Wenhai Wang. Enze Xie, Xiang Li, Deng- Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. arXiv preprint arXiv:2102. 12 /22, 2021.
Zhenlin Xu, Deyi Liu, Junlin Yang, Colin Raffel, and Mare Niethammer. Robust and generalizable visual representation learning via random convolutions. arXiv preprint arXiv:2007. 13003, 2020.
Sangdoo Yun, Dongyoon Han, Scong Joon Oh, Sanghyuk Chun, Junsuk Choe. and Youngjoon Yoo.
Cutmix: Regularization strategy to train strong classifiers with localizable features. Proceed- ings of the IEEE/CVF International Conference Computer Vision, pp. 6023-6032. 2019.
12

---

Published as a conference paper at ICLR 2022
Amir R. Zamir, Alexander Sax, Nikhil Cheerla, Rohan Suri, Zhangjie Cao, Jitendra Malik, and Leonidas J. Guibas. Robust learning through cross-task consistency. In CVPR, 2020.
Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empiri- cal risk minimization. In ICLR, 2018.
13

---

Published as a conference paper at ICLR 2022
A APPENDIX
The Appendix is organized as follows. Appendix extends the discussion on learning objective. Appendix A.I presents more figures visualize the attention, the decoded images, and the decoding failure cases. Appendix A. compares designs for combing the pixel and discrete embeddings. In the end, Appendix A.11 discusses the implementation details.
Inageset
Jellyties
Blank
Jak-ol-Jasters
Rendition
sketch
Imagetet
Imagetet-C
vallin (mstras))
www. L-P seprrstte
(Caeation
Caberr
DECEED
HEG S BEREE
I-
Figure 7: Visualization eight evaluation benchmarks. Each image consists of the original test image (Left) and the decoded image (Right) from the finetuned discrete embeddings. The encoder and decoder are trained only on ImageNet 2012 data but generalize on out-of-distribution datasets.
A.1 VISUALIZATION
A.1.1 RECONSTRUCTED IMAGES FROM DISCRETE REPRESENTATION
Fig. 7 shows more examples of reconstructed images decoded from the finetuned discrete embed- ding. 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 bench- marks, e.g., the objects in ImageNet-A learned only on the ImageNet 2012 training dataset (the first row in Fig. can generalize to the other -itt---distrbuution test datasets.
much smaller. Note that while the discrete representation
14

---

Published as a conference paper at ICLR 2022
Blarh
Iobach oon
16
perallel complex Lextures
small objectx parts
Figure 8: Visualization of failure cases for the decoded images. Each image consists of the original test image (Left) and the decoded image (Right) from the finetuned discrete embeddings.
A.1.2 LIMITATIONS DISCRETE REPRESENTATION
We find four failure cases that the discrete representation are unable to capture: text, human faces, parallel lines, and small objects parts. The example images are illustrated in Fig. 8 Without prior knowledge, the discrete representation has difficulty to reconstruct text in image. On the other hand, it is interesting find the decoder can model animal faces but not human faces. This may be result of lacking facial images in the training data of ImageNet and ImageNet-21IK.
The serious problem found for recognition is its failure capture small objects parts. Though this can be beneficial sometimes, e.g. force the model to recognize the tire without using the "BMW" logo (see the first image in the last nizing the small whales in the last image. As result, our proposed model leverages the power of both embeddings promote the interaction between modeling global and local features.
of Fig. (), it can cause problems in many cases, e.g. recog-
A.1.3 FILTERS
In Fig. 3b of the main paper, illustrate the visual comparison of the learned pixel embeddings between the standard ViT and our best-performing model. In this subsection, we extend the visual- ization our models with varying number of pixel dimensions when concatenating the discrete and pixel embedding, where their classification performances are compared in Table [3] As the dimension pixel embedding grows, the filters started pick up more high-frequency signals. which
However, the default setting in the paper using only 32 pixel dimensions seems induce an inductive bias by limiting the budget pixel representation, turns out to be the best performing model.
1.1.4 ATTENTIONS
We visualize the attention following the steps in (Dosovitskiy et al. 2020). To be specific, com- pute maps of the attention from the output token to the input space, we use Attention Rollout (Abnar & Zuidema, 2020). Briefly, average attention weights of VIT-B and ViT-B+Ours across all heads and then recursively multiply the weight matrices all layers. Both VIT and ViT+Ours are trained on ImageNet.
In Fig. [I Fig. 12 and Fig. 1 we visualize the attention of individual images in the first mini-batch of each test set. As shown, model attends image regions that are semantically relevant for
15

---

Published as a conference paper at ICLR 2022
Tendon be betton be
adlction crickling
nor codaskaa Tlop
enieeday ben proppe conuetent
aod medeciii Steo
not
ViT (all pixels)
Ours (pixel_dim=32)
Ours (pixel_dim=64)
Ours (pixel_ dim=128)
Figure 9: Comparing visualized pixel embeddings of the ViT and our model. The row shows the randomly selected filters and the Bottom shows the first principal components. Our models with varying pixel dimensions are shown, where their classification performances are compared in Ours (pixeL.dim=22) works the best and is used as the default model the main paper.
Table 13
classification and captures more global contexts relative to the central object. This can be better from Fig. [0 where the heat map averages the attention values of all images in the first mini- batch of each test set. As found 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 exhibit the same defect.
18
(a) ViT (ImageNet)
(b) VIT (ImageNet-R)
(c) ViT (Stylized)
(d) VIT (ObjectNet)
(c) Ours (ImageNet)
(f) Ours (ImageNet-R)
(g) Ours (Stylized)
(h) Ours (ObjectNet)
Figure 10: Comparison of average attention of the ViT (top row) and the proposed model (bottom
row) on four validation datasets: ImageNet 2012, ImageNet-R, Jyyizee-Imaageett, and ObjectNet. The heat map averages the attention values of images in the first mini-batch of each test set. The results show attention capture more global context relative to the central object.
A.2 DETAILS ON LEARNING OBJECTIVE
Let x denote an image with the class label y, and z represent the discrete tokens z E Z. For notational convenience, we use single random variable represent the discrete latent variables whereas our implementation actually represents an image as flattened I-D sequence following the raster scan.
(2(xz) denotes the VQ- VAE encoder, parameterized by d. to represent an image x into the discrete token
16

---

Published as a conference paper at ICLR 2022
(2(22) is the decoder that models the distribution over the RGB image generated from discrete tokens.
Pv (y/ [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 using the factoriza- tion p(r, y, =) = Pa (ar (pp(2l =)p(z). Our overall procedure can be viewed as maximizing the evidence lower bound (EI (ELBO) the joint likelihood, which yields:
p(2 v)
Dkulgs(2|x), p(= v)) = [2(22)z) log
(9)
(0(2(z)
p( v, p( (00
Σ (z(z))) log
(10)
Since the f-divergence is non-negative, we have:
(=(2)) log p(r, y, (6(2(2)
+ log logp(z, v\ > (0(2|x))] 2 0
(11)
logp(z, v) z Σ 2o (z((x) ((alax))0o log p(a, 2, 4al (z(
(12)
Using the factorization, have:
logp(z,)) ga)(=/ [log Po(ul z z)po (=)=))(()))
(13)
(0(2(2)
Equation 1 yields the ELBO:
logp(æ, v) 2 Bgg() Pog + E (=|x) flog Pe (yla,, z)p(z)
(14)
(=(2(a))
2 B pa (x|=)] (x (02x,,))p ()(()]
(15)
In the first stage, maximize the ELBO with respect & and which corresponds learning the VQ-VAE encoder and decoder. We assume a uniform prior for both (vl(,,r; and p(=). Given that (2(1)) predicts 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 deterministie. We note similar assumption 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 pv (y)u,,
is in the DALL-E model (Ramesh et al., 202T) used
In addition to the reconstruction loss, the VQ-VAE's training objective adds a dictionary learning loss and commitment loss, calculated from:
LvQ-VAE log ((x)2,(())
(5)(=(())) w1| Bl]z sg|u]]l.
(16)
where ze (ar) and n, (z) represent the encoder output and the decoder input, respectively. v de- notes the discrete embedding for the image x. sg(.) is the stop gradient function. The training used straight- ttrough estimation which just copies gradients from (x) to (z). Notice that our implementation uses VQ-GAN Esser et al. 202T) which adds an additional GAN loss:
LVo-GAN = log (r) (x)) + (x)) w1 3 ( sg[wi]| + ALGAN,
(17)
where B 0.25 and À 0.1 are used to balance different loss terms. LVQ.GAA also has perceptual loss (Johnson et al. 2016).
In the second stage of finetuning, optimize while holding and 0 fixed, which corresponds to learning a vision transformer for classification and finetuning the discrete embedding V. By rearranging Equation 14 we have:
logp(a, v) 2 Egg(s ( 2)] (=)o) [log po(x)2)p(2)
(18)
(2(2)))
17

---

Published as a conference paper at ICLR 2022
where Pe (y/ar, 2) is the proposed ViT model illustrated in Fig. of the main paper. The first term denotes the likelihood for classification that is learned by minimizing the multi-class cross-entrops loss. The second term becomes constant given the fixed 0 and Φ.
In the finetuning stage, two sets of parameters are updated: the vision transformer &b is learned from scratch and the discrete embedding V is finetuned. We fix the encoder and stop the gradient prop- agating 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- finnttnniinnuunggggg strategy with the joint learning strategy in which all parameters are optimized simultaneously. We find the joint learning significantly underperforms the proposed learning strategy.
A.3 COMPARING TO DATA AUGMENTATIONS
Since our model only changes the token representation in the VIT architecture, it is conceptually different and complementary to data augmentation. Nevertheless, this subsection compares our method with seven combinations of recent data augmentation strategies, including CutOut (De Taylor 2017), CutMix (Yun et al. 2019), RandAug (Cubuk et al. 2020), and Mixup (Zhang al,, 2018). As discussed following (Steiner et al. 2021), the VIT bascline and model both use the default augmentation (RandAug Mixup).
As shown in Table 7 our method achieves the best performance on the robustness benchmarks while is marginally worse than the best baseline with state- of-th--attt data augmentation strategy (CutMix RandAug) on the ImageNet in-distribution validation set. Incorporating state-of-the tht-hrr-r data augmentation in our model is worthy of future research.
Table 7: Model performance trained on ImageNet. We compared with several recent data augmentations, including CutOut (DeVries & Taylor (017)) and CutMix Yun et al. (2019). on the ViT-B model. All columns indicate the top-1 classification accuracy except the last column ImageNet-C which indicates the mean Cor- ruption Error (lower the better). The bold number indicates higher accuracy than the corresponding baseline. The box highlights best accuracy. While the state-of-the-ee data augmentation (CurMix+ (CunMix--andAug) can improve Imagenet in-distribution accuracy, they are worse than ours on the OOD test accuracy.
Out of Distribution Rohustness Test
Model
ImageNet
Real
Rendition
Stylized
Sketch
ObjectNet
V2
A
C L
Augmentation Ours (discrete only)
72.47
13.44
59.32
5,23
73.73
79.63
34.61
8.94
23,66
20.84
60.30
6,45
74.82
VIT-B CutQut
72.27
77.91
25.74
7.89
16.50
17.82
58.05
10.23
69.01
VIT-B CurMix
75.51
80.53
28.45
7,89
17.15
21.62
62.36
14.72
64.07
VIE-B CutOut+ Mixup
16.28
56.84
8.33
V CuiMix Mixup
71.66
76.82
23.46
4.92
13.91
16.99
$6.98
10.23
80.54
CutQut RandAug
79.07
28.32
66.88
15.77
54,52
VIT-B CulMix RandAug
79.63
85.24
38.34
10.94
27.46
29.80
68.05
18.80
$1.63
VIT-B + RandAug + Mixup
78.73
84.85
38.15
10.39
28.60
28.71
67.34
16.92
53.51
Our VT..B RandAug Mixup
79.48
84.86
44.77
19.38
34.59
30.55
68.05
17.20
46.22
A.4 THE IMPORTANCE OF DISCRETE INFORMATION
This subsection compares to new baseline that directly concatenates the global feature extracted by CNN to the input token of VIT. The only difference between this baseline and our model is that ours concatenates discrete embeddings rather than global CNN features. In Table g we extensively search the concatenation dimension of the CNN features in [64, 128, 256, 384, 512, 640]. How- ever, none of the models improve robustness. Our results show that concatenating the global CNN feature performs no better than the VIT baseline and significantly worse than concatenating discrete representations, demonstrating the necessity discrete representations for robustness.
A.5 TEXTURE VS. SHAPE STUDY
In Fig. of of the main paper, we performs the "shape vs. texture biases" analysis following (Geirhos) al. 2019 (see Figure in their paper). It uses the score "fraction of shape decision" quantify
18

---

Published as a conference paper at ICLR 2022
Table 8: replace our discrete token representation with global continuous (GC) features from CNN model that has the same CNN architecture as our VQGAN encoder. We denote the dimension for the global features Dim. We vary the dimension of the pixel token to concatenate with the global continuous CNN features. Simply concatenating the global feature extracted from CNN does not improve robustness.
GF.
Out of Distriburion Robustness Test Sketch
Model
Dim
ImageNet
Real
Rendition
Stylized
ObjectNet
V2
A
C 4
GC CNN + VIT
640
78.48
83.27
34.75
9.
24.85
28.37
65.65
16.13
58.71
GC CNN VIT
512
78.59
83.24
34.31
25.50
27.50
65.35
16.48
57.66
GC CNN + VII
384
78.51
83:44
34.56
10:47
25.14
28.35
65.86
15.97
57.64
GC VIT
58.35
GC CNN VIT
128
78.29
82.96
34.32
9.22
24.90
26.91
65,28
15.73
58.71
GC CNN + VIT
78.34
82.98
34.35
9.19
25.06
27.07
65.48
15.40
58.58
VIT
78.73
84.85
38.15
10.39
28.60
28.71
67.34
16.92
53.51
Ours: (discrete only)
78.67
84.28
48.82
22.19
39.10
30.27
66.52
14.77
55,21
Ours
79.48
84.86
44.77
19.38
34.59
30.55
68.05
17.20
46.22
model's 's shape-biat between 0 and 1. To compute the score, we follow the steps on their Githut as follows: evaluate the model on all 1280 images Stylized ImageNet; 2) map the model decision to 16 classes; 3) exclude images without cue conflict; 4) take the subset of "correctly" classi- fied images (either shape or texture category correctly predicted); (5) compute "shape bias" as the fraction between correct shape decisions and (correct shape decisions correct texture decisions).
Fig. 5a in the main paper presents the scores on categories along their average denoted by the colored horizontal line. We compute the scores for three models trained on ImageNet, i.e. ResNet- 50, ViT-B/16, and ours, and quote human scores from (Geirhos et al. 2019). We also compute the score for the concatenating global CNN features baseline which is about 0.3. Note that the OOD generalization setting "IN-SIN" (Geirhos et al., 2019) is used, where the model trained only on the ImageNet (IN) and tested the Stylized ImageNet (SIN) images with confficting shape and texture cue.
A.6 THE EFFECT OF CODEBOOK ENCODER CAPACITY ON ROBUSTNESS
Table 9: The impact of codebook encoder capacity for robustness.
CodeBook
Out of Distribution Robustness Test
Size
ImageNet
Real
Rendition
Stylized
Sketch
ObjectNot
C
Small VQ Encoder
76.63
77.69
45.92
21.95
35.02
26.03
64.18
10.68
58.64
Encoder
78.67
84.28
48.82
22.19
39.10
30.27
14.77
55.21
We analyze how the capacity of the discrete encoder affects the model's robustness. We train a small encoder and standard encoder and compare their performance. We describe the VQ encoder's architecture configuration in Table 17] where the small encoder has only 13% of the #parameters of the standard VQ encoder. We run experiments Table Our results show that using the standard VQ-encoder is better but the small VQ-encoder comparatively yields reasonable results considering the reduced model capacity.
A.7 THE EFFECT OF PRETRAINING CODEBOOK WITH MORE DATA ON ROBUSNTESS
Table 10: The impact of pretraining codebook on sufficiently large data for robustness. The code- book is pretrained on the standard ImageNet ILSVRC 2012 or the ImageNet21K dataset (about 1 larger). Using the codebook, we train ViT models on ImageNet21K. KK.
Codebook
Out of Distribution Robustnoss Test
Pretrained
ImageNet
Real
Rendition
Stylized
ObjectNet
V2
ImageNet
83.32
51.56
14.77
41.10
42.50
72.57
33.51
55.31
ImageNNe21K
83.40
88.44
55.26
19.69
44.72
43.92
72.68
36.64
45.86
We analyze the effect of pre-training codebook on sufficiently large data. We use the same model
architecture, but pretrain two codebooks (and encoders) the ImageNet ILSVRC 2012 and the
https://githhb..co http://c...eeeeeeeeeeeettrreeeee httes/////tttoooo////////hhhh/ee
shape
19

---

Published as a conference paper at ICLR 2022
ImageNet2IK, (about 11x larger), respectively. Using the pre-trained codebook, train ViTs on ImageNet21K and compare the results in Table where it shows pretraining on large data improves the robustness.
A.8 ADDITIONAL MODELS FOR LEARNING DISCRETE REPRESENTATION
The results show that the ability of discrete representations to improve VIT's robustness is general which is not limited to specific VQ-GAN models.
Table 11: Model architecture for discrete representation. We bold the numbers if they improve robustness over the baseline ViT model. The numbers are boxed where VQ-VAE further improves robustness than VQ-GAN under apple-to-apple comparison.
Codebook
Out of Distribution Robustness Test
Pretrained Model
ImageNet
Real
Rendition
Stylized
Sketch
ObjectNet
V2
C
Baseline (No discrete)
78.73
84.85
38.15
10.39
28.60
28.71
67.34
16.92
53.51
VQ-VAE VAE (Discrete Only)
78.36
84.35
46.22
23.36
35.17
29.34
66.24
13.61
52.20
VQ-GAN (Discrete Only)
78.67
84,28
48.82
22.19
39.10
30.27
66.52
14.77
55.21
VQ-VAE
78.51
83.68
41.54
17.50
30.91
27.43
65.74
15.61
50.47
VQ-GAN
A.9 COMPARING DESIGNS FOR COMBINING EMBEDDINGS
In this subsection, we compare different designs combine the pixel and discrete embeddings. See the combination operation in Fig- T and in Equation the main paper. The results are used verify design of using simple concatenation presented in the paper. Four designs are considered and discussed below. Their performances are compared Table 12
Table 12:
The accuracy using different combination methods. While directly adding pixel token to discrete token improves the most ImageNet and ImageNet-A accuracy, concatenation method improves the overall robustness.
Combining
Out of Distribution Robustness Test ObjectNet
Method
ImageNNet
Real
Rendition
Stylized
Sketch
V2
A
C.
Addition
79.67
84,89
40.68
13.59
30.89
29.18
67.19
17.68
50.85
Concatenation
34.59
30.55
68.05
Residual Gating
74.59
79.66
41,06
17.03
29.56
25.34
62.37
9,33
59.88
Cros-Atiention 79.18
84.53
43.75
18.67
34.19
29.71
66.73
15.81
50.06
Addition. first 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
د X (V + xx E2
(19)
where Vz, and x,, E represent the discrete and pixel embeddings. respectively, and E₂ 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 dimension of for the pixel embedding, dimension of 256 for the discrete embedding, thus we pad 0 to the vector if the input dimension the ViT is higher than 288.
f(V N. x V2.; E;
(20)
where 55"" indicates the concatenation operation.
Residual Gating. As pixel embeddings may contain nuisances, inspired by (Vo et al. 2019). we leam gate from both embeddings, and then multiply it with the pixel embeddings filter out details that are unimportant to classification. Specifically, we calculate the gate by 2-layer MLP, and apply softmax the output from the last layer.
20

---

Published as a conference paper at ICLR 2022
= Softmax(MLP( 2. X))
(21)
Then the pixel embeddings are gated by:
P= P=Goxp
(22)
Then we concatenate the gated embedding with the discrete embedding:
J(V2, S.. x, E) Vz,ii P;
(23)
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).
A = MCA(V, K2z X, E, Xp Ep
(24)
We then concatenate the attended feature output with the discrete embedding. We also pad 0 if the VIT requires larger input dimensionality.
f(V ne X₂
Vz.i Rai A
(25)
A.9.1 THE ROBUSTNESS EFFECT OF USING CONTINUOUS REPRESENTATION WITH DISCRETE REPRESENITION.
The above study shows that concatenation the most effective way to combine the discrete and continuous representations. As we fixed the dimensionality for the discrete token to be 256, we can study the optimal ratio between continuous and discrete representations by changing the dimension- ality on the pixel token.
The robust performance under different pixel token dimension in concatenation combin- ing 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 dis- tribution robustness performance. The more continuous representations the model use, the lower robustness the out-of-diitribution set that depicting object shape, but also aciheves higher curacy 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 choice to use dimension 32, which balance the trade-off and achieves the best overall robustness.
Table 13:
Pixel Token
Out Distribution Rohustness Test Sketch
Dimension
ImageNet
Real
Rendition
Stylized
ObjectNet
V2
A
C
78.67
84.28
22.19
39.10
30.27
66.52
14.77
55.21
8
79.05
84,50
46.08
21.17
34.67
30.20
66.71
15.15
48.42
16
79.45
84.84
46.00
20.47
34.88
29.42
67.09
15.81
46.91
32
79.48
84.86
44.77
19.38
34.59
30.55
68.05
17.20
46.22
64
79.71
84.98
43.97
18.75
33.89
29.91
67.98
17.47
46.10
128
80.04
85.07
41,96
15.47
31.65
30.05
68.03
17.80
49.58
All Pixel
78.73
84,85
38.15
10.39
28.60
28.71
67.34
16.92
53.51
A. THE IMPORTANCE OF SPATIAL INFORMATION FOR DIFFERNT ARCHITECTURES
In addition to the analysis the main paper, we also investigate whether 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 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.
21

---

Published as a conference paper at ICLR 2022
Table 14: Contribution of position embedding robust recognition as measured by the relative performance drop when the position embedding is removed. We experiment on VIT. Ours, Hybrid- ViT. and Ours with discrete token only. Note that Hybrid uses the continuous, global feature from ResNet-50 CNN as the input token. A larger relative drop indicates the model relies spatial information more. Adding discrete token input representation drives position information and spatial structure more crucial in model.
Out of Distribution Robustness Test ObjectNet
ImageNet
Rendition
Stylized
Sketch
C4
VITBB
78.73
84.85
38.15
10.39
28.6
28.71
67.34
16.92
53.51
-/o. PosEmb
76.51
77.01
28.25
5.86
15.17
24.02
63.22
13.13
69.99
Relative drop((8)
2.8%
9.2%
26.0%
43.6%
47.0%
16.3%
6.1%
22.4%
30.8%
Hybrid VIT-B
74.94
80.54
33.03
7.5
25.33
23.08
61.30
7.44
69.61
Hybrid-ViT-B- w/o. PosEmb
75.13
33.11
7.5
25.15
22.78
61.75
68.51
relative drop (8)
-0.3
-0.1
-0.2
0
0.7
3.0
0.7
2.1
1,6
(discrete only)
48.82
22.19
30.27
66.52
55.21
Ours (discrete only) w/o. PosEmnh
16.34
5.09
1.95
2.09
1.20
1.75
11.75
1.29
124
Relative Drop (%)
96.0
90.5
96.9
82.3
124.6
Ours
79.48
84.86
44.77
19.38
34.59
30.55
68.05
17.20
46.22
Ours- w/o. PosEmb
56.27
59.06
17,24
4.14
6.57
11.36
42.98
3.51
89.76
Relative (%) (%)
29.2%
30.4%
61.5%
78.6%
81.0%
62.8%
36.8%
79.6%
94.2%
Table 15: Model configuration for the Transformer Encoder.
Model
Layers
Hidden size
MLP size
Heads
#Params
Tiny (Ti)
12
192
768
3
5.8M
Small (S)
384
1546
22.2M
Base (B)
12
768
3072
12
86M
Large (L)
24
1024
4096
16
307M
A.11 IMPLEMENTATION DETAILS
A.11.1 ARCHITECTURE
We use three variants of ViT transformer encoder backbone in our experiments. We show the config- uration details and the number of parameters in Table [5] We also use hybrid model with ResNet50 as the tokenization backbone. We show the configuration in Table T6
We use B-16 for the MLPMixer Tolstikhin et al. 2021)
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. (021). For Small VQ Encoder, we decrease both the number the resisual blocks and the number for filter channel by half, which is lightweight model that decrease the model size by around 8 times. We show the configuration detail Table 17
A.12 TRAINING
Implementation details We implement our model in Jax and optimize all models with Adam (Kingma & Ba 2014). Unless specified otherwise, the input images are resized to 224x224, trained with batch size of 4,096, with a weight decay of 0.1. We use a linear learning rate warm- up and cosine decay. On ImageNet, the models are trained for 300 epochs using three ViT model variants, i... Ti, S, and B, from small to large. The VQ-GAN model also trained on ImageNet for 100K steps using batch size of 256. The codebook size is K 1024, and the codebook embeds the discrete token into d, = 256 dimensions. For discrete ViT-B, we use average pooling on the final features. On ImageNeet22K., we train 90 epochs. We finetune with Adam and Ic-4 learning rate, batch size 512, for 20K iterations. The quantizer model a VQ- ImageNet-21K only with codebook size K 8, 192. All models including
higher resolution ImageNet,
GAN and is trained 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.
ImageNet Training We train our model with the Adam (Kiggma & Ba 2014) optimizer. Our batch- size is 4096 and trained on 64 TPU cores. We start from a learning rate of 0.001 and train 300 epoch.
22

---

Published as a conference paper at ICLR 2022
Table 16: Model configuration for ResNettVViT hybrid models.
Model
Resblocks
Patch-size
#Params
Hybrid-S
[3,4,6.3]
46.1
Hybrid-B
[3,4,6,3]
111
Table 17: Model configuration for the VQ encoders.
Model
Resblocks
Filter Channel
Embedding dimension
#Params
Small VQ Encoder
[1,1,111,1]
64
256
3.0M
VQ Encoder
[2,2,2,2,2]
128
256
23.7M
We use linear warm up for 10k iterations and then cosine annealing for the learning rate schedule. We use 224 x 224 resolution for the input image. We conduct experiment on three variants of the VIT model, S, and B, from small to large. addition, we also experiment on the newly released MLPMixer, which also uses image patch as the token embeddings. We denote our approach as us- ing both representations, and we also run 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, concatenate both representations, and pad zeros if additional dimension required for the transformer. We use the same augmentation and RandAug with hyper-parameter (2,15)
(Steiner et al. 2021), where and mixup with a = 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 en- coder only with a codebook size of 1,024. The codebook integer is embedded into a 256 dimension vector. For the ViT-B discrete, we find training 800 epoch instead of 300 epoch can give another 0.1-1% gain over the test set, while training others for longer results in decreased performance.
model regularization
22K Training Our models use learning rate of 0.001, weight decay of 0.03, and train 90 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 use the same augmentation as (Steiner et al. 2021). For the quantizer model, we use VQGAN that is trained on unlabeled ImageNet-21K only. We use codebook size 192, and the codebook integer embedded into a 256 dimension vector. To evaluate on the standard test set, we further finetune 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 larger resolution 384 x 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. adversarial loss with weight 0.1. We use L1 gradient penalty of 10 during the optimization. The model is trained with resolution of 256 x 256. The model can scale to different image resolution without finetuning.
23

---

Published as a conference paper at ICLR 2022
VTT-B/16
(b) Ours ViT-B/16
Figure 11: Attention comparison of the ViT and the proposed model on ImageNet 2012.
24

---

Published as a conference paper at ICLR 2022
VTT-B/16
(b) Ours ViT-B/16
Figure 12: Attention comparison of the ViT and the proposed model on ImageNet-R.
25

---

Published as a conference paper at ICLR 2022
(a) ViT-B/16 (Sketch)
(b) Ours ViT-B/16 (Sketch)
(c) VIT-B/16 (ObjectNet)
(d) Ours VIT-B/16 (ObjectNet)
Figure 13: Attention comparison of the VIT and the proposed model on ImageNet Sketch and ObjectNet.
26


Reviews:
[{"review_id": "PkSlYmJ_N43", "reviewer": "Reviewer_rENv", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Authors state that ViTs are overly reliant on local features and fail to make adequate use of global context. To me, this is counter-intuitive and is not supported by concurrent recent findings. \n\nInspired by https://openreview.net/forum?id=Bygh9j09KX that shows CNNs are biased towards texture; [1] conducts an extensive study, and demonstrates that ViTs are less biased towards texture, compared with their CNN counterparts. The authors show that ViTs have shape-bias, comparable to humans. See Fig5 of the paper. \n\n[1] “Intriguing Properties of ViTs” Neurips’21 https://arxiv.org/pdf/2105.10497.pdf\n\nTo make a claim on texture vs Shape bias, the authors must conduct a principled study as in “ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness”  https://openreview.net/forum?id=Bygh9j09KX\n\nIn recent works, ViTs demonstrate better out-of-domain generalization, as demonstrated in extensive experiments conducted in Fig14 of https://arxiv.org/pdf/2105.10497.pdf  and Fig5 of  https://arxiv.org/abs/2106.09785\n", "summary_of_the_paper": "Authors propose discrete tokens (instead of the standard continuous pixel-values projected and fed as tokens to ViTs), to enhance the shape learning capability of ViTs, and thus make them robust against out-of-distribution data. The underlying assumption is that discrete tokens preserve the global structure and lose local details, and are therefore better for robustness.\n", "summary_of_the_review": "In the paper, the experiments are conducted primarily on ImageNet robustness datasets. It will be interesting to see, out-of-domain generalization of the features learned by the proposed approach e.g., ImageNet pre-trained features on fine-grained datasets (Flowers, Pets, Fungi, Sketches, Birds, etc).\n\nI am not entirely convinced by the shape-vs-texture bias claims made in the paper, due to lack of principled study as in  https://openreview.net/forum?id=Bygh9j09KX\n\nAlso, it will be interesting to see if the proposed discrete tokens enhance ViTs robustness against common corruptions (e.g., artifacts introduced by rain, haze etc), and adversarial perturbations. The considered robustness analysis is too restrictive in my opinion.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "7vldYHiSLhP", "reviewer": "Reviewer_YfcP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pro:\n- How to improve the robustness of ViTs is rarely studied in the literature. The authors proposed a novel approach to improve the ViT's robustness, and made a comprehensive study on it.\n- The paper is well written.\n- Results showed the effectiveness of the proposed approach on improving the robustness of ViTs.\n\nCon:\n- The authors claimed that ViTs' fail to make adequate use of global context. I am not entirely convinced by this claims. From results of the paper, I believe that the proposed method can improve the robustness of ViTs, but it can not support the claim. The experiment of discarding position embedding can only prove that positional information contribute small information for the classification. It risk to claim that this observation shows that ViTs do not fully use the global context information. ", "summary_of_the_paper": "The authors have an observation that ViTs trained on ImageNet heavily depend on local features, but fail to use the global features (shape or structure). To address this issue and improve the robustness of ViTs, the authors proposed to replace the linear embedding layer by the a vector-quantized encoder. The authors claims that it can push the ViTs to learn the global information by this replacement. Experiments are conducted on ImageNet and other ImageNet variant datasets. Results show that the proposed method can improve ViTs' robustness on various benchmarks.", "summary_of_the_review": "The authors proposed an novel approach to improve the ViT's robustness. Comprehensive experiments are conducted on various robustness benchmark. Results showed the proposed method can improve the robustness of ViTs. My main concern is that the major claim of the paper, that ViT's fail to make adequate use of global information, is not supported by the experiments. I'd like to accept the paper if the authors can make the claim more supportive. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "5e96LD-A2xA", "reviewer": "Reviewer_mNoS", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1. The proposed method share a similar idea with BEiT using discrete VAE to \"tokenize\" image while apply the idea to improve ViT in different ways. To the best of my knowledge, concatenating discrete representation with the continuous representations for ViT has not been studied in previous work.\n2. The proposed method is technically sound and the writing is properly formatted, well organized and easy-to-follow.\n3. The proposed method is simple yet effective as shown in the experiments. It is a general approach to improve the robustness of ViT. The experiments show that the approach consistently outperforms baseline methods on the benchmarks and achieve SoTA on ImageNet-Rendition, Stylized-ImageNet and ImageNet-Sketch.\n4. The authors present the observation that discrete representations derived from VQ-VAE preserve shape and structure object information which can be integrated into ViT to improve the robustness. They also conduct insightful experiments to demonstrate this observation qualitatively. The quantitative evaluation on the ImageNet benchmarks seems to support this observation.\n5. The authors present extensive ablation studies as listed below. The results are convincing.\n - As the proposed method capture shape information, positional embeddings play important role for their proposed models comparing to plain ViT\n - the improvement does not trivially come from using larger models\n - analysis on codebook size is reasonable and outcome is consistent to ones would expect.\n\n\nWeaknesses:\n1. In Figure 3(b). the visualization shows that discrete representations promote ViT to learn structured filters. Why the proposed modification is able to achieve this behavior is unclear. According to the rest of the paper, the discrete representation itself carries structure information to improve the robustness. Ideally the filters should pick up the complementary information that are more related to style and texture.\n2. There are high attention scores on the corners for ViT In Figure 4 (ii). Why is this the case?\n3. The description of the optimization on \\phi in equation (7) is not precise. A term E_q(z|x)[logp(y|x,z)] can be factored out from the KL term and the gradient of this term w.r.t. \\phi seems not considered in learning VQ_VAE. As a result, the optimization is not exactly maximizing the ELBO described in the equation. However, this term seems impossible to address using pre-training and fine-tuning approach given that p(y|x, z) is not properly modeled during pre-training.\n4. The comparison to the SoTA methods presented In Table 3 might be unfair. The proposed method is based on ViT-B (86M parameters) which is bigger than ResNet-50 (25M parameters) on which the other competing methods are based. The competing methods might achieve similar performance using ViT-B. Some of the improvement might come from a bigger model and it would be great if the authors can clarify this and add the discussion in the paragraph.\n5. Minor issues:\n - Notation L seems used to denote two different things in Section 3.1. It denotes the length of the sequence in the third sentence (L=HxW/P^2) while it also denotes number of transformer layers in equation (2) and (3).\n- In the paragraph above equation (6), z_d takes value from {1, 2, .., K}. Maybe it is better to write z_d \\in {1, 2, ..., K}^L.\n - In the paragraph above section 4.3, the authors refer ImageNet-Rendition as ImageNet-R. It would be great to stay consistent using ImageNet-Rendition as ImageNet-R only appears here.\n", "summary_of_the_paper": "The author present an observation in this paper that discrete image token representations derived from a vector quantized image encoder is able to preserve shape and structure object information. Inspired by this observation, the authors propose a modification to ViT architectures that appending these token representations to the input and show that the resultant models generalize better for out-of-distribution data on ImageNet classification. The modification to the input is simple and can be integrated into variants of ViT. Experiments show that the proposed method consistently outperforms baseline models on ImageNet, ImageNet-Real and seven out-of-distribution datasets derived from ImageNet. The margin of improvement is especially larger for the tasks where textures and image style change significantly while object structure is more discriminative. The main contributions of this paper are two-fold: 1) a novel approach to enrich existing ViT architectures with shape and structure information derived from discrete representations and 2) improve the robustness over the baseline ViT models and achieve SoTA results on ImageNet-Rendition, Stylized-ImageNet and ImageNet-Sketch.\n", "summary_of_the_review": "The submission is pretty solid and the claims made in the paper are convincing. To the best of my knowledge, the proposed method is novel and the contributions are significant. The authors conduct detailed ablation study which are pretty interesting to learn. There are some observations unclear to me, but they are not critical issues given the solid experiments presented in the paper. Overall I believe the quality of this submission is great. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "0-YpZn0GxeM", "reviewer": "Reviewer_a2pv", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros.\n1. The start point to use discrete information to handle robustness is novel.\n2. The authors  provide some interesting opinions, such as \"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.\"\n3. This paper is written distinctly, especially Introduction and Related work section.\n\nCons.\n1. #Abstract: (1) The claim ViTs \"are overly reliant on local features\" rather than \"make adequate use of global context\" leads ViTs to fail to generalize OOD, real-world data. #Introductioin: (2) ViTs's robustness comes from capturing globally-contextualized inductive bias than CNNs. (3) Spatial structure contributes less than 3% of ViT's performances. The (1)(2)(3) seems to be contradictory.\n2. The number of baselines is very limited. Missing many important methods, such as CutMix, Mosaic, Puzzle Mix, Cutout, Manifold.\n3. The reviewer is confused about whether the global information is the key weapon of your model or the discrete information.\n4. Missing enough discussion of the method itself. (1) After reading the paper, the reviewer still gets less information on the \"codebook\". (2)  Is VQ-GAN the chosen one, how about using other technologies, and so forth.\n5. Missing some related works such as Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions; Twins: Revisiting the design of spatial attention in vision transformers. \n\nOthers issues.\n1. In Figure 3(b), the author noticed that the \"Ours\" model captures more structural and shape patterns. And in 4.2 Qualitative results, authors hypothesize the robustness improvement stems from better capture of global contexts. So is that mean:  the ViT (or related models) learns higher frequencies is bad? \n2. Why using discrete representation can improve robustness? However, such an operation results in information loss. If it works, how about directly concat a global feature extracted by CNN?", "summary_of_the_paper": "This paper introduces Dr. ViT, which is conducted on 7 ImageNet robustness benchmarks. The motivation of this work is to demonstrate the robustness can be enhanced when adding discrete representation. Sufficient experiments demonstrate the effectiveness of this paper.\n", "summary_of_the_review": "Suggestion.\n1. The authors organize splendid discussions in the introduction and related work section, and some discussions look interesting. But the latter paragraphs could not work well in concert with the previous writing.  It is suggested that supplement enough in-depth analysis to demonstrate the mentioned illustrations. After reading this paper, the reviewer could just observe \"ViT + some global information would be better\" rather than your method is necessary.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Summary: Authors present an approach to improve the robustness of vision transformers by mapping standard tokens into discrete tokens that are invariant to small perturbations. Method is applied to a variety of backbone architectures and evaluated on a range of out of distribution forms of ImageNet test set. Significant performance gains are measured across many of these tasks.

Pros:
- Novel, simple, effective approach
- General approach applicable across model variants, complimentary to other methods to improve robustness.
- Comprehensive study, evaluated on many ImageNet robustness benchmarks
- Well written overall

Cons:
- Biggest issue: 3 reviewers point out concerns about validity of claims that ViT architecture is more reliant on local patterns and less on global context. This seems mostly a semantic issue around conjectures about why the method works – it does not invalidate the value of the new approach or its solid results.  Authors have responded to reviewer concerns by changing wording in paper to relax the claims, specifying “shape information” rather than “global information”.  They have also added experiments to measure shape bias, as defined in prior art, to backup these claims. 
- Paper missing baselines of data augmentation strategies. Authors have responded by including such comparative experiments. 
- Paper is missing ablation studies on changing the type of codebook. Authors have responded by including multiple variations of codebooks, and varying the codebook size.

This paper was a close call based on the reviews. However, in AC opinion, the critiques have been adequately addressed by the authors. This is confirmed by adding an extra expert reviewer to the pool, who agreed with some earlier critiques, and was satisfied with the changes and additional experiments presented by the authors. AC recommendation is to accept.
Decision: Accept (Poster)