SoLA-Vision: Fine-grained Layer-wise Linear Softmax Hybrid Attention

TL;DR

SoLA-Vision introduces a fine-grained layer-wise hybrid attention backbone that combines linear attention (with exponential decay) and a small number of global softmax layers to balance accuracy and efficiency in vision tasks. Through analytical and empirical studies, the authors show that distributed softmax layers act as layer-wise shortcuts to reestablish long-range interactions, while a Hidden State Bridge preserves high-resolution context for dense predictions. The proposed architecture achieves strong ImageNet-1K performance and superior results on COCO object detection and ADE20K semantic segmentation with fewer softmax layers and linear scaling in high-resolution stages. These findings offer a practical, scalable path toward efficient, high-accuracy attention mechanisms in vision models and suggest future work on learnable layer-placement policies.

Abstract

Standard softmax self-attention excels in vision tasks but incurs quadratic complexity O(N^2), limiting high-resolution deployment. Linear attention reduces the cost to O(N), yet its compressed state representations can impair modeling capacity and accuracy. We present an analytical study that contrasts linear and softmax attention for visual representation learning from a layer-stacking perspective. We further conduct systematic experiments on layer-wise hybridization patterns of linear and softmax attention. Our results show that, compared with rigid intra-block hybrid designs, fine-grained layer-wise hybridization can match or surpass performance while requiring fewer softmax layers. Building on these findings, we propose SoLA-Vision (Softmax-Linear Attention Vision), a flexible layer-wise hybrid attention backbone that enables fine-grained control over how linear and softmax attention are integrated. By strategically inserting a small number of global softmax layers, SoLA-Vision achieves a strong trade-off between accuracy and computational cost. On ImageNet-1K, SoLA-Vision outperforms purely linear and other hybrid attention models. On dense prediction tasks, it consistently surpasses strong baselines by a considerable margin. Code will be released.

Figures (7)

• Figure 1: Performance and efficiency of SoLA-Vision. (a) ImageNet-1K Top-1 accuracy versus compute (GFLOPs per image at 224×224), compared with representative CNN, Transformer, Linear, and Hybrid backbones. (b) Compute scaling of SoLA-Vision with input resolution from $16^2$ to $1280^2$, compared against Swin (linear-complexity windowed attention), DeiT (quadratic global attention), and a SoLA-FullSoftmax variant. Patch size is 4.
• Figure 2: Illustration of (a) Intra-Block Hybrid and (b) Layer-wise Hybrid attention.
• Figure 3: Feature activation maps of pretrained linear and softmax attention layers.
• Figure 4: Overall architecture of SoLA-Vision. (a) The SoLA-Vision backbone (SoLA-T shown) with Hidden State Bridge (HSB) connections indicated. Stages 1–2 use only linear attention, whereas Stages 3–4 combine linear attention with a small budget of global softmax self-attention layer via fine-grained, layer-wise placement. (b) The Hidden State Bridge mechanism. (c) The WKV linear attention layer architecture.
• Figure 5: Linear and Hybrid Stacking Performance$K$ denotes the total layers of model. Layers are same as \ref{['tab:1.1']} and are trained for 300 epochs on the ImageNet-1K 100-class subset.