💻 computer science

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

本文介绍了 HSMLA，这是一种结合了基于 ReLU 的线性注意力、选择性 Softmax 精炼以及多尺度深度卷积的新型注意力机制，旨在克服标准视觉 Transformer 的二次方复杂度，在保持医疗图像分割和病理分析等密集预测任务高准确度的同时，实现显著的推理加速。

原作者： Dong Liu, Yanxuan Yu, Renata Borovica-Gajic, Tong Geng, Ying Nian Wu

发布于 2026-08-12

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Dong Liu, Yanxuan Yu, Renata Borovica-Gajic, Tong Geng, Ying Nian Wu

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图绘制一幅巨大且细节极其丰富的城市壁画，但你手里只有一个小桶颜料，而且完成任务的时间非常短。如果你试图同时观察每一栋建筑里的每一块砖头，以决定使用什么颜色，那么在你还没开始动笔之前，时间就会耗尽。这就是一种被称为“视觉 Transformer”（Vision Transformer）的计算机大脑所面临的问题。这些数字艺术家在理解图像方面非常出色，但当图像变得巨大时——比如高分辨率的医学扫描图或自动驾驶汽车拍摄的街景——它们就会感到力不从心。它们试图将每一个像素与每一个其他像素进行比较，这就像是试图让体育场里的每一个人都同时与其他人互相介绍一样。这太慢了，而且消耗太多能量。

为了解决这个问题，科学家们尝试了一种名为“线性注意力机制”（linear attention）的捷径。与其让每个人都互相介绍，不如让每个人向全场喊出一个总体的摘要。这非常快，但缺点是细节会变得模糊。这就像是听到人群在欢呼，但你却听不清每个人具体在喊什么词。对于需要寻找人体扫描中肿瘤确切边缘或面部细微线条的任务来说，这种模糊是无法接受的。因此，一个大问题一直是：我们能否既拥有捷径带来的速度，又拥有细致观察带来的清晰度？

于是，一项名为 HSMLA（分层 Softmax 多尺度线性注意力机制）的新发明诞生了，它由一组研究人员团队创造。把 HSMLA 想象成一位超级聪明的艺术总监，他知道什么时候该打破规则。HSMLA 不会对整幅壁画采取同样的处理方式：对于图像中那些无聊、空旷的部分（比如天空或平整的墙壁），它使用高效的方法；而一旦它发现了一些棘手的东西——比如一根参差不齐的树枝、复杂的建筑边缘或医学图像中可疑的点——它会立即切换到“超级模式”。它会放大并针对那个微小且重要的点，进行缓慢且细致的逐像素比较。

论文显示，这种“混搭”方法是一个游戏规则的改变者。通过仅对大约 30% 的图像（即真正需要关注的部分）进行繁重、缓慢的工作，而将其余部分交给快速方法，HSMLA 在超分辨率（使模糊图像变清晰）等任务上的运行速度比标准方法快了高达 4.2 倍。在医学成像领域，结果更加令人印象深刻。在用于寻找器官的 CT 扫描中，这种新方法实现了 87.3% 的 Dice 分数（衡量计算机勾勒器官轮廓完美程度的指标），同时比之前的标准方法快了 3.2 倍。在用于检测癌症的病理切片中，它达到了 94.2% 的 AUC（衡量检测准确度的指标），并实现了 4.1 倍的加速。

研究人员在从识别城市街道上的汽车到发现微小肺结节等各种任务上进行了测试，结果是一致的：你不再需要在速度和准确度之间做选择。该系统足够聪明，知道何时可以滑行，何时需要冲刺。这不仅仅是一个理论构想；团队构建了它，并在用于自动驾驶汽车和医疗设备的芯片等真实硬件上对其进行了测试，证明了它的有效性。他们发现，通过使用一个“门控”系统来决定图像的哪些部分需要缓慢、仔细的观察，他们可以在保持全局图像清晰的同时，精准地锐化局部细节。这有点像拥有一支画家团队：背景艺术家以闪电般的速度工作，而几位大师级的细节画家只在处理精细的花朵和面孔时被请出来，从而确保整幅壁画在快速完成的同时，不会丢失任何一笔细节的质量。

技术摘要：HSMLA（分层 Softmax 多尺度线性注意力机制）

1. 问题陈述

视觉 Transformer (ViTs) 在各种视觉任务中取得了最先进的性能，但在高分辨率密集预测中面临着显著的计算瓶颈，这是由于自注意力机制具有二次方复杂度 ( $O(N^2)$ )。虽然线性注意力机制提供了 $O(N)$ 的复杂度，但它们通常会产生较平坦的注意力分布，这削弱了对局部上下文的建模能力。这种局限性在需要锐利、局部化注意力的应用中尤为关键，例如医学成像（器官边界、肿瘤边缘）以及高分辨率密集预测任务（如语义分割和超分辨率）。现有的高效架构通常统一地应用线性注意力，或依赖于旨在用于分类而非密集预测的令牌剪枝策略，这导致无法在不牺牲效率的情况下恢复必要的局部表达能力。

2. 方法论

作者提出了 HSMLA（分层 Softmax 多尺度线性注意力），这是一种旨在平衡全局效率与局部精度的混合注意力机制。该架构的工作原理是将昂贵的完整 Softmax 计算仅分配给一小组关键的空间区域，而对其余部分则依赖于高效的线性注意力。

核心组件：

多尺度 ReLU 线性注意力（全局骨干网络）：
受 EfficientViT 的启发，HSMLA 使用基于 ReLU 的线性注意力来高效计算全局上下文。输入令牌通过多尺度深度卷积（使用 3×3、5×5 和 7×7 的卷积核大小）进行投影和增强，以捕捉多个感受野下的局部上下文。线性注意力机制利用结合律将复杂度降低到 $O(Nd^2)$ ，在每层计算一次全局聚合项（ $Z$ 和 $D$ ）并将其应用于所有查询。
分层块级门控（Hierarchical Block-Level Gating）：
为了解决线性注意力缺乏局部锐度的问题，HSMLA 引入了一种在块级别（连续的空间块）而非令牌级别运行的学习门控机制。
- 头级评分（Head-Level Scoring）： 每个注意力头根据多尺度特征计算门控分数。
- 块级聚合（Block-Level Aggregation）： 这些分数被聚合以产生每个空间块共享的门控值 $M_b \in [0,1]$ 。
- 路由（Routing）： 只有门控值超过阈值 $\tau$ （通常选择约 30% 的块）的块会被路由到局部 Softmax 精细化路径。这种块级方法确保了连续的内存访问，并避免了逐令牌的分支分歧，使其对硬件更加友好。
双路径残差内核（Two-Path Residual Kernel）：
HSMLA 并没有采用用线性注意力直接替换 Softmax 的二元替换方案，而是采用了残差精细化策略：
- 全局路径： 每个令牌都通过线性注意力路径保留其全局上下文（ $O_{linear}$ ）。
- 局部路径： 对于选定的块，在小窗口（ $w \times w$ ）内计算局部 Softmax 注意力。
- 残差修正： 最终输出是全局线性输出与稀疏残差修正（ $R_i = O_{local} - O_{lin-local}$ ）之和。这种公式确保了全局上下文永远不会被丢弃，且昂贵的 Softmax 操作被限制在局部邻域内，将复杂度从 $O(\alpha N^2 d)$ 降低到 $O(\alpha N w^2 d)$ 。
硬件感知设计：
训练过程包含两个正则化项，以强制执行硬件友好的激活模式：
1. 预算项（Budget Term）： 惩罚偏离目标 Softmax 预算比例 ( $\rho$ ) 的行为，以确保可预测的 FLOPs 和内存流量。
2. 空间平滑项（Spatial Smoothness Term）： 惩罚相邻块门控之间的突变，鼓励形成连续的瓦片簇而非分散的激活。这优化了内存层次结构的使用（HBM $\to$ L2 $\to$ SRAM）并减少了内核启动开销。
块组合（Block Composition）：
HSMLA 被集成到一个“三明治”架构中（类似于 EfficientViT），将注意力模块与深度卷积和前馈网络 (FFN) 交替排列。

3. 主要贡献

局限性分析： 本文指出标准线性注意力无法建模密集预测所需的锐利局部结构，特别是在医学成像领域。
HSMLA 架构： 一种新型的分层注意力模块，结合了多尺度线性注意力与内容感知、查询稀疏的选择性 Softmax 门控。
残差精细化策略： 一种双路径公式，将全局线性聚合与稀疏局部修正解耦，在保持全局上下文的同时增加了局部精度。
硬件协同设计： 引入了特定的正则化损失和瓦片紧凑型推理流水线，以确保学习到的门控掩码与 GPU 内存层次结构及执行模式相匹配。

4. 实验结果

作者在四种视觉场景下评估了 HSMLA：语义分割、图像分类、超分辨率和生物医学成像。

语义分割 (Cityscapes & ADE20K)：
- HSMLA-B2 在 Cityscapes 上实现了 81.8% mIoU，比 SegFormer-B2 快 3.7 倍。
- 在 ADE20K 上，它达到了 45.9% mIoU，同时保持 3.7 倍的加速。
- 它在准确率（+1.5 mIoU）和延迟方面均优于 EfficientViT-B2。
图像分类 (ImageNet-1K)：
- HSMLA-B2 实现了 82.1% Top-1 准确率，与 ConvNeXt-B 持平，并优于 EfficientViT-B2 (+0.3%)，且比 Swin Transformer-B 快 2.3 倍。
生物医学成像：
- CT 器官分割 (BTCV)： 实现了 87.3% Dice 和 8.6 mm HD95，速度提升 3.2 倍，在准确率和边界描绘方面显著优于 TransUNet 和 SwinUNet。
- 病理 WSI (CAMELYON16)： 达到了 94.2% AUC，速度提升 4.1 倍。
- 手术器械分割 (EndoVis 2017)： 在 Jetson AGX Orin 上实现了 89.7% Dice，帧率为 42 FPS，且降低了内存占用。
- 肺结节检测 (LUNA16)： 达到了 95.0% 敏感度，速度提升 2.9 倍。
超分辨率 (DIV2K)：
- HSMLA-B1 实现了 31.92 dB PSNR，比 Restormer 快 4.2 倍。

5. 重要性与主张

本文声称 HSMLA 成功弥合了线性注意力的效率与 Softmax 注意力的局部表达能力之间的差距。通过仅将稀疏的查询子集路由到全量 Softmax 注意力，该方法在保持或提高准确率的同时，在各项密集预测任务中实现了高达 4.2 倍的推理时间加速。

作者强调，HSMLA 对于高分辨率和医学成像应用特别重要，因为在这些应用中：

模型必须在边缘设备（如 Jetson 类 GPU）上受限于严格的延迟和内存预算运行。
保留精细结构细节（器官边界、肿瘤边缘）至关重要。
计算可以动态分配到具有诊断相关性的区域，而不牺牲全局上下文。

这项工作表明，将多尺度线性注意力与内容感知、查询稀疏的 Softmax 精细化相结合，是实现高效密集预测的一个实用且有效的方向，为资源受限环境下的标准二次方注意力机制提供了一个可行的替代方案。