TransNuSeg:一种用于细胞核分割的轻量级多任务Transformer
TransNuSeg: A Lightweight Multi-Task Transformer for Nuclei Segmentation
AI总结:
本文提出纯Transformer多任务框架TransNuSeg,通过三解码器、自蒸馏一致性损失、注意力共享和token MLP瓶颈提升细胞核分割性能并降低复杂度。
摘要:
细胞核尺寸较小,但在实际临床实践中,全局空间信息以及细胞核与背景之间颜色或亮度对比度的相关性,一直被认为是实现准确细胞核分割的关键组成部分。然而,自动细胞核分割领域由卷积神经网络(CNNs)主导,同时,近来流行的Transformer在捕捉局部—全局相关性方面能力强大,但其潜力尚未得到充分探索。为此,我们首次尝试构建用于细胞核分割的纯Transformer框架,称为TransNuSeg。与以往工作不同,我们将具有挑战性的细胞核分割任务解耦为一个内在的多任务学习任务,其中采用三解码器结构,分别用于细胞核实例、细胞核边缘和簇状边缘分割。为消除以往工作中不同分支之间存在分歧的预测,我们引入了一种新颖的自蒸馏损失,以显式地在各分支之间施加一致性约束。此外,为刻画各分支之间的高度相关性并减少参数数量,我们提出了一种高效的注意力共享方案,即在三个解码器之间部分共享自注意力头。最后,使用token MLP瓶颈替代参数过多的Transformer瓶颈,以进一步降低模型复杂度。在包括MoNuSeg在内的两个不同模态数据集上的实验表明,我们的方法能够以少30%的参数量,在Dice指标上优于CA2.5-Net等最先进同类方法2%至3%。总之,TransNuSeg证实了Transformer在细胞核分割场景中的优势,因此可作为实际临床实践中的高效解决方案。代码可在https://github.com/zhenqi-he/transnuseg获取。
英文摘要:
Nuclei appear small in size, yet, in real clinical practice, the global spatial information and correlation of the color or brightness contrast between nuclei and background, have been considered a crucial component for accurate nuclei segmentation. However, the field of automatic nuclei segmentation is dominated by Convolutional Neural Networks (CNNs), meanwhile, the potential of the recently prevalent Transformers has not been fully explored, which is powerful in capturing local-global correlations. To this end, we make the first attempt at a pure Transformer framework for nuclei segmentation, called TransNuSeg. Different from prior work, we decouple the challenging nuclei segmentation task into an intrinsic multi-task learning task, where a tri-decoder structure is employed for nuclei instance, nuclei edge, and clustered edge segmentation respectively. To eliminate the divergent predictions from different branches in previous work, a novel self distillation loss is introduced to explicitly impose consistency regulation between branches. Moreover, to formulate the high correlation between branches and also reduce the number of parameters, an efficient attention sharing scheme is proposed by partially sharing the self-attention heads amongst the tri-decoders. Finally, a token MLP bottleneck replaces the over-parameterized Transformer bottleneck for a further reduction in model complexity. Experiments on two datasets of different modalities, including MoNuSeg have shown that our methods can outperform state-of-the-art counterparts such as CA2.5-Net by 2-3% Dice with 30% fewer parameters. In conclusion, TransNuSeg confirms the strength of Transformer in the context of nuclei segmentation, which thus can serve as an efficient solution for real clinical practice. Code is available at https://github.com/zhenqi-he/transnuseg.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。