【UNet 改进 | 注意机制篇】UNet引入STA超级令牌注意力机制（CVPR 2023），稀疏关联采样打破高分计算瓶颈，二次创新

本文教的是方法，也给出几种改进方法，二次创新结构，百变不离其宗，一文带你改进自己模型，科研路上少走弯路。

前言

在医学图像分割任务中，病灶区域往往形态各异、边界模糊，且经常与周围组织的对比度较低，这要求模型具备极强的特征提取和细节辨别能力。传统的U-Net网络虽然通过跳跃连接融合了多尺度特征，但在处理复杂语义和精细边缘时仍有不足。尤其是在早期的浅层网络中，如果使用局部卷积或局部自注意力，往往难以捕获长距离的全局依赖；而如果直接使用全局自注意力，又会因为图像分辨率过高（Token数量巨大）导致难以承受的计算开销。为了解决这些痛点，本文引入了在计算机视觉领域表现出色的 STA（Super Token Attention，超级令牌注意力机制）。STA 巧妙地借鉴了超像素（Superpixel）的思想，通过稀疏关联学习对视觉内容进行语义上有意义的聚合，大幅减少了参与自注意力计算的 Token 数量，同时完美保留了全局上下文建模的能力。这种创新设计使得网络能够以极高的计算效率在浅层就开始捕获全局依赖，让 U-Net 模型在医学图像病灶分割上取得显著突破！

论文发表信息

• 论文标题：STViT: Vision Transformer with Super Token Sampling
• 发表会议/期刊：CVPR 2023（计算机视觉顶会）
• 作者单位：中国科学院自动化研究所（CASIA）、中国科学技术大学等
• 核心贡献：提出了一种超级令牌注意力（STA）机制。该机制首先通过稀疏关联学习从原始视觉 Token 中采样出数量较少的超级令牌（Super Tokens），随后在超级令牌上执行自注意力计算，最后再将更新后的特征映射回原始 Token 空间。STA 成功将传统全局注意力分解为稀疏关联图和低维注意力的乘积，极大地提升了全局依赖建模的效率。
• 论文地址：https://arxiv.org/abs/2211.11167
• 代码地址：https://github.com/hhb072/STViT

STA 核心模块理论

STA（超级令牌注意力）机制的核心思想在于通过语义上有意义的局部聚合来大幅降低自注意力的计算复杂度，同时保持全局建模能力。传统的 ViT 在高分辨率图像上计算全局自注意力会导致二次方增长的计算量，STA 通过以下创新实现了破局：

1. 超级令牌采样（Super Token Sampling）：受传统超像素算法的启发，STA 引入了稀疏关联学习。模块将输入特征划分为多个局部区域，并在每个区域内生成超像素中心（超级令牌）。通过计算原始像素与超级令牌之间的亲和度矩阵（Affinity Matrix），实现特征的高效聚合。
2. 超级令牌上的自注意力（Self-Attention on Super Tokens）：经过采样后，Token 的数量大幅减少。STA 直接在这些高度浓缩了语义信息的超级令牌上执行多头自注意力（MHSA）计算。由于参与计算的节点数量极少，这一步骤即使在浅层高分辨率特征图上也能以极低的计算成本捕获全局长距离依赖。
3. 特征映射还原（Token Upsampling）：在完成全局信息交互后，STA 利用之前计算得到的亲和度矩阵的转置，将更新后的超级令牌特征反向投影回原始像素空间。这不仅恢复了特征图的空间分辨率，还将全局上下文信息精准地传递给了每一个局部像素。

作为一个兼顾计算效率与全局感受野的先进模块，STA 能够让模型在不增加过高计算负担的前提下，从早期阶段就具备捕获长距离依赖和复杂病灶特征的能力，非常适合应用于对精度和感受野要求极高的医学图像分割任务。

训练代码参考和下载：手把手带你完成基于Unet的医学图像分割系统设计与实现，Unet网络架构讲解，Unet复现，使用vibe coding工具，开发一个完整系统，包含完整源码

🐴一、实战细节

⚡⚡实验结果画图

画图效果如下，代码可一键运行

画图代码：

# -*- coding: utf-8 -*-""" @Auth ：落花不写码 @File ：画图.py @IDE ：PyCharm @Motto :学习新思想，争做新青年 """importmatplotlib.pyplotaspltimportpandasaspdimportos plt.rcParams