TMSS: 一个端对端的用于分割以及生存预测的基于transformer的多模态网络

github:GitHub - ikboljon/tmss_miccai: TMSS: An End-to-End Transformer-based Multimodal Network for Segmentation and Survival Prediction

概要

当肿瘤学家估计癌症患者的生存期时，他们依赖于多模态数据。尽管文献中已经提出了一些多模态深度学习方法，但大多数方法依赖于拥有两个或多个独立的网络，在整个模型的后期阶段共享知识。

这项工作提出了一种深度学习方法，模拟肿瘤学家在量化癌症和估计患者生存率时的分析行为。我们提出了TMSS，一个基于transformer的端到端多模态网络分割和生存预测，利用transformer的优势，处理不同模态。

在HEad & NeCK TumOR segmentation和PET/CT图像挑战(HECKTOR)中结果预测的训练数据集上训练并验证了该模型的分割和预后任务。
我们所提出的预后模型显著优于最先进的方法，其一致性指数（concordance index）为0.763±0.14，同时实现了与独立分割模型相当的dice score为0.772±0.030。

介绍

背景

以往研究使用基于cnn的模块提取医学图像的特征，并与电子健康记录连接。输出通过全连接层传递到多任务逻辑回归(MTLR)（multi-task logistic regression）模型。同时，将电子健康记录输入Cox比例风险(Cox - proportional hazard, Cox - ph)模型来预测患者的风险评分。最后，通过取MTLR和CoxPH模型输出的平均值来计算风险预测。该集成模型实现了一致性指数(C-index：一种评估预后准确性的常用指标)在HECKTOR测试集上为0.72，优于其他提出的解决方案。虽然这项工作使用了多模态数据，但医学图像和电子病历的学习是不相关的，这可能会导致深度学习模型学习到的特征辨别能力较差，从而影响最终结果。

贡献

TMSS 结合CT、PET扫描及患者EHR(electronic health record），预测肿瘤的分割掩膜及患者的生存风险评分。标准卷积神经网络主要关注成像模态，不能使用其他输入特征。为了解决这个问题，我们提出了一个基于Transformer的编码器，它能够处理可用的多模态输入数据以及它们之间的交互。

-提出了TMSS，优于在相同数据集上训练的SOTA模型。

-展示了一个视觉转换器编码器可以关注多模态数据来预测分割和疾病结果，其中多模态数据被投影到相同的嵌入空间。

-我们提出了一个用于分割掩码和风险评分预测的组合损失函数。

*Pet扫描

PET是唯一可在活体上显示生物分子代谢、受体及神经介质活动的新型影像技术，现已广泛用于多种疾病的诊断与鉴别诊断、病情判断、疗效评价、脏器功能研究和新药开发等方面。

多用于：肿瘤，神经系统，心血管疾病

PET扫描用于训练不同流行的CNN架构，如AlexNet [16]， GoogleLeNet[25]和ResNet[12]，所有这些架构都在ImageNet[4]上进行了预训练，以将其与在临床记录上训练的传统方法进行比较。通过比较所有四种不同的方法，他们得出结论，使用PET扫描诊断DL模型可以预测无进展生存期和治疗结果。

*多任务逻辑回归(MTLR)（multi-task logistic regression）

多任务逻辑回归(MTLR)是一种用于同时解决多个相关分类问题的机器学习模型。在MTLR中，目标是学习一组参数，这些参数可以准确地预测一组相关任务的标签。MTLR假设每个任务都有自己的特性集，但是这些特性在所有任务中都是共享的。该模型由多个逻辑回归分类器组成，每个任务一个，具有共享的特征表示。共享的特征表示允许模型在所有任务中学习特征和标签之间的关系，提高其对每个单独任务做出准确预测的能力。

网络结构

 所提出的TMSS体系结构和多模态训练策略的说明。TMSS将EHR和多模态图像线性投影到特征向量中，并将其馈送到变压器编码器中。CNN解码器使用输入图像进行馈电，不同层的跳过连接输出，最后一层输出进行分割，而预测端则使用编码器最后一层的输出进行风险评分预测。

Transformer编码器 Transformer Encoder

该网络的主要优点是编码器本身嵌入CT/PET和EHR数据，并相应地编码它们的位置，同时提取不同模态之间的依赖关系(即注意力)。

维数的三维图像重塑变成一个平面二维的序列

，其中H、W、D分别为三维图像的高度、宽度和深度，C为通道数， P × P × P表示各patch的尺寸，是提取的patch的数量。然后将这些patch投影到嵌入维数h，形成一个矩阵。

同时，EHR数据也被投影到一个维度

 图像投影与EHR相连接，形成一个矩阵

 将具有相同维度的位置编码添加到每个patch和EHR投影中作为可学习参数。

我们把类Token 从Vitsion Transformer 中删除，因为我们的解决方案没有处理分类任务。

生成的embeddings被馈送到由12层组成的变压器编码器，遵循与原始ViT相同的pipline，具有归一化、多头注意和多层感知器。

使用自注意的目的是学习n + 1个嵌入之间的关系，包括图像和EHR。多头注意中的自注意可以写成:

 

分割端 Segmentation End

分割端是一个基于cnn的解码器，类似Unetr中的解码器。原始图像与从ViT层Z3、Z6、Z9和Z12(最后一层)传递的跳过连接一起被馈送到解码器。

只有图像潜在表征通过这些跳跃连接然后输入的cnn解码器，

 上采样阶段使用卷积、反卷积、批归一化和整流线性单元(ReLU)激活。

医学预测端 Prognostic End

预测路径接收具有维度的编码器的输出，并计算其平均值，将维数降至然后，这个潜在向量被转发到两个完全连接的层，分别将维数从h降低到512和128。然后将得到的特征图馈送到MTLR模型，用于最终的风险预测。MTLR模块将 Future Horizon 划分为不同的time bins，设置为超参数，对于每个time bins，使用逻辑回归模型来预测事件是否发生。

*future horizon

Future horizon指的是在未来的一段时间内所做的预测或计划。它是当前和未来时间点之间的时间间隔，与特定问题或决策过程相关。未来的范围可能因所考虑问题的背景和性质而有所不同。例如，在金融预测中，未来的范围可能是几个月或几年，而在气候建模中，未来的范围可能是几十年甚至几百年。一般而言，在作出预测或计划时，未来视界是一个重要的考虑因素，因为较长的未来视界可能与较大的不确定性和潜在风险有关。因此，重要的是要仔细考虑用于预测或计划的数据和模型的准确性和可靠性，并纳入应急计划或替代方案，以防意外的发展或变化。

损失函数 Loss Function

由于网络同时执行两个任务，三种损失的组合被表述为最终的目标函数。

分割端

由dice loss(公式2)和focal loss(公式3)之和支持，其中N为样本量，p^为模型预测，y为ground truth， α为焦点损失中精度和召回率权衡的权重(设为1)，γ为聚焦参数(经验设为2)。

 预测端

具有负对数似然损失(NLL)，如式(4)所示。这里，NLL loss 中的第一行对应未约束数据，第二行对应约束数据，第三行是归一化常数。

是模型预测， bk为偏差项，yk为ground truth。

最终损失

我们的网络训练使用的最终loss在公式5中，作为三种损耗的组合。超参数β，提供了模型路径两侧的权重，按照经验设置为0.3。

实验设置

数据集

HECKTOR挑战平台1提供了PET和CT图像的多中心数据集、它们的分割掩码和电子健康记录。

训练和测试的记录分别为224和101份。对于分割和预测任务，测试集的ground truth被隐藏，因此不能用于验证我们的方法。因此，在训练集上进行k-fold(其中k=5)交叉验证。

EHR包括与性别、体重、年龄、肿瘤分期(N-、M-和t -期)、烟酒消费、化疗经历、人乳头瘤病毒(HPV)、TNM,TNM组版本等相关的数据
影像数据包括CT、PET、肿瘤分割mask;样本切片分别如图2所示

 

 图2:来自成像数据集的样本。(a)为原始PET扫描图;(b)为原始CT扫描图，外加真实值蒙版;(c)为80 × 80 × 48裁剪PET图;(d)为80 × 80 × 48裁剪CT图，外加真实值蒙版。

预处理

CT和PET图像都重新采样到各向同性体素间距1.0 mm^3。它们的强度值在被馈送到网络之前被标准化。对CT图像的HU值窗口进行经验剪辑(- 1024,1024)，之后图像在(- 1,1)之间归一化。另一方面，PET图像采用Z-score归一化。此外，图像裁剪到80 × 80 × 48mm3，如[21]，主要有两个目的:

第一个是公平地将我们的结果与[21]中的最先进技术进行比较，后者也使用了这些尺寸的图像。

其次，这种图像维度的减少，反过来又加快了训练和推理过程，并允许运行多个实验。

EHR是多中心的，大多数中心缺少一些关于烟草、酒精消费量、性能状况、HPV状况和SUV估计重量的数据;因此，他们被放弃了。75%的数据被审查，被认为停止了对医院的追踪。

消融实验

在我们的实验中，我们使用了一台NVIDIA RTX A6000 (48GB)。

我们使用PyTorch来实现网络，并训练模型50个epoch。

batch size: 16，

学习率: 4e-3，

weight decay设置为1e-5,

步进衰减学习率策略(step decay learning rate strategy)用于在35个epoch后将学习率降低10倍。
 

图像(scans)被修补成16 × 16 × 16的大小，投影到768的嵌入维度。encoder中使用的总层数为12，每层有12个attention heads。

损失函数中的β设为0.3。所有超参数均采用OPTUNA[1]框架进行经验选择。

预后风险的评价指标为concordance index(C-index)，分割的评价指标为dice similarity coefficient(DSC)。

*OPTUNA

Optuna是一个自动超参数优化软件框架，专为机器学习而设计。它具有命令式、 运行时定义的用户 API。Optuna的用户可以动态地构建超参数的搜索空间。

实验结果

我们使用HECKTOR数据集用于头颈部癌症患者的诊断和预后。我们在内部进行了几个实验，所有实验都使用挑战中的训练数据集进行5倍交叉验证。所有实验都使用相同的设置进行训练和交叉验证。

table1： 不同模型在HECKTOR数据集上的预测性能。
报告的是5折交叉验证的平均值和标准偏差。

所有实验结果如表1所示。我们从生存分析的常用算法开始。CoxPH, MTLR和Deep 采用MTLR作为基线。考虑到它们的体系结构性质，这三种方法都只使用EHR数据来计算风险。MTLR+cnn和EHR+CoxPH的集合在测试集上获得了最高的c-index（0.704），该测试集也实现了与原始工作相同的训练和验证方式。最后，我们的模型在输入中embed了EHR信息，并使用了不同于ensemble的transformer，优于所有其他模型，平均c-index为0.763。

我们只使用其中一个fold来优化超参数，然后使用整个数据集进行k-fold交叉验证。然而，对于特定的fold有遗漏的风险，我们用hold-out test set重新做了测试集。

为了排除模型训练的统计依赖性，我们将数据集随机分为两个子集，训练和测试，分别占80%和20%的比例。模型超参数使用训练集的一个小子集进行优化，并使用hold-out test set进行测试。我们在测试集上的Cindex分数是0.74。

在hold-out test set上的生存预测（prognosis task score）任务得分略低于k-fold交叉验证得分0.76，但它更可靠，并且大于之前的最佳得分0.70。

为了进行细分比较，我们实现了UNETR，这是一个细分独立网络。我们的模型实现了
DSC为0.772±0.03，仅比分割优化后的UNETR网络DSC为0.774±0.01低0.002。

*hold-out set：就是我们传统意义上的测试集（test set）的别称

*UNETR

讨论

传统的癌症患者自动化诊断和预后的方法通常分两个阶段进行;要么是一个提取肿瘤放射组学(如肿瘤体积[15])并将其提供给预后模型的独立网络，要么是在SOTA[21]中使用一组cnn来提取扫描特征并与EHR连接，然后将它们提供给另一个网络进行风险预测。

然而，我们的方法在端到端网络中同时解决了这两个问题，使其更简单、更容易训练。我们展示了如何我们的方法优于其他模型的良好边际使用Vit。将EHR数据编码到网络中是新引入的，以模仿医生审查患者数据的方式。

这有效提高了预后的准确性，如表1所示。上述结果表明了Transformer 在处理多模态数据方面的优越性。我们假设，嵌入在transformer block 中的注意力，以及它们容纳多模态数据的能力，允许它们在模态之间和中介表示中找到关系。这可以帮助他们更好地处理手头的任务。多重损失的使用提高了模型在给定数据中更好地内插的能力，并有望在遇到未见过的数据时变得更加健壮。权重变量β的引入值为0.3，对模型的预后错误惩罚更大，迫使它更好地学习特征，并相应地调整其权重以获得准确的预后。

虽然我们的模型的主要目标是预后，而不是分割，但我们获得了与UNETR相当的结果，后者对分割进行了优化。
这进一步证明了我们的假设，即这两项任务相互补充和帮助，以获得更好的表现。它还阐明了如何改进分割任务进而完善预后结果，并帮助模型更好地学习图像和EHR数据的表示。

总结

在这项工作中，我们提出了一个头颈部肿瘤诊断和预后的端到端多模式框架。该模型利用了Transformer在处理多模态数据方面的优势，以及它在模态内部和跨模态之间找到长期关系的能力，以获得更好的模型性能。我们在头颈部CT/PET图像和患者EHR上训练和验证我们的模型，并将我们的结果与当前最先进的预后和分割方法进行比较。对于未来的工作，可以探索网络的自我监督学习和预训练。事实证明，它们可以帮助模型更好地学习，尤其是在数据有限的情况下，就像我们的情况一样。此外，目前的网络可以应用于具有不同数据集的类似任务，以测试模型的泛化性。