D YNAMIC T OKEN N ORMALIZATION I MPROVES V ISION T RANSFORMER
Anonymous authors
Paper under double-blind review
A BSTRACT
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 difﬁcult 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 uniﬁed 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 .
9% mCE in robustness experiments on ImageNet-C, and by 0 .
8% accuracy in Long ListOps on Long-Range Arena.
I NTRODUCTION
Vision Transformers (ViTs) have been employed in various tasks of computer vision, such as image classiﬁcation (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-efﬁcient 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.
(a) LN Tokens
Tokens
Reduce Differenc e H e a d s before LN after LN B samples C embeds
Intra-token
T Tokens (b) IN
H e a d s
Tokens Tokens
Preserve before IN after IN
Inter-token
Tokens Tokens
Preserve (c) DTN before DTN after DTN
Intra-and Inter-token
Position-aware
H e a d s
Differenc e Differenc e
In this work, we investigate the problem of inductive bias in vision transformers from the perspective of the normalization method.
Speciﬁcally, 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 ﬁnd 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 beneﬁts.
(1) DTN is built on a uniﬁed 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.
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 classiﬁcation 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.
9 Layers
M e a n A tt n .
D i s t a n c e
LN  9 Layers
IN  9 Layers
DTN
Head 1 Head 2 Head 3 Head
Head 5 Head
R ELATED W ORK
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 classiﬁcation, 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-efﬁcient 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 classiﬁcation, 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.
M ETHOD
This section ﬁrstly 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).
3. R EVISITING L AYER N ORMALIZATION
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 ∈ R B × 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 afﬁne transformation to obtain the output tokens.
The formulation of LN is written by
[Equation]
where t and c are indices of tokens and embeddings of a token respectively, (cid:15) is a small positive constant to avoid zero denominator, and γ c , β c are two learnable parameters in afﬁne 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
[Equation]
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 ﬁnd 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 afﬁne parameters { γ c , β c } Cc = 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 speciﬁc 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 deﬁnition of IN is the same with LN in Eqn.(1) except for the acquisition of normalization constants,
[Equation]
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
[Figure]
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 ﬁnd that ViT-S with IN obtains much lower top-1 accuracy ( 77 . 7 %) on ImageNet than its counterpart with LN ( 79 . 9 %).
Section 3.
resolves these challenges by presenting dynamic token normalization.
3. D YNAMIC T OKEN N ORMALIZATION (DTN)
DTN encourages attention heads in MHSA to model both global and local context by inheriting the advantages of LN and IN.
Deﬁnition.
DTN is deﬁned by a uniﬁed formulation.
Given the feature of tokens x ∈ R T × C , DTN normalizes it through
[Equation]
where γ , β are two C-by-1 vectors by stacking all γ c and β c into a column, and µ h ∈
R T × CH , ( σ 2 ) h ∈ R T × CH are normalization constants of DTN in head h where H 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 speciﬁc 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 difﬁculties, DTN obtains normalization constants by trading off intra- and intertoken statistics as given by
[Equation]
[Equation]
where ( µ ln ) h ∈ R T × CH , (( σ 2 ) ln ) h ∈ R T × CH are intra-token mean and variance obtained by stacking all µ ln t in Eqn.(2) into a column and then broadcasting it for C/H columns, x h ∈ R T × CH represents token embeddings in the head h of x .
In Eqn.(5), P h x h ∈ R T × CH , [ P h ( x h (cid:12) x h ) − (
[Table]
)] are expected to represent inter-token mean and variance respectively.
Towards this goal, we deﬁne
P h as a T-by-T learnable matrix satisfying that the sum of each row equals 1 .
For example, when
P h
= 1 T 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.
Xu et al.
, 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 uniﬁed 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
= 1 T 1 , we have µ h = 1 T 1 x h and ( σ 2 )
h
= 1 T 1 ( x h (cid:12) x h ) − ( 1 T 1 x h (cid:12) 1 T 1 x 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 ﬂexibility 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 signiﬁcantly 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.
Speciﬁcally, 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 h by 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
[Equation]
where R ∈ R T × T × 3 is a constant tensor representing the relative positional embedding.
To embed the relative position between image patches, we instantiate R ij as written by
R ij =
[( δ xij ) 2 + ( δ yij ) 2 ,δ xij ,δ yij ] T where δ xij and δ yij 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, a h ∈ R 3 × 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 a h as equation below, P h gives larger weights to tokens in the neighbourhood of size √ H × √ H relative to the underlying token
[Equation]
where ∆ h =
[∆ h 1 , ∆ h 2 ] 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 that P 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 .
3. D ISCUSSION
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
[Table] [Table]
implemented in a forward pass.
The introduced learnable parameters by DTN are weight ratio λ h and a h 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 a h 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 (cid:28) 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 efﬁcient, we adopt a global average pooling operation with pooling size s on the token level, which reduces the length of tokens to T/s 2 .
Hence, the computational complexity of DTN decreases to O ( BCT 2 /s 4 ) .
We also investigate the FLOPs of vision transformers with DTN in Table 1 and Table 2.
EXPERIMENT
This section extensively evaluates the effectiveness of the proposed Dynamic Token Normalization (DTN) in various computer vision tasks, including classiﬁcation 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 efﬁcient 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.
4. R ESULTS ON I MAGE N ET
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 ×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 ∗ .
[Table] [Table]
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 beneﬁts 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 ﬂexible 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.
Speciﬁcally, 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.
4. R ESULTS ON L IST O PS OF L ONG -R ANGE A RENA
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 efﬁcient 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 difﬁcult.
We compare DTN with other normalizations using vanilla Transformer (Vaswani et al., 2017) and two recently proposed efﬁcient 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.
As shown, DTN yields consistent performance gains ( > 0 .
compared to other listed normalizations, including the SOTA method Reformer.
For the other normalizations, we observe a signiﬁcant 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. A BLATION STUDY
[Table]
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 components by considering ﬁve variants of DTN.
(a) λ h
= 1 .
In this case, DTN becomes LN; (b) λ h
=
and P h = 1 T 1 .
For this situation, DTN turns into IN; (c) λ h
= 0 and P h is learnable.
(d) λ h = 0 .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 ﬁnd that both intra-token statistics from LN and positional attention matrix P 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 a h .
We investigate the effect of different initializations of a h using ViT-T ∗ with DTN.
To compare with the initialization in Eqn.(7), we also use a truncated normal distribution to initialize a h .
In experiment, we ﬁnd DTN with the initialization of a h followed truncated normal distribution achieves 72.7% top-1 accuracy, which is worse than DTN with a h 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.
C ONCLUSION
In this work, we ﬁnd 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
A M ORE D ETAILS ABOUT A PPROACH A. A RCHITECTURAL D ETAIL
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 ﬁrst l dtn transformer encoders and leave the last L − l dtn 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 l dtn -th transformer encoder by following d’Ascoli et al.
; Wang et al.
(2021b).
By default, we set l dtn = 10 for ViTs and l dtn = 3 for PVT which indicates that DTN is used in the ﬁrst 3 blocks of
PVT.
Since Swin do not utilize class token, we insert DTN into all the blocks of Swin.
Note that l dtn determines how many DTNs are plugged into the transformer.
We investigate the inﬂuence of the number of DTN used in transformers by the ablation study on l dtn as shown in Sec.
B.4.
L a y e r
L a y e r L a y e r L a y e r
L a y e r L a y e r
L a y e r L a y e r
L a y e r L a y e r
L a y e r L a y e r
L a y e r L a y e r
L a y e r L a y e r
L a y e r L a y e r
L a y e r L a y e r
Algorithm 1 Forward pass of DTN.
, ( σ 2 ) lnt = 1 C (cid:80) Cc =1 ( x btc − µ lnt ) 2 ,t ∈ [ T ] .
(cid:46) intra-token statistics.
9: calculate: ( µ it ) h
= P h ( x b ) h , (( σ 2 ) it ) h
=
[ P h (( x b ) (cid:12) ( x b )) − ( P h ( x b ) (cid:12) P h ( x b ))] .
(cid:46) inter-token statistics.
10: calculate: µ h = λ h ( µ ln ) h
+ (1 − λ h )( µ it ) h ,h ∈
[ H ] .
(cid:46) mean of DTN.
11: calculate: ( σ 2 ) h
= λ
h (( σ 2 ) ln )
h
+ (1 − λ h )(( σ 2 ) it ) h ,h ∈
[ H ] .
(cid:46) variance of DTN.
12: concatenate statistics in all heads: µ = Concate h ∈
[ H ] { µ h } , σ 2 =
Concate h ∈
[ H ] { ( σ 2 ) h } . 13: calculate DTN output:
˜ x
= γ (cid:12) ( x − µ ) / √ σ 2 + (cid:15) + β . 14: end for
B M ORE EXPERIMENTAL RESULTS B. I MPLEMENTATION DETAIL
ImageNet.
We evaluate the performance of our proposed DTN using ViT models with different sizes on ImageNet, which consists of 1 .
M training images and 50 k validation images.
Top-1 and
[Table]
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 .
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 ﬁrst 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.
Speciﬁcally, all evaluated models have an embedding dimension of 512 , with 8 heads and 6 layers.
The models are trained for 5 K steps on 2 K length sequence with batch size 32 .
Deﬁnition of mean attention distance in Fig.2.
The mean attention distance is deﬁned by d =
T (cid:80) Ti =1 d i , d i = (cid:80)
Tj =1 A ij δ ij where A ij 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 A ij is 1 , d i 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. R ESULTS ON I MAGE N ET -C AND I MAGE N ET -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 speciﬁc 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. R ESULTS ON S ELF - 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 ﬁne-tune the classiﬁcation head by 100 epochs for
[Table] [Table]
[Figure] [Figure] [Figure] r a t i o s o f v a r h
[Figure]
both ViT-S with LN.
Please see the Appendices for more implementation details.
Compared to the training procedure without SAM, we ﬁnd 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. M ORE 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, l dtn determines how many DTNs are plugged into the transformer.
We investigate the inﬂuence of the number of DTN used in transformers by the ablation study on l dtn .
We set l dtn = 4 , 6 , 8 , 10 for VIT-S ∗ which means LN in the ﬁrst 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 l dtn ’s on ImageNet.
As shown in Table, more DTN layers bring greater improvement on classiﬁcation accuracy. [Table] [Table]
Performance of DTN on larger models.
we employ DTN on some large-scale models, including
Swin-S, PVT-L [Table] and PVTv2-B3 (Wang et al., 2021a).
Note that these models either have large sizes ( > 45 M) or attain outstanding top-1 accuracy ( > 82 .
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 inﬂuence 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.
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 ﬁrst two principle components are visualized in Fig.9.
Results are obtained on the ﬁrst 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) DTN can consistently improve both AP and box AP by a large margin [Table] [Figure]
[Table]
[Figure]
are visualized before and after LN (rows 1-2), and before and after DTN (rows 3-4).
We see 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.
Results are obtained on the ﬁrst layer of trained ViT-S with LN and DTN, respectively.
[Table]
could be easier than learning diverse directions (C-1 Dimension).
Hence, our DTN can reduce the optimization difﬁculty in learning diverse token representations.
In the experiment, we also ﬁnd 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
[-Head 
-]