DYNAMIC TOKEN NORMALIZATION IMPROVES VISION TRANSFORMER
Year: 2021

Abstract:
Vision Transformer (ViT) and its variants (e.g., Swin, PVT) have achieved great success in various computer vision tasks, owing to their capability to learn longrange contextual information. Layer Normalization (LN) is an essential ingredient in these models. However, we found that the ordinary LN makes tokens at different positions similar in magnitude because it normalizes embeddings within each token. It is difficult for Transformers to capture inductive bias such as the positional context in an image with LN. We tackle this problem by proposing a new normalizer, termed Dynamic Token Normalization (DTN), where normalization is performed both within each token (intra-token) and across different tokens (intertoken). DTN has several merits. Firstly, it is built on a unified formulation and thus can represent various existing normalization methods. Secondly, DTN learns to normalize tokens in both intra-token and inter-token manners, enabling Transformers to capture both the global contextual information and the local positional context. Thirdly, by simply replacing LN layers, DTN can be readily plugged into various vision transformers, such as ViT, Swin, PVT, LeViT, T2T-ViT, BigBird and Reformer. Extensive experiments show that the transformer equipped with DTN consistently outperforms baseline model with minimal extra parameters and computational overhead. For example, DTN outperforms LN by 0.5% 1.2% top-1 accuracy on ImageNet, by 1.2 1.4 box AP in object detection on COCO benchmark, by 2.3% 3.9% mCE in robustness experiments on ImageNet-C, and by 0.5% 0.8% accuracy in Long ListOps on Long-Range Arena.

text: Vision Transformer (ViT) and its variants (e.g., Swin, PVT) have achieved great success in various computer vision tasks, owing to their capability to learn longrange contextual information. Layer Normalization (LN) is an essential ingredient in these models. However, we found that the ordinary LN makes tokens at different positions similar in magnitude because it normalizes embeddings within each token. It is difficult for Transformers to capture inductive bias such as the positional context in an image with LN. We tackle this problem by proposing a new normalizer, termed Dynamic Token Normalization (DTN), where normalization is performed both within each token (intra-token) and across different tokens (intertoken). DTN has several merits. Firstly, it is built on a unified formulation and thus can represent various existing normalization methods. Secondly, DTN learns to normalize tokens in both intra-token and inter-token manners, enabling Transformers to capture both the global contextual information and the local positional context. Thirdly, by simply replacing LN layers, DTN can be readily plugged into various vision transformers, such as ViT, Swin, PVT, LeViT, T2T-ViT, BigBird and Reformer. Extensive experiments show that the transformer equipped with DTN consistently outperforms baseline model with minimal extra parameters and computational overhead. For example, DTN outperforms LN by 0.5% - 1.2% top-1 accuracy on ImageNet, by 1.2 - 1.4 box AP in object detection on COCO benchmark, by 2.3% - 3.9% mCE in robustness experiments on ImageNet-C, and by 0.5% - 0.8% accuracy in Long ListOps on Long-Range Arena.

heading: 1 INTRODUCTION
text: Vision Transformers (ViTs) have been employed in various tasks of computer vision, such as image classification (Dosovitskiy et al., 2020; Yuan et al., 2021), object detection (Wang et al., 2021b; Liu et al., 2021) and semantic segmentation (Strudel et al., 2021). Compared with the conventional Convolutional Neural Networks (CNNs), ViTs have the advantages in modeling long-range dependencies, as well as learning from multimodal data due to the representational capacity of the multi-head self-attention (MHSA) modules (Vaswani et al., 2017; Dosovitskiy et al., 2020). These appealing properties are desirable for vision systems, enabling ViTs to serve as a versatile backbone for various visual tasks.
However, despite their great successes, ViTs often have greater demands on large-scale data than CNNs, due to the lack of inductive bias in ViTs such as local context for image modeling (Dosovitskiy et al., 2020). In contrast, the inductive bias can be easily induced in CNNs by sharing convolution kernels across pixels in images (Krizhevsky et al., 2012). Many recent advanced works attempt to train a data-efficient ViT by introducing convolutional prior. For example, by distilling knowledge from a convolution teacher such as RegNetY-16GF (Radosavovic et al., 2020), DeiT (Touvron et al., 2021) can attain competitive performance when it is trained on ImageNet only. Another line of works alleviates this problem by modifying the architecture of ViTs. For example, Swin (Liu et al., 2021) and PVT (Wang et al., 2021b) utilize a hierarchical representation to allow for hidden features with different resolutions, preserving the details of locality information. Although these works present excellent results in the low-data regime, they either require convolution or require tedious architectural design.
In this work, we investigate the problem of inductive bias in vision transformers from the perspective of the normalization method. Specifically, it is known that layer normalization (LN) (Ba et al., 2016) dominates various vision transformers. However, LN normalizes the embedding within each token, making all the tokens have similar magnitude regardless of their spatial positions as shown in Fig.1(a). Although LN encourages transformers to model global contextual information (Dosovitskiy et al., 2020), we find that transformers with LN cannot effectively capture the local context in an image as indicated in Fig.2(a), because the semantic difference between different tokens has been reduced.
To tackle the above issue, we propose a new normalizer for vision transformers, termed dynamic token normalization (DTN). Motivated by Fig.1(b), where normalization in an inter-token manner like instance normalization (IN) (Ulyanov et al., 2016) can preserve the variation between tokens, DTN calculates its statistics across different tokens. However, directly aggregating tokens at different positions may lead to inaccurate estimates of normalization constants due to the domain difference between tokens as shown in Fig.3. To avoid this problem, DTN not only collects intra-token statistics like LN, but also employs a position-aware weight matrix to aggregate tokens with similar semantic information, as illustrated in Fig.1(c). DTN has several attractive benefits. (1) DTN is built on a unified formulation, making it capable of representing various existing normalization methods such as LN and instance normalization (IN). (2) DTN learns to normalize embeddings in both intra-token and inter-token manners, thus encouraging transformers to capture both global contextual information and local positional context as shown in Fig.2(c). (3) DTN is fully compatible with various advanced vision transformers. For example, DTN can be easily plugged into recently proposed models such as PVT (Wang et al., 2021b) and Swin (Liu et al., 2021) by simply replacing LN layers in the original networks.
The main contributions of this work are three-fold. (1) From the perspective of the normalization, we observe that LN reduces the difference in magnitude between tokens regardless of their different spatial positions, making it ineffective for ViTs to induce inductive bias such as local context. (2) We develop a new normalization technique, namely DTN, for vision transformers to capture both long-range dependencies and local positional context. Our proposed DTN can be seamlessly plugged into various vision transformers, consistently outperforms its baseline models with various normalization methods such as LN. (3) Extensive experiment such as image classification on ImageNet (Russakovsky et al., 2015), robustness on ImageNet-C (Hendrycks & Dietterich, 2019), self-supervised pre-training on ViTs (Caron et al., 2021), ListOps on Long-Range Arena (Tay et al., 2021) show that DTN can achieve better performance with minimal extra parameters and marginal increase of computational overhead compared to existing approaches. For example, the variant of ViT-S with DTN exceeds its counterpart of LN by 1.1% top-1 accuracy on ImageNet under the same amount of parameters with only 5.4% increase of FLOPs.

heading: 2 RELATED WORK
text: Vision transformers. Transformer models have been widely utilized in the task of computer vision in recent years. For example, the vision transformer (ViT) (Dosovitskiy et al., 2020) splits the source image into a sequence of patches, which is then fed into transformer encoders. Although ViT presents excellent results on image classification, it requires costly pre-training on large-scale datasets (e.g., JFT-300M (Sun et al., 2017)) as ViT models are hard to capture inductive bias. Many works tackle this problem by introducing convolutional priors into ViTs because convolution inherently encodes inductive bias such as local context. For example, DeiT Touvron et al. (2021) uses a token-based strategy to distill knowledge from convolutional teachers. Moreover, PVT (Wang et al., 2021b), Swin (Liu et al., 2021), T2T (Yuan et al., 2021), CVT (Wu et al., 2021) and et al. introduces hierarchical representation to vision transformers by modifying the structure of transformers. Different from these works, our DTN trains data-efficient transformers from the perspective of normalization. We show that DTN enables vision transformers to capture both global context and local positional locality. More importantly, DTN is a generally functional module and can be easily plugged into the aforementioned advanced vision transformers.
Normalization methods. Normalization techniques have been extensively investigated in CNN. For different vision tasks, various normalizers such as BN (Ioffe & Szegedy, 2015), IN (Ulyanov et al., 2016), LN (Ba et al., 2016), GN (Wu & He, 2018), SN (Luo et al., 2019a) and et al. are developed. For example, BN is widely used in CNN for image classification, while IN performs well in pixellevel tasks such as image style transfer. To our interest, LN outperforms the above normalizers in transformers and has dominated various transformer-based models. Although ScaleNorm (Nguyen & Salazar, 2019) and PowerNorm (Shen et al., 2020) improves LN in language tasks such as machine translation, it does not work well in vision transformer as shown in Table 3. Instead, we observe that vision transformer models with LN cannot effectively encode local context in an image, as shown in Fig.2. To resolve this issue, we propose a new normalizer named DTN, which can capture both global contextual information and local positional context as shown in Fig.2(c).
Dynamic Architectures. Thanks for the suggestion. Our work is related to dynamic architectures such as mixture-of-experts (MoE) (Shazeer et al., 2018; Eigen et al., 2013) and dynamic channel gating networks (Hua et al., 2018). Similar to these works, DTN also learns weight ratios to select computation units. In DTN, such a strategy is very effective to combine intra- and inter-token statistics. In addition, DTN is substantially different from DN (Luo et al., 2019b), where dynamic normalization can be constructed as the statistics in BN and LN can be inferred from the statistics of IN. However, it does not hold in the transformer because LN normalizes within each token embedding. Compared to SN (Luo et al., 2019a), we get rid of BN, which is empirically detrimental to ViTs. Moreover, we use a position-aware probability matrix to collect intra-token statistics, making our DTN rich enough to contain various normalization methods.

heading: 3 METHOD
text: This section firstly revisits layer normalization (LN) (Ba et al., 2016) and instance normalization (IN) (Ulyanov et al., 2016) and then introduces the proposed dynamic token normalization (DTN).

heading: 3.1 REVISITING LAYER NORMALIZATION
text: In general, we denote a scalar with a regular letter (e.g. ‘a’) and a tensor with a bold letter (e.g. a). Let us begin by introducing LN in the context of ViTs. We consider a typical feature of tokens x ∈ RB×T×C in the transformer where B, T , and C denote batch size, token length, embedding dimension of each token respectively, as shown in Fig.1(a). Since LN is performed in a data-independent manner, we drop the notation of ‘B’ for simplicity.
LN standardizes the input feature by removing each token’s mean and standard deviation and then utilizes an affine transformation to obtain the output tokens. The formulation of LN is written by
x̃tc = γc xtc − µlnt√ (σ2)ln + + βc (1)
where t and c are indices of tokens and embeddings of a token respectively, is a small positive constant to avoid zero denominator, and γc, βc are two learnable parameters in affine transformation. In Eqn.(1), the normalization constants of LN µlnt and (σ
2)lnt are calculated in an intra-token manner as shown in Fig.1(a). Hence, for all t ∈ [T ], we have
µlnt = 1
C C∑ c=1 xtc and (σ 2)lnt = 1 C C∑ c=1 (xtc − µlnt )2 (2)
Previous works show that LN works particularly well with the multi-head self-attention (MHSA) module to capture long-range dependencies in vision tasks, as can also be seen from Fig.2(a) where most heads in MHSA after LN attend to most of the image. However, we find that LN reduces the difference in magnitude between tokens at different positions, preventing the MHSA module from inducing inductive bias such as local context. To see this, Eqn.(2) shows that the mean and variance are obtained within each token, implying that each token would have zero mean and unit variance. Further, the standardized tokens are then operated by the same set of affine parameters {γc, βc}Cc=1 through Eqn(1). Therefore, all the tokens returned by LN would have a similar magnitude regardless of their positions in the image.
This fact can also be observed by visualizing token magnitude in Fig.1(a). As we can see, the difference between tokens after LN is reduced. However, since tokens are generated from image patches at different spatial locations, they should encode specific semantic information to embed the local context in an image. As a result, MHSA module after LN cannot effectively capture local context as presented in Fig.2(a) where only a few heads have a small attention distance.
Recent works tackle the issue of inductive bias in ViTs by combining the design of convolutional neural network (CNN) and vision transformers (Wang et al., 2021b; Liu et al., 2021). Although these methods have achieved good performance on various vision tasks, they still require complicated architecture design. In this paper, we aim to improve vision transformers by designing a new normalization method.
Instance Normalization (IN). IN provides an alternative to normalizing tokens while preserving the variation between them. In CNN, IN learns invariant features to a pixel-level perturbation, such as color and styles. Hence, IN can be employed to learn features of tokens with the local context. The definition of IN is the same with LN in Eqn.(1) except for the acquisition of normalization constants,
µinc = 1
T T∑ t=1 xtc (σ 2)inc = 1 T T∑ t=1 (xtc − µinc )2. (3)
Since IN obtains its statistics in an inter-token manner, the tokens returned by IN still preserves the variation between tokens in each head as shown in Fig.1(b). In addition, As we can observe from Fig.2(b), MHSA in the transformer with IN have more heads with a small mean attention distance than that of LN, showing that IN encourages MHSA to model local context.
However, IN has a major drawback in the context of vision transformers. From Fig.3(a), the normalization constants in IN are acquired by considering all tokens in the image. However, tokens
encoded from different image patches in the vision transformer may come from different semantic domains. For example, the top patch presents a tree, while the bottom is a snow picture. If we take an average over all the tokens to calculate mean and variance, it would lead to inaccurate estimates of normalization constants as shown in Fig.3(b). Empirically, we find that ViT-S with IN obtains much lower top-1 accuracy (77.7%) on ImageNet than its counterpart with LN (79.9%). Section 3.2 resolves these challenges by presenting dynamic token normalization.

heading: 3.2 DYNAMIC TOKEN NORMALIZATION (DTN)
text: DTN encourages attention heads in MHSA to model both global and local context by inheriting the advantages of LN and IN.
Definition. DTN is defined by a unified formulation. Given the feature of tokens x ∈ RT×C , DTN normalizes it through
x̃ = γ x− Concateh∈[H]{µh}√ Concateh∈[H]{(σ2)h}+ + β (4)
where γ, β are two C-by-1 vectors by stacking all γc and βc into a column, and µh ∈ RT× CH , (σ2)h ∈ RT× CH are normalization constants of DTN in head hwhereH denotes the number of heads in transformer. The ‘Concate’ notation indicates that DTN concatenates normalization constants from different heads. This design is motivated by two observations in Fig.2. First, attention heads attend to patches in the image with different attention distances, encouraging diverse contextual modeling for different self-attention heads. Second, by obtaining statistics in two different ways (intra-token and inter-token), LN and IN produce different patterns of attention distance in MHSA. Hence, we design DTN by calculating normalization constants specific to each head.
Normalization constants in DTN. As aforementioned in Sec.3.1, LN acquires its normalization constants within each token, which is helpful for global contextual modeling but fails to capture local context between tokens. Although IN can achieve self-attention with locality, it calculates the normalization constants across tokens, resulting in inaccurate mean and variance estimates. To overcome the above difficulties, DTN obtains normalization constants by trading off intra- and intertoken statistics as given by
µh = λh(µln)h + (1− λh)P hxh, (σ2)h = λh((σ2)ln)h + (1− λh)[P h(xh xh)− (P hxh P hxh)]
(5)
where (µln)h ∈ RT× CH , ((σ2)ln)h ∈ RT× CH are intra-token mean and variance obtained by stacking all µlnt in Eqn.(2) into a column and then broadcasting it for C/H columns, xh ∈ RT× C H represents token embeddings in the head h of x. In Eqn.(5), P hxh ∈ RT× CH , [P h(xh xh)− (P hxh P hxh)] are expected to represent inter-token mean and variance respectively. Towards this goal, we define P h as a T-by-T learnable matrix satisfying that the sum of each row equals 1. For example, when P h = 1T 1 and 1 is a T-by-T matrix with all ones, they become mean and variance of IN respectively. Moreover, DTN utilizes a learnable weight ratio λh ∈ [0, 1] to trade off intra-token and inter-token
statistics. By combining intra-token and inter-token normalization constants for each head, DTN not only preserves the difference between different tokens as shown in Fig.1(c), but also enables different attention heads to perform diverse contextual modelling in MHSA as shown in Fig.2(c). In DTN, the weight ratios for µh and σh can be different as the mean and variance in the normalization plays different roles in network’s training Luo et al. (2018); Xu et al. (2019), but they are shared in Eqn.(5) to simplify the notations.
Representation Capacity. By comparing Eqn.(5) and Eqn.(2-3), we see that DTN calculates its normalization constant in a unified formulation, making it capable to represent a series of normalization methods. For example, when λh = 1, we have µh = (µln)h and (σ2)h = ((σ2)ln)h. Hence, DTN degrades into LN. When λh = 0 and P h = 1T 1, we have µ h = 1T 1x h and (σ2)h = 1T 1(x h xh)− ( 1T 1x h 1T 1x h) which are the matrix forms of mean and variance in Eqn.(3). Therefore, DTN becomes IN in this case. Moreover, when λh = 0 and P h is a normalized k-banded matrix with entries in the diagonal band of all ones, DTN becomes a local version of IN which obtains the mean and variance for each token by only looking at k neighbouring tokens. Some examples of k-banded matrix is shown in Fig.4 of Appendix.
The flexibility of λh and P h enables DTN to mitigate the problem of inaccurate estimates of statistics in IN in two ways. Firstly, by giving weight ratio λ a large value, DTN behaves more like LN, which does not suffer from inaccurate estimates of normalization constants. Secondly, by restricting the entries of P h such as k-banded matrix, DTN can calculate the statistics across the tokens that share similar semantic information with the underlying token. As shown in Fig.3(b), the relative variation between tokens and the corresponding statistics (mean) in DTN is significantly reduced compared with that of IN.
However, it is challenging to directly train a huge matrix with a special structure such as k-banded matrix through the vanilla SGD or Adam algorithm. Moreover, the extra learnable parameters introduced by P h in DTN are non-negligible. Specifically, every P h has T 2 variables, and there are often hundreds of heads in a transformer. Therefore, directly training P h would introduce millions of parameters. Instead, we design P hby exploring the prior of the relative position of image patches which leads to marginal extra parameters.
Construction of P h. We exploit the implicit visual clue in an image to construct P h. As shown in Fig.3(a), the closer two image patches are to each other, the more similar their semantic information would be. Therefore, we employ positional self-attention with relative positional embedding (d’Ascoli et al., 2021) to generate positional attention matrix P h as given by
P h = softmax(Rah) (6)
where R ∈ RT×T×3 is a constant tensor representing the relative positional embedding. To embed the relative position between image patches, we instantiate Rij as written by Rij = [(δ x ij) 2 + (δyij) 2, δxij , δ y ij ] T where δxij and δ y ij are relative horizontal and vertical shifts between patch i and patch j (Cordonnier et al., 2019), respectively. An example of the construction of R is illustrated in Fig.5 of Appendix. Moreover, ah ∈ R3×1 are learnable parameters for each head. Hence, Eqn.(6) only introduces 3 parameters for each head, which is negligible compared to the parameters of the transformer model. In particular, by initializing ah as equation below, P h gives larger weights to tokens in the neighbourhood of size √ H × √ H relative to the underlying token,
ah = [−1, 2∆h1 , 2∆h2 ]T (7) where ∆h = [∆h1 ,∆h2 ] is each of the possible positional offsets of the neighbourhood of size√ H × √ H relative to the underlying token. It also denotes the center of attention matrix P h, which means thatP h would assign the largest weight to the position ∆h. For example, there are 4 optional attention centres (denoted by red box) when H = 4 as shown in Fig.6. Since the weights of each P h concentrates on a neighbourhood of size √ H × √ H Eqn.(6-7), DTN can aggregates the statistics across tokens with similar semantic information through positional attention matrix P h.

heading: 3.3 DISCUSSION
text: Implementation of DTN. DTN can be inserted extensively into various vision transformer models by replacing LN in the original network. In this work, we verify the effectiveness of DTN on ViT, PVT, and Swin. However, it can also be easily applied to recently-proposed advanced vision transformers such as CVT (Wu et al., 2021). As shown in Algorithm 1 of Appendix, DTN can be
implemented in a forward pass. The introduced learnable parameters by DTN are weight ratio λh and ah by Eqn.(7) which are initialized through line 2 of Algorithm 1. Note that All the computations involving in DTN are differentiable. During training, λh and ah can be optimized jointly with parameters in transformer models by automatic differentiation, making it easy for DTN to be implemented in advanced deep learning platforms such as PyTorch.
Complexity analysis. The computation and parameter complexity for different normalization methods are compared in Table 7b. Note that the number of heads in the transformer is far less than the embedding dimension i.e., H C. DTN introduces almost the same learnable parameters compared with IN and LN. Moreover, our DTN calculates inter-token normalization constants for each token embedding through positional attention matrix P h in Eqn.(5), resulting in O(BCT 2) computational complexity. The computational overhead proportional to T 2 is nonnegligible when the token length T is large. To make the computation efficient, we adopt a global average pooling operation with pooling size s on the token level, which reduces the length of tokens to T/s2. Hence, the computational complexity of DTN decreases to O(BCT 2/s4). We also investigate the FLOPs of vision transformers with DTN in Table 1 and Table 2.

heading: 4 EXPERIMENT
text: This section extensively evaluates the effectiveness of the proposed Dynamic Token Normalization (DTN) in various computer vision tasks, including classification on ImageNet, robustness on ImageNet-C and ImageNet-R (Sec.B.2 of Appendix), and self-supervised pre-training (Sec.B.3 of Appendix). Besides, we also test DTN on the task of ListOps using transformers with different efficient self-attention modules. An ablation study is provided to analyze our proposed DTN comprehensively (see more details in Sec.B.4 of Appendix). The training details can be found in Sec.B.1.

heading: 4.1 RESULTS ON IMAGENET
text: DTN on ViT models with different sizes. We evaluate the performance of the proposed DTN on ViT models with different sizes, including ViT-T, ViT-S, and ViT-B. Note that DTN obtains inter-token normalization constants by positional attention matrix P h in Eqn.(5), which gives larger weights to entries in the neighbourhood of size √ H × √ H where H denotes the number of heads. To exploit this property, we increase the number of heads from 6 to 9 in order to produce 3 × 3 neighborhood. To avoid the soaring increase of the model size, we also decrease the embedding dimension in each head from 64 to 48. The resulting models are denoted by ViT-T∗, ViT-S∗, and ViT-B∗.
As shown in Table 1, ViT models with DTN can consistently outperform the baseline models using LN with a marginal increase of parameters and FLOPs on the ImageNet dataset, demonstrating the effectiveness of our DTN. For example, DTN surpasses LN by a margin of 1.1% top-1 accuracy on ViT-S∗ with almost the same parameters and only 5.4% increase of FLOPs compared with LN. On the other hand, we see the margin of performance becomes larger when ViT models are instantiated with a larger number of heads. For instance, DTN exceeds LN by 0.8% top-1 accuracy on ViT-B∗ while DTN outperforms LN by 0.5% top-1 accuracy on ViT-B. It demonstrates that our proposed DTN benefits from a larger head dimension in transformers as implied by Eqn.(6) and Eqn.(7).
DTN on various vision transformers. We investigate the effect of DTN on various vision transformers on ImageNet. Since normalization is an indispensable and lightweight component in transformers, it is flexible to plug the proposed DTN into various vision transformers by directly replacing LN. In practice, two representative vision transformers, including PVT Wang et al. (2021b) and Swin Liu et al. (2021) are selected to verify the versatility of our DTN. We use their respective open-sourced training framework. As shown in Table 2, DTN improves the performance of PVT and Swin by a clear margin in terms of top-1 and top-5 accuracies compared with baseline using LN with a marginal increase of FLOPs. Specifically, PVT-Tiny, PVT-Small, and Swin-T with DTN surpass their baseline with LN by 1.2%, 0.9%, and 0.7% top-1 accuracy, respectively.
DTN versus other normalization methods. We show the advantage of our DTN over exiting representative normalizers that are widely utilized in computer vision models such as BN, IN, and GN, and in language models such as LN, ScaleNorm, and PowerNorm. The results are reported in Table 3. It can be seen that our DTN achieves competitive performance compared with other normalization methods on both ViT-S and ViT-S∗. In particular, we see that vanilla BN and IN obtain the worst performance, which is inferior to LN by 2.6% and 2.2% top-1 accuracy, respectively. This may be caused by inaccurate estimates of mean and variance in IN and BN. For example, PowerNorm improves BN by providing accurate normalization constant through moving average in both forward and backward propagation, achieving comparable top-1 accuracy compared with LN. Our proposed DTN further calculates the statistics in both intra-token and inter-token ways, consistently outperforming previous normalizers.

heading: 4.2 RESULTS ON LISTOPS OF LONG-RANGE ARENA
text: We further evaluate our DTN on the ListOps task of Long-Range Arena benchmark (Tay et al., 2021) to validate the proposed methods on non-image sources. The Long-Range Arena (LRA) benchmark is designed explicitly to assess efficient Transformer models under the long-context scenario. To demonstrate the effectiveness of DTN, we choose the Long ListOps task, one of the most challenging tasks in the LRA benchmark. The Long Listops task is a longer variation of standard ListOps task (Nangia & Bowman, 2018) with up to 2K sequence lengths, which is considerably difficult.
We compare DTN with other normalizations using vanilla Transformer (Vaswani et al., 2017) and two recently proposed efficient Transformer models including BigBird (Zaheer et al., 2020) and Reformer (Kitaev et al., 2020). All of the experiments are conducted with default settings in (Tay et al., 2021) and the results are reported in Table. 4. As shown, DTN yields consistent performance gains (> 0.5%) compared to other listed normalizations, including the SOTA method Reformer. For the other normalizations, we observe a significant accuracy drop on vanilla Transformer models with IN, BN, and GN, which is possibly caused inaccurate estimates of statistics (as mentioned in Sec.3.1).
4.3 ABLATION STUDY
Effect of each component. As shown in Eqn.(5), DTN obtains its normalization constants by three crucial components including inter-token statistics from LN, intra-token statistics controlled by positional attention matrix P h, and learnable weight ratio λh. We investigate the effect of each compo-
nents by considering five variants of DTN. (a) λh = 1. In this case, DTN becomes LN; (b) λh = 0 and P h = 1T 1. For this situation, DTN turns into IN; (c) λ
h = 0 and P h is learnable. (d) λh = 0.5 and P h is learnable; (e) Both λh and P h are learnable which is our proposed DTN. The results are reported in Table 5. By comparing (b) and (a & c), we find that both intra-token statistics from LN and positional attention matrixP h can improve IN by a large margin, showing the reasonableness of our design in Eqn.(5). It also implies that both intra-token and inter-token statistics are useful in normalization. By comparing (e) and (c & d), we see that a learnable λh can better trade off inter-token and intra-token statistics for each head, which further improves the performance of DTN.
Different initilization of ah. We investigate the effect of different initializations of ah using ViT-T∗ with DTN. To compare with the initialization in Eqn.(7), we also use a truncated normal distribution to initialize ah. In experiment, we find DTN with the initialization of ah followed truncated normal distribution achieves 72.7% top-1 accuracy, which is worse than DTN with ah initialized by Eqn.(7) (73.2% top-1 accuracy). The effectiveness of initialization in Eqn.(7) is further revealed by visualization in Fig.6. We can see that P h generated by Eqn.(7) gives larger weights to its neighbouring tokens before training while P h obtained by initialization with truncated normal distribution has uniform weights, indicating that Eqn.(7) helps DTN aggregate statistics over tokens with similar local context. After training, the initialization in Eqn.(7) can also better capture local positional content than initialization using truncated normal as shown in Fig.7.

heading: 5 CONCLUSION
text: In this work, we find that layer normalization (LN) makes tokens similar to each other regardless of their positions in spatial. It would result in the lack of inductive bias such as local context for ViTs. We tackle this problem by proposing a new normalizer named dynamic token normalization (DTN) where normalization constants are aggregated on intra- and inter-token bases. DTN provides a holistic formulation by representing a series of existing normalizers such as IN and LN. Since DTN considers tokens from different positions, it preserves the variation between tokens and thus can capture the local context in the image. Through extensive experiments and studies, DTN can adapt to ViTs with different sizes, various vision transformers, and tasks outperforming their counterparts. In particular, DTN improves the modeling capability of the self-attention module by designing a new normalizer, shedding light on future work on transformer-based architecture development. For example, DTN could be combined with a sparse self-attention module because it encourages selfattention with a small attention distance, as shown in Fig.2.
Ethics Statement. We improve the performance of vision transformers by the proposed dynamic token normalization (DTN). We notice that DTN can be plugged into various deep architectures and applied in a wide range of learning tasks. Hence, our work and AI applications in different tasks would have the same negative impact on ethics. Moreover, DTN may have different effects on
different image sources, thus potentially producing unfair models. We will carefully study our DTN on the fairness of the model output.
Reproducibility Statement. To reproduce the experimental results, we provide training details and hyper-parameters in Appendix Sec.B.1. Moreover, we will also make our code available by a link to an anonymous repository during the discussion stage.

heading: A MORE DETAILS ABOUT APPROACH
text: A.1 ARCHITECTURAL DETAIL
We integrate DTN into various vision transformers such as ViT, PVT and Swin which is instantiated by stacking multiple transformer encoders. Note that our proposed DTN utilizes the relative positional encoding, implying that the class token can not be directly concatenated with the tokens encoded from image patches. To tackle this issue, we employ DTN in the first ldtn transformer encoders and leave the last L − ldtn transformer unchanged where L denotes the number of transformer encoders. In this way, the class token can then appended to the output tokens of the ldtn-th transformer encoder by following d’Ascoli et al. (2021); Wang et al. (2021b). By default, we set ldtn = 10 for ViTs and ldtn = 3 for PVT which indicates that DTN is used in the first 3 blocks of PVT. Since Swin do not utilize class token, we insert DTN into all the blocks of Swin. Note that ldtn determines how many DTNs are plugged into the transformer. We investigate the influence of the number of DTN used in transformers by the ablation study on ldtn as shown in Sec.B.4.
Algorithm 1 Forward pass of DTN. 1: Input: mini-batch inputs x ∈ RB×T×C , where the b-th sample in the batch is xb ∈ RT×C , b ∈ {1, 2, ..., B};
learnable parameters ωh to generate λh, h ∈ [H]; relative positional embeddingR ∈ RT×T×3 and ah ∈ R3×1 in Eqn.(6) in Eqn.(7); scale and shift parameters γ,β ∈ RC×1.
2: Initialization: initialize ωh = 0 and ah = [−1, 2∆h1 , 2∆h2 ]T before training. 3: Hyperparameters: . 4: Output: the normalized tokens {x̃n, n = 1, 2, ..., N}. 5: calculate: λh = sigmoid(ωh). 6: calculate: P h = softmax(Rah). 7: for b = 1 to B do 8: calculate: µlnt = 1C ∑C c=1 x b tc, (σ 2)lnt = 1 C ∑C c=1(x b tc − µlnt )2, t ∈ [T ]. . intra-token statistics.
9: calculate: (µit)h = P h(xb)h, ((σ2)it)h = [P h((xb) (xb))− (P h(xb) P h(xb))]. . inter-token statistics. 10: calculate: µh = λh(µln)h + (1− λh)(µit)h, h ∈ [H]. . mean of DTN. 11: calculate: (σ2)h = λh((σ2)ln)h + (1− λh)((σ2)it)h, h ∈ [H]. . variance of DTN. 12: concatenate statistics in all heads: µ = Concateh∈[H]{µh},σ2 = Concateh∈[H]{(σ2)h}. 13: calculate DTN output: x̃ = γ (x− µ)/√σ2 + + β. 14: end for

heading: B MORE EXPERIMENTAL RESULTS
text: B.1 IMPLEMENTATION DETAIL
ImageNet. We evaluate the performance of our proposed DTN using ViT models with different sizes on ImageNet, which consists of 1.28M training images and 50k validation images. Top-1 and
Top-5 accuracies are reported on ImageNet. By default, we set the patch size of ViT models as 16. Moreover, we utilize the common protocols, i.e., number of parameters and Float Points Operations (FLOPs) to obtain model size and computational consumption. We train ViT with our proposed DTN by following the training framework of DeiT (Touvron et al., 2021) where the ViT models are trained with a total batch size of 1024 on all GPUs. We use Adam optimizer with a momentum of 0.9 and weight decay of 0.05. The cosine learning schedule is adopted with the initial learning rate of 0.0005. We use global average pooling for PVT and Swin, where pooling sizes for the first two blocks are set to 4 and 2, respectively.
ListOps on LRA. For all experiments on the LRA benchmark, we follow the open repository released by (Tay et al., 2021) and keep all of the settings unchanged. Specifically, all evaluated models have an embedding dimension of 512, with 8 heads and 6 layers. The models are trained for 5K steps on 2K length sequence with batch size 32.
Definition of mean attention distance in Fig.2. The mean attention distance is defined by d = 1 T ∑T i=1 di, di = ∑T j=1Aijδij where Aij and δij indicate the self-attention weight and Euclidean distance in 2D spatial between token i and token j, respectively. We calculate the mean attention distance for each head by averaging a batch of samples on the ImageNet validation set. When computing the attention weight between token i and other tokens, we deem token i as the attention center. Since the sum over j of Aij is 1, di indicates the number of tokens between the attention center token i and other tokens. Therefore, a large mean attention distance implies that self-attention would care more about distant tokens relative to the center token. In this sense, self-attention is thought to model global context. On the contrary, a small mean attention distance implies that self-attention would care more about neighboring tokens relative to the center token. In this case, self-attention can better capture local context.
B.2 RESULTS ON IMAGENET-C AND IMAGENET-R
DTN can calculate the normalization constants for each token embedding. Thus, DTN is expected to be robust to input perturbations as the mean and variance are specific to each token. To verify this, we compare the robustness to input perturbations of ViT models with LN (baseline) and our DTN. To capture different aspects of robustness, we rely on different, specialized benchmarks ImageNet-C (Hendrycks & Dietterich, 2019) and ImageNet-R (Hendrycks et al., 2021), which contains images in the presence of ‘natural corruptions’ and ‘naturally occurring distribution shifts,’ respectively. For the experiment on ImageNet-C, we can see from Table 6 that DTN can surpass the baseline LN by 2.2% mCE metric on ViT-S and 3.9% mCE metric on ViT-S∗. Moreover, DTN is also superior to LN in the presence of various natural corruptions such as ‘Nosie’ and ‘Blur’ as shown in Table 6. For the experiment on ImageNet-R, DTN obtains 30.7% and 32.7% top-1 accuracy on ViT-S and ViT-S∗ respectively, which are superior to LN by 1.5% and 3.0% top-1 accuracy. Therefore, our DTN is more robust to natural corruptions and real-world distribution shifts than LN.
B.3 RESULTS ON SELF-SUPERVISION
A special way of improving a neural model’s generalization is supervised contrastive learning [9, 11, 25, 31]. We couple DTN with the supervised contrastive learning on ViT-S∗ for 100 epoch pre-training, followed by a lineal evaluation. We fine-tune the classification head by 100 epochs for
both ViT-S with LN. Please see the Appendices for more implementation details. Compared to the training procedure without SAM, we find considerable performance gain as our DTN can encourage transformers to capture both the long-range dependencies and the local context, improving the ImageNet top-1 accuracy of ViT-S∗ from 74.7% to 75.2%.
B.4 MORE ABLATION STUDIES
Learning dynamics of λh. In the implementation, we treat the weight ratios λh of mean and variance in Eqn.(5) differently because mean and variance in normalization play a different role in network’s training. For DTN layers of ViT-T∗, we plot the learning dynamics of λh for (a) mean and (b) variance. As implied by Eqn.(5), the smaller λh is, the more important inter-token statistics would be. As shown in Fig.8, we have three observations. First, the weight ratio of mean and variance have distinct learning dynamics. λh of the mean for different heads are more diverse than that of variance. Second, different DTN layers have different learning dynamics, which are smoothly converged in training. Third, multiple heads in shallow layers prefer inter-token statics. Whereas larger λh are typically presented in higher layers. It is consistent with Fig.2 where some heads in shallow layers have a small attention distance while most heads in higher layers have a large attention distance.
Effect of the number of DTN layers. As discussed in Sec.A.1, ldtn determines how many DTNs are plugged into the transformer. We investigate the influence of the number of DTN used in transformers by the ablation study on ldtn. We set ldtn = 4, 6, 8, 10 for VIT-S∗ which means LN in the first 4, 6, 8, 10 transformer encoders are replaced by our DTN while leaving the remaining transformer encoder layers unchanged. The results are obtained by training VIT-S∗ with different ldtn’s on ImageNet. As shown in Table, more DTN layers bring greater improvement on classification accuracy.
Performance of DTN on larger models. we employ DTN on some large-scale models, including Swin-S, PVT-L, and PVTv2-B3 (Wang et al., 2021a). Note that these models either have large sizes (> 45M) or attain outstanding top-1 accuracy (> 82.5%) on ImageNet. As shown in Table 9, our DTN achieves consistent gains on top of these large models. For example, DTN improves the plain Swin-S and PVT-L by 0.5% and 0.6% top-1 accuracy, respectively. For the improved version of the PVT model, e.g., PVTv2-B3(Wang et al., 2021a), we also observe the performance gain (+0.5% top-1 accuracy).
How does RPE (i.e. R in Eqn.(6) affects the performance of DTN. We clarify that our DTN is a normalization component that helps a variety of transformer models (ViT, PVT, Swin, T2T-ViT, BigBird, Reformer, etc.) learn better token representation, even some of the models use a relative positional encoding. To separate the effect of RPE in Eqn.(6) from DTN, we conduct two ablation studies on Imagenet as shown in Table 10. (a) We put RPE in Eqn.(6) into the self-attention module of ViT and PVT and use LN to normalize the tokens. Note that neither ViT nor PVT uses RPE in their implementation. Hence, ablation (a) can evaluate how RPE in Eqn.(6) affects the performance of the original model. We also investigate the effect of RPE in Eqn.(6) by removing it from DTN; (b) We plug DTN into vision transformers using carefully-designed RPE such as LeViT(Graham et al., 2021) and SWin (Liu et al., 2021), which demonstrates that DTN can still improve the models with RPE. Results. From ablation (a) in Table 10, we can see that a naive RPE used in Eqn.(6) has a marginal improvement when directly put into the MHSA module. However, the performance boosts a lot when such simple RPE is used in our DTN as did in Eqn.(6). We also see that removing RPE in Eqn.(6) has limited influence on the performance of DTN. Moreover, our DTN can still improve those models with carefully-designed RPE as shown by ablation (b) in Table 10, demonstrating the versatility of DTN as a normalization technique.
Results on downstream task. We assess the generalization of ourDTN on detection/segmentation track using the COCO2017 dataset ( Lin et al. (2014)). We train our model on the union of 80k training images and 35k validation images and report the performance on the mini-val 5k images. Mask R-CNN and RetinaNet are used as the base detection/segmentation framework. The standard COCO metrics of Average Precision (AP) for bounding box detection (APbb) and instance segmentation (APm) is used to evaluate our methods. We use MMDetection training framework with PVT models as basic backbones and all the hyper-parameters are the same as Wang et al. (2021b). Table 11 shows the detection and segmentation results. The results show that compared with vanilla LN, our DTN block can consistently improve the performance. For example, our DTN with PVT-Tiny is 1.4 AP higher in detection and 1.2 AP higher in segmentation than LN. To sum up, these experiments demonstrate the generalization ability of our DTN in dense tasks that require local context.
The effect of DTN on the token direction and magnitude. In our DTN, we have an interesting observation about the learning of direction and magnitude. We perform PCA visualization of tokens before and after normalization. The projected tokens in first two principle components are visualized in Fig.9. Results are obtained on the first layer of trained ViT-S with LN and DTN, respectively. It can be seen that the PCA projected tokens after DTN are closer to each other than LN. It implies that DTN encourages tokens to learn a less diverse direction than LN since tokens normalized by DTN have already presented diverse magnitudes. As we know, learning diverse magnitudes (1 Dimension)
could be easier than learning diverse directions (C-1 Dimension). Hence, our DTN can reduce the optimization difficulty in learning diverse token representations. In the experiment, we also find that ViT models with DTN can converge much faster than LN. Understanding the role of token magnitude and direction in self-attention modules would be a meaningful future research direction.
DTN is also effective on those models with local context already induced. Our DTN as a normalization technique can be seamlessly applied to various transformer models, even those with local context already encoded. For example, we replace LN with our DTN on T2T-ViT (Yuan et al., 2021), LeViT (Graham et al., 2021), and Swin. These models introduce positional information, convolutions, or window attention to induce inductive bias. As shown in Table 12, DTN can consistently outperform LN on these models. For example, DTN improves plain Swin-T and Swin-S with LN
by 0.7% and 0.5% top-1 accuracy, respectively. Also, DTN achieves 0.7% higher top-1 accuracy on both T2T-ViTt-14 and LeViT-128S

Review summaries:
This work presents a token normalization method in replacement of Layer Norm (LN) and Instance Norm (IN) for vision transformers (ViTs). The motivation is that the authors find that the common normalization used in most existing ViTs, LN, will reduce the difference in token magnitude, and this may lead to failure of capturing positional context and other inductive bias. Hence, they propose a Dynamic Token Normalization (DTN) component by combining LN and relative positional embedding based transformation. DTN can be plug in varying ViTs e.g., ViT, Swin, PVT. The experiments are done on ImageNet, ImageNet-C, ImageNet-R, and ListOps in supervised learning setting, as well as self-supervised pre-training. 

This paper introduces a dynamic normalization, named Dynamic Token Normalization (DTN), to replace the vanilla layer norm in ViT. It learns to normalize tokens in both intra-token and inter-token manners, enabling Transformers to capture both the global contextual information and the local positional context. Experimental results show that DTN can improve some Vision Transformers in ImageNet Classification and Long ListOps tasks. 

This paper studies properties in layer normalization and instance normalization, and states that the two normalization operations have their own drawbacks in vision transformers: LN suffers from lacking inductive bias while IN may be affected by the different semantics within the tokens. Accordingly, this paper proposes a new normalization method called DTN and considers both inter- and intra- token normalization into vision transformers. According to the experiments in the paper, the proposed DTN normalization boosts the performance of various vision transformer models on classification with ImageNet dataset as well as robustness with IMAGENET-C and IMAGENET-R.

The paper first analyzes the limitation of LN in Transformers and then proposes DTN to capture both long-range dependencies and local positional context. DTN is a unified version that balances LN and IN. Extensive experiments show the effectiveness of the proposed DTN with some small/middle-scale Transformers on the ImageNet.