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