AI Research Portal

This paper introduces M²-CoFS, a unified model designed to break task boundaries between 3D medical image fusion (MIF) and segmentation (MIS). The core idea is a "network-guided network learning" paradigm that models the weight spaces of MIF and MIS as high-dimensional manifolds and uses a lightweight neural network to implicitly construct a shared manifold, yielding unified weights for both tasks. Experimental results demonstrate that M²-CoFS consistently outperforms state-of-the-art comparison models on both MIF and MIS.

The paper addresses several challenges in integrating 3D medical image fusion (MIF) and segmentation (MIS):

Existing methods often treat MIF and MIS independently, lacking inter-task collaboration and limiting overall performance.
Current MIF approaches frequently focus on 2D slices, neglecting the spatial continuity inherent in 3D volumetric data, which hinders integration with 3D MIS methods.
The differing optimization objectives of MIF (capturing complementary information) and MIS (delineating local boundaries) can lead to gradient conflicts during joint optimization.

A collaborative MIF-MIS network is designed to enable mutual promotion between the two tasks, overcoming the limitations of single-task approaches.
A "network-guided network learning" strategy is proposed, utilizing a lightweight Weight Generation Network (WGN) to implicitly construct a shared manifold for unified, task-order-irrelevant weights.
A tailored three-stage training strategy is introduced to progressively integrate weights from task-specific spaces into a shared space, ensuring advanced performance for both MIF and MIS.

The M²-CoFS model employs a three-stage training strategy to achieve joint optimization of MIF and MIS.

Stage I: Task-Specific Weights Obtaining In this initial stage, MIF and MIS are trained independently to obtain high-quality task-specific weights. This provides a foundation for subsequent collaborative optimization. The optimization objectives are defined by minimizing the respective task losses, LF and Ls. The encoder utilizes a 3D CNN for local features and a Swin Transformer for global context, while feature-fusion modules and task-specific heads handle multimodal integration and segmentation.

Stage II: Cross-Task Weights Adaptation based on Continual Learning To reduce discrepancies between the weight distributions of the two tasks, a cross-task weights adaptation mechanism is employed. This involves initializing one task's weights with the other's pre-trained weights and continuing training. To mitigate catastrophic forgetting, Synaptic Intelligence (SI) is incorporated as a regularizer, quantifying weight importance and allowing adaptation while preserving critical knowledge.

Stage III: Network-Guided Network Learning - A Unified Task Order-Irrelevant Weights Generation Network based on Implicitly Constructing Shared Manifold The primary goal of this stage is to obtain unified, task-order-irrelevant weights (θu) that can be used for inference in both MIF and MIS. The weight spaces of MIF and MIS are treated as high-dimensional manifolds. A lightweight MLP-based Weight Generation Network (WGN) implicitly constructs a shared manifold and generates unified weights. To preserve the geometric structure of the individual task manifolds, a manifold distance-based loss function is embedded into the WGN's objective. This involves projecting weights into a shared space, constructing a KNN-based graph, and computing distance matrices to form a regularizer.

Loss Function The overall loss function in Stage III combines the task-specific losses (LF' and Ls') with a manifold-constraint term (LM). The task-specific losses are adapted from Stage I, and the manifold-constraint term uses computed distance matrices to regularize the WGN.

Architecture: The encoder combines a 3D CNN and a Swin Transformer. Feature-fusion modules (FF and Fs) use 3D CNNs. The fusion head uses stacked 3D CNNs, and the segmentation head leverages the Swin UNETR architecture. A lightweight MLP-based Weight Generation Network (WGN) is used in Stage III.
Key Equations:
- Inference: Yf, Ys = G(Φ(A, B); θu) (Equation 1)
- Stage I Optimization: arg min LF (ψ ◦ FF (Φ(A, B); θF)) and arg min Ls (φ ◦ Fs (Φ(A, B); θs)) (Equations 2 & 3)
- Stage II Adaptation: Os→F = arg min LF (ψ ◦ FF (Φ(A, B); θs)) and θF→s = arg min Ls (φ ◦ Fs (Φ(A, B); θF)) (Equations 4 & 5)
- Synaptic Intelligence (SI): Ωk = Σ(t=1 to T) ||θ(t) - θ(t-1)||² / (||θ(t)||² + ξ) (Equation 6)
- Stage II Loss with SI: L'F = LF + λLSI and L's = Ls + µLsI (Equations 18 & 19)
- Stage III Optimization: arg min θu [LF'(ψ ◦ FF(Φ(A, B); θu)) + Ls'(φ ◦ Fs(Φ(A, B); θu)) + LM(θu, θF→s, θs→F)] (Equation 8)
- WGN Output: G(θF→s, θs→F) = θu (Equation 10)
- Manifold Distance Computation: LM = DUS + DU→F (Equation 15)
- Stage I Loss (MIF): LF = LL₁ + α LG (Equation 16)
- Stage I Loss (MIS): Ls = 1 - (2 * Σ(i=1 to J) Σ(j=1 to I) si,j * Li,j) / (Σ(i=1 to J) Σ(j=1 to I) si,j + Σ(i=1 to J) Σ(j=1 to I) Li,j) (Equation 17)
- Stage III Loss: L = L'F + ηL's + γLM (Equation 20)
Loss/Objective: The model uses task-specific losses (e.g., L1 loss, Dice loss) and incorporates Synaptic Intelligence (SI) for regularization against catastrophic forgetting. In Stage III, a manifold distance-based loss (LM) is added to guide the WGN.
Optimization: AdamW optimizer is used with a learning rate of 3 × 10⁻⁴.
Augmentation: Not stated in the paper.
Data Handling: The MRBrainS MRI dataset is used with a 4:1 training/test split. The BraTS19 dataset is used with a 9:1 training/test split (100 random cases).
Implementation Details: Experiments are conducted on Ubuntu 22.04 with an Intel i9-14900K CPU, RTX 4090 GPU, and PyTorch 2.3.0.
- Stage I: α = 0.1, 200 epochs.
- Stage II: λ = 0.1, μ = 0.1, 100 epochs.
- Stage III: n = 2, γ = 0.1, 200 epochs.

Datasets: MRBrainS MRI dataset and BraTS19 dataset.
Metrics:
- MIF: QMI, QTE, QNICE, QG, MI, VIFp.
- MIS: Dice Score (%), Hausdorff Distance (mm).
Baselines:
- MIF: BSAFusion, C2RF, EMMA, FILM, GIFNet, RMR-Fusion.
- MIS: 3D UX-Net, TransBTS, RMR-Fusion, Swin UNETR, UNETR, nnFormer.
Numeric Results:
- MIF: The paper reports quantitative results in Table 1, showing M²-CoFS achieving the best or second-best performance across multiple metrics on both MRBrainS MRI and BraTS19 datasets. For example, on BraTS19, M²-CoFS achieved QMI of 0.858 and QTE of 0.604.
- MIS: The paper reports quantitative results in Table 1, showing M²-CoFS achieving the best or second-best performance across multiple metrics on both MRBrainS MRI and BraTS19 datasets. For example, on BraTS19, M²-CoFS achieved a Dice Score of 81.96% and a Hausdorff Distance of 4.376 mm.
- Ablation Study (Table 2): Removing Stage II or Stage III led to performance drops. Removing SI (c1) or LM (c2) also negatively impacted results.
- MIS Improvement (Table 3): Replacing backbones with M²-CoFS strategy showed improvements over original models for TransBTS, 3D UX-Net, and UNETR.
Runtime/Training Cost: Not stated in the paper.

The paper notes that the explicit manifold distance computation is "not perfectly precise" but argues that embedding it into the loss function still preserves manifold geometry by maintaining relative distances.

Method evidence: "We propose M²-CoFS, a unified model capable of jointly handling both tasks. Our core contribution is a "network-guided network learning" paradigm designed to break the task boundaries." (Abstract)
Result evidence: "Experimental results show that M²-CoFS consistently outperforms state-of-the-art comparison models on both MIF and MIS." (Abstract)
Limitation evidence: "Note that while this explicit manifold distance computation is not perfectly precise, embedding it into the loss function still preserves manifold geometry." (Section: Method - The Process of Manifold Distance Computation)

Breaking Task Boundaries: A Unified Model for 3D Medical Image Fusion and Segmentation Guided by Manifold Perspective

Summary

Problem

Key Contributions

Method

Technical Details

Results

Weaknesses / Limitations

Evidence

Additional Figures