CodeBrain: Imputing Any Brain MRI via Modality- and Instance-Specific Codes

TL;DR

CodeBrain addresses unified brain MRI imputation under varying modality availability by reframing inter-modality synthesis as a two-stage latent code prediction problem. It constructs a latent space of scalar-quantized modality- and instance-specific codes (dimension $d$, quantized to $L$ levels) and learns modality-agnostic common features in Stage I, then Stage II uses a prior encoder to predict full-modality codes from incomplete data, enabling high-fidelity synthesis without modality-specific modules. Evaluations on IXI and BraTS 2023 show CodeBrain outperforms four SOTA methods across one-to-one and many-to-one scenarios and improves downstream brain tumor segmentation on BraTS 2023. The approach simplifies generalization, reduces model complexity, and offers a scalable solution for unified brain MRI imputation, with future work addressing hallucination and integrating MRI physics.

Abstract

Unified MRI imputation, which can adapt to diverse imputation scenarios, is highly desirable as it reduces scanning costs and provides comprehensive MRI information for improved clinical diagnosis. Existing unified MRI imputation methods either rely on specific prompts to guide their transformation network or require multiple modality-specific modules. However, these approaches struggle to capture large modality and instance variations or become too complex to generalize effectively. To address these limitations, we propose CodeBrain, a fundamentally different pipeline for unified brain MRI imputation. Our key idea is to reframe various inter-modality transformations as a full-modality code prediction task via a two-stage framework. In the first stage, CodeBrain reconstructs a target modality from any other modalities by learning a compact scalar-quantized code for each instance and modality. Any target modality can then be reconstructed with high fidelity by combining the corresponding code with shared features extracted from any available modality. In the second stage, a projection encoder is trained to predict full-modality compact codes from any incomplete MRI samples, effectively simulating various imputation scenarios. We evaluate our CodeBrain on two public brain MRI datasets (i.e., IXI and BraTS 2023). Extensive experiments demonstrate that CodeBrain outperforms state-of-the-art methods, setting a new benchmark for unified brain MRI imputation. Our code will be released.

Figures (6)

• Figure 1: Illustration of different MRI imputation settings. Given multiple MRI modalities, there are various one-to-one and many-to-one transformation scenarios for synthesizing a missing one. Our proposed CodeBrain model aims to achieve a unified imputation by leveraging modality- and instance-specific prompts (i.e., learnable codes) to generate any missing brain MRI modality.
• Figure 2: Pipeline of our proposed CodeBrain framework for unified Brain MRI imputation. It consists of two training stages: Stage I constructs a latent quantized space, where any target anchor is projected to its particular scalar code by $E_{posterior}$ and then reconstructed by $D_a$ using this code along with the common features $F_c$ extracted from any input incomplete modalities by $E_s$. Stage II aims to predict full-modality scalar-quantized codes from any available source modalities by $E_{prior}$ with a grading loss. During inference, the decoder $D_a$ synthesizes the missing anchor modality using the predicted anchor code by $E_{prior}$ and the extracted common features from input incomplete modalities via $E_s$. Note that all modules here are modality-agnostic.
• Figure 3: Exemplar synthesized brain MRI scans of our CodeBrain on the IXI dataset. 1st and 3rd columns: O$\rightarrow$O imputed results; 5th column: M$\rightarrow$O imputed results; 7th column: the original modality; Each synthesized scan is equipped with its corresponding error map on the right. The three bits in an annotation "010" correspond to T1, T2, and PD, respectively, where "1" indicates that modality is available.
• Figure 4: Exemplar comparison results between four public models and our proposed CodeBrain in the T2$\rightarrow$T1 (Top) and T1, PD$\rightarrow$T2 (Bottom) scenarios on the IXI dataset, along with corresponding error maps.
• Figure 5: Ablation studies with different $d$ in our CodeBrain on the IXI dataset. Here, common features $F_c$ enhance the reconstruction performance, in Stage I, see "Stage I (wo $F_c$) vs. (w/ $F_c$)", while the Stage II performance is improved by using a grading loss, see "Stage II (w/ CLS) vs. (w/ Grad.)".