Source-linked AI summary
Interlaced Sparse Self-Attention for Semantic Segmentation
Lang Huang, Yuhui Yuan, Jianyuan Guo, Chao Zhang, Xilin Chen, Jingdong Wang
TL;DR
Dense self-attention is costly for high-resolution vision inputs, motivating a more efficient mechanism. The paper factorizes its dense affinity matrix into successive long-range and short-range sparse attentions, achieving competitive semantic-segmentation performance with substantially lower efficiency costs.
Problem
High-resolution vision tasks make dense self-attention expensive in computation and memory, limiting its practical application.
Method
The method factorizes the dense affinity matrix into two sparse affinity matrices and applies long-range and short-range attention successively.
Results
The approach achieves competitive performance across semantic-segmentation datasets while being much more efficient than conventional self-attention.
Takeaways & Limitations
Interlaced sparse self-attention captures dense long-range dependencies more efficiently for high-resolution vision processing.
Takeaways & Limitations
The formulation assumes an input size satisfying N = P × Q.
Abstract
from arXiv · showhide
In this paper, we present a so-called interlaced sparse self-attention approach to improve the efficiency of the \emph{self-attention} mechanism for semantic segmentation. The main idea is that we factorize the dense affinity matrix as the product of two sparse affinity matrices. There are two successive attention modules each estimating a sparse affinity matrix. The first attention module is used to estimate the affinities within a subset of positions that have long spatial interval distances and the second attention module is used to estimate the affinities within a subset of positions that have short spatial interval distances. These two attention modules are designed so that each position is able to receive the information from all the other positions. In contrast to the original self-attention module, our approach decreases the computation and memory complexity substantially especially when processing high-resolution feature maps. We empirically verify the effectiveness of our approach on six challenging semantic segmentation benchmarks.
1. Introduction
The paper targets the high computation and memory cost of dense self-attention for high-resolution vision tasks. It introduces interlaced sparse self-attention, which factorizes dense affinities into two sparse stages while preserving information propagation between all positions.
- Motivation: Stacking nearly hundreds of 3 × 3 convolutions may be needed to capture dependencies between any positions in 256 × 256 inputs.The paper motivates self-attention as a more direct way to model long-range dependencies.
- Motivation: Self-attention models dependencies between any positions in one layer, but its computation complexity is about O(N^2).This cost becomes especially problematic for high-resolution inputs used in object detection and semantic segmentation.
- Motivation: More than 64 GB of GPU memory can be required to train self-attention models with batch size 8 on semantic segmentation inputs.The paper therefore identifies reducing self-attention computation and memory as having practical value.
- Approach: The proposed scheme factorizes dense affinity computation into two sparse affinity matrices, A_L and A_S, whose product is dense.Computing the sparse block affinity matrices is substantially cheaper than computing the conventional dense affinity matrix.
- Approach: Long-range attention groups positions with long spatial interval distances, while successive short-range attention groups originally nearby positions.The two stages use sparse connections and together allow every output position to receive information from all input positions.
- Evaluation: The approach is empirically evaluated across vision tasks and reports similar or better semantic-segmentation performance than conventional self-attention.The introduction also reports comparisons with mechanisms including CGNL and RCCA.
2. Related Work
Related work improves or approximates self-attention through alternative affinity mechanisms and interlacing-based architectures. This paper instead applies interlacing to factorize self-attention affinities and organize long-range pixel groups.
- Self-Attention and Non-local: Self-attention and non-local mechanisms model relations between all positions and have been applied to video, detection, segmentation, and re-identification.These methods provide the broader context for the proposed attention design.
- Efficient Attention: CGNL approximates pairwise similarities with a Taylor expansion, while RCCA uses consecutive criss-cross attention to approximate self-attention.Other related methods use global representations or alternative sparse structures to improve efficiency.
- Interlacing: Prior interlacing methods include interleaved group convolution, ShuffleNet, channel local convolution, and space-to-channel mechanisms.These works establish interlacing as a design pattern for network architectures.
- Position of This Work: Unlike those methods, this work uses interlacing to decompose dense self-attention into two sparse affinity matrices and group pixels with long spatial intervals.The paper also notes a concurrent Sparse Transformer using a similar factorization for sequential tasks.
3. Approach
The approach decomposes self-attention into successive long-range and short-range sparse attention operations, whose sparse affinity matrices together reproduce dense information propagation. This design reduces computation for high-resolution inputs while retaining a straightforward implementation.
- Self-Attention: Self-attention forms a dense affinity matrix from transformed input features, with θ and φ mapping X to lower-dimensional representations before similarity computation.The affinity matrix records similarities between positions, while g learns the output embedding.
- Interlaced Sparse Self-Attention: The proposed method factorizes the dense affinity matrix into long-range and short-range sparse block affinity matrices, AL and AS.Their product represents the dense affinity matrix while each component is cheaper to compute.
- Long-range Attention: Long-range attention permutes X, partitions it into groups of Q positions, and independently applies self-attention within each group to estimate AL.Each subset is sampled from separated spatial regions, and the resulting affinity matrices are merged after attention.
- Short-range Attention: Short-range attention permutes the long-range output, partitions it into groups of P neighboring positions, and independently estimates AS within each group.The updated representations are merged after local attention; combining both stages lets each output position receive information from all input positions.
- Complexity and Implementation: The approach is more efficient than conventional self-attention for high-resolution inputs, while its PyTorch implementation uses permutation, reshaping, and 1 × 1 convolution-based transforms.Figure 4 compares computation cost in GFLOPs as feature-map resolution increases.
4. Experiments
The experiments evaluate interlaced sparse self-attention across six semantic segmentation benchmarks and additional COCO detection and instance-segmentation tasks. The approach generally matches or exceeds competing methods while substantially reducing computational cost.
- Semantic Segmentation: The approach is evaluated on six semantic segmentation benchmarks using mIoU and pixel accuracy, plus COCO object detection and instance segmentation.The segmentation datasets are Cityscapes, ADE20K, LIP, PASCAL VOC 2012, PASCAL-Context, and COCO-Stuff.
- Semantic Segmentation: 80.3% on Cityscapes outperforms AAF by 1.2% while requiring much less computation than DANet.The comparison uses the Cityscapes test set with multi-scale and flip testing.
- Semantic Segmentation: 45.04% mIoU on ADE20K improves the same-backbone GCU result by 0.2%.The evaluation uses a ResNet-101 backbone and multi-scale testing.
- Semantic Segmentation: The method reaches 55.07% on LIP, 83.2% mIoU on PASCAL VOC 2012, and 54.1% mIoU on PASCAL-Context.It is reported as state of the art on LIP, slightly better than DANet on PASCAL VOC 2012, and better than other listed methods on PASCAL-Context.
- Detection and Instance Segmentation: On COCO, the module improves Mask-RCNN by about 1% across detection and instance-segmentation metrics, reaching 39.7 box AP and 35.7 mask AP with a 2× schedule.Its performance is comparable to a non-local block while reducing computation complexity significantly.
- Efficiency and Ablation: The method is much more efficient than PPM, SA, DANet, RCCA, and CGNL in GPU memory, GFLOPs, and inference time.With Ph = Pw = 8, the authors use a fixed setting for all experiments; the approach also consistently outperforms PPM on ADE20K, Cityscapes, and LIP.
5. Conclusion
The paper presents interlaced sparse self-attention as an efficient alternative that factorizes dense affinity computation into two sparse affinity matrices while retaining competitive semantic-segmentation performance.
- Interlaced sparse self-attention factors the dense affinity matrix into the product of two sparse affinity matrices.The approach is designed as a simple extension of existing self-attention implementations.
- The method achieves competitive performance across various semantic segmentation datasets.
- The proposed mechanism is substantially more efficient than conventional self-attention.