HSMLA: Hierarchical Softmax Multi-scale Linear Attention for Efficient Vision Transformers

2026-08-06

AI Summary

This work addresses the computational bottleneck of vision Transformers in high-resolution dense prediction tasks, where the quadratic complexity of self-attention impedes scalability and existing linear attention methods struggle to capture local context effectively. To overcome these limitations, the authors propose HSMLA, a novel framework that uniquely integrates ReLU-based linear attention for global context modeling, a selective Softmax mechanism to refine critical local features, and depthwise convolutional multi-scale token representations within a hierarchical attention architecture. This design achieves strong local modeling capability while preserving linear computational complexity. Extensive experiments demonstrate that HSMLA delivers state-of-the-art efficiency–accuracy trade-offs, achieving up to 4.2× faster inference—alongside competitive performance—on tasks such as CT organ segmentation (87.3% Dice, 3.2× speedup) and pathological whole-slide image classification (94.2% AUC, 4.1× speedup).

Abstract

Vision transformers face significant computational overheads in high-resolution dense prediction due to the quadratic complexity of self-attention. Linear attention offers efficiency but sacrifices local context modeling. We propose \textbf{HSMLA (Hierarchical Softmax Multi-scale Linear Attention)}, which combines ReLU-based linear attention for global context, selective softmax refinement for critical local features, and multi-scale token representations via depthwise convolutions. HSMLA achieves superior accuracy-efficiency trade-offs: up to $4.2\times$ inference-time speedup across dense prediction tasks, $87.3%$ Dice with $3.2\times$ speedup on CT organ segmentation, and $94.2%$ AUC with $4.1\times$ speedup on pathology WSI.

Problem

• vision transformers
• computational overhead
• self-attention
• linear attention
• local context modeling

Innovation

• Linear Attention
• Multi-scale Representation
• Efficient Vision Transformers
• Hierarchical Attention
• Dense Prediction