Dynamic Token Normalization improves Vision Transformers

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 . 1 INTRODUCTION . 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 can not 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 . 2 RELATED WORK . 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 can not 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 . 3 METHOD . 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 ) . 3.1 REVISITING LAYER NORMALIZATION . 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 can not 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 .

Summary: 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.