PaperlayerView exact revision on arXiv

Source-linked AI summary

LViT: Language meets Vision Transformer in Medical Image Segmentation

Zihan Li, Yunxiang Li, Qingde Li, Puyang Wang, Dazhou Guo, Le Lu, Dakai Jin, You Zhang, Qingqi Hong

arXiv:2206.14718v4cs.CV

TL;DR

Limited high-quality labeled medical images motivate LViT, a text-augmented segmentation model that uses medical annotations, EPI, and LV loss. Across constructed multimodal CT and X-ray datasets, the paper reports superior fully supervised and semi-supervised segmentation performance.

  • Problem

    Medical image segmentation is limited by insufficient high-quality images and the high cost of expert annotation.

  • Method

    LViT incorporates medical text into a hybrid CNN-Transformer model and uses EPI for pseudo-label refinement with LV loss for unlabeled-image supervision.

  • Results

    LViT achieves superior segmentation performance across fully supervised and semi-supervised experiments on multimodal medical-image datasets.

  • Takeaways & Limitations

    Medical text can compensate for image-data quality deficiencies and support improved pseudo labels in semi-supervised medical image segmentation.

Abstract

from arXiv · show

Deep learning has been widely used in medical image segmentation and other aspects. However, the performance of existing medical image segmentation models has been limited by the challenge of obtaining sufficient high-quality labeled data due to the prohibitive data annotation cost. To alleviate this limitation, we propose a new text-augmented medical image segmentation model LViT (Language meets Vision Transformer). In our LViT model, medical text annotation is incorporated to compensate for the quality deficiency in image data. In addition, the text information can guide to generate pseudo labels of improved quality in the semi-supervised learning. We also propose an Exponential Pseudo label Iteration mechanism (EPI) to help the Pixel-Level Attention Module (PLAM) preserve local image features in semi-supervised LViT setting. In our model, LV (Language-Vision) loss is designed to supervise the training of unlabeled images using text information directly. For evaluation, we construct three multimodal medical segmentation datasets (image + text) containing X-rays and CT images. Experimental results show that our proposed LViT has superior segmentation performance in both fully-supervised and semi-supervised setting. The code and datasets are available at https://github.com/HUANGLIZI/LViT.

I. INTRODUCTION

LViT addresses limited medical-image quality and costly annotation by incorporating complementary medical text for segmentation and semi-supervised pseudo-label refinement.

  • Motivation: High-quality medical images and expert annotations are difficult or expensive to obtain, limiting medical image segmentation performance.The paper motivates using accompanying medical records as complementary information because they are generated with patients without extra access cost.
  • Challenges: The introduction frames two challenges: using existing image-text information to improve segmentation and using text to improve pseudo-label quality.These challenges motivate both the multimodal model and semi-supervised learning mechanisms.
  • Approach: LViT combines text embeddings with a hybrid CNN-Transformer structure and Pixel-Level Attention Modules to merge text, global features, and local image features.The CNN component retains local-feature extraction while the Transformer encodes global information.
  • Semi-supervised learning: EPI progressively refines pseudo labels through exponential moving average, while LV loss directly supervises unlabeled images using text information.Together, these mechanisms address pseudo-label quality and direct use of text during semi-supervised training.
  • Evaluation: The study constructs multimodal CT and X-ray segmentation datasets and reports superior performance in fully supervised and semi-supervised settings.The datasets include MosMedData+, ESO-CT, and QaTa-COV19.

B. Vision-language model

The paper situates LViT within vision-language segmentation and presents a Double-U architecture that merges image and text features while retaining local and global image representations.

  • Prior vision-language models: Vision-language segmentation methods use multimodal interactions such as pixel-word attention to integrate linguistic and visual features.Prior work includes VLT and LAVT, while medical-image analysis has also begun using text assistance.
  • LViT design: LViT uses a CNN-ViT Interaction Module to combine convolutional image processing with Transformer-based multimodal feature interaction.The design aims to retain local image features while encoding global features.
  • Semi-supervised learning: The LV loss uses text information to supervise training on unlabeled medical images.This mechanism is introduced to address scarcity of medical-image labels.
  • LViT design: The model is a Double-U structure formed from a U-shaped CNN branch and a U-shaped Transformer branch.The CNN branch provides image input and prediction output, while the ViT branch merges image and text information.

A. LViT Model

The LViT model uses a U-shaped CNN branch for image encoding and prediction, paired with a U-shaped ViT branch that merges text and image features before returning them to the CNN pathway.

  • U-shape CNN Branch: The U-shaped CNN branch receives image information and outputs the segmentation mask through convolution, batch normalization, ReLU, pooling, and concatenation operations.Downsampling occurs between DownCNN modules, while UpCNN modules use concatenation.
  • CNN-ViT interaction: The CNN-ViT interaction module upsamples and aligns ViT features, then combines them with CNN features through residual connections.PLAM is placed at skip connections to improve local-feature segmentation capability.
  • U-shape ViT Branch: The U-shaped ViT branch merges image features from DownCNN with text features and processes the combined representation with a Transformer encoder.Text features are supplied by BERT-Embed, and image features are formed through PatchEmbedding.
  • U-shape ViT Branch: Corresponding ViT features are transferred back through UpViT and merged with CNN features to extract global information while reducing sensitivity to inaccurate text annotations.The fusion is performed at corresponding feature sizes.

3) Pixel-Level Attention Module (PLAM):

PLAM combines local image-feature preservation with semantic text fusion, while EPI iteratively stabilizes pseudo labels for unlabeled pixels using an EMA-style update.

  • Pixel-Level Attention Module (PLAM): PLAM preserves local image features and merges them with semantic features from text.It also enhances convolutional representations of local features.
  • Exponential Pseudo-label Iteration mechanism (EPI): EPI iteratively updates pseudo labels using an EMA-style mechanism rather than directly reusing one model generation’s prediction.The momentum parameter β is set to 0.99.
  • Exponential Pseudo-label Iteration mechanism (EPI): EPI gradually optimizes each unlabeled pixel’s segmentation prediction and improves robustness to noisy labels.The method avoids sharp deterioration caused by directly using one generation’s pseudo label as the next target.
  • Exponential Pseudo-label Iteration mechanism (EPI): Under the stated assumption that recent model weights fluctuate around the optimum, recent pseudo-label predictions likewise fluctuate around the mask.This assumption motivates averaging predictions across generations.
  • Exponential Pseudo-label Iteration mechanism (EPI): EPI assigns gradient-step coefficient 1−β^(n−i), which decreases with iteration index and stabilizes pseudo-label changes.The authors state that this produces a high-confidence pseudo label.

C. LV (Language-Vision) Loss

LV loss uses structured text to identify approximate anatomical masks and directly supervise unlabeled-image training, helping prevent severe pseudo-label deterioration without refining boundaries.

  • LV (Language-Vision) Loss: LV loss uses structured text describing organ locations to form corresponding contrastive masks.The method computes cosine similarity between text feature vectors to relate pseudo-label text with contrastive text.
  • LV (Language-Vision) Loss: The highest-similarity contrastive text determines the corresponding segmentation mask for comparison with the predicted pseudo label.The comparison uses cosine similarity between pseudo-label and contrastive-label feature vectors.
  • LV (Language-Vision) Loss: Contrastive labels provide approximate location information rather than boundary refinement.Accordingly, LV loss primarily targets mis-segmentation or mislabeling with substantial differences.
  • LV (Language-Vision) Loss: LV loss is applied only to unlabeled data, where it complements consistency supervision to avoid sharp pseudo-label quality deterioration.The authors state that contrastive labels provide little performance improvement when data is labeled.

D. Proof of CNN-Transformer structure superiority

LViT combines CNN and ViT processing without a text encoder, aiming to retain CNN-like local features and ViT-like global features for medical image segmentation.

  • CNN-Transformer structure: LViT uses interaction between CNN and ViT to extract image and text features without a text encoder.This is presented as the model’s distinction from previous vision-language work.
  • CNN-Transformer structure: CNN convolution produces channel-wise outputs and is described as suitable for learning shallow, affine-invariant features.Shared convolution kernels also support whole-image filtering focused on boundary features.
  • CNN-Transformer structure: ViT self-attention computes similarities among transformed inputs and is described as focusing on input-feature invariance.Because ViT processes the whole image, the authors characterize it as more suited to robust global-feature learning than CNN.

A. Setup

The experiments use three medical segmentation datasets, extend text annotations for multimodal training, and evaluate models with Dice and mIoU under labeled and unlabeled settings.

  • Datasets: Three datasets are used to evaluate LViT, including MosMedData+, QaTa-COV19, and ESO-CT.MosMedData+ contains CT lung-infection slices, while QaTa-COV19 contains COVID-19 chest X-ray radiographs with manual lesion annotations.
  • Datasets: The authors extend text annotations for training, describing infection laterality, lesion count, and approximate infected-area locations.The MosMedData+ text structure is described as similar to QaTa-COV19’s.
  • Training and evaluation: The training loss uses Dice and cross-entropy terms, with an additional LV-loss term for unlabeled data and α = 0 for labeled data.For unlabeled data, α = 0.1.
  • Training and evaluation: The cross-entropy term is defined over pixels and categories using prediction probabilities pij and binary targets yij.In the experiments, the number of categories C is set to 1.
  • Training and evaluation: Dice score and mIoU are used to evaluate segmentation performance.The metrics are applied to LViT and other state-of-the-art methods.

C. Implementation Details

The experiments compare LViT with CNN- and Transformer-based segmentation models across two datasets, including fully and semi-supervised settings. LViT achieves strong segmentation performance while reducing mis-segmentation and preserving finer boundaries.

  • LViT-T improves Dice by 3.24% and mIoU by 4.3% over nnUNet on QaTa-COV19.
  • LViT-T remains superior to other state-of-the-art methods when only 1/4 of the training labels are used.
  • On MosMedData+, LViT-T improves Dice by 2.15% and mIoU by 1.15% compared with GLoRIA.
  • LViT-TW and LViT-T (1/4) achieve performance comparable to nnUNet and UCTransNet on MosMedData+.
  • Qualitative comparisons report less mis-segmentation and more delicate segmentation boundaries for LViT than for UNet++, nnUNet, and TransUNet.The comparison uses qualitative results from QaTa-COV19 and MosMedData+.

E. Ablation Study

Ablation studies evaluate supervised components, semi-supervised components, and training configurations. The results support contributions from text information, EPI, LLV, and combined attention modules.

  • 1) Effectiveness of Proposed Components:: Text produces the most significant performance improvement among the evaluated supervised components.The supervised ablation examines DownViT, UpViT, PLAM, and Text.
  • 1) Effectiveness of Proposed Components:: Combining GAP and GMP in PLAM performs better than using either attention module alone.GAP/PLAM also outperforms GMP/PLAM, possibly because global average pooling integrates diverse information better.
  • 1) Effectiveness of Proposed Components:: EPI, Text, and LLV each improve semi-supervised performance in the ablation experiments.EPI progressively refines pseudo labels, while LLV supports continued improvement using unlabeled data.
  • The ablations include text-encoder and text-embedding-layer comparisons and hyperparameter studies of batch size and learning rate.
  • 1) Effectiveness of Proposed Components:: Adding EPI increases Dice by 0.26% with text annotations and 0.21% without text annotations for LViT (1/4).

2) Ablation Study on Model Size:

The ablations examine model size, hyper-parameters, text representation, and semi-supervised performance. Text annotation improves segmentation substantially with only modest added computational cost, while EPI and LV loss support strong semi-supervised results.

  • Model size: Adding text annotation improves segmentation significantly over LViT-W with only 1.7M additional parameters and 0.1G additional computation.The comparison spans LViT variants across QaTa-COV19 and MosMedData+.
  • Model size: Larger models can improve performance on difficult, heterogeneous datasets, while increasing model size reduces performance jitter and improves robustness.The authors relate size-dependent gains to dataset distribution and segmentation difficulty.
  • Hyper-parameters: LViT performs optimally with batch size 24, using learning rate 3e-4 on QaTa-COV19 and 1e-3 on MosMedData+.These settings are reported from the batch-size and learning-rate ablations.
  • Text representation: A text encoder requires nearly three times the parameters and nearly twice the computation of a text embedding layer without improving performance in well-structured reports.Performance even decreases for the encoder in the well-structured-text experiments, supporting the embedding-layer design.
  • Semi-supervision: Across 25% and 50% label ratios, LViT achieves superior semi-supervised segmentation performance relative to traditional and multimodal comparison methods.The reported advantage is attributed to the EPI mechanism and LV loss, with or without text information in the pipeline.

F. Practical

The practical studies test LViT on esophageal CT data and inspect its activation patterns with GradCAM. Text annotations provide approximate tumor location and are associated with stronger cross-dataset segmentation performance, while inference currently requires text input.

  • Practical application: The ESO-CT dataset contains 286 esophageal cancer cases with masks and radiologist-annotated clinical information.Clinical information divides the esophagus into four sections and records the tumor’s rough location.
  • Practical application: On ESO-CT, LViT-TW is comparable to nnUNet, while text-augmented LViT-T outperforms the other compared methods by a large margin.The comparison includes UNet, UNet++, nnUNet, and TransUNet; nnUNet uses complex image preprocessing that LViT does not require.
  • Practical application: The ESO-CT text input encodes the tumor’s approximate location within four anatomical esophageal sections.The sections are Cervical, Upper Thorax, Middle Thorax, and Lower Thorax.
  • Interpretability: GradCAM is used to compare attention regions and assess whether LViT notices lesions and whether text enhances lesion-focused attention.The interpretability study is conducted on QaTa-COV19 across multiple approaches and LViT layers.
  • Scope and limitations: In the current LViT version, text inputs are required during inference, motivating future automatic generation of annotations from images.The paper also identifies extension from 2D to 3D segmentation as future work.
Loading 2206.14718v4…