Source-linked AI summary
BiFormer: Vision Transformer with Bi-Level Routing Attention
Lei Zhu, Xinjiang Wang, Zhanghan Ke, Wayne Zhang, Rynson Lau
TL;DR
Vision-transformer attention captures long-range dependencies but is costly because it computes pairwise interactions across all spatial locations. The paper introduces query-adaptive bi-level routing and builds BiFormer around it, reporting strong results across four vision tasks and 83.8% top-1 accuracy at 4.6G FLOPs on ImageNet-1K. Its GPU throughput can remain lower than similar-FLOPs models because routing adds kernel-launch and memory-transaction overheads.
Problem
Global attention captures long-range dependencies but has high computational complexity and heavy memory footprints, while prior sparse methods are static or query-agnostic.
Method
Bi-level routing filters irrelevant key-value pairs at a coarse region level, then applies token attention to the union of routed regions using dense matrix multiplications.
Results
BiFormer shows superior performance across image classification, object detection, instance segmentation, and semantic segmentation; BiFormer-T reaches 83.8% top-1 accuracy at 4.6G FLOPs on ImageNet-1K.
Takeaways & Limitations
Query-adaptive routing gives BiFormer a better computation-performance trade-off by attending to small subsets of relevant tokens.
Takeaways & Limitations
Routing introduces GPU kernel-launch and memory-transaction overheads, giving BiFormer lower throughput than some similar-FLOPs models.
Abstract
from arXiv · showhide
As the core building block of vision transformers, attention is a powerful tool to capture long-range dependency. However, such power comes at a cost: it incurs a huge computation burden and heavy memory footprint as pairwise token interaction across all spatial locations is computed. A series of works attempt to alleviate this problem by introducing handcrafted and content-agnostic sparsity into attention, such as restricting the attention operation to be inside local windows, axial stripes, or dilated windows. In contrast to these approaches, we propose a novel dynamic sparse attention via bi-level routing to enable a more flexible allocation of computations with content awareness. Specifically, for a query, irrelevant key-value pairs are first filtered out at a coarse region level, and then fine-grained token-to-token attention is applied in the union of remaining candidate regions (\ie, routed regions). We provide a simple yet effective implementation of the proposed bi-level routing attention, which utilizes the sparsity to save both computation and memory while involving only GPU-friendly dense matrix multiplications. Built with the proposed bi-level routing attention, a new general vision transformer, named BiFormer, is then presented. As BiFormer attends to a small subset of relevant tokens in a \textbf{query adaptive} manner without distraction from other irrelevant ones, it enjoys both good performance and high computational efficiency, especially in dense prediction tasks. Empirical results across several computer vision tasks such as image classification, object detection, and semantic segmentation verify the effectiveness of our design. Code is available at \url{https://github.com/rayleizhu/BiFormer}.
1 City University of Hong Kong 2 SenseTime Research
The section lists contact email addresses for the paper’s authors.
- The listed City University of Hong Kong contacts are lzhu68-c and zhanghake2-c.
- The listed SenseTime contacts are wangxinjiang and wayne.zhang.
- The contact addresses use my.cityu.edu.hk and sensetime.com domains.
1. Introduction
The introduction motivates query-aware sparse attention to reduce attention’s computational and memory costs, then presents BRA and BiFormer as an efficient solution with strong vision-task performance.
- Global attention captures long-range dependencies but computes pairwise affinities across all spatial locations, creating high complexity and heavy memory footprints.
- Existing sparse-attention methods use handcrafted patterns or query-agnostic key-value subsets, despite queries attending to different pairs across semantic regions.
- Per-query selection can retain vanilla attention complexity, while local offset prediction can hinder long-range dependency modeling.
- BRA prunes a region-level affinity graph to top-k routed regions before applying token-to-token attention over their union.
- BiFormer uses BRA to attend content-aware subsets of relevant tokens and achieves a better computation-performance trade-off across vision applications.
- 4.6G FLOPs yields 83.8% top-1 accuracy for BiFormer-T on ImageNet-1K without external data or distillation.
2. Related Works
Related work covers transformer architectures and efficient attention methods, including sparse patterns and hierarchical alternatives with different efficiency trade-offs.
- Vision transformers use channel-wise MLP blocks for channel mixing and attention blocks for cross-location spatial mixing.
- Efficient-attention research reduces vanilla attention’s computation and memory bottlenecks through sparse connections, low-rank approximations, or recurrent operations.
- Quadtree Transformer uses recursive multiscale processing that hurts parallelism, whereas BiFormer gathers key/value tokens and applies dense matrix multiplications more efficiently.
3. Our Approach: BiFormer
BiFormer introduces Bi-level Routing Attention (BRA), which first routes among semantically related regions and then performs token-level attention on the selected regions. This query-adaptive sparsity preserves long-range dependency while reducing attention complexity and using GPU-friendly dense operations.
- Bi-Level Routing Attention: BRA addresses MHSA’s O(N^2) scalability problem by replacing all-to-all token interactions with dynamic, query-adaptive sparse attention.Each query attends to a small subset of relevant key-value tokens rather than all tokens.
- Region partition and routing: BRA partitions the feature map into regions, averages queries and keys within each region, and constructs a region-level affinity graph.The region-level queries and keys are produced by per-region averaging after linear projection.
- Region partition and routing: Row-wise top-k pruning retains the most relevant regions for each region, producing routing indices used to select candidate key-value pairs.The ith routing-index row contains the k most relevant regions for the ith region.
- Token-to-token attention: Token-to-token attention is then applied to all key-value pairs in the union of the routed regions, with gathered tensors enabling efficient matrix multiplication.The implementation uses dense batch matrix multiplication and adds a depthwise-convolution local context enhancement term.
- Complexity analysis: BRA provides direct long-range dependency modeling with O((HW)^(4/3)) complexity, compared with O((HW)^2) for vanilla attention.The complexity depends on choosing an appropriate region partition factor S.
- BiFormer architecture: BiFormer uses BRA in a four-stage pyramid backbone with progressively reduced spatial resolution and increased channel dimensions.Three model sizes are created by scaling network width and depth, with task-specific routing configurations.
4. Experiments
BiFormer is evaluated across image classification, object detection, instance segmentation, and semantic segmentation, with strong results under comparable computational settings. Ablations and visualizations support the effectiveness of bi-level routing attention and its content-aware selection of relevant regions.
- Image Classification: BiFormer-S achieves 83.8% top-1 accuracy at approximately 4G FLOPs, reported as the best result without extra training data or training tricks.
- Object Detection and Instance Segmentation: BiFormer shows a clear advantage on COCO object detection and instance segmentation across bounding-box and mask AP metrics at IoU thresholds of 50% and 75%.
- Semantic Segmentation: With Semantic FPN on ADE20K, BiFormer-S/B achieves 48.9/49.9 mIoU, improving CSWin-T/S by 0.7 mIoU; similar gains occur with UperNet.
- Ablation Study: Under matched Swin-T architecture designs, bi-level routing attention outperforms existing sparse attention mechanisms on image classification and semantic segmentation.
- Attention Visualization: Visualizations show routed regions covering semantically related entities and nonadjacent objects, indicating that the mechanism captures long-range inter-object relationships.
5. Limitation and Future Work
BiFormer has lower GPU throughput than some similarly computationally sized models because routing adds kernel-launch and memory-transaction overheads. The authors identify kernel fusion and hardware-aware sparse attention as future directions.
- Routing overheads from extra GPU kernel launches and memory transactions reduce BiFormer's throughput relative to some models with similar FLOPs.
- The authors suggest GPU kernel fusion and hardware-aware sparse attention as ways to mitigate this efficiency limitation.
6. Conclusion
The paper introduces bi-level routing attention (BRA) for dynamic, query-aware computation allocation and uses it as the core of BiFormer. BRA routes at the region level before applying token-level attention, while regional representations are formed by average pooling.
- Contribution: BRA dynamically allocates computation by filtering irrelevant key-value pairs at the region level before fine-grained attention.It builds and prunes a region-level graph, then attends within the union of routed regions.
- Contribution: BiFormer uses BRA as its core building block and achieves strong performance across image classification, object detection, instance segmentation, and semantic segmentation.
- Regional representations: Regional representations Qr and Kr are derived with average pooling for region-to-region routing.
- Regional representations: Average-token affinity is equivalent to average token-to-token affinity between two regions, supporting the routing criterion.
B. Throughput Comparison
BiFormer-STL sacrifices throughput relative to Swin-T because routing adds kernel-launch and memory-transaction overheads, but remains substantially faster than QuadTree-STL. The comparison uses training and inference benchmarks on a Tesla V100 under FP32 and AMP modes.
- Experimental setup: The benchmark compares Swin-T, QuadTree-STL, and BiFormer-STL using training and inference throughput on a 32GB Tesla V100.Results are reported under FP32 and automatic mixed precision modes.
- Throughput: BiFormer-STL training and inference throughput decreases by approximately 30% and 40%, respectively, versus Swin-T.The reduction is attributed to routing-related GPU kernel launches and memory transactions.
- Throughput: BiFormer-STL remains 3×–6× faster than QuadTree-STL.BRA uses token gathering followed by GPU-friendly dense matrix multiplications, whereas QuadTree attention uses recursive and sparse operations.
- Limitation: Routing overheads may be reduced through engineering techniques such as GPU kernel fusion.The paper leaves this optimization for future work.
C. Choices of top-k and partition factor S
The choices of partition factor S and top-k are guided mainly by engineering constraints and stage-dependent complexity. Ablations indicate that increasing attended tokens can hurt accuracy, suggesting that explicit sparsity may reduce background distraction.
- Parameter choices: S is chosen as a feature-map size divisor, while larger S values are used in dense prediction tasks to balance routing and token-attention complexity.For 224 × 224 classification, S = 7 matches the stage feature-map dimensions.
- Parameter choices: k is gradually increased in later stages as region sizes become smaller, maintaining a reasonable number of attended tokens.
- Ablation: Increasing the number of attended tokens may hurt IN-1K accuracy, implying that explicit sparsity can regularize against background distractions.
D. Adapting Pretrained Plain ViT with BRA
The paper adapts pretrained DeiT-B to semantic segmentation by replacing full attention with BRA and transferring pretrained projection weights. On ADE20K, BRA outperforms local-window attention in both decoder settings reported.
- Adaptation: Pretrained DeiT-B is adapted to ADE20K by replacing all or part of its full multi-head self-attention modules with BRA.BRA projection weights are initialized from the original MHSA weights.
- Results: Without propagation blocks, BRA surpasses local-window attention by 2.4 mAP.
- Results: With four global propagation blocks, both architectures improve, while BRA retains a 0.2 mAP advantage.
- Experimental setup: The ADE20K setup uses a Simple Feature Pyramid followed by a Upernet head.
E. More Visualization Results
The paper presents additional visualization results in Figure 6 to further demonstrate how bi-level routing attention (BRA) works.
- Figure 6 contains additional visualization results for BRA.
- The visualizations are intended to demonstrate how BRA works.
- This section extends the paper’s visualization evidence beyond earlier results.