PaperlayerView exact revision on arXiv

Source-linked AI summary

MCSeg: Pre-training and Fine-tuning Volumetric Pyramid Transformer for Multi-modal Cardiac Image Segmentation

Zhiyu Ye, Hairong Zheng, Tong Zhang

arXiv:2608.30371v1cs.CV

TL;DR

Cardiac segmentation must generalize across modalities despite scarce finely annotated 3D data. MCSeg combines self-supervised ViT pre-training, an SFP bridge to a CNN decoder, and RMI-based fine-tuning, achieving SOTA results across four datasets and strong few-shot performance.

  • Problem

    Scarce finely annotated 3D data, rare-disease imbalance, and CNN domain-generalization limits hinder robust cardiac segmentation across diverse datasets.

    Source: §I · p. 1

  • Method

    MCSeg pre-trains a volumetric ViT with masked autoencoders, remaps its output through an SFP to a CNN decoder, and uses RMI loss during fine-tuning.

    Source: §I · p. 2, 2) Loss Function · p. 4, §V · p. 8

  • Results

    MCSeg achieves SOTA performance across four cardiac datasets and demonstrates cross-modality adaptability and few-shot learning potential.

    Source: §I · p. 2

  • Takeaways & Limitations

    MCSeg provides a cardiac-specific pipeline that transfers anatomical priors across CT and MRI tasks and adapts to scarce annotated data.

    Source: E. Discussion · p. 8, §V · p. 8

    Large-scale 3D MAE pre-training requires substantial computational and memory resources, including hundreds of GPU hours on advanced hardware.

    Source: E. Discussion · p. 8

Abstract

from arXiv · show

Automatic cardiac image segmentation is pivotal for diagnosing and treating cardiac diseases. In this work, we introduce MCSeg, a volumetric transformer-based network tailored for multi-modal cardiac segmentation. To overcome the architectural mismatch inherent in existing hybrid networks, we propose a novel Scaling Feature Pyramid (SFP). Unlike conventional skip connections, the SFP effectively bridges the single-scale 3D Vision Transformer (ViT) encoder and the multi-scale CNN decoder by transforming the ViT's output into a hierarchical feature pyramid, ensuring that global contextual information is effectively leveraged. For the training paradigm, the ViT encoder first undergoes self-supervised pre-training via masked image modeling. Subsequently, the network is fine-tuned on downstream tasks, during which a regional mutual information (RMI) loss is integrated to improve boundary segmentation accuracy. In experiments, MCSeg consistently outperforms eleven SOTA methods on CT dataset ImageCHD, multi-modal dataset MM-WHS, MRI dataset HVSMR-2.0 and MSD Heart, highlighting the effectiveness of our MCSeg for multi-modal cardiac segmentation tasks. Furthermore, MCSeg's superior performance in few-shot experiment showcases its significant potential in adapting to limited data scenarios. Codes and pre-trained ViT-B weights are open-sourced at https://openi.pcl.ac.cn/OpenMedIA/MCSeg

I. INTRODUCTION

MCSeg targets cardiac segmentation across modalities by combining volumetric transformer representations with self-supervised pre-training and specialized fine-tuning. The approach addresses limited annotated 3D data and seeks robust adaptation across diverse cardiac tasks.

  • Cardiac image segmentation supports non-invasive assessment, diagnosis, treatment planning, disease monitoring, and prognosis.Source: §I · p. 1
  • Limited finely annotated 3D data, rare-disease imbalance, and CNN domain-generalization limits complicate training universally applicable cardiac segmentation models.Source: §I · p. 1
  • Self-supervised masked image modeling and ViTs are used to exploit unlabeled images and model long-range dependencies for diverse 3D cardiac tasks.Source: §I · p. 2
  • MCSeg bridges a single-scale ViT encoder and multi-scale decoder with a Scaling Feature Pyramid to exploit pre-trained 3D representations.Source: §I · p. 2
  • The two-stage training paradigm combines self-supervised 3D MAE pre-training, downstream fine-tuning, and RMI loss for boundary segmentation accuracy.Source: §I · p. 2
  • MCSeg achieves SOTA performance across ImageCHD, HVSMR-2.0, MM-WHS, and MSD, with reported cross-modality adaptability and few-shot learning potential.Source: §I · p. 2

II. RELATED WORKS

Prior cardiac segmentation methods include CNNs, pure transformers, hybrid networks, and self-supervised vision-transformer approaches. MCSeg builds on these directions by combining transformer pre-training with a feature-pyramid interface for CNN decoding.

A. Network Architecture

MCSeg uses a non-hierarchical volumetric ViT encoder, an SFP that remaps its single-scale output into multiple resolutions, and a CNN decoder for segmentation.

B. Fine-tuning for Cardiac Segmentation

During fine-tuning, the frozen pre-trained ViT encoder feeds the SFP and decoder, which are trained with a compound loss combining Dice, cross-entropy, and regional mutual information. RMI adds local spatial constraints and is smoothly introduced during training.

  • The pre-trained ViT backbone remains frozen, while only the SFP and decoder are trained to produce segmentation results.Source: 1) Training Strategy · p. 3
  • The compound objective combines soft Dice, cross-entropy, and RMI losses.Source: 2) Loss Function · pp. 3–4
  • RMI complements Dice and cross-entropy by enforcing local consistency between predicted and ground-truth regions for smoother boundaries.Source: 2) Loss Function · p. 4
  • The number of input regions is determined by the down-sampling factor ds and square region size r.Source: 2) Loss Function · p. 4
  • The RMI contribution is regulated by a smooth step function using the current and maximum training epochs.Source: 2) Loss Function · p. 4

IV. EXPERIMENTS

The experiments use ImageCAS for pre-training and ImageCHD for downstream congenital-heart-disease segmentation, with dataset-specific annotations and splits.

  • ImageCAS contains 1,000 CTA images and is used entirely for pre-training because complete heart structures are visible.Source: 1) ImageCAS · p. 4
  • ImageCHD contains 110 3D CT images spanning 16 congenital heart diseases, with seven annotated cardiac substructures.Source: 2) ImageCHD · p. 4
  • ImageCHD is split into 77 training, 11 validation, and 22 test cases, while training and validation images also support pre-training.Source: 2) ImageCHD · p. 4

2) ImageCHD:

The broader evaluation covers cardiovascular MRI, multimodal CT/MRI, and left-atrium MRI segmentation datasets, alongside implementation and evaluation details.

  • HVSMR-2.0 provides 60 cardiovascular MRI scans with masks for four chambers and four great vessels.Source: 3) HVSMR-2.0 · p. 4
  • MM-WHS contains 60 CT and 60 MRI images, with 20 annotated cases per modality and remaining unlabeled images used for pre-training.Source: 4) MM-WHS · p. 4
  • MSD Task02 evaluates left-atrium segmentation using 20 labeled MRI training cases.Source: 5) MSD Task02 Heart · p. 4
  • The model is implemented in PyTorch and MONAI, pre-trained on four 80G A100 GPUs, and fine-tuned on a single 32G V100 GPU.Source: B. Implementation Details · p. 4
  • Experimental results are measured with the Dice coefficient.Source: B. Implementation Details · p. 4

C. Comparative Experiments

MCSeg is compared with eleven volumetric segmentation networks across full-data and few-shot cardiac experiments. It achieves strong overall performance and benefits substantially from pre-training under limited data.

Comparative setup

  • The comparison includes pure CNN, pure transformer, and hybrid CNN–Transformer networks, with selected models using original or shared pre-trained weights.Source: C. Comparative Experiments · pp. 4–5

Full-data evaluation

Few-shot evaluation

2) Whole-heart Segmentation on HVSMR-2.0 Dataset:

On HVSMR-2.0, MCSeg achieves strong segmentation performance and shows qualitative advantages over competing models, particularly on MRI images.

Qualitative comparison

Cross-modality comparison

  • MCSeg shows a greater advantage over other models on MRI images than on CT images.This comparison is explicitly reported after evaluating the HVSMR-2.0 results.Source: 2) Whole-heart Segmentation on… · p. 6

Evaluation design

  • The evaluation includes five-fold cross-validation on multimodal datasets using models trained from scratch and with pre-trained weights.The referenced cross-validation results compare eight models across various datasets.Source: 3) Whole-heart Segmentation on MM-WHS… · p. 6

5) Computational Efficiency:

The reported experiments examine cross-validation, qualitative results, computational costs, and architectural ablations for MCSeg and competing models.

Cross-validation

  • MCSeg achieves the highest Dice scores and lowest intra-group variance in nearly all MM-WHS cross-validation folds.This result is reported for both CT and MRI evaluations, with models trained from scratch and with pre-trained weights.Source: Figure 6 · p. 7

Computational efficiency

  • MCSeg requires fewer trainable parameters, uses less GPU memory, and achieves the shortest inference time among compared methods.Measurements were obtained on ImageCHD with an input size of (128, 128, 128).Source: 5) Computational Efficiency · p. 7, Table II · p. 7

Ablation studies

  • The ablation experiments assess SFP configuration, pre-training data scale and modality, fine-tuning strategy, and loss-function combinations.These experiments target the main architectural and training choices in MCSeg.Source: D. Ablation Studies · p. 7

SFP ablation

  • The proposed SFP outperforms a conventional FPN by 1.5%.Additional SFP analysis finds that constructing the pyramid from the deepest ViT block z12 with the selected scaling configuration gives optimal performance.Source: 1) SFP Design · p. 7

Pre-training ablation

  • Pre-trained weights improve Dice scores by over 2.3% on ImageCHD and by 7.0% to 9.5% on two MRI datasets compared with training from scratch.The reported comparison also considers pre-training computational costs and identifies ViT-B223 as the optimal trade-off.Source: 2) Data Scale and Modalities for… · p. 7

2) Data Scale and Modalities for Pre-training:

The ablations compare downstream training strategies and loss combinations while also evaluating pre-training data scales, modalities, and computational costs.

Fine-tuning strategy

  • Frozen-encoder fine-tuning achieves the highest Dice scores on both ImageCHD and HVSMR-2.0.Among the three compared training strategies, it is reported as the most effective for downstream adaptation.Source: 3) Fine-tuning Strategy · p. 7

Loss functions

  • The Dice, CE, and RMI combination achieves a Dice score of 0.892 and an HD95 of 13.4.It is compared with alternative loss-function combinations on ImageCHD.Source: 4) Loss Function · p. 7, Table V · p. 8
  • RMI loss refines boundary accuracy beyond the standard Dice and CE combination.The reported evidence links the combined loss to improved boundary-sensitive performance.Source: 4) Loss Function · pp. 7–8

E. Discussion

MCSeg consistently outperforms existing methods across CT and MRI cardiac segmentation, with transfer benefits from CT pre-training and strong few-shot performance. Its gains combine large-scale cardiac pre-training with SFP and RMI, but deployment remains constrained by computational demands and cardiac-specific scope.

Performance

  • MCSeg outperforms SOTA methods across four cardiac datasets and consistently surpasses existing methods on both CT and MRI datasets.The evaluated datasets include ImageCHD, HVSMR-2.0, MM-WHS, and MSD.Source: E. Discussion · p. 8, §V · p. 8

Cross-modality transfer

  • CT-only pre-training substantially improves both CT and MRI performance because the modalities share anatomical structures.The benefit is especially pronounced for MRI tasks affected by scanner and protocol variability.Source: E. Discussion · p. 8

Ablation and architecture

  • The SFP module’s effectiveness is isolated by applying identical pre-trained ViT weights to UNETR, after which MCSeg remains consistently superior.This comparison directly evaluates the encoder-decoder bridging design rather than changing the pre-training variant.Source: E. Discussion · p. 8

Limitations and accessibility

  • Large-scale 3D MAE pre-training requires hundreds of GPU hours and advanced hardware, limiting practical accessibility despite released weights and code.The authors open-source pre-trained and fine-tuned weights plus the complete source code to support deployment and reproducibility.Source: E. Discussion · p. 8

Design implications

  • MCSeg combines pre-trained global anatomical representations with SFP and RMI loss to improve segmentation precision and boundary detail.Large-scale pre-training raises the performance floor, while SFP and RMI loss further improve precision.Source: §V · p. 8

Data efficiency

  • Few-shot experiments demonstrate MCSeg’s adaptability when annotated cardiac data are scarce.This supports the model’s use in limited-data segmentation settings.Source: §V · p. 8

Related papers

Loading 2608.30371v1…