Blending Anti-Aliasing into Vision Transformer

1 Introduction . Transformers have led to impressive breakthroughs in language understanding , dominating a wide range of natural language processing ( NLP ) tasks . Meanwhile , continuous researches strive to leverage transformers for vision tasks , revolutionizing the conventional inductive bias in convolutional neural networks ( CNNs ) . After a big leap made by vision transformer ( ViT ) [ 1 ] , promising results on vision tasks have recently emerged [ 2–4 ] , demonstrating a possibility of using transformers as a primary backbone for vision applications . While appealing advantages such as long-range context modeling and parameter efficiency are introduced by self-attention , this mechanism also brings an inevitable aliasing issue to vision transformers . Aliasing [ 5 ] traditionally refers to the phenomenon that high-frequency signals become indistinguishable when undersampled [ 6 ] . This effect occurs when discrete patterns are utilized to capture a more continuous signal , resulting in frequency ambiguity . In terms of ViT [ 1 ] , images are split into non-overlapping patches during tokenization , which are then fed into transformer blocks . Compare to the more “ continuous ” representation of image , tokenization and self-attention performed on its discontinuous patch embeddings can be regarded as subsampling operations . While alleviating computational costs , these operations leave a side-effect which leads to possible aliasing . A simple solution to mitigate aliasing phenomenon is to increase the sampling rate . Similar properties emerge in vision transformers , where overlapped tokens [ 7 ] and smaller patch sizes [ 1 , 8 ] lead to improved performance . As increased sampling rates lead to quadratic computation costs , we hereby conjecture that proper anti-aliasing filters that integrated into the “ attending ” process could also provide a fix . A potential concern is that the general purpose of anti-aliasing : to introduce ∗Work done during an internship at Amazon . 1Code is made available at https : //github.com/amazon-research/anti-aliasing-transformer . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . “ smoothness ” to signals , might be contradictory with the goal of self-attention : to capture more significant and high-frequency features in contrast . How would traditional anti-aliasing techniques influence vision transformers ? And what if we equip these state-of-the-art transformer architectures with suitable anti-aliasing options ? To this end , we design the Aliasing Reduction Module ( ARM ) for vision transformers , which consists of an anti-aliasing filter and an external modulation module . Extensive investigations are conducted primarily , including the filtering choices such as the basic Gaussian blurring filters in [ 10 ] and learnable convolutional filters , as well as different placements of these filters . Notably , based on the observations that the “ jagged ” phenomenon varies across images and locations , we propose an adaptive filtering method which draws inspirations from dictionary learning [ 11 ] . Specifically , a lightweight aliasing estimation network is constructed to predict the combination coefficients of the pre-defined low pass filter bank adaptively . Furthermore , the position of integrating our Aliasing Reduction Module in highly-modularized transformers is crucial . The model benefits from an early anti-aliasing operation right after the self-attention computation , which is distinctive from the observations in CNNs [ 10 , 12 ] . As elaborated in Figure 1 , the module brings perceptually smoother attention , which leads to better feature localization and eventually correct prediction . By integrating ARM to state-of-the-art architectures , we consistently find a considerable margin of improvement . To summarize , our contributions are : • We investigate anti-aliasing techniques for vision transformers , and propose the Aliasing Reduction Module ( ARM ) , which is compatible with most existing architectures . • We tailor various anti-aliasing strategies and different placements of aliasing reduction , which are distinctive from the observations in CNN structures . • Experiments show that our simple yet effective design boosts sophisticated state-of-the-art vision transformers . Furthermore , we observe stronger generalization , robustness , and data-efficiency brought by our anti-aliasing strategy . 2 Related Works . Vision Transformers . The transformer structure was introduced in [ 13 ] for machine translation , which further became a general purpose model for many NLP tasks . The Vision Transformer ( ViT ) [ 1 ] firstly closes the gaps on image classification with previous CNN models using pure transformer . DeiT [ 9 ] further improves training data efficiency of ViT , utilizing an additional distillation token for teacher-student training as well as stronger data augmentations [ 14 , 15 ] . Given its superior performance on classification , recent works apply transformer architecture to various vision tasks , including object detection [ 16 , 17 ] , segmentation [ 3 , 18 , 19 ] , point cloud processing [ 20 ] , image generation [ 21–24 ] and so on . In order to obtain better transformer architectures that generalize well to general vision tasks , some concurrent advances have introduced hand-crafted designs . For example , CPVT [ 25 ] and CoaT [ 26 ] improves the original positional encoding . T2T [ 7 ] improves the naive tokenization process . Multi-scale information [ 27 , 28 ] , convolutional designs [ 29–31 ] , and pyramid structures [ 32 , 33 ] are also utilized in vision transformer extensions . Due to the inefficiency and unnecessity of self-attention , some researchers also explore approximated attention [ 34 ] , local attention [ 31 , 32 ] or even replacing it with MLP structure [ 35 ] . Different from these designs , our work dive into the “ feature-resampling ” property of the self-attention mechanism , which leads to aliasing and performance degradation . To this end , we investigate a lightweight anti-aliasing module to help mitigate the effect . This design is general and compatible with most vision transformer architectures . Anti-aliasing in CNNs . The definition of aliasing was originally covered in signal processing [ 6 ] , which causes the distortion on sampled signals to reconstruct original signals . The aliasing issue occurs in CNNs with any strided layers , such as max pooling and strided convolution . While early smoothing methods like average pooling [ 36 ] degrades performance on modern models and benchmarks , the elegant Blurpool [ 10 ] introduces both anti-aliasing properties and increased accuracies to strong baselines . DDAC [ 12 ] further assigns adaptive filters for each spatial location and channel group . These anti-aliased CNNs demonstrate both better recognition , generalization to downstream tasks and robustness towards corruptions [ 37 , 38 ] . Differs from CNN structures , vision transformers do not have explicit “ strided ” subsampling except for the tokenization step , i.e . patch embedding . Furthermore , their potential downsampling [ 32 , 33 ] choices are mostly densely-connected layers . Thus the observations and techniques for CNNs may not generalize well to transformers . In contrast , we argue that in transformer , self-attention on split patch tokens can be viewed as a “ resampling ” step which tries to reconstruct the informative part on the original continuous pixel representation of images . To this end , we investigate on anti-aliasing self-attention operation , in order to learn a more continuous and smoothed attention representation . 3 Methodology . We propose to blend the anti-aliasing property into vision transformers by processing the inner representation with a lightweight Aliasing Reduction Module ( ARM ) . ARM consists of an antialiasing filter as well as an external smoothing module . The anti-aliasing operator can be chosen flexibly among a traditional low-pass filter , e.g . Gaussian filter , a learnable convolutional filter , or a pre-defined filter bank . In this section , we will first explain what raises the issue of aliasing in vision transformers , as well as its difference compared with the twin problem in CNNs . Secondly , we provide the details about how we design the proposed ARM and the crucial choice of where to integrate it . 3.1 What makes for aliasing in vision transformer ? . Frequency aliasing is the phenomenon occurs when performing subsampling on any type of signal , for instance an audio , an image or a video . When the sampling rate is lower than bandwidth of the original signal and does not meet the Nyquist rate [ 39 ] , the resampled signal will be aliased where its high frequency patterns interlace with the low frequency components additively . As demonstrated in Figure 2 , this effect usually results in visible corruptions and artifacts . As discussed in [ 10 ] , convolutional networks also suffer from aliasing which is caused by strided operations , i.e. , strided convolution or max pooling , that commonly exist in widely adopted CNN structures . The subsequent layers following the strided operations are likely to extract features from the distorted inputs . Such artifacts might affect the final prediction , as well as model ’ s robustness towards natural corruptions . Self-attention mechanism allows transformers to integrate information across all spatial locations , by resampling the importance of feature representations based on self-affinities . However , with quadratic cost in terms of large number of pixels , naive application of “ per-word ” tokenization in NLP does not scale to realistic image sizes . A widespread practice is to process each input image into a sequence of patches as in ViT [ 1 ] . The patch tokenization and embedding indeed alleviate computational costs while inevitably bring in the problem of aliasing due to the discontinuous processing . When there exists an aliasing phenomenon in very early layers , the subsequent transformer blocks will be overwhelmed by the misled signals . Similar to CNNs , vision transformers have another potential aliasing sources from the downsampling layers . Note that there exists no progressive downsampling in ViT-based architectures , we further analyze whether the feature map downsampling in recent transformer architectures [ 32 , 30 , 33 ] influences the model in Section 3.2 and Section 4.2 . In conclusion , the discontinuous tokenization process and its succeeding self-attention make for the aliasing phenomenon in vision transformer , further limiting their modeling power . Meanwhile , it ’ s computationally infeasible to simply extend the finite token sets through smaller resolutions or even pixel-wise processing . To this end , we aim to develop a generic anti-aliasing module that mitigates this omitted phenomenon in vision transformers . 3.2 Incorporating anti-aliasing module into vision transformer . In this section , we answer two questions raised above : how we design the structure of ARM and where to place it inside the highly-modularized transformer architecture . We first re-visit the implementation of a transformer block for vision transformer . Given an input x ∈ RC×H×W , the block first splits x into a sequence z ∈ RN×d with a length of N where the dimension of patch embedding is d. The sequence representation is then passed into a fully connected layer . The FC layer projects the sequences into query , key and value with the same hidden dimension [ q , k , v ] ∈ RD×3Dh , where D and Dh represent the input and hidden dimension . Then self-attention SA can be formulated as : A = softmax ( qkT / √ Dh ) ( 1 ) SA ( z ) = Av ( 2 ) Then with the projection and normalization , the attention map x̂ is obtained . The attentive signal is accumulated back to the original input x . As we ’ ve analyzed in Section 3.1 , the self-attention operation is conducted on split patch tokens from the original feature maps . This process “ attends ” to finite sets of discrete tokens and identifies relevant features from them , collecting stronger significance while bringing in aliasing . In other words , the attention map x̂ is a resampled signal from the original representation x , which is more discrete and sparse . To this end , we choose to directly act on the aliased attention map and apply anti-aliasing operation on it . The general purpose is to make the attention signal more continuous and robust . A recent work FNet [ 41 ] , which replaces self-attention in BERT [ 42 ] with Fourier Transforms , may also share the same spirit . Besides , we ’ ve also tried to put the smoothing filter to the different positions , such as tokenization process or downsampling step in Swin Transformer [ 32 ] . Detailed discussion about the placement can be found in Section 4.2 . Figure 3 provides an overview of the proposed Aliasing Reduction Module ( ARM ) . The module adaptively smooths aliased attention maps computed from the discontinuous token sets . We tailor three valid choices of the anti-aliasing filter as visualized in Figure 3 : Gaussian Filter . In order to investigate whether the model suffers from aliasing , the simplest approach is to smooth the attention map with a traditional low pass filter . Hereby we choose the Gaussian blur , which is effective at reducing image noise and leaving fewer sharp edges . The filter is kept frozen during training , introducing roughly zero computation cost to training and testing . Learnable Convolution Filter . While Gaussian blurring servers as an efficient operator at removing noises , applying such fixed smoothing to the entire feature map inevitably results in some valuable high frequency information loss . As has been discussed in [ 12 ] , it ’ s sub-optimal to apply the same filter across the varying content . To this end , we also propose a learnable convolution-based filter . The content-aware filter adaptively learns to produce varying anti-aliasing effect . On the contrary , this design incurs additional trainable parameters and extra GPU memory . Pre-defined Filter Bank . Due to the unconstrained nature of attention representations , the training of shallow Anti-Aliasing convolution filter is unstable as the gradients oscillate between varying attention maps . While other strategies such as residual learning and encoder-decoder structures may alleviate these issues , this line of method demands even more computation costs . To balance the trade-off , we predefine a filter bank that contains disparate degrees of smoothness . A lightweight convolutional layer is optimized to estimate the combination coefficient of the atomic filters . Specifically , The filter Fi ∈ Rk×k , which denotes a filter in the i-th channel of the input attention maps , is constructed as a linear combination of n pre-defined filter dictionary D from the coefficient Φ : Fi = N∑ j=1 ΦjDj ( 3 ) Based on the predicted coefficient from the weight head , a complex filter can be derived from the pre-defined filter bank . Since the estimated filter evolves within the linear space constructed from the atomic filters in the dictionary , it is feasible and fast to optimize . The linear assembling strategy has been proven effective in computational photography [ 43 ] , parametric human body representation [ 44 ] , network compression [ 45 ] and acceleration [ 46 ] . Specifically , we compose our pre-defined filter bank with Gaussian filters as well as a small portion of difference of Gaussians ( DoG ) , which can endow the feasibility when edge-preserving properties are required . As visualized in Figure 4 , the filter bank is constructed by altering the covariance matrix of multivariate Gaussian kernels , reflecting in varying scale , rotation , and ellipticalness . Each filter is normalized to the summation of 1 . Thus the weights of a k × k sized kernel is defined based on the probability density function ( PDF ) of the Gaussian distribution with a covariance matrix Σ : k ( x− x0 ; Σ ) = 1 2π|Σ| 12 e− 1 2 ( x−x0 ) T Σ−1 ( x−x0 ) T ( 4 ) Beyond the “ internal ” smoothing operation within each attention map , we also explore the effect brought by “ external ” modulation approach . Using buffered frames , temporal AntiAliasing ( TAA ) [ 47 ] in computer graphics tackles the aliased patterns that the edge-smoothing filters handle poorly . Inspired by this , we propose to ensemble our anti-aliasing filter with an external regularization which reduces the internal shift in attention maps using batched samples . We operationalize the module using a batch norm [ 48 ] after the filtered attention maps . As the normalization operation has been analyzed to bring smoothness into the optimization process of neural networks [ 49 ] , it is non-trivial to be ensembled into our Aliasing Reduction Module for being able to mitigate separate patterns and oscillations , of which the filter-based methods are incapable . In summary , our anti-aliasing component is composed of an anti-aliasing filter , which can either consist of a Gaussian filter , the learnable convolution , and the proposed filter bank . Subsequently , an external modulation module connects to the filters , rescaling the filtered features . The module directly applies smoothness to the attention maps , thus being pluggable to most vision transformation architectures .

Summary: This paper proposes a plug-and-play anti-aliasing reduction method to solve the aliasing problem in the vision transformer. It analyzes the effectiveness and generalization ability in multiple tasks and various vision transformer methods. Experimental results demonstrate the effectiveness of the proposed method.