arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2408.05889cs.CV

通过Token级表示学习增强医学图像3D Transformer分割模型

Enhancing 3D Transformer Segmentation Model for Medical Image with Token-level Representation Learning

University of Notre Dame(圣母大学)

Xinrong Hu, Dewen Zeng, Yawen Wu, Xueyang Li, Yiyu Shi

AI总结:

针对医学图像3D分割预训练难题,提出token级表示学习损失及“旋转-恢复”防坍塌机制,改进对比损失,在两个公开数据集上性能优于现有最优预训练方法。

摘要:

在医学图像领域,尽管已有诸多研究证实Swin Transformer在逐像素稠密预测任务中表现出色,但在不使用额外数据集的情况下对这类模型进行预训练,能否进一步提升下游语义分割任务的性能,目前仍有待探索。此前的表示学习方法受限于3D体数据数量有限且计算成本高昂,难以广泛应用。此外,大多数专为Transformer设计的预训练任务并不适配Swin Transformer的层级结构。为此,本研究提出一种token级表示学习损失,该损失分别最大化不同增强视图下token嵌入之间的一致性,而非基于体级全局特征进行计算。此外,我们发现了一种仅由该新损失导致的潜在表示坍塌问题。为防止坍塌,我们设计了一种简单的“旋转-恢复”机制:对输入体数据的一个增强视图进行旋转和翻转,随后恢复特征图中token的顺序。我们还对对比损失进行了改进,以区分不同体数据中相同位置的token。我们在两个公开医学分割数据集上测试了所提预训练方案,下游分割任务的结果表明,与其他现有最优预训练方法相比,我们的方法带来了更显著的性能提升。

英文摘要:

In the field of medical images, although various works find Swin Transformer has promising effectiveness on pixelwise dense prediction, whether pre-training these models without using extra dataset can further boost the performance for the downstream semantic segmentation remains unexplored.Applications of previous representation learning methods are hindered by the limited number of 3D volumes and high computational cost. In addition, most of pretext tasks designed specifically for Transformer are not applicable to hierarchical structure of Swin Transformer. Thus, this work proposes a token-level representation learning loss that maximizes agreement between token embeddings from different augmented views individually instead of volume-level global features. Moreover, we identify a potential representation collapse exclusively caused by this new loss. To prevent collapse, we invent a simple "rotate-and-restore" mechanism, which rotates and flips one augmented view of input volume, and later restores the order of tokens in the feature maps. We also modify the contrastive loss to address the discrimination between tokens at the same position but from different volumes. We test our pre-training scheme on two public medical segmentation datasets, and the results on the downstream segmentation task show more improvement of our methods than other state-of-the-art pre-trainig methods.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑