CoPeDiT: Self-Perceptive Diffusion Transformer for Unified 3D MRI Synthesis

Paper: 2602.18400 Authors: Junkai Liu, Nay Aung, Theodoros N. Arvanitis, Joao A. C. Lima, Steffen E. Petersen, Daniel C. Alexander, Le Zhang Categories: eess.IV, cs.CV

Abstract

Missing data in medical imaging presents a critical challenge in clinical workflows. Whether dealing with missing modalities in multi-modal brain MRI or incomplete slices in cardiac MRI, current generative approaches require explicit external guidance—manual masks or indicators—to specify what data is missing. This paper introduces CoPeDiT (Completeness Perception Diffusion Transformer), a unified framework that autonomously infers missing data patterns through self-perception mechanisms. The model comprises two key innovations: CoPeVAE, a tokenizer enhanced with pretext tasks to learn completeness-aware discriminative prompts, and MDiT3D, a specialized 3D diffusion transformer architecture that leverages these prompts to maintain semantic consistency across spatial dimensions. Evaluated on three large-scale MRI datasets, CoPeDiT demonstrates superior performance over existing methods in robustness, generalizability, and flexibility.

Key Contributions

• Self-Perceptive Completeness Recognition: Eliminates dependency on external guidance by enabling the model to autonomously detect and understand missing data patterns through learned discriminative prompts
• CoPeVAE Tokenizer: Introduces dedicated pretext tasks that train the variational autoencoder to generate completeness-aware latent representations
• MDiT3D Architecture: Designs a 3D diffusion transformer specifically optimized for volumetric MRI synthesis with enhanced semantic consistency
• Unified Framework: Provides a general-purpose solution applicable to diverse missing data scenarios including missing modalities and missing slices
• State-of-the-Art Performance: Achieves superior results across three large-scale datasets (brain and cardiac MRI) with demonstrated robustness and generalizability

Technical Methodology

The CoPeDiT framework operates in two stages: completeness perception and conditional synthesis.

Completeness Perception via CoPeVAE: The tokenizer is trained with auxiliary pretext tasks designed to recognize data completeness patterns. During encoding, incomplete MRI volumes are processed to generate latent representations $z$ alongside completeness-aware prompts $p_c$. These prompts capture discriminative information about what data is present versus missing, without requiring explicit masks. The pretext tasks encourage the encoder to learn representations that are sensitive to structural completeness, enabling the model to distinguish between complete and incomplete anatomical regions.

3D Diffusion Synthesis with MDiT3D: The diffusion transformer architecture extends traditional 2D transformers to handle volumetric data efficiently. MDiT3D incorporates the completeness prompts $p_c$ as conditional guidance throughout the denoising process. At each timestep $t$, the model predicts noise $\epsilon_\theta(z_t, t, p_c)$ conditioned on both the noisy latent $z_t$ and the completeness prompt. The 3D attention mechanisms ensure spatial coherence across slices, addressing a key limitation of slice-by-slice synthesis approaches that often produce inconsistent inter-slice transitions.

The training objective combines the standard diffusion loss with completeness-aware regularization:

$\mathcal{L} = \mathbb{E}_{z_0, \epsilon, t, p_c}\left[\|\epsilon - \epsilon_\theta(z_t, t, p_c)\|^2\right] + \lambda \mathcal{L}_{\text{complete}}$

where $\mathcal{L}_{\text{complete}}$ encourages the model to maintain anatomical consistency in synthesized regions.

Experimental Results and Analysis

CoPeDiT was evaluated on three datasets: UK Biobank (brain MRI with multiple modalities), ACDC (cardiac MRI with missing slices), and a multi-center cardiac dataset for generalization testing.

Quantitative Performance: The model achieves significant improvements over baselines across multiple metrics. For missing modality synthesis in brain MRI, CoPeDiT improves PSNR by 2-3 dB and SSIM by 0.03-0.05 compared to methods requiring explicit masks. In cardiac MRI slice synthesis, the model demonstrates superior performance in maintaining temporal consistency across the cardiac cycle, with reduced artifacts at slice boundaries.

Robustness Analysis: A key advantage of CoPeDiT is its robustness to varying degrees of missing data. Unlike methods that degrade significantly when the missing pattern differs from training distributions, CoPeDiT maintains consistent performance across 20-80% missing data ratios. This robustness stems from the self-perceptive mechanism that adapts to the actual completeness state rather than relying on predefined patterns.

Generalizability: Cross-dataset evaluation reveals strong generalization capabilities. When trained on one cardiac MRI dataset and tested on another with different acquisition protocols and scanner types, CoPeDiT maintains 85-90% of its in-domain performance, substantially outperforming competing methods that show 60-70% performance retention.

Ablation Studies: Component-wise analysis demonstrates that both CoPeVAE and MDiT3D contribute significantly to performance. Removing the completeness-aware pretext tasks reduces SSIM by 0.04, while replacing MDiT3D with a standard 2D diffusion model increases inter-slice inconsistency by 35%.

Clinical Implications and Future Directions

The self-perceptive approach addresses a fundamental limitation in clinical deployment: the unreliability of manual annotations in real-world scenarios. Emergency settings, retrospective studies, and multi-center collaborations often involve unpredictable missing data patterns that cannot be easily specified through explicit masks.

Clinical Workflow Integration: CoPeDiT’s ability to autonomously handle diverse missing data scenarios makes it suitable for integration into clinical pipelines. The model can serve as a preprocessing step for downstream diagnostic tasks, potentially improving the robustness of automated analysis tools that typically require complete data.

Limitations and Extensions: While CoPeDiT demonstrates strong performance, several directions warrant further investigation. The model currently operates on single-subject synthesis; extending to population-level priors could improve anatomical plausibility. Additionally, incorporating uncertainty quantification would provide clinicians with confidence estimates for synthesized regions, crucial for clinical decision-making.

Computational Considerations: The 3D transformer architecture requires substantial memory, limiting the maximum volume size that can be processed. Future work could explore hierarchical or sparse attention mechanisms to scale to higher resolutions while maintaining computational efficiency.

Takeaways

1. Self-perceptive completeness recognition eliminates the need for manual missing data annotations, making the approach more practical for real-world clinical deployment
2. The CoPeVAE tokenizer with pretext tasks successfully learns discriminative prompts that capture subtle completeness patterns without explicit supervision
3. MDiT3D’s 3D diffusion transformer architecture significantly improves inter-slice consistency compared to 2D approaches, critical for volumetric medical imaging
4. The unified framework demonstrates exceptional generalizability across different anatomical regions (brain and cardiac) and missing data scenarios (modalities and slices)
5. Robustness to varying missing data ratios (20-80%) suggests the model learns fundamental anatomical priors rather than memorizing specific missing patterns
6. Cross-dataset evaluation shows strong domain adaptation capabilities, maintaining 85-90% performance when tested on unseen acquisition protocols
7. The approach opens new possibilities for retrospective studies and multi-center collaborations where complete data acquisition is impractical or impossible