Source-linked AI summary
Unpaired Multi-modal Segmentation via Knowledge Distillation
Qi Dou, Quande Liu, Pheng Ann Heng, Ben Glocker
TL;DR
Unpaired CT/MRI segmentation lacks the paired, registered inputs commonly used in multimodal learning. The paper shares convolutional kernels, uses modality-specific normalization, and adds a KL-divergence knowledge-distillation loss; across cardiac and abdominal tasks, the scheme consistently outperforms single-model and prior multimodal approaches.
Problem
Existing multimodal segmentation typically uses paired, co-registered images, while unpaired CT/MRI data present distinct distributions and limited pixel-to-pixel correspondence.
Method
The method shares all convolutional kernels across CT and MRI, uses modality-specific internal normalization, and distills semantic knowledge through a KL-divergence-based loss.
Results
Across cardiac and abdominal segmentation with 2D and 3D networks, the method consistently outperforms single-model training and previous multimodal approaches.
Takeaways & Limitations
A single parameter-efficient network can integrate unpaired CT and MRI segmentation, with reported benefits across both evaluated tasks and network architectures.
Takeaways & Limitations
The relatively basic 2D dilation network and 3D U-Net may limit segmentation accuracy compared with more complex multimodal architectures.
Abstract
from arXiv · showhide
Multi-modal learning is typically performed with network architectures containing modality-specific layers and shared layers, utilizing co-registered images of different modalities. We propose a novel learning scheme for unpaired cross-modality image segmentation, with a highly compact architecture achieving superior segmentation accuracy. In our method, we heavily reuse network parameters, by sharing all convolutional kernels across CT and MRI, and only employ modality-specific internal normalization layers which compute respective statistics. To effectively train such a highly compact model, we introduce a novel loss term inspired by knowledge distillation, by explicitly constraining the KL-divergence of our derived prediction distributions between modalities. We have extensively validated our approach on two multi-class segmentation problems: i) cardiac structure segmentation, and ii) abdominal organ segmentation. Different network settings, i.e., 2D dilated network and 3D U-net, are utilized to investigate our method's general efficacy. Experimental results on both tasks demonstrate that our novel multi-modal learning scheme consistently outperforms single-modal training and previous multi-modal approaches.
I. INTRODUCTION
The paper addresses unpaired CT/MRI segmentation and proposes a compact multimodal scheme that shares convolutional kernels while using modality-specific normalization and knowledge distillation. It validates the approach across two segmentation tasks and CNN settings.
- Research gap: Unpaired multimodal segmentation remains unclear because existing methods typically require paired, co-registered images from the same patient.The paper also identifies limited exploration of CT/MRI combinations beyond multi-sequence MRI.
- Research gap: CT and MRI have distinct acquisition physics and visual distributions, requiring new ways to exchange cross-modality information.The paper frames semantic relationships, rather than pixel-to-pixel coherence, as the relevant connection for unpaired images.
- Proposed approach: The proposed compact model shares all convolutional kernels across CT and MRI while using modality-specific internal normalization parameters.This design explicitly addresses distribution shift while reusing network parameters.
- Proposed approach: A knowledge-distillation mechanism extracts semantic knowledge from high-level CNN representations and derives a loss function to guide multimodal learning.The method explicitly leverages shared information across modalities through the derived loss.
- Validation: The method is evaluated on two multi-class segmentation tasks using both 2D and 3D CNN architectures.The experiments are intended to demonstrate general effectiveness across cardiac and abdominal segmentation settings.
A. Independent normalization of CT and MRI
The method uses modality-specific normalization to handle CT/MRI statistical differences while sharing convolutional kernels, building on prior multimodal and knowledge-distillation ideas. Its normalization layer computes activation statistics and applies trainable scale and shift parameters separately by modality.
- Background: Prior CT/MRI multimodal segmentation methods commonly use independent modality-specific encoders and decoders, with an X-shaped architecture reported as strongest among those designs.The paper motivates a more compact alternative to these dual-stream architectures.
- Background: Knowledge distillation transfers learned information from a larger model to a smaller model using soft-label supervision.The paper draws on this concept to motivate semantic knowledge transfer across modalities.
- A. Independent normalization of CT and MRI: The proposed model shares the same CNN kernels across modalities and separates internal normalization parameters to address CT/MRI distribution differences.This design aims to retain shared representations while calibrating modality-specific activations.
- A. Independent normalization of CT and MRI: The normalization layer computes modality-specific activation statistics and then applies trainable scale γ and shift β parameters.The activation set can be defined using Batch, Instance, Layer, or Group Normalization.
- A. Independent normalization of CT and MRI: Separate normalization avoids defective features that can arise when CT and MRI statistics are combined before normalization.The paper illustrates that opposing modality means could average to zero, producing a meaningless shared statistic.
B. Knowledge distillation loss
The method distills semantic knowledge from pre-softmax activations and aligns CT and MRI prediction distributions with a symmetric KL-divergence loss. Shared convolutional kernels are therefore trained to capture cross-modality patterns while modality-specific normalization handles differing statistics.
- Shared convolutional kernels are expected to extract universal representations capturing patterns common to CT and MRI.The paper identifies shape as an example of a potentially robust, discriminative cross-modality pattern.
- The method averages pre-softmax activations over locations belonging to each ground-truth class to form a C-dimensional semantic vector.Each vector element records activation behavior across the class channels for pixels of a particular class.
- Temperature-scaled softmax converts each class-specific semantic vector into a probability distribution describing predictions across classes.The temperature is set to T =2 in the experiments, while T = 1 gives the ordinary softmax.
- The KD-loss aligns CT and MRI distilled knowledge by minimizing their relative-entropy divergence between class-specific prediction vectors.The two modality distributions are updated together during dynamic training, using a symmetric KL-divergence formulation.
C. Overall loss function and training procedure
The compact network combines ordinary segmentation losses with the KD-loss and an L2 regularizer. Training feeds CT and MRI through shared convolutions and modality-specific normalization layers while updating both modalities jointly.
- The overall objective combines per-modality segmentation losses, weighted KD-loss, and L2 regularization.The segmentation term combines Dice loss with pixel-wise weighted cross-entropy; α is generally set as 0.5 and η as 1e−4.
- Each training batch contains half CT and half MRI samples processed by shared convolutional kernels and their respective internal normalization layers.The KD-loss is computed from activation tensors of samples from the two modalities.
- The multimodal scheme is architecture-independent and can be integrated into existing 2D and 3D CNN models.
IV. EXPERIMENTS
The experiments evaluate the proposed multimodal learning approach on cardiac structure and abdominal multi-organ segmentation. Different 2D and 3D CNN architectures are used to assess its flexibility and general efficacy.
- The evaluation covers two multi-class tasks: cardiac structure segmentation and abdominal multi-organ segmentation.
- The study implements 2D and 3D models with different network architectures to demonstrate the method’s flexibility and general efficacy.
A. Datasets and networks
The study uses unpaired CT and MRI datasets for cardiac and abdominal segmentation, pairing a 2D dilated CNN with the cardiac task and a 3D U-Net with the abdominal task. The abdominal setting also tests performance when MRI has substantially fewer samples than CT.
- Datasets: The cardiac dataset contains 20 unpaired CT and 20 unpaired MRI images covering four cardiac structures.The structures are LVM, LAC, LVC, and AA.
- Datasets: Cardiac images are resampled to approximately 1.0×1.0×1.0 mm3, cropped to 256×256 in the coronal plane, and split by modality into training, validation, and testing sets.The split is 70% training, 10% validation, and 20% testing, with per-modality zero-mean, unit-variance intensity normalization.
- Networks: The cardiac task uses a 2D dilated CNN with three adjacent input slices and the middle slice supervised by a ground-truth mask.Its architecture is illustrated in Fig. 2.
- Datasets: The abdominal task segments liver, spleen, right kidney, and left kidney using 30 CT patients and 9 available MRI cases.One CT case was excluded for low image quality, creating a setting where one modality has far fewer samples.
B. Experimental settings
The experiments compare seven fixed-architecture settings for unpaired CT/MRI segmentation, ranging from separate or fully shared models to the proposed kernel-sharing architecture with knowledge distillation. Performance is assessed against single-modality baselines and prior multi-modal approaches using Dice and Hausdorff metrics.
- Experimental settings: Seven experimental settings compare individual, jointly shared, fusion-based, and proposed multi-modal architectures under fixed network architectures and hyper-parameters.The settings include Individual, Joint, Joint+KD, “Y”-shaped, “X”-shaped, “Chilopod”-shaped, and Ours.
- Experimental settings: The “Chilopod”-shaped architecture shares all CNN kernels while retaining modality-specific internal feature normalization layers.This is the proposed architecture component used to address modality differences while reusing network parameters.
- Experimental settings: The full “Ours” scheme combines the “Chilopod”-shaped architecture with the proposed knowledge-distillation loss.This setting evaluates the architecture and distillation loss together.
- Evaluation: Segmentation performance is evaluated using volume Dice coefficient (%) and surface Hausdorff distance (mm), reporting class-wise means and standard deviations.Mean values over each modality and over both modalities are also presented for direct comparison.
- Evaluation: Individual models serve as single-modality baselines, while comparisons include two state-of-the-art multi-modal approaches and available challenge-winning results.The evaluation is designed to demonstrate the effectiveness of multi-modal learning.
1) Results on multi-modal cardiac segmentation:
The proposed compact multimodal scheme improves segmentation across 2D cardiac and 3D abdominal tasks, with shared convolutional kernels, modality-specific normalization, and KD-loss contributing to strong results.
- 2D cardiac segmentation: 89.0% Dice on CT and 84.5% on MRI follow adding KD-loss to the 2D cardiac Joint model.The Joint model without KD-loss achieves 88.6% on CT and 84.0% on MRI.
- 2D cardiac segmentation: 88.8% overall Dice and 2.10mm overall mean Hausdorff distance make the proposed cardiac scheme outperform Individual training and the cited challenge winner.The overall Dice comprises 91.7% on CT and 86.0% on MRI; the challenge winner achieves 85.5% overall Dice.
- 3D abdominal segmentation: 81.0% average Dice results after KD-loss improves the 3D abdominal Joint model by 2.0 percentage points from 79.0%.Without KD-loss, the Joint model falls from Individual training’s 87.1% to 79.0%.
- 3D abdominal segmentation: 90.8% average Dice and 2.34mm average Hausdorff distance are achieved by the full scheme, outperforming Individual, X-shaped, and Y-shaped models.The Chilopod model shares convolution kernels and uses modality-specific BN layers, while KD-loss supplies the final improvement.
- 3D abdominal segmentation: MRI abdominal segmentation improves by 5.6%, from 83.7% to 89.3%, relative to Individual training.The reported gains mainly involve spleen and left kidney segmentation, which are challenging under MRI data scarcity.
3) Statistical analysis on significance:
Statistical tests report significant improvements for the proposed method across all evaluated settings, while ablations examine parameter compactness and normalization choices.
- Statistical significance: p < 0.05 is obtained in all settings when jointly testing CT and MRI results, indicating statistically significant improvement.Student’s t-tests are reported for both 2D cardiac and 3D abdominal segmentation.
- Model compactness: 39.32M parameters are used by the 2D Joint model, compared with 1916K parameters for the 3D Joint model.The paper presents compactness as a benefit of sharing parameters with modality-specific internal normalization.
- Comparison: The paper reports superior performance for the proposed method relative to existing multimodal schemes.This comparison is stated alongside the compactness analysis.
- Normalization ablation: Instance, Layer, and Group normalization are evaluated in five ablation settings on the cardiac dataset.The proposed method consistently outperforms Individual learning across these normalization methods.
2) Different weights of KD-loss:
Varying the KD-loss weight shows that activating the loss stabilizes cross-modality semantic alignment, while the method remains applicable across tested architectures and challenging abdominal cases.
- KD-loss behavior: The KD-loss weight α is evaluated from 0.1 to 0.9 in increments of 0.1 using mean segmentation Dice on cardiac CT and MRI.The α = 0 setting corresponds to specified columns in Table I, providing the unconstrained comparison.
- KD-loss behavior: KD-loss activation stabilizes the loss and improves alignment between CT and MRI class-probability distributions.Without constraints (α = 0), distilled knowledge diverges; activating KD-loss stabilizes L_kd and better aligns the distributions.
- Semantic alignment: Confusion matrices for CT and MRI evolve during training, with their absolute difference planes used to visualize alignment.The visualization compares the modality-specific matrices from initialization through convergence; difference values are amplified fivefold in Fig. 8.
- Method scope: The approach is designed for unpaired CT and MRI segmentation without relying on image alignment, addressing distribution shift and absent registration.Its compact architecture shares convolutional kernels while using modality-specific normalization and knowledge distillation to leverage shared semantic information.
- Robustness: On an abdominal CT outlier case with artefacts, the proposed approach achieves higher Dice and smaller Hausdorff distance than comparison methods.The training cases were good-quality images without artefacts, making this a lower-quality test case.
- Limitations: A limitation is that the cardiac 2D dilation network and abdominal 3D U-Net are relatively basic compared with more complex multimodal architectures.The authors state that these plain architectures may limit segmentation accuracy and plan to integrate the scheme into better-designed networks.
VI. CONCLUSION
The paper presents a parameter-efficient scheme for unpaired CT and MRI segmentation using shared convolutional kernels, modality-specific normalization, and a new KD-loss. It demonstrates effectiveness across two tasks and both 2D and 3D architectures, with potential deployment and generalization benefits.
- The proposed scheme combines shared convolutional kernels with modality-specific normalization and a new KD-loss for unpaired CT and MRI segmentation.
- The method is effective across two segmentation tasks and both 2D and 3D network architectures.
- A single parameter-efficient network may ease multimodal model deployment and improve clinical usability.
- The KD-loss has potential to support model generalization in medical image segmentation applications.