方向:深度学习、计算机视觉、Transformer、医学图像
一、信息
期刊
Medical Image Analysis Q1
时间
2024
单位
Department of Computer Science, Johns Hopkins University, Baltimore, MD 21218, USA (约翰霍普金斯大学)
Department of Computer Science and Engineering, University of California, Santa Cruz, CA 95064, USA (加州大学圣克鲁兹分校)
Department of Radiation Oncology, Stanford University, Stanford, CA 94305, USA (斯坦福大学 放射肿瘤学系)
Shanghai AI Lab, Xuhui District, Shanghai, 200000, China (上海人工智能实验室)
The School of Medicine, Stanford University, Stanford, CA 94305, USA (斯坦福大学 医学院)
DAMO Academy, Alibaba Group, New York, NY 10014, USA (阿里巴巴达摩院)
The East China Normal University, Shanghai 200062, China (华东师范大学)
The Australian Institute for Machine Learning, University of Adelaide, Australia (澳大利亚阿德莱德大学 机器智能研究院)
数据集
BTCV 多器官分割 (Multi-organ Segmentation)
13 个腹部器官(如脾脏、胰腺、肾脏、肝脏等)。类别多,器官间对比度差异大。
BraTS2021 脑肿瘤分割 (Brain Tumor Segmentation)
脑部病变的不同子区域:增强肿瘤(ET)、坏死核心(NCR)和水肿(ED)。不同的肿瘤子区域可能在空间上重叠,需要模型具有处理非互斥类别的能力。
MSD HepaticVessel 肝脏血管分割
肝脏中的门静脉和肝静脉。血管是细小且分支复杂的结构,对模型的细节捕捉和连接性要求极高。
大规模内部胰腺肿块数据集 (Large-scale in-house Pancreatic Mass Dataset)
胰腺肿瘤/肿块。胰腺肿瘤在 CT/MR 图像中对比度通常较低,且目标体积小,边界模糊,极易被背景噪声干扰。
评估方法
Dice 相似系数 (Dice Similarity Coefficient, DSC)
敏感度 (Sensitivity)
Sensitivity = \frac{TP}{TP + FN}
临床/召回
召回率。衡量模型成功识别出所有阳性实例(前景)*的能力。在医学中,高敏感度意味着*漏诊率低(避免漏掉肿瘤)。
特异性 (Specificity)
Specificity = \frac{TN}{TN + FP}
临床/精确度
衡量模型成功识别出所有阴性实例(背景)*的能力。高特异性意味着*误诊率低(避免将健康组织误判为病变)。
95% Hausdorff 距离 (95% HD)
95\% HD = \text{max}_{95\%} ( \sup_{a \in A} \inf_{b \in B} d(a, b), \sup_{b \in B} \inf_{a \in A} d(a, b))
边界距离
衡量预测边界(A)与真实边界(B)之间的最大距离(取 95% 分位数以减少异常值影响)。HD 值越小,边界吻合度越高,分割越精细。
二、摘要
背景:长期以来在计算机视觉领域,卷积神经网络都是大家常用的方法,在医学图像分割领域以U-Net这样基于卷积的方法被持续沿用。但卷积方法在建模长距离依赖关系方面存在局限性。因此人们引用Transformer架构优化U-Net,但人们缺乏全面理解。
本研究的贡献:提出TransUNet通用框架,将Transformer的编码器解码器分别封装到两个模块中(1)编码器对卷积神经网络的特征图的图像块进行Token化处理,促进全局上下文提取;(2)解码器,通过候选区域与U-Net特征之间的交叉注意力机制细化候选区域。配置很灵活,形成三种配置:仅编码器、仅解码器、编码器+解码器。
实验结果:编码器建模在多个腹部器官之间的相互作用方面效果显著,解码器在处理肿瘤等小目标方面表现出色。在多器官分割与胰腺肿瘤分割中处于领先位置。
三、引言
U-Net
U-Net的架构有两个主要组成部分:CNN编码器、CNN解码器,
CNN编码器
作用:负责收缩路径,特征提取和捕捉上下文信息。
结构:由重复的卷积层和下采样(如最大池化)操作组成。
过程: 随着网络层级的深入,特征图的空间分辨率(宽和高)逐渐减小,而特征通道数逐渐增加(深度增加)。这使得网络能够学习到图像的更抽象、更全局的特征。
CNN解码器
作用:负责恢复空间分辨率,并利用编码器捕获的特征来精确定位分割目标。
结构:由上采样(如转置卷积或插值后接卷积)操作和重复的卷积层组成。
过程:通过上采样,特征图的空间分辨率逐渐增加,恢复到接近输入图像的大小。
由于下采样时会丢失图像细节和位置信息。U-Net提出Skip Connections的跳跃连接方式,如图所示,编码器的每一层直接连接到解码器的对应层,以保证高分辨率的图像信息。
局限性
U-Net 基于标准卷积操作,其固有的局部性和对局部纹理的归纳偏置,使其难以高效且鲁棒地捕获和建模跨越大区域的全局结构信息和长距离空间依赖关系。因此,当不同患者的病灶在形状和大小上表现出巨大差异时,U-Net 难以准确适应这些高变异性的几何变化,导致分割精度下降。
Transformer
现在研究界越来越关注Transformer,2021年首次提出的TransUNet,这种方法同时利用U-Net编码器高分辨率空间细节,同时借助Transformer的全局上下文建模能力。尽管如此,对于Transformer的注意力机制在U-Net不同组件中的作用,仍缺乏全面理解。
本文贡献
1.在医学图像分割的序列到序列预测中融入了自注意力机制和交叉注意力机制。
2.将逐像素分割重新定义为掩码分类。提出从粗到细的注意力优化方法,以提升小目标/肿瘤的分割效果。
3.首次对Transformer在U-Net的编码过程中的战略性调整进行全面研究。
4.各个医学分割任务上遥遥领先。
四、相关工作
基于卷积神经网络
基于Transformer架构
用于分割的掩码分类
五、方法
3.1作为编码器的Transformer
1.图像序列化。
将输入图像重塑为扁平化的3D补丁序列,图像的输入维度是x\in\mathbb{R}^{D\times{H}\times{W}\times{C}},空间分辨率为D\times{H}\times{W},通道数为C。其中每个Patch大小为P\times{P}\times{P},N=DWH/P^3。使用可训练的线性投影将向量化补丁x^P映射到潜在的d_enc维度嵌入空间
LN()层归一化
3.2作为编码器的Transformer
3.2.1粗略候选估计
传统U-Net主要按照逐像素分割任务,每个像素都可能被分为K个类别之一,这通常使用逐像素交叉熵损失训练分割模型。
本文将医学分割视为一种掩码分类问题。引入器官查询概念,即d_{dec}维特征,它表示图像中每个器官向量,通过一组预定义的N个器官查询,对于包含K个分类的图像,我们要将图像分割为N个候选区域,随后为每个区域分配标签。
N的值不一定和类别数量一致,有意将N设置为大于K,以最小化假阴性风险。
粗预测的分割图可以通过初始化器官查询P^0\in\mathbb{R}^{N\times{d_{dec}}}与U-Net最后一个块特征F\in\mathbb{R}^{D\times{H}\times{d_{dec}}}(CNN编码器最后一层)的嵌入之间的点积来计算:
g(·)是sigmoid激活函数,阈值设置为0.5
3.2.2Transformer解码器
Transformer目的是优化器官查询,从而伴随由粗到细的策略改进Z。每层中的自注意力机制将使Transformer解码器能够全面处理图像特征并捕捉器官之间的关系。
同时训练CNN解码器和Transformer解码器,经过优化的器官查询表示为p^t\in\mathbb{R}^{N\times{d_{dec}}}。同时,中间的U-Net特征被映射到一个d_{dec}维的特征空间记为F。当上采样块数和解码器层数相同时,每一层都可以进行一次优化的查询,符合从粗-细的策略。
这是一个注意力机制的计算公式,第t个查询经过权重矩阵形成下一层查询。U-Net特征F经过权重矩阵变成K和V。同时使用残差连接。
3.2.3从粗到细的注意力优化
这里当z不为1的时候代表是背景噪声,在运算的时候给到负无穷在softmax之后就变为0,让它不参与注意力运算,把目光注意在前景(目标)。
3.3 TransUNet
流程
Input:
权重矩阵\mathbf{w}_q,\mathbf{w}_k,\mathbf{w}_v
初始器官嵌入,最后一层U-Net特征F
U-Net的第t层特征
最大迭代次数
Output:
Z_t:最终精细分割掩码 维度R^{N\times(D·H·W)},经过T次修正后的结果,每一行代表一个查询识别出的器官。1表示器官,0表示背景。
\hat{y} :预测类别标签 维度R^N。
计数器t置0
将初始查询变量赋予P_0
第一次尝试,直接用P_0与最深层的图像特征F做点积
循环
根据这一次的注意力与掩码更新下一次的查询
根据查询更新掩码
t=t+1
t=T
用 P_T 生成最终掩码 Z_T。
同时,通过一个线性层(公式 8、9)给这个掩码贴上标签(比如:这是“肝脏”)。
3.3.1 仅编码器
仅加入Transformer编码器部分,优点:
1.它允许我们在解码路径中利用中间的高分辨率CNN特征图
2.混合CNN-Transformer编码器的性能优于单纯使用纯Transformer作为编码器
使用:混合分割损失,该损失由逐像素交叉熵和骰子损失组成
因为Transformer编码器擅长捕捉全局上下文信息,因此它在多器官分割方面应该最为有效。
3.3.2 仅解码器
仅加入Transformer解码器部分,优点:
使用匈牙利匹配损失
Transformer解码器采用从粗到细的注意力机制来细化小而难得目标,应该更适合肿瘤分割。
匈牙利损失算法
第一步:构建代价矩阵 (Cost Matrix)
我们计算每一个预测结果 i 与每一个真实标签 j 之间的“距离”或“代价” C_{i,j}。这个代价通常由两部分组成:
分类代价:预测类别与真实类别是否一致。
分割代价:预测掩码与真实掩码的重叠度(通常用 Dice Loss 或 Focal Loss 计算)。
第二步:寻找最优分配 (\sigma)
通过匈牙利算法(Hungarian Algorithm),在矩阵中找到一种“一对一”的匹配方案,使得总的匹配代价达到最小。
第三步:计算损失函数 (Loss Calculation)
一旦匹配完成(例如:预测 A 匹配到了肝脏,预测 B 匹配到了脾脏,预测 C 匹配到了“空/背景”):
对于匹配成功的对:计算分类损失+分割损失。
对于匹配到“空”的预测:只计算分类损失(标签为背景)
3.3.3 编码器+解码器
六、实验与讨论
数据集:BTCV多器官分割数据集,BraTS2021脑肿瘤分割挑战数据集、医学分割十项全能(MSD)肝血管数据集、以及一个大规模的内部胰腺肿块数据集。
BTCV数据集严格遵循文献中的分割方式,其余四个数据集均进行无折交叉验证验证。
6.1数据集与评估
1. BTCV 多器官分割数据集 (Synapse)
来源:MICCAI 2015 挑战赛。
数据量:30 例腹部 CT 扫描(共 3779 张轴向切片)。
图像规格:分辨率为 512 \times 512,切片数 85 \sim 198 层。
空间分辨率:轴向分辨率约 0.54 \text{mm},层厚 2.5 \sim 5.0 \text{mm}。
分割目标(8 类):主动脉、胆囊、脾脏、左肾、右肾、肝脏、胰腺、胃。
数据划分:18 例训练(2212 张切片),12 例验证。
2. BraTS2021 脑肿瘤分割数据集
来源:脑肿瘤分割挑战赛,目前规模最大的脑肿瘤数据集。
数据量:1251 例多模态 MRI 扫描。
模态(4 种):T1、T1Gd(增强)、T2、T2-FLAIR。
预处理:已进行颅骨去除(Skull-stripping)、各向同性 1 \text{mm}^3 重采样及空间对齐。
图像规格:240 \times 240 \times 155。
分割目标:增强肿瘤 (ET)、周围水肿/侵犯组织 (ED)、坏死肿瘤核心 (NCR)。
数据划分:采用 5 折交叉验证。
3. MSD HepaticVessel 肝脏血管数据集
来源:Medical Segmentation Decathlon (MSD) 挑战赛任务。
数据量:443 例门静脉期 CT 扫描(包含各类原发性和转移性肝肿瘤患者)。
特点:肝血管呈管状且相互连接,紧邻异质性肿瘤,分割难度极大。
分割目标:肝脏内的血管和肿瘤。
数据划分:采用 5 折交叉验证。
4. 大规模内部胰腺肿块数据集 (Pancreas)
来源:美国某高容量医院收集的静脉期 CT。
数据量:2930 例 CT 扫描(规模极大的胰腺肿瘤专项数据集)。
临床背景:重点关注胰腺导管腺癌 (PDAC),其 5 年生存率仅约 10%。
分割目标:胰腺 (Pancreas)、导管腺癌 (PDAC)、囊肿 (Cyst)。
数据划分:2123 例用于训练(含 PDAC、囊肿及正常胰腺),807 例用于测试。
6.2实现细节
1. 架构配置与主干网络
主干网络:采用3D nn-UNet作为基础架构。TransUNet 展现出极强的适应性,会根据不同数据集自动调整下采样层数和各阶段的通道数(具体参考 Table 1)。
Transformer 编码器:实验了1 层和12 层 ViT两种方案。
12 层版本使用了在 ImageNet21k 上预训练的模型。
引入了LayerScale技术以稳定训练。
隐藏层维度:* 编码器维度 (d_{enc}):768
解码器维度 (d_{dec}):192
2. 训练设置与损失函数
硬件与批次:使用单张Nvidia RTX 8000 GPU,Batch Size 设为2。
数据增强:完全沿用 nn-UNet 的标准数据增强流程(包括旋转、缩放、翻转、噪声等)。
匈牙利匹配损失系数:* \lambda_0 = 0.7 (通常对应分类损失)
\lambda_1 = 0.3 (通常对应掩码分割损失)
3. 推理与测试策略
3D 优先:由于 3D 基准性能显著优于 2D(参考 Table 2),研究重心主要集中在3D 实验。
滑窗预测 (Sliding-window):处理大型扫描件时,将图像切块并滑动窗口进行推理。
结果聚合 (Aggregation):* 对于重叠区域的每个体素 (i, j, s),将所有 N 个预测块的概率向量进行累加:\sum_{n=1}^N (Z_{T n, ijs})。
最终决策:经过argmax操作,从聚合的概率向量中选出最大值,得到最终的“硬预测”标签。
6.3. 分析性研究
本节验证了关于 Transformer 编码器和解码器分工的核心假设:
三种配置对比:
Encoder-only:在BTCV 多器官分割上表现卓越。特别是使用 ImageNet 预训练的 12 层 ViT,Dice 分数提升了 0.8%(达到 88.11%),证明其擅长捕获全局器官关系。
Decoder-only:在MSD 血管/肿瘤分割上优势巨大,增益达 1.63%。证明“由粗到精”的机制非常适合细小且具有挑战性的目标。
Encoder+Decoder:综合表现最稳健。在多器官任务上甚至微超单编码器版本(88.39%)。
默认选型:由于 1 层编码器+解码器的组合在性能上与 12 层组合相当,且计算量更低,因此Encoder(1层)+Decoder被设定为默认架构。
泛化性:在内部胰腺肿瘤数据集上,Decoder-only 架构再次胜出,验证了其在处理小目标(PDAC 和 Cyst)时的普适性。
最优选择:Encoder+Decoder,因为它们在不同任务上取得了最佳结果,且网络参数数量相当少。
6.4. Transformer 解码器深度分析
对解码器的关键组件进行了剥离实验:
查询向量数:性能对 Query 数量不敏感。即便从 20 个变动到 40 个,结果依然稳定,说明模型鲁棒性强。
多尺度 CNN 特征:将 CNN 解码器不同分辨率的特征引入更新过程至关重要,相比单尺度特征有显著提升。
由粗到精的细化:
掩码机制有效性:通过公式 (7) 将注意力限制在前景,有效排除了背景噪声。
视觉验证:图 2 展示了预测掩码从第 1 次迭代到第 3 次迭代如何显著改善了肿瘤的边界质量。
位置编码:去除后性能仅微降 0.1%,说明卷积层本身已经隐含捕捉了足够的空间位置信息。
6.5. 与 SOTA 模型对比 (Comparison with SOTA)
TransUNet 在各维度均超越了现有最先进技术:
2D/3D 框架通用性:无论基于 U-Net 还是 3D nnU-Net,TransUNet 均表现优异。
超越 Transformer 竞品:在挑战性的胆囊分割任务中,比 CoTr、nnFormer 等模型提升了约10%。
顶尖赛事水平:性能超越了BraTS2021挑战赛的冠军方案(nnUNet-Large)
6.6. 小肿瘤检测分析 (Small Tumor Detection)
本节通过量化数据证明了 TransUNet 处理“极端情况”的能力:
极小肿瘤 (< 20mm):在 PDAC 肿瘤分割中,比 nnU-Net 基准高出9.7%DSC。
微小囊肿 (< 10mm):比基准高出4.3%DSC。
结论:虽然在大肿瘤上也有提升,但 TransUNet 在小目标上的精度增益最为显著。
6.7. 效率分析 (Efficiency Analysis)
在性能提升的同时,保持了工程上的可行性:
参数量与显存:Encoder+Decoder 版本(41.4M 参数)比 SwinUNETR-v2 等模型更轻量。显存占用< 12GB,支持在普通消费级显卡(如 Titan-XP)上训练。
计算开销:虽然由于注意力机制,耗时略高于纯卷积模型,但在推理速度、训练时间、FLOPs 等 5 个维度上均优于 3D UX-Net 等近期 Transformer 模型。
七、总结
7.1创新点
1.掩码分类
传统做法:像 U-Net 这种像素分类,容易出现“孤岛像素”或边界断裂,因为它不理解物体的整体性。
TransUNet 做法:它通过Organ Query,强制模型以“物体”为单位进行思考。模型输出的是 N 个完整的掩码方案。
意义:这赋予了分割结果更好的拓扑结构一致性。这也是为什么它在处理细长血管和复杂肿瘤时,比传统卷积方法更稳健的原因。
2.“残差式”由粗到细的闭环修正
Z_t(预测)\rightarrow 约束 Attention(缩小视野)\rightarrow 提取精细特征 \rightarrow Z_{t+1}(修正)。
它利用Masked Attention解决了 Transformer 容易在全局信息中“迷失细节”的问题。这种“带枷锁的注意力”强迫模型在迭代中不断“磨掉”背景噪声,只打磨目标边界。
3.证明了 Transformer 在 U-Net 中的“模块化分工”
Encoder适合多器官分割
Decoder适合肿瘤、血管分割