IA-RED$^2$: Interpretability-Aware Redundancy Reduction for Vision Transformers

1 Introduction . Transformer , a self-attention-based architecture processing sequential input without any recurrent or convolutional operations , has set off a storm in the computer vision literature recently . By dividing the input image into a series of patches and then tokenizing them with linear transformation , the transformer can effectively process the visual data in different modalities [ 13 , 53 , 54 , 28 , 3 , 17 , 66 ] . Despite its versatility , the transformer is always deeply troubled with inefficient computation and its vague interpretability . The vision transformer suffers heavy computational costs , especially when the input sequence is long . As the attention module in the vision transformer computes the fullyconnected relations among all of the input patches , the computational cost is then quadratic with regard to the length of the input sequence . On the other hand , previous works [ 6 , 8 ] have already shown the vulnerable interpretability of the original vision transformer , where the raw attention comes from the architecture sometimes fails to perceive the informative region of the input images . Recently , more designs of vision transformer architecture [ 34 , 65 , 18 , 56 , 14 , 9 , 3 ] are proposed to get higher accuracy with less computational cost . Although these methods anchor good trade-offs between efficiency and accuracy , their compression makes the vision transformer even more lack interpretability . Most of these methods assume that the input sequences are sampled from a regular visual input in a fixed shape rule , and thus their network architectures are not flexible as well , which makes the vision transformer ( 1 ) no longer able to process the input sequence with arbitrary length as the architecture is designed for a specific input shape ; ( 2 ) neither model-agnostic nor task agnostic anymore ; or ( 3 ) neglect the fact that the model redundancy is also input-dependant . We yet argue that 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . there is no inherent tension between efficiency and interpretability , and achieving them both does not have to pay design flexibility as a price . Indeed , starting from the philosophy of Occam ’ s razor , the law of parsimony , or always pursuing more compact solutions when possible , is always treated as a rule-of-thumb for pursing interpretability , especially in complicated fitting problems [ 21 ] . This paper aims to seek the win-win between efficiency and interpretability while keeping the flexibility and versatility of the original vision transformer . We propose a novel Interpretability-Aware REDundancy REDuction ( IA-RED2 ) framework for reducing the redundancy of vision transformers . The key mechanism that IA-RED2 uses to increase efficiency is to dynamically drop some less informative patches in the original input sequence so that the length of the input sequence could be reduced . While the original vision transformer tokenizes all of the input patches , it neglects the fact that some of the input patches are redundant and such redundancy is input-dependant ( see from Figure 1 ) . As the computational complexity of the attention module is quadratically linear to the input sequence length , the effect of reducing input sequence length would be magnified in the amount of the computation . Motivated by this , we leverage the idea of dynamic inference [ 39 , 37 , 38 , 61 , 57 ] , and adopt a policy network ( referred to as multi-head interpreter ) to decide which patches are uninformative and then discard them . Our proposed method is inherently interpretability-aware as the policy network learns to discriminate which region is crucial for the final prediction results . To summarize , the main contributions of our work includes : ( 1 ) We propose IA-RED2 , the first interpretability-aware redundancy reduction framework for vision transformer . ( 2 ) Our IA-RED2 framework is one of the first input-dependent dynamic inference framework for vision transformer , which adaptively decides the patch tokens to compute per input instance . ( 3 ) IA-RED2 is both model-agnostic and task-agnostic . We conduct experiments with IA-RED2 framework spanning different tasks , including image recognition and action recognition , and different models , including DeiT [ 53 ] , TimeSformer [ 3 ] . ( 4 ) We attain promising interpretable results ( shown in Figure 3 ) over baselines , with a 1.4× acceleration over DeiT on image recognition tasks , and a 4× acceleration over TimeSformer on video action recognition task while largely maintaining the accuracy . We also provide both qualitative results regarding interpretability with heatmaps by our method and those from other baseline methods like raw attention , MemNet [ 29 ] ; as well as the quantitative comparison with current state-of-the-art model interpretability methods , such as GradCAM [ 44 ] , on ImageNet-Segmentation [ 16 ] dataset with the weakly-supervised image segmentation task . 2 Related Work . Interpretability of Neural Networks . Besides improving the discrimination power of deep neural networks , model interpretability has recently raised another significant and popular research question . One of the important goals is to predict the heatmap visualization that precisely indicates the objects or contexts of relevance . Simonyan et . al . [ 45 ] attempts to maximize the class score that generates a saliency map for the given inputs . Dabkowski et al . [ 12 ] mask the salient parts of the inputs to manipulate the scores of the classifier , which generalizes well to unseen images and enables fast saliency detection . Khosla et al . [ 29 ] provide the largest annotated image memorability dataset to benchmark the visualization and explanation of natural images . After that , gradientbased methods [ 49 , 46 , 47 ] are proposed to generate precise heatmaps , by computing the gradient with respect to input during the backpropagation . While all the above approaches are studying the interpretability of convolutional neural networks ( CNNs ) , only a few works contribute to the visualization of the vision transformer . Chefer et al . shed some light on its visualization by assigning the local relevance on Transformer layers . Caron et al . [ 6 ] demonstrate that a self-supervised trained ViT produces explicit representation about the semantic location of a given object in natural images . Different from all of them , our approach starts with a novel multi-head interpreter which is supervised by an efficiency-driven signal , and then benefits from this powerful interpreter by reducing the redundancy of transformer , achieving a “ win-win ” between interpretability and efficiency . Dynamic Networks . Neural networks are found as redundant regarding their huge computation cost [ 19 , 23 , 35 ] . To overcame this issue , many adaptive computation methods are explored during the inference stage [ 1 , 2 , 57 , 15 , 26 , 58 , 20 ] . These adaptive computation strategies help speed up the inference time of convolutional neural networks ( CNNs ) [ 33 ] , recurrent neural network ( RNNs ) [ 15 ] , and also self-attention based methods ( BERT ) [ 25 ] . Besides the model-level adaptation , others further extend this idea to data-level adaptation , by either reducing the spatial redundancy [ 63 ] or focusing on key area [ 59 ] . However , those methods are limited by the convolutional structure , where only 2D data can be taken as input . Different from those approaches , our methods naturally benefit from the unstructured input taken by vision transformer , and thus can provide a much more precise glance at the target object with the affection of background being eliminated . Vision Transformer . Transformer , as a self-attention based model , has been widely adopted in natural language processing area before . The recent advance [ 13 ] shows that the transformer can also achieve incredible performance on computer vision tasks . While vision transformer suffers from the necessity large-scaled dataset [ 48 ] , many recent works try to encode strong inductive prior by either combining it with convolutional layer [ 60 , 32 , 64 , 62 ] or introducing 2D-hierarchical structure to vision transformer [ 34 , 56 , 14 , 9 ] . Besides , transformer also shows strong power in other vision tasks , including semantic segmentation [ 67 ] , object detection [ 5 , 68 ] , image processing [ 10 ] , and image generation [ 28 , 27 ] . These successes further suggest the potential of transformer to become the universal model for general vision tasks . Some other works [ 40 , 11 , 51 ] also make meaningful efforts on vision transformer efficiency . Different from those methods , the proposed method achieves a “ win-win ” on both efficiency and interpretability . 3 Proposed Method . Our main goal is to reduce the redundancy in vision transformers by dynamically dropping less informative patches in the original input sequence while classifying it correctly with the minimum computation . Our method is built on top of vision transformer ( ViT ) [ 13 ] . We start from presenting a brief overview of ViT , including the computational complexity of each module regarding the input sequence length . We then describe our proposed IA-RED2 framework for hierarchically reducing the redundant patch tokens at different layers of the vision transformer . 3.1 Overview of Vision Transformer . Vision transformer mainly consists of three main modules : ( 1 ) Multi-head Self Attention layer ( MSA ) to learn relationships between every two different patches among all the input tokens . There are h self-attention heads inside the MSA . In each self-attention head , the input token Xi is first projected to a query Qi , a key Ki , and a value Vi by three different linear transformations . Then , the query Qi computes the dot products with all the keys K and these dot products will be scaled and normalized by the softmax layer to get the attention weights . After that , it outputs the token Yi by weighted sum up all the values V with the obtained attention weights . Finally , the outputs from all heads are concatenated and re-projected by a linear layer into an output token . ( 2 ) Feed-Forward Network ( FFN ) which consists of two linear layers which are connected by the GeLU activation [ 24 ] function . For each output token Yi ∈ RD from the precedent MSA layer , FFN processes it individually . The first linear layer upgrades its dimension from D to 4D , and the second linear layer downgrades its dimension from 4D to D. Both MSA and FFN are functioning as residual connection [ 22 ] . ( 3 ) Linear Patch Embedding and Positional Encoding : For an image or a video clip , ViT first splits it into several fixed-size patches and embeds them into input tokens with a linear layer . After transforming the original image and video into a series of tokens , the network is no longer capable of being aware of the positional information of the input tokens . Thus the positional embeddings are added to the input tokens right after the patch embedding to learn the positional information of each token . Computational Complexity . For an input sequence N × D , where N is the length of the input sequence and D is the embedding dimension of each input token . The computation complexity of the MSA is O ( 4ND2 + 2N2D ) . While for the FFN , the computational complexity is O ( 8ND2 ) . As the computational complexity of patch embedding can be neglected compared with the MSA and FFN , the total computational complexity of the ViT is O ( 12ND2 + 2N2D ) .

Summary: This paper proposes the IA-RED2 framework, which tries to reduce the redundant tokens in vision transformer in a sequential manner. The redundant tokens are determined by the proposed multi-head interpreter (as a policy network). The multi-head interpreter is trained with REINFORCE, utilizing a reward function that considers the sparsity/accuracy tradeoff. The training of the full network follows a complicated curriculum learning training pipeline, which contains multiple stages and first training multi-head interpreter and then finetuning the rest of the networks in each stage. The authors conduct experiments on ImageNet1k classification with DeiT models and Video action recognition with the TimeSformer model.  This paper is among the first wave of sparse vision transformers with data-dependent sparsity, together with several other concurrent works like "Not All Images are Worth 16x16 Words: Dynamic Vision Transformers with Adaptive Sequence Length". It attains promising interpretable results, as the authors claimed.