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: The paper proposes an Aliasing Reduction Module (ARM) consisting of an anti-aliasing filter and an external smoothing module. The paper further iterates that the module is lightweight and compatible with most existing transformer architectures. Furthermore, the method shows comparisons between a naive Gaussian filter, a Learnable Convolution Filter, and the proposed Pre-defined Filter Bank. Experiments are performed to show that placing the anti-aliasing module after the Self-Attention block yields the best results. Additionally, the method illustrates that placing the anti-aliasing module in earlier layers translates to the performance boost while inserting the module at deeper stages leads to reduced performance.