Slim UNETR面向计算资源受限场景的混合 Transformer 轻量化三维医学图像分割网络TMI 2024论文代码链接GitHub - aigzhusmart/Slim-UNETR: An Efficient, High-Quality 3D Segmentation for Medical Image Analysis with Constrained Computational Resources · GitHubhttps://github.com/aigzhusmart/Slim-UNETRAbstractHybrid transformer‑based segmentation approaches have shown great promise in medical image analysis. However, they typically require considerable computational power and resources during both training and inference stages, posing a challenge for resource‑limited medical applications common in the field. To address this issue, we present an innovative framework called Slim UNETR, designed to achieve a balance between accuracy and efficiency by leveraging the advantages of both convolutional neural networks and transformers. Our method features the Slim UNETR Block as a core component, which effectively enables information exchange through self‑attention mechanism decomposition and cost‑effective representation aggregation. Additionally, we utilize the throughput metric as an efficiency indicator to provide feedback on model resource consumption. Our experiments demonstrate that Slim UNETR outperforms state‑of‑the‑art models in terms of accuracy, model size, and efficiency when deployed on resource‑constrained devices. Remarkably, Slim UNETR achieves 92.44% dice‑accuracy on BraTS2021 while being 34.6× smaller and 13.4× faster during inference compared to Swin UNETR.翻译基于混合 Transformer 的分割方法在医学图像分析领域展现出巨大潜力。然而这类方法在训练与推理阶段均需要消耗巨大的算力与硬件资源这对于医疗领域普遍存在的资源受限应用场景构成巨大挑战。针对该问题本文提出创新网络框架 Slim UNETR该网络结合卷积神经网络与 Transformer 两者的优势实现分割精度与运行效率之间的平衡。该方法的核心是 Slim UNETR 模块通过对自注意力机制进行分解、低成本特征表征聚合实现高效的特征信息交互。除此之外本文将吞吐量throughput作为效率评价指标用来反映模型的资源消耗情况。实验结果表明在资源受限设备上部署时Slim UNETR 在分割精度、模型参数量、推理效率上均优于现有 SOTA 方法。值得注意的是在 BraTS2021 数据集上 Slim UNETR 取得 92.44% 的 Dice 精度对比 Swin UNETR模型参数量仅为其 1/34.6推理速度提升 13.4 倍。创新点1提出面向资源受限场景的轻量化混合 U 型三维医学分割框架Slim UNETR。融合卷积神经网络与 Transformer 各自优势在保证分割精度的同时大幅压缩模型参数量、提升推理速度解决现有 CNN‑Transformer 混合模型算力开销大、难以在低算力医疗设备部署的痛点。2设计核心Slim UNETR Block 模块。内部构建上下文集成流程拆解传统自注意力机制由局部特征聚合 LRC、LineConv、全局稀疏 Transformer GST、局部反向扩散 LRD 四部分组成。通过稀疏 token 策略避免对全部 token 执行注意力运算以较低计算成本完成局部与全局特征之间的信息交互。3提出将吞吐量throughput作为模型实际效率评价指标。指出传统 FLOPs 无法反映内存访问开销、硬件并行特性吞吐量可以更加真实地衡量资源受限硬件上模型的实际推理性能为轻量化医学分割模型提供更贴合落地场景的评估方式。4在 BraTS2021、MSD2019 脑肿瘤、MSD2021 肝血管三套大规模三维医学数据集开展充分实验。Slim UNETR 的整体网络框架该模型采用带有跳跃连接的 U 型编码器-解码器架构用于生成多尺度分层特征图。网络左侧为编码器一共包含 4 个阶段Stage1~Stage4每一个阶段均由 3D 深度卷积 DW Conv 与 Slim UNETR Block 顺序堆叠而成输入是维度为 H×W×D×4 的四维脑肿瘤多模态体数据随着层级逐级下采样特征图的空间尺寸不断缩小依次变为 H/4×W/4×D/4 、 H/8×W/8×D/8、H/16×W/16×D/16、H/32×W/32×D/32通道数同步提升在这一过程中网络逐层提取由细到粗的多尺度上下文特征。图右下角放大框展示了核心的 Slim UNETR Block也就是每个阶段内完成特征交互的基础单元模块内部依次执行局部特征聚合 LRC、线性卷积 LC、全局稀疏 Transformer GST、局部反向扩散 LRD最终采用残差连接输入特征与模块输出相加完成局部细节与全局长距离依赖的低成本信息交互规避传统完整自注意力带来的二次复杂度开销。网络右侧为解码器分支采用 3D 转置卷积 TS Conv 搭配 Block 完成逐级上采样逐步恢复特征图的空间分辨率编码器每一级输出的特征图会通过跳跃连接直接传递至解码器对应的层级与上采样后的特征逐元素相加弥补下采样过程丢失的细节信息。在完成全部解码阶段之后特征送入分割头 Seg Head最终输出 H×W×D×3 的分割预测结果对应脑肿瘤的三类分割子区域。整体架构继承 UNETR 系列 U 型编解码的范式但依靠轻量化 Slim UNETR Block在构建分层特征的同时大幅降低了混合 Transformer 架构的计算开销适配资源受限设备的三维医学图像分割任务。Slim UNETR Block 的内部结构核心该模块构建了一个低成本瓶颈结构内部设计上下文集成流程实现特征在局部表征与全局表征之间的转换以分解自注意力的方式用较低计算开销完成三维医学体数据的全局 - 局部特征交互。模块输入特征为X_in整体是多段残差串联的结构每一个子模块的输出都和输入特征进行逐元素相加也就是残差连接有效缓解深层网络训练时的梯度消失问题。模块首先进入LRC 局部表征聚合Local Representations Congregation分支对应图左侧放大子模块。LRC 模块由 BatchNorm、逐点 3D 卷积 Pointwise Conv3D、3D 深度卷积 Depthwise Conv3D、BatchNorm 以及另一层逐点卷积串联构成。其前向计算原文公式写为LRC 依靠深度卷积捕捉三维空间内局部邻域细节逐点卷积完成通道维度的特征融合专门负责提取医学图像的局部纹理、边界等细粒度信息计算完成后使用残差相加保留原始输入特征避免局部信息丢失。经过 LRC 之后特征送入LineConvLineConv 采用 Pointwise Conv3D、GELU 激活、Pointwise Conv3D 的结构完成通道间非线性变换对应公式LineConv 在不破坏三维空间结构的前提下做通道映射对局部聚合后的特征做进一步平滑与表征增强。接下来特征经过 Patch Partition 操作划分特征块后送入GST 全局稀疏 TransformerGlobal Sparse Transformer也就是图右下放大的子模块这是整个 Block 最核心的创新部分用于替代传统完整自注意力。GST 首先使用步幅为r的逐点卷积对特征做稀疏采样得到 Q、K、V 三组特征仅在稀疏采样后的 token 上执行多头自注意力而不是对全部 token 做全局计算以此降低 O(N^2)的二次复杂度。完成多头自注意力计算后再通过转置卷积恢复到原始分辨率最后搭配逐点卷积完成特征映射。GST 的计算过程原文表示为GST 只在少量 token 之间建模长距离依赖兼顾全局上下文建模能力同时大幅削减计算量与参数量解决传统 Transformer 在三维体图像上开销巨大的问题。GST 输出的特征送入LRD 局部反向扩散Local Reverse Diffusion模块。LRD 的作用是将 GST 学到的全局上下文信息扩散回原始稠密特征网格把全局语义信息重新注入到每个局部空间位置弥补稀疏采样造成的局部细节丢失计算公式补充反向扩散GST 阶段用步幅 r 的逐点卷积把一整块窗口内很多个体素压缩为 1 个中心 token在 LRD 中使用相同步幅 r 的 3D 转置卷积做反向映射该卷积核会把这一个中心 token 存储的全局特征值按照卷积核权重重新分配、加权叠加回它原本对应的整个窗口所有体素位置上相当于把单个 token 携带的全局信息扩散给到窗口里面每一个空间坐标完成信息的 “辐射分发”分发、辐射本质就是转置卷积的加权上采样运算将特征图分辨率还原为送入 GST 之前的稠密尺寸完成上采样之后再将得到的特征张量和 GST 输入的原始稠密特征做残差相加原始特征中 LRC 提取的局部纹理信息得以完整保留不会被覆盖最终每一个体素既保留原有的局部细节又融合了从稀疏中心 token 传递过来的全局上下文以此弥补 GST 稀疏采样仅对少量 token 做注意力造成的空间信息缺失同时保证张量维度匹配后续网络的计算要求。LRD 完成全局信息向局部空间的回传实现全局表征到局部表征的转换。经过 LRD 之后再次接入一层 LineConv做最后的通道非线性特征校准得到模块最终输出特征X_out完整模块整体前向过程整体来看Slim UNETR Block 这套上下文集成流程完成了“局部提取→通道映射→稀疏全局建模→全局信息回传”的闭环在同一个模块内部交替完成局部特征与全局特征之间的转换。和 UNETR、Swin UNETR 中直接对全部 token 做自注意力的方式不同Slim UNETR Block 通过分解自注意力、稀疏 token 采样在保留全局上下文建模能力的前提下显著降低计算开销是整个 Slim UNETR 轻量化、高精度的核心基础单元该模块可以堆叠在编码器的每一个 Stage 中分层提取多尺度三维医学特征。补充Slim‑UNETR 对自注意力机制的分解原理传统完整的多头自注意力需要全部 token 两两之间计算相关性复杂度是O(N^2)三维医学体数据 token 数量巨大直接计算显存和算力开销极高。Slim‑UNETR 并不是修改注意力数学公式而是把完整的 “局部提取‑全局注意力‑全局信息回写” 这一个整体自注意力流程拆解成三个子阶段用卷积与稀疏注意力分工协作以此实现对自注意力机制的分解。第一步由LRC 局部表征聚合先用卷积在各个局部窗口内完成局部信息的聚合浓缩把稠密的原始 token 集合提炼出一批代表局部窗口信息的中心 token不再拿全部原始 token 进入注意力计算。第二步送入GST 全局稀疏 Transformer只对上面得到的少量中心代表 token 执行标准多头自注意力完成跨窗口的全局依赖建模这一步只在稀疏采样后的少量 token 之间做(Q,K,V)计算避开全部 token 两两交互带来的平方复杂度。第三步依靠LRD 局部反向扩散把 GST 在少量中心 token 上学到的全局语义通过转置卷积反向扩散把全局信息重新分发回原始全部稠密空间位置。简单来说它把原本一步完成的 “全部 token 同时做全局自注意力”拆解成卷积做局部浓缩 → 少量代表 token 做稀疏全局注意力 → 再把全局信息扩散回全部位置。用卷积承接局部运算仅用少量 token 承担全局注意力运算实现分解。Slim UNETR Block 内部的上下文集成流程按照论文原文描述该流程旨在提升大规模分割模型性能的同时降低计算开销整体由 LRC 局部表征聚合、GST 全局稀疏 Transformer、LRD 局部反向扩散三个核心子过程构成完成局部特征提取→全局上下文建模→全局信息回灌到局部的完整闭环。图 (a) 为局部表征聚合 LRCLocal Representations Congregation。该模块将输入三维体特征切分为多个 3D 窗口嵌入块在每个窗口内部利用卷积对窗口内的局部 Patch 做特征聚合提取图像局部纹理、边界这类细粒度信息把原始稠密特征浓缩为局部集中特征。该步骤将原始输入拆分成更小的子集降低后续模块的处理难度专门负责捕捉三维医学图像上的邻域局部信息。图 (b) 是全局稀疏 Transformer GSTGlobal Sparse Transformer。它接收 LRC 输出的局部集中特征选取中心注意力 Patch 作为计算基准仅在稀疏采样得到的少量 Patch 之间执行自注意力计算学习跨局部窗口的高层全局依赖。不同于传统 Transformer 对全部 token 做注意力GST 只对稀疏选取的 Patch 建模长距离关联以此规避原生自注意力 O(N^2) 的高计算复杂度在大幅减少计算量的前提下学习全局语义模式。图 (c) 为局部反向扩散 LRDLocal Reverse Diffusion。该模块的功能是把 GST 学到的全局语义知识重新扩散、融合回每一个局部表征当中。它将全局注意力得到的特征反向映射回原始稠密的 3D 窗口网格生成局部扩散特征把全局上下文信息注入各个局部位置补偿稀疏注意力采样过程丢失的空间细节完成从全局表征到局部表征的转换。模型效率评估论文指出传统 FLOPs 仅代表理论计算量无法反映内存读写、硬件并行等实际硬件约束不能代表低算力设备上真实推理性能。因此论文引入吞吐量 Throughput作为核心效率评价指标单位为 vol/s代表每秒能够处理的三维体数量。吞吐量直接在目标硬件上实测得到可以真实体现模型在资源受限设备上的推理速度。作者同时结合参数量、Dice 指标做综合评估。实验结果显示Slim UNETR 相比 Swin UNETR在 Dice 更高的情况下参数量缩小 34.6 倍吞吐量提升 13.4 倍证明该轻量化方案可以在不损失精度的前提下大幅降低部署成本。消融实验论文设计多组消融实验验证上下文集成流程中 LRC、GST、LRD 各个组件的有效性同时探究稀疏比例、窗口大小等超参对性能与效率的影响。实验逐步移除 LRC、GST、LRD 模块观察分割 Dice 指标的下降幅度。实验结果证明移除任意一个子模块分割精度都会明显下降证明 LRC 局部特征聚合、GST 稀疏全局注意力、LRD 局部反向扩散三者缺一不可协同完成局部与全局特征交互。除此之外作者针对 GST 的稀疏采样比例做对照权衡稀疏度带来的精度 - 速度取舍同时测试不同窗口尺寸分析窗口大小对全局建模能力与显存占用的影响选取最优配置。消融实验完整验证了 Slim UNETR Block 每一个设计的必要性。