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

TL;DR

This paper introduces a novel token-level representation learning method with a rotate-and-restore mechanism to improve 3D Transformer-based medical image segmentation, achieving superior results without extra data pre-training.

Contribution

It proposes a new token-level loss and a rotate-and-restore strategy to enhance pre-training of Swin Transformers for medical image segmentation.

Findings

Outperforms existing pre-training methods on medical segmentation datasets.

Effectively prevents representation collapse during training.

Improves downstream segmentation accuracy significantly.

Abstract

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.