TransNuSeg: A Lightweight Multi-Task Transformer for Nuclei Segmentation
Zhenqi He Mathias Unberath Jing Ke Yiqing Shen
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.