QT-UNet: A Self-Querying All-Transformer U-Net for 2D and 3D Segmentation Augmented by Self-Supervised Learning

Abstract

I 2017 revolusjonerte Transformer-modellen naturlig språkbehandling ved å tilgjengeliggjøre store modeller som var i stand til å forstå komplekse sammenhenger over store avstander i tekst med en håndterbar beregningskostnad. I 2021 brakte introduksjonen av Vision Transformere med seg en liknende revolusjon i bildebehandling, med liknende skaleringsfordeler.

Utviklingen av effektive Vision Transformere med lineær tidskompleksitet som Swin Transformeren bidrog til ytteligere opptak i bruk av Vision Transformere, spesielt i felter som Autonomt Syn og Medisinsk Bildeanalyse. Modeller som VT-UNet smeltet sammen tradisjonelle UNet med Swin Transformere for å skape en ny sterk volumetrisk segmenteringsmodell for hjernesvulster, med et nyskapende Enkoder-Dekoder Cross-Attention konsept.

Parallelt med disse revolusjonene opplevde Self-Supervised Learning en lignende revolusjon og økning i bruk innen flere datasynsdomener, spesielt domener der det er knapt med treningsdata. Swin-UNETR forhåndstrente en sterk Swin-basert koder med et stort CT-datasett og kontrast-, rekonstruksjons- og rotasjonsbaserte treningsoppgaver, og viste sterk nedstrøms ytelse i Medical Segmentation Decathlon (MSD) og Beyond The Cranial Vault (BTCV).

Dette prosjektet smelter disse fremskrittene sammen med modellen Querying Transformer UNet (QT-UNet): Et all-Swin Transformer UNet med Enkoder-Dekoder Cross-Attention, forsterket med Self-Supervised Learning. QT-UNet testes med flere Medical Image Computing datasett for å evaluere modellens effektivitet som en generell volumetrisk segmenteringsmodell. Vi samler et stort datasett kalt CT-SSL med 3.597 CT-skanninger til pretrening. En 2D-versjon, QT-UNet-2D, spinnes ut av hovedmodellen for å evaluere effektiviteten til teknikkene i en 2D Autonomt synskontekst.

Vår beste modell er konkurransedyktig med "state of the art" i BraTS2021 med 40% færre FLOPs enn vår baseline VT-UNet, med en gjennomsnittlig Dice score på 88,61 og Hausdorff Distance på 4,85 mm. Vi finner mindre gode resultater med BTCV og MSD, men demonstrer effektiviteten til både vår nye Cross-Attention mekanisme og vår SSL-pipeline ved pretrening på CT-SSL. Vi overfører også teknikkene til en 2D-kontekst med CityScapes, og finner at vår Cross-Attention mekanisme og SSL-pipeline er effektiv uten endringer.

In 2017, the Transformer model revolutionised the Natural Language Processing field, bringing large-scale models capable of understanding complex long-range dependencies in text at a manageable computational cost. In 2020, the Vision Transformer brought a similar revolution to Computer Vision, with similar scaling benefits.

The development of linear time complexity Vision Transformers like the Swin transformer further aided the adoption of Vision Transformers, leading to a large number of applications in Autonomous Driving and Medical Image Computing. Models like VT-UNet melded traditional UNets with Swin transformers to create a strong volumetric segmentation model for brain tumour segmentation, introducing a novel Encoder-Decoder Cross-Attention concept.

Parallel to these revolutions, Self-Supervised Learning saw a similar revolution and uptake in use within several Computer Vision subdomains, particularly Medical Image Computing where training data is often scarce. Notably, Swin-UNETR pre-trained a strong Swin-based encoder with a large CT dataset utilising contrastitive, reconstructive, and rotation tasks, demonstrating strong performance in downstream Medical Segmentation Decathlon (MSD) and Beyond The Cranial Vault (BTCV) tasks.

Our research melds these advances together to produce the Querying Transformer UNet (QT-UNet): A all-Swin Transformer UNet with Encoder-Decoder Cross-Attention, enhanced by Self-Supervised Learning (SSL). QT-UNet is tested with several Medical Image Computing datasets to evaluate its efficacy as a general volumetric segmentation model. We also collect a large CT pretraining dataset dubbed CT-SSL with 3,597 CT scans. A 2D version, QT-UNet-2D, is spun out to evaluate the techniques in a 2D Autonomous Driving context.

Our best model is competitive with State of the Art in BraTS2021 despite a 40% reduction in FLOPs against our baseline VT-UNet, with an average Dice score of 88.61 and Hausdorff Distance of 4.85mm. We find weaker results with BTCV and Medical Segmentation Decathlon, but demonstrate the effectiveness of both our new Cross-Attention mechanism, and our SSL pipeline when pre-training with our CT-SSL dataset. We transfer the model to a 2D context with CityScapes, finding that our new Cross-Attention mechanism and SSL pipeline are effective without modification.