[论文解读] SegTransVAE: Hybrid CNN -- Transformer with Regularization for medical image segmentation

Quan-Dung Pham, Hai Nguyen-Truong+2 more|arXiv (Cornell University)|Jan 21, 2022

AI in cancer detection参考文献 12被引用 5

一句话总结

SegTransVAE 提出了一种混合 3D CNN-Transformer 架构，配备变分自编码器（VAE）分支用于医学图像分割，结合了通过 3D CNN 提取的局部特征、通过 Transformer 实现的全局上下文建模，以及通过 VAE 实现的自监督重建，以提升泛化能力。该方法在 BraTS 和 KiTS19 数据集上实现了最先进性能，Dice 分数和 95% Hausdorff 距离指标均优于 3D U-Net、UNETR 和 SegresnetVAE，同时推理速度与轻量级 CNN 相当。

ABSTRACT

Current research on deep learning for medical image segmentation exposes their limitations in learning either global semantic information or local contextual information. To tackle these issues, a novel network named SegTransVAE is proposed in this paper. SegTransVAE is built upon encoder-decoder architecture, exploiting transformer with the variational autoencoder (VAE) branch to the network to reconstruct the input images jointly with segmentation. To the best of our knowledge, this is the first method combining the success of CNN, transformer, and VAE. Evaluation on various recently introduced datasets shows that SegTransVAE outperforms previous methods in Dice Score and $95\%$-Haudorff Distance while having comparable inference time to a simple CNN-based architecture network. The source code is available at: https://github.com/itruonghai/SegTransVAE.

研究动机与目标

• 解决 CNN 在捕捉长距离空间依赖性方面的局限性，以及 Transformer 在缺乏归纳偏置和在低数据环境下泛化能力不足的问题，以提升医学图像分割性能。
• 通过联合训练分割任务与基于 VAE 分支的自重建任务，提升分割精度，减少过拟合并增强特征学习能力。
• 在统一的编码器-解码器框架中整合 3D CNN（局部不变性）、Transformer（全局注意力）和 VAE（正则化）的优势。
• 在具有挑战性的医学影像基准上实现最先进性能，同时保持与现有模型相当的推理速度。

提出的方法

• 编码器使用堆叠的改进型 ResNet 块，配合实例归一化和 Leaky ReLU，从输入体数据中提取分层 3D 局部特征。
• 将 3D 特征图展平并使用可学习的线性投影映射到 d 维嵌入空间，同时添加可学习的位置嵌入以保留空间结构。
• 通过多头注意力与前馈网络层堆叠形成 Transformer 编码器，实现对所有 patch 的自注意力机制，支持全局特征建模。
• 通过特征映射模块将 Transformer 输出重新塑形为 3D 特征图，并通过卷积块减少通道维度，作为解码器的输入。
• 在最终的 Transformer 层上附加一个 VAE 分支，用于重建输入图像，提供正则化并提升在小样本数据上的泛化能力。
• 解码器使用 3D 转置卷积，并通过来自编码器的跳跃连接重建高分辨率分割掩码。

实验结果

研究问题

• RQ1 基于 VAE 的自监督重建的混合 CNN-Transformer 架构是否能提升医学影像分割任务的性能？
• RQ2 VAE 正则化的集成是否能减少过拟合并提升泛化能力，尤其是在小样本医学数据集上？
• RQ3 结合局部（CNN）与全局（Transformer）特征学习的方法是否能在分割精度上超越纯 CNN 或纯 Transformer 模型？
• RQ4 与现有最先进模型相比，该架构在推理速度和参数效率方面表现如何？

主要发现

• 在 BraTS 2021 数据集上，SegTransVAE 在增强肿瘤（ET）上的 Dice 分数达到 85.48%，全肿瘤（WT）为 90.52%，肿瘤核心（TC）为 92.60%，优于 3D U-Net、UNETR 和 SegresnetVAE。
• SegTransVAE 在 95% Hausdorff 距离上分别达到 2.89 mm（ET）、3.57 mm（WT）和 5.84 mm（TC），显著优于基线方法。
• 在 KiTS19 数据集上，SegTransVAE 在肾脏分割中 Dice 分数达到 95.28%，在肿瘤分割中达到 66.31%，两项指标均优于 3D U-Net、UNETR 和 SegresnetVAE。
• 肾脏分割的 95% Hausdorff 距离为 3.28 mm，显著低于 3D U-Net（6.32 mm）和 UNETR（8.84 mm）。
• SegTransVAE 拥有 4470 万个参数，推理时间为 0.45 秒，与 3D U-Net（560 万个参数，0.45 秒）相当，且快于 SegresnetVAE（750 万个参数，0.55 秒）。
• 尽管参数量多于 3D U-Net，SegTransVAE 的 GFLOPs（607.5）低于 3D U-Net（1000+），表明其具有更高的计算效率。