Source-linked AI summary
IA-RED$^2$: Interpretability-Aware Redundancy Reduction for Vision Transformers
Bowen Pan, Rameswar Panda, Yifan Jiang, Zhangyang Wang, Rogerio Feris, Aude Oliva
TL;DR
Vision transformers are powerful but computationally expensive and difficult to interpret. IA-RED2 dynamically removes less informative patches through an interpretable module and extends this reduction hierarchically. Across image and video settings, it provides model- and task-agnostic acceleration while retaining accuracy and producing human-understandable evidence.
Problem
Vision transformers suffer heavy computation and limited interpretability, especially when long patch sequences make attention expensive.
Method
IA-RED2 uses multi-head interpreters to dynamically discard less informative patch tokens and hierarchically reduce computation across transformer stages.
Results
IA-RED2 works across image classification and video understanding as a model-agnostic and task-agnostic framework, with reported acceleration and accuracy benefits.
Takeaways & Limitations
The framework jointly provides computational reduction and human-understandable trajectories for vision-transformer inference.
Abstract
from arXiv · showhide
The self-attention-based model, transformer, is recently becoming the leading backbone in the field of computer vision. In spite of the impressive success made by transformers in a variety of vision tasks, it still suffers from heavy computation and intensive memory costs. To address this limitation, this paper presents an Interpretability-Aware REDundancy REDuction framework (IA-RED$^2$). We start by observing a large amount of redundant computation, mainly spent on uncorrelated input patches, and then introduce an interpretable module to dynamically and gracefully drop these redundant patches. This novel framework is then extended to a hierarchical structure, where uncorrelated tokens at different stages are gradually removed, resulting in a considerable shrinkage of computational cost. We include extensive experiments on both image and video tasks, where our method could deliver up to 1.4x speed-up for state-of-the-art models like DeiT and TimeSformer, by only sacrificing less than 0.7% accuracy. More importantly, contrary to other acceleration approaches, our method is inherently interpretable with substantial visual evidence, making vision transformer closer to a more human-understandable architecture while being lighter. We demonstrate that the interpretability that naturally emerged in our framework can outperform the raw attention learned by the original visual transformer, as well as those generated by off-the-shelf interpretation methods, with both qualitative and quantitative results. Project Page: http://people.csail.mit.edu/bpan/ia-red/.
1 Introduction
Vision transformers face high computation and limited interpretability, motivating IA-RED2, which dynamically removes less informative patches while preserving flexibility across tasks and models.
- Motivation: Vision transformers incur quadratic attention costs as input patch sequences grow, while their interpretability remains vague.The attention module computes fully connected relations among input patches.
- Motivation: Existing efficient designs improve accuracy-cost trade-offs but further reduce interpretability and often assume fixed-shape input sequences.These assumptions also limit architectural flexibility for arbitrary-length inputs.
- IA-RED2: The framework seeks efficiency and interpretability together without sacrificing the flexibility and versatility of the original vision transformer.The paper connects compact solutions with interpretability through a parsimony-based motivation.
- IA-RED2: IA-RED2 dynamically identifies and discards less informative, input-dependent patches to shorten sequences and reduce attention computation.A multi-head interpreter policy network makes the patch-selection decisions.
- Contributions: IA-RED2 is presented as model-agnostic and task-agnostic, with experiments spanning image recognition, action recognition, DeiT, and TimeSformer.The introduction reports 1.4× acceleration over DeiT and 4× acceleration over TimeSformer while largely maintaining accuracy.
2 Related Work
Related work covers interpretability methods, dynamic networks, and vision-transformer architectures, positioning IA-RED2 as an efficiency-driven interpreter for reducing transformer redundancy.
- Interpretability of Neural Networks: Neural-network interpretability methods commonly seek heatmaps that identify input regions relevant to predictions.Prior approaches include saliency maximization, input masking, and other visualization strategies.
- Interpretability of Neural Networks: IA-RED2 uses an efficiency-driven signal to obtain interpretable patch selection while reducing transformer redundancy.The paper characterizes this combination as a win-win between interpretability and efficiency.
- Dynamic Networks: Dynamic-network research adapts computation during inference across CNNs, RNNs, self-attention models, and data-level spatial redundancy.These approaches include model-level and input-level computation adaptation.
- Vision Transformer: Vision transformers extend transformer architectures to computer vision, while later designs add convolutional or hierarchical inductive structure.The related work also notes their reliance on large-scale datasets.
3 Proposed Method
IA-RED2 reduces vision-transformer computation by hierarchically evaluating and dropping less informative patch tokens before MSA and FFN blocks. Its multi-head interpreters are trained to balance efficiency and accuracy while producing informative scores for interpretability.
- Computational motivation: ViT complexity depends on sequence length, with total computation O(12ND^2 + 2N^2D) for sequence length N and embedding dimension D.The quadratic term motivates reducing the number of retained patch tokens.
- Framework overview: The framework partitions ViT blocks into D groups, each containing a multi-head interpreter followed by L MSA-FFN blocks.Patch tokens are evaluated before entering each group’s MSA-FFN blocks.
- Multi-head interpreter: The multi-head interpreter uses a policy token and multiple heads to estimate each input patch’s importance and drop uninformative tokens.The interpreter is modeled after MSA and operates on position-aware patch-token sequences.
- Token selection: Tokens are sampled as keep-or-discard Bernoulli actions, where keeping corresponds to u_i = 1 and discarding corresponds to u_i = 0.The action probability is parameterized by the interpreter’s predicted importance score I_ij.
- Training objective: The reward combines prediction correctness with the percentage of patches kept, using τ to control the efficiency–accuracy trade-off.This objective encourages correct predictions with as few patch tokens as possible.
- Hierarchical training and interpretability: Interpreters are trained hierarchically in a curriculum: earlier interpreters are trained first, then fixed while subsequent MSA-FFN modules are optimized.Informative scores from different groups are visualized after removed patches receive zero scores and the sequences are interpolated back to image or video size.
4 Experiments
Experiments evaluate IA-RED2 for interpretability and redundancy reduction across image and video models. The method localizes informative regions, reduces input-token redundancy, and maintains competitive recognition performance while enabling efficient inference.
- 4.1 Emergence of Interpretability: IA-RED2 better localizes objects of interest, especially their part-level regions, than memorability maps and raw attention in DeiT-S visualizations.The comparison uses heatmaps from MemNet, raw attention, and the proposed multi-head interpreter.
- 4.1 Emergence of Interpretability: IA-RED2 significantly outperforms memorability maps, raw attention, LIME, GradCAM, and LRP on weakly-supervised ImageNet image segmentation.The evaluation uses binary object masks and reports pixel accuracy, mean accuracy, and mean IoU.
- 4.2 Redundancy Reduction on Recognition Tasks: 79.1% top-1 accuracy is achieved by IA-RED2 on ImageNet-1K, versus 76.0% for a DeiT-S teacher-student model with around 70% of the teacher’s FLOPs.IA-RED2 also provides visual evidence for the informative regions supporting classification.
- 4.2 Redundancy Reduction on Recognition Tasks: IA-RED2 outperforms the attention baseline but trails random and temporal-difference token dropping on video action recognition.The authors suggest that higher video redundancy may make the model more robust to random patch dropping.
- 4.2 Redundancy Reduction on Recognition Tasks: 1360 fps and 79.1% top-1 accuracy are obtained by IA-RED2 on DeiT-S, compared with 720 fps and 77.9% for Sinkhorn Transformer and 663 fps and 77.7% for Routing Transformer.The comparison is conducted on ImageNet-1K using an NVIDIA Tesla V100 GPU.
5 Conclusions
IA-RED2 hierarchically reduces vision-transformer computation while producing human-understandable trajectories across image and video tasks. The framework is model-agnostic, task-agnostic, and complementary to model-compression approaches.
- IA-RED2 hierarchically reduces computational cost and speeds up vision transformers with human-understandable trajectories.
- Experiments cover image classification and video understanding tasks, demonstrating that IA-RED2 is model-agnostic and task-agnostic.
- IA-RED2 is complementary to model-compression approaches such as weight pruning.
- The work acknowledges IBM's donation of the Satori GPU cluster and support from the MIT-IBM Watson AI Lab, Nexplore, and Woodside.
A Pseudo Code of Our Training Process
The training pseudocode optimizes multi-head interpreters and MSA-FFN blocks for DeiT-S using token sequences and labels as inputs.
- For DeiT-S with D = 3, each group trains its multi-head interpreter for 10 epochs, followed by 20 epochs for the remaining MSA-FFN blocks.
- Algorithm 1 optimizes multi-head interpreters and MSA-FFN blocks on DeiT-S.
- The training procedure requires a token sequence X after positional embedding and its label Y.
- The pseudocode uses policy-gradient computation for interpreter parameters Wp and gradient computation for MSA-FFN parameters Wb.
B Discussion on the Training Time of Our Method
Training DeiT-S takes around 4.5 hours on 24 NVIDIA Tesla V100-32GB GPUs, with REINFORCE used for interpreter training during one third of the epochs.
- 4.5 hours on 24 NVIDIA Tesla V100-32GB GPUs are required for 90-epoch DeiT-S training.
- REINFORCE trains the multi-head interpreters for one third of the epochs without requiring backbone gradients, saving computation.
C Random Baseline with Different Seeds
Across four random seeds, the learned dropping policy consistently outperforms random dropping with DeiT-S. The comparison also examines straight-through Gumbel, which tends to highlight background regions rather than foreground objects.
- Across four random seeds, the learned dropping policy consistently outperforms the random baseline with DeiT-S.
- The experiment compares random dropping with dropping based on the learned policy.
- At similar FLOPs, the study compares straight-through Gumbel with REINFORCE on DeiT-S.
- Straight-through Gumbel often highlights background regions instead of foreground objects, leading to discarding tokens with relatively higher softmax values.
E Effect of Threshold in Discarding Tokens
The DeiT-B token-discarding threshold controls the trade-off between efficiency and accuracy: higher thresholds improve efficiency, while lower thresholds improve accuracy.
- Higher thresholds produce a more efficient DeiT-B model, while lower thresholds produce a more accurate model.The evaluated thresholds are 0.48, 0.49, 0.50, 0.51, and 0.52.
F Ablation Study on Square Reward and Insights on τ
The ablation study examines squared versus linear rewards and varying τ, finding that τ controls accuracy-efficiency trade-offs and that removing the squared reward increases sensitivity to τ.
- Changing τ produces different accuracy-efficiency trade-offs.
- Without the squared reward, accuracy-efficiency trade-offs become more sensitive to changes in τ.
G More Interpretability Results and Demo Tool
The paper provides additional interpretability visualizations for image and video tasks and releases a tool for generating interpretability results from an image.
- Additional visualizations cover interpretability results for image tasks, hierarchical redundancy reduction, and video action recognition.The video experiments use JointST TimeSformer on Kinetics-400.
- The interpretation tool can be run with `python interpreter.py -p {image_path} -o {output_dir}`.The documented environment requires Python==3.6 or above and torch==1.7 or above.
H Broader Impact
The paper presents IA-RED^2 as reducing transformer computation while producing visual interpretability evidence, with broader impacts including potential energy savings and possible AI misuse.
- The authors state that reduced computational cost could save energy and reduce industry carbon emissions, while easier AI use may enable abuse.
- The method’s heatmaps highlight informative image regions for MemNet, raw attention, and the DeiT-S model.
- Hierarchical redundancy reduction visualizations show the ratio of remaining patches for the DeiT-S model.
- TimeSformer visualizations show redundant patches being filtered while informative patches are retained for final prediction.