X-Diffusion: Generating Detailed 3D MRI Volumes From a Single Image Using Cross-Sectional Diffusion Models

## Abstract

Abstract In this work, we present X-Diffusion , a cross-sectional diffusion model tailored for Magnetic Resonance Imaging (MRI) data.
X-Diffusion is capable of generating the entire MRI volume from just a single MRI slice or optionally from few multiple slices, setting new benchmarks in the precision of synthesized MRIs from extremely sparse observations. The uniqueness lies in the novel view-conditional training and inference of X-Diffusion on MRI volumes, allowing for generalized MRI learning.
Our evaluations span both brain tumour MRIs from the BRATS dataset and full-body MRIs from the UK Biobank dataset. Utilizing the paired pre-registered Dual-energy X-ray Absorptiometry (DXA) and MRI modalities in the UK Biobank dataset, X-Diffusion is able to generate detailed 3D MRI volume from a single full-body DXA. Remarkably, the resultant MRIs not only stand out in precision on unseen examples (surpassing state-of-the-art results by large margins) but also flawlessly retain essential features of the original MRI, including tumour profiles, spine curvature, brain volume, and beyond. Furthermore, the trained X-Diffusion model on the MRI datasets attains a generalization capacity out-of-domain ( e.g . generating knee MRIs even though it is trained on brains). The code is available on the project website https://emmanuelleb985.github.io/XDiffusion/ .

## 1 Introduction

Medical imaging stands as a cornerstone in modern healthcare, its innovations playing a critical role in disease diagnosis and treatment planning. Traditional MRI scans, though detailed, are often time-consuming and come with significant economic implications []. The urgency to tackle these impediments has propelled research endeavors in the past, but the quest for a cost-efficient, rapid, and precise alternative persists[, , ].
A rapid and affordable MRI process would catalyze early disease detection, potentially saving countless lives. Moreover, by reducing the barriers to access, we would be ensuring a more holistic healthcare approach, promptly addressing diseases before they escalate.

Figure: Figure 1: X-Diffusion. We present X-Diffusion, a method that can generate detailed and dense MRI volumes from (i) a single MRI slice, or (ii) from a single DXA scan. This is the first work that we know of in the medical imaging literature to generate MRIs that preserve key properties from extremely sparse inputs. X-Diffusion can optionally accept multiple MRI slice inputs as well for full MRI generation from few MRI slices.

Traditionally, inverse 2D or 3D fast Fourier transform (FFT) [] on k-space data with full Cartesian sampling are used to reconstruct MR images from raw data or with the help of machine learning models [, , , , , , ].
Recent years saw a pivot towards machine learning-based frameworks such as Generative Adversarial Networks (GANs) and Convolutional Neural Networks (CNNs), harnessing the power of deep neural networks to enrich MRI reconstruction [, ].
However, a pervasive challenge remained: the synthesis of high-resolution MRIs from extremely limited observations (or even a single 2D image). The previous works either target compressive sensing for increasing the frequency resolution of the MRI [, , , ] or target increasing the slice density when a sufficient number of slices is available (*e.g*. more than 30) []. These existing gaps in the MRI reconstruction landscape underscore the significance of our X-diffusion in MRI reconstruction from an extremely small number of observations.

Recently, the use of diffusion-based models for image inverse problems has shown great success [, , , , , ].
This motivated our X-Diffusion to investigate learning volumes instead of images.
In this light, our X-Diffusion proposes a novel architecture to allow learning on 3D volumetric data by view-dependent cross-sections. This allows for full MRI generation with unprecedented accuracy from a single MRI slice, multiple slices, or even from DXA image if paired data is available (see Figure [1](#S1.F1)). To the best of our knowledge, X-Diffusion is the first work to successfully generate detailed MRI volumes from a single DXA scan, bridging the gap between two common data modalities in medical imaging. It is important to note that the generated MRIs are not clinical replacements for true MRIs, but could provide a quick, affordable, and informative “pseudo-MRI" before conducting a full MRI examination.

Contributions:
(i) We introduce X-Diffusion, a cross-sectional diffusion model that generates MRI slices conditioned on a single input MRI slice or multiple slices. The proposed X-Diffusion achieves state-of-the-art results on MRI reconstruction and super-resolution compared to recent methods on BRATS, a large public dataset of annotated MRIs for brain tumours.
(ii) We adapt our X-Diffusion to leverage paired and registered full-body MRI and DXA images from UK Biobank dataset to generate full-body 3D MRI from a single DXA for the first time in the literature.
(iii) We validate the generated MRIs on a wide range of tasks that ensure the generated MRIs retain important features of the original MRIs, including tumor profiles, spine curvature, brain volume, and more, without using this meta-information in the generation process.
(iv) We showcase the generalization of trained X-Diffusion on different datasets (knee MRIs) illustrating the potential of X-Diffusion to be the first 3D volumetric foundation model in medical imaging.

## 2 Related Work

### 2.1 3D Understanding and Generation

Multi-View 3D Reconstruction.
Multi-view 3D reconstruction predicts 3D from 2D RGB images captured from different views [, ].
Recently, Neural Radiance Fields (NeRF) [, ] marked a significant shift towards 3D radiance, enabling realistic view synthesis [, ]. Current research optimizes NeRF for both few-shot and one-shot settings [, , , ].
One important difference between these few-shot NeRF works and our X-Diffusion is that we learn the explicit 3D prior on volumes by cross-sectional conditioning while those other works rely on 2D priors to enhance the neural rendering of mostly surfaces, making them less suitable for 3D volumes.

Single-View 3D Reconstruction.
Previous efforts leveraged CLIP’s capabilities [] for tasks like 3D modeling [, , , ].
The seminal work of DreamFusion [] distilled a ready-made diffusion mechanism [] into NeRF [, ]. This methodology ignited a myriad of new techniques, both in converting text to 3D (*e.g*. [, ]) and transitioning visuals to 3D forms (*e.g*. [, , , ]).
These frameworks were considerably improved by Zero-123  [], explicitly conditioning on camera-views while finetuning Stable Diffusion on the large 3D CAD dataset Objaverse []. While Zero-123 learns to generate surface renderings of a target view given a single image, X-Diffusion learns to generate a cross-sectional slice, conditioned on the angle and depth index of the slice, allowing for dense 3D volume generation and targeting MRI medical imaging.

### 2.2 MRI Analysis and Reconstruction

Full-Body MRI Analysis.
Most methods on automatic MRI analysis focused on developing methods for local segmentation of organs or tumours [, , , ]. Relatively few studies looked at whole-body scans. Most of them were developed to detect and segment the spine in tasks such as scoliosis detection[, , , , ].
We leverage the MRI analysis techniques for validating the viability of the generated MRIs for tumor, spine, and other discriminative features of interest from the medical imaging community.

MRI Reconstruction.
With the recent rise of foundation models in computer vision [, , ], several attempts have shown promise in steering these models for medical imaging domain [, ]. However, this is mainly limited to discriminative tasks such as segmentation, classification and detection. For Medical imaging inverse problem tasks, mostly classical methods were employed for incensing the resolution of the reconstruction [, ], or adopt diffusion models without great leverage of image pretraining [, , , ]. ScoreMRI and TPDM [, ] make use of diffusion probabilistic model (DPM) performing conditional sampling-based inverse problem. TPDM [] proposed to overcome the limitation of ScoreMRI being an image-to-image model and leveraged the 3D prior distribution of the data using a product of two 2D diffusion models. Although this approach enables 3D generation, it only samples from two fixed canonical planes from the 3D MRI and does not work for sparse input. On the other hand, X-Diffusion leverages the full 3D volume by sampling the brain in all directions and leverages the Stable Diffusion huge image pretraining for 3D MRI volumes from a single MRI slice or an aligned DXA image.

Figure: Figure 2: X-Diffusion Pipeline. A single or multi-slice input is fed into the Latent Diffusion U-Net conditioned on the target slice index $d$ and target rotation from 360^∘ slicing. The 3D volume is reconstructed by vertical stacking of the slices from a fixed axis of rotation. The final volume $\mathcal{X}$ is obtained after averaging the $N$ realigned view-dependant volumes $R_{i}^{\intercal}\mathcal{X}_{R_{i}}$ from a set of predefined target rotations $R_{i}$.
Refer to caption: /html/2404.19604/assets/x2.png

## 3 Methodology

In this section, we present the methodology underpinning X-Diffusion. Our approach can be delineated into three primary aspects (as shown in Figure [2](#S2.F2)):

- 1.
Synthesizing MRI volumes using X-Diffusion models conditioned on a single MRI slice or multiple slices, with view direction and slice of volume indexing control.
- 2.
Leveraging the multi-view capability of the X-Diffusion to aggregate multiple view-conditioned volumes to generate the final MRI output.
- 3.
Adopting X-Diffusion to work on converting Dual X-ray Absorptiometry (DXA) to MRI volumes by utilizing a registered DXA-MRI paired dataset to create the first DXA to 3D MRI model in the medical imaging literature.

### 3.1 Diffusion Models Preliminaries

In previous works on view-conditional diffusion [, , ],
the diffusion model $\epsilon_{\theta}$ is trained based on the objective:

$$ $\min_{\theta}\;\mathbb{E}_{z\sim\mathcal{E}(x),t,\epsilon\sim\mathcal{N}(0,1)}||\epsilon-\epsilon_{\theta}(z_{t},t,c(x,R,\tau))||_{2}^{2}$ (1) $$

In Equation[1](#S3.E1), $\theta$ denotes the model parameters that are being optimized. The latent variable $z$ is sampled from a distribution $\mathcal{E}(x)$, where $x$ indicates the input data and $\mathcal{E}$ is a fixed encoder. $t\sim[1,2,...,T]$ specifies a particular time step during the diffusion process with maximum $T$ steps. The term $\epsilon$ is a noise variable, which is sampled from a standard normal distribution, $\mathcal{N}(0,1)$. The function $\epsilon_{\theta}$ is representative of the model’s prediction for a given $z_{t},t$, and transformation $c(x,R,\tau)$, where $R$ and $\tau$ are rotation and translation parameters, respectively.

Proceeding, the gradient of the Score Jacobian Chaining (SJC) loss, which approximates the score towards the non-noisy input as described in [, ], is given by:
$\nabla\mathcal{L}_{SJC}=\nabla_{I_{\pi}}\log p_{\sqrt{2}\epsilon}(x_{\pi})$. The term $\nabla_{I_{\pi}}$ specifies the gradient with respect to the image $I_{\pi}$. The expression $p_{\sqrt{2}\epsilon}(x_{\pi})$ denotes the probability distribution of the transformed image $x_{\pi}$ under noise level $\sqrt{2}\epsilon$.

In our setup, $\tau$ is replaced with the index $d$ of the slice of the MRI volume and $R$ is the rotation applied on the MRI volume for the cross-sectional processing.

### 3.2 X-Diffusion for Cross-Sectional MRI Synthesis

Upon acquiring the MRI slice $x\in\mathbb{R}^{H\times W}$ , we seek to synthesize the entire MRI volume $\mathcal{X}\in\mathbb{R}^{H\times W\times D}$ . For this, we employ X-Diffusion $\epsilon_{\theta}$, a cross-sectional diffusion model. The fundamental idea stems from the analogy that a 3D volume can be built crosswise by stacking slices from a certain direction, just like a loaf of bread. The full target volume $\bar{\mathcal{X}}$ can be reconstructed from limited slices by generating target slices indexed by their depth $d\in[1,2,...,D]$ in the MRI volume conditioned on a certain direction $R$ where the volume is oriented. This simplifies the learning of cross-sections since the rotated MRI volume $R\mathcal{X}$ will have the same size $H\times W\times D$
as the original volume where zero padding is used. For simplicity of the processing of the data, we use the same dimensions for all directions ($H=W=D$). This allows varying the depth after rotating the ground truth MRI $\bar{\mathcal{X}}$ volume by simply indexing by the depth index $d$, and hence the slice that is used for training will be $\bar{x}_{d}=(R\mathcal{X})_{d,:,:}$. The full objective of training X-Diffusion is as follows.

$$ $\displaystyle\min_{\theta}$ $\displaystyle\mathbb{E}||\epsilon-\epsilon_{\theta}(z_{t},t,c(x,d,R))||_{2}^{2}$ (2) s.t. $\displaystyle z\sim\mathcal{E}(\bar{x_{d}}),~{}~{}t\sim[1,2,...,T]$ $\displaystyle\epsilon\sim\mathcal{N}(0,1),~{}~{}d\sim[1,2,...,D],~{}~{}R\sim SO(3)$ $$

The X-Diffusion model is trained with cross-sections from all different directions $R$ and all different depths $d$, which allows it to generate the target from any arbitrary rotation and depth.
At inference , X-Diffusion is applied $D$ times with $d\in[1,2,..,D]$ from an arbitrary orientation $R_{i}$ to obtain the view-conditional volume $\mathcal{X}_{R_{i}}$.

$$ $\displaystyle\mathcal{X}_{R_{i}}=\begin{bmatrix}\epsilon_{\theta}(z_{t},t,c(x,1,R_{i}))\\ \vdots\\ \epsilon_{\theta}(z_{t},t,c(x,D,R_{i}))\end{bmatrix}~{}~{},~{}~{}t=1,2,...,T.$ (3) $$

The volume is then rotated back by $R_{i}^{\intercal}\mathcal{X}_{R_{i}}$ to the Canonical orientation in order to proceed for validation.

Multi-slice input.
While the pipeline described above is effective, it relies on heavy diffusion operations for each slice input and output. Adding more slices by simply inflating the network will create computational and memory difficulties. Therefore, to efficiently allow X-Diffusion’s pipeline to accept $K$ slices as input while maintaining the same original weights structure of Stable Diffusion [], we perform a cumulative sum operation on the dot product of consecutive slices to reduce to a single slice input. The reduction operation of the $K>1$ input slices $(x_{1},x_{2},...,x_{K})$ is similar to what is followed in TPDM [] in the conditioning volume, and it can be described as follows. $x=\frac{1}{K-1}\sum_{j=1}^{K}x_{j}\cdot x_{j+1}$.

Figure: Figure 3: Qualitative Results of DXA to 3D MRI Generation with X-Diffusion. We show a single DXA example, Two corresponding ground-truth MRI slices (index 68 and 100), the corresponding generated MRI slice, and a difference map to qualitatively measure the error between generated and ground-truth MRI. Note that when stitching the MRIs in the UKBioBank dataset, there is a disparity in the bias field effect which is strongest at the knee region (brighter pixels). The same pattern is present on all samples in the dataset for a fixed depth, and hence it is learned by X-Diffusion as well.
Refer to caption: /html/2404.19604/assets/x3.png

Figure: Figure 4: Visualisations of 3D Brain Generation. For the input slice (slice index 76), we show examples of slices from generated 3D brain MRI volumes with varying slice index (top) and its ground-truth brain slices (bottom). We show the tumour profile segmentation map in all output and ground truth slices to highlight the differences and show the 3D tumor in the generated MRI and ground truth MRI in the most right column. Red is used for non-enhancing and necrotic tumor core, green for the peritumoral edema, and blue for the enhancing tumor core.
Refer to caption: /html/2404.19604/assets/images/MultipleSlices/Brains/76_sample_input_76.png

### 3.3 Multi-View MRI Volume Generation

One advantage of our cross-sectional diffusion is that it can learn and generate the volume $\mathcal{X}_{R_{i}}$ from any arbitrary view direction $R_{i}$ (as in Equation [3](#S3.E3)). In training, this allows X-Diffusion to train on MRIs from all types of cross-sections, unlike the typically followed common 3 planes (coronal, sagittal, and axial) [, , ], which allows the model to generalize much better. At inference, we leverage this power to generate $N$ volumes from $N$ different views predefined as equally distributed views around the $360^{\circ}$ around the azimuth horizontal rotations $R_{i}\in\{R_{\text{azim}}(\frac{i\times 360^{\circ}}{N})\}_{i=1}^{N}$, where $R_{\text{azim}}(r)$ is the rotation matrix defined by rotating by $r$ degrees around the vertical axis (0,1,0). The final MRI volume output $\mathcal{X}$ is then obtained by averaging the view-conditional volumes at inference after rotating back to the canonical orientation of the output as follows.

$$ $\mathcal{X}=\frac{1}{N}\sum_{i=1}^{N}R_{i}^{\intercal}\mathcal{X}_{R_{i}}$ (4) $$

This multi-view aggregation is inspired by how typically multi-view discriminative methods learn a global representation by average/max pooling multiple views features[, , ].
We show in Section [6.1](#S6.SS1) the utility of the volume averaging compared to a single volume.

### 3.4 DXA to MRI Volume Generation

DXA is a single image data modality that is similar to X-ray but includes other non-bony information such as tissue mass []. It measures bone mineral density
and body fat composition. The radiation level is low enough that it is acceptable for conducting studies of healthy participants, such as the UK Biobank. In order to leverage X-Diffusion to synthesize the MRI volume $\mathcal{X}_{R_{i}}$ from a single DXA $u\in\mathbb{R}^{H^{\prime}\times W^{\prime}}$ (as in Figure [3](#S3.F3)), they have to be aligned, and registered. Note that, the size of the DXA does not match the MRI ($H^{\prime}\times W^{\prime}\neq H\times W$) and the scans are not registered. The two modalities in UK BioBank were not taken simultaneously but close in time hence why we believe registration is feasible for these two sequences as illustrated extensively in [].

In order to tackle this domain gap, we leverage a registration network [] paired with a X-Diffusion to achieve DXA to MRI slice generation.
[] introduced a multi-modal image-matching contrastive framework, that is able to learn correspondences between DXA and middle coronal MRI slices.
These networks extrapolate the DXA scan by a transform $\mathbf{A}$ to the coronal MRI slice by harnessing the embedded patterns and features of the DXA and the coronal MRI mid-slice.
X-Diffusion is then trained on the registered DXA ($\mathbf{A}(u)\in\mathbb{R}^{H\times W}$) and corresponding MRI slices in ($\bar{x}\in\mathbb{R}^{H\times W}$) in the target MRI volumes and is able to produce precise MRI volumes $\mathcal{X}$ that align with the DXA scans (see Figure [5](#S3.F5) left).
The other details are similar to Section [3.2](#S3.SS2) and Section [3.3](#S3.SS3).

Figure: Figure 5: Alignment and Error Spread of MRI Generation. (left) We show from left to right: input DXA, ground-truth MRI, generated MRI, and overlay of the two modalities to test the alignment. The 3D PSNR for this example is 26.38 dB. (right) We show the Mean Square Error of the generated MRI from a single MRI slice (index 78) as a function of the output depth index. It shows that the error is not equally distributed in the generated MRI.
Refer to caption: /html/2404.19604/assets/x17.png

## 4 Experiments

### 4.1 Datasets

BRATS. The largest public dataset of brain tumours consisting of 5,880 MRI scans from 1,470 brain diffuse glioma patients, and corresponding annotations of tumours[, , ]. All scans were skull-stripped and resampled to 1 mm isotropic resolution. All images have a resolution of 240 $\times$ 240 $\times$ 155, and we use the flair T2 sequence. Tumours are annotated for 3 classes: Whole Tumour (WT), Tumour Core (TC), and Enhanced Tumour Core (ET).

UK Biobank. A more comprehensive dataset of 48,384 full-body MRIs from more than 500,000 volunteers[], capturing diverse physiological attributes across a broad demographic spectrum. These Dixon MRIs do not come stitched, the scans are scanned axially and there is a disparity in the bias field effect (a common artifact of MRI machines) which is strongest at the knee region.
The same knee pattern is present on all samples in the dataset. UK Biobank MRIs are resampled to be isotropic and cropped to a consistent resolution (501 $\times$ 160 $\times$ 224). 48,384 whole-body MRIs are paired with antero-posterior (AP) DXA scans of the same subjects.

IXI. A dataset of T1-weighted 1.5 Tesla brain MRI images of 582 healthy subjects, freely available online [].

Knee fastMRI. A public dataset of raw k-space data from NYU Langone[, ]. We use the test set provided (n=109) of fastMRI single coil, dimensions 640x372x30. These are center-cropped to 320x320x30.

### 4.2 Evaluation Metrics

We use the standard 3D PSNR [] and 2D SSIM [] metrics to evaluate 3D MRI reconstruction and the following metrics for the validation experiments:

Dice Score. It is used to evaluate the performance of our model at segmenting the brain tumours [].
Dice Score = $\frac{2|Y\cap\hat{Y}|}{n(|Y|+|\hat{Y}|)}$, where $Y$ is the prediction, $\hat{Y}$ is the ground-truth label and $n$ the total number of slices.

Brain Volume.
We measure brain volume in $mm^{3}$ by counting the non-zero voxels in the volume multiplied by the voxel spacing [].

Spine Curvature.
Let $\gamma(t)=(x(t),y(t))$ be the equation of a twice differentiable plane curve parametrized by $t\in[0,209]$. We measure the spine curvature $\kappa$ similar to []: $\kappa=(y^{{}^{\prime\prime}}x^{{}^{\prime}}-x^{{}^{\prime\prime}}y^{{}^{\prime}})/(x^{{}^{\prime}2}+y^{{}^{\prime}2})^{\frac{3}{2}}$.

**Table 1: Model Performance on Test Brain Data and Whole-Body MRIs. We compare the MRI reconstruction for baselines ScoreMRI [], TPDM [] and our X-Diffusion model for varying input slice numbers in training and inference. We report the mean 3D test PSNR on BRATS (BR) brain dataset and the UK Biobank body dataset (UK). The results showcase huge improvement over the baselines, especially on the small number of input slices (particularly at 1). The parameters count and inference time of a single 3D MRI. On a single GPU a6000, 48GB of RAM for the three methods are as follows: ScoreMRI (860M, 139.1s), TPDM (1720M, 149.5s), and X-Diffusion (990M, 141.5s).**
|  | Test 3D PSNR $\uparrow$ |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | 1 slice | 2 slices | 3 slices | 5 slices | 10 slices | 31 slices |  |  |  |  |  |  |
|  | BR | UK | BR | UK | BR | UK | BR | UK | BR | UK | BR | UK |
| ScoreMRI [] | 9.37 | 8.54 | 10.25 | 9.16 | 10.68 | 10.42 | 12.37 | 11.88 | 14.31 | 13.24 | 29.24 | 19.01 |
| TPDM [] | 10.48 | 9.29 | 10.86 | 9.99 | 11.33 | 11.09 | 14.13 | 12.62 | 16.65 | 15.88 | 31.48 | 21.70 |
| X-Diffusion (ours) | 23.1 | 22.42 | 25.2 | 23.04 | 29.43 | 25.26 | 31.25 | 26.85 | 33.27 | 27.44 | 35.48 | 29.01 |

Figure: Figure 6: Visual Comparison of MRI Brain Reconstruction. We benchmark different methods of reconstructed 3D brains on test set with multi-slice inputs. We show a generated slice from 3D brain generated from ScoreMRI [] (top), TPDM[] (middle), and X-Diffusion (bottom) conditioned on a varying number of input slices. The red zoomed crop is placed in the exact location in all images to highlight the differences.
Refer to caption: /html/2404.19604/assets/images/Brains_Multi/Slide5.png

Figure: Figure 7: Effect of the Number of Input Slices. We plot the test PSNR *vs*. number of input slices for X-Diffusion and our baselines i.e. TPDM [] and ScoreMRI [] on the brain MRI dataset. $N$ is the number of averaged view-dependent volumes. We show the standard deviation of each run to account for potential randomness.
Refer to caption: /html/2404.19604/assets/x22.png

### 4.3 Baselines

We compare X-Diffusion’s performance against state-of-the-art MRI generation techniques, namely ScoreMRI [] and Two-Perpendicular-Diffusion-Models TPDM [] using ncsnpp model [].
For the multiple slice input (nx256x256) in X-Diffusion, we aggregated the multiple inputs to form a single batch (1x256x256). For comparison with Score-MRI, being an image-to-image model, we uniformly sampled $n$ slices along the z-axis. As for TPDM, we conditioned on $n$ slices from the full volume after the fusion of the two diffusion models.

### 4.4 Implementation Details

To facilitate using the pretrained weights of Zero-123 [] (based on Stable Diffusion []), we use the same channel size in the input 3, repeating the grayscale images. For the size of the MRI volumes, we used $H=W=D=155$, as originally the sizes in the dataset were 155 slices.
For model training, we use a base learning rate of $1.0e^{-06}$, LambdaLinearScheduler with warm-up every 100 steps. Batch size is set to 32. In the diffusion sampling, we used $T=1000$ time steps and an ETA of 1.0. More details about the datasets, metrics, and setup are provided in supplementary material.

## 5 Results

### 5.1 Main Results

Our results unequivocally highlight the superior performance of X-Diffusion in terms of both qualitative and quantitative metrics. Representative MRI volumes generated by our pipeline, when juxtaposed with ground-truth images, showcased remarkable similarity, with even intricate physiological features like tumor information, spine curvature, and fat distribution being accurately captured.

Notably, X-Diffusion achieves state-of-the-art $PSNR>30$ dB for a few input slices while baselines require more than 60 input slices to achieve similar performance (Figure [7](#S4.F7)). The margin is more than 12 dB PSNR for the 1-slice input in both the BRATS and the UK Biobank benchmarks (see Table [1](#S4.T1) and Figure [6](#S4.F6)). For reference, two randomly sampled MRIs from UK Biobank would have a PSNR of 15.95 dB $\pm$ 0.36 (on 4800 randomly sampled examples). Omitting the preprocessing step of alignment DXA to MRI, leads to a drop of PSNR on average by 2.87 dB (29.01 dB $\rightarrow$ 26.14 dB). The slices from 3D reconstructed volumes at varying depths and axis of rotation, visually match the ground truths for both brain and whole-body scans (see Figures [4](#S3.F4) and [5](#S3.F5) left). We also plot the error map (Figure [3](#S3.F3)) and the spread of the error (Figure [5](#S3.F5) right) of such X-Diffusion generations to highlight the differences with the ground truth MRIs.

### 5.2 MRI Validation Results

Brain Volumes Preservation.
The generated MRIs by our X-Diffusion retain almost the exact same average brain volume $1.28e^{6}$ $mm^{3}$ *vs*. $1.31e^{6}$ $mm^{3}$ of the real MRIs.

Tumour Information Preservation.
For the brain tumor segmentation, we use a Swin UNETR model[, ], trained with random rotation, and intensity as data augmentation.
On the test set with human ground-truth annotations ($n=333$), the brain volumes generated from single slice input preserve the volume of the different tumour components (paired t-test, $p-value<0.05$ for all 3 classes). In Figure [4](#S3.F4), we highlight the tumor profiles of the generated MRIs compared to the ground truth tumour profile. The real MRI Dice score in the test set is 85.15 while the generated MRIs from a single slice have a dice score of 83.09. This shows how the generated MRIs indeed preserve the tumor information and can act as an affordable and informative pseudo-MRI, before conducting an actual costly MRI examination in hospitals. More detailed results are provided in supplementary material.

Preservation of Spine Curvature.
For the spine segmentation on UK Biobank, we use a UNet++ model [] with Dice Loss. We use a model trained to predict curves on DXA on UK Biobank []). We measure the Pearson correlation factor [] of spine curvature measured on the generated MRIs where the input is a single MRI coronal slice, a single sagittal slice, or from the paired DXA, against the curvature of reference real MRIs of the same samples. The correlation coefficients are 0.89 for the coronal MRIs, 0.88 for the sagittal MRIs, and 0.87 for the DXAs on the test set of 308 human-annotated angles. We can then bin the curvature of the spines under different scoliosis categories based on human-annotated angles: mild: $0.06<\kappa<0.12$, moderate: $0.12\leq\kappa<0.15$, and severe $\kappa\geq 0.15$. We show the results in Figure [8](#S5.F8). This illustrates that the generated MRIs preserve the spine curvature from normal to severe scoliosis cases. Additional details about spine curvature are provided in supplementary material.

Figure: Figure 8: Scoliosis Categories of Generated MRIs. We show spine curvature predicted *vs*. reference curvature and human annotated angles for scoliosis categories in section [5.2](#S5.SS2). The barplot indicates that our generated MRIs maintain almost the same distribution of scoliosis categories for then set of 308 patients annotated in the UK Biobank.
Refer to caption: /html/2404.19604/assets/x23.png

### 5.3 Out-of-Domain Generalisation

One way to test the generalization capability of the trained X-Diffusion is to test it on a completely different domain from an MRI dataset not seen during training.
We report the single-slice results on NYU fastMRI [, ], a knee MRI dataset, using the X-Diffusion trained on the BRATS brain MRIs. The results are shown in Figure [9](#S5.F9) and Table [2](#S6.T2). It shows how successfully X-Diffusion is able to generate knee MRIs from a single image, despite not seeing knees at all in training. To qualitatively assess how realistic our generated 3D volumes were (produced from a single slice), we gave 20 generated examples alongside their real MRI counterparts to an expert orthopaedic surgeon. He was then asked to identify the real example from a given pair. The surgeon identified with certainty only 10 real knee MRIs out of 17, while could not decide on the remaining 3 of the 20 MRI pairs. This further validates the generated out-of-domain MRIs.

Figure: Figure 9: Out-of-Domain Generations of X-Diffusion. We show an example of knee 3D MRI generation using X-Diffusion from the single input slice on the left. We show (top): different slices of the generated 3D MRI, (bottom): ground truth slices of the same sample as reference. Our X-Diffusion can generate high-fidelity 3D MRIs of knees, even though it is trained on BRATS brain MRI dataset, illustrating its potential as a foundation model for 3D MRI generation.
Refer to caption: /html/2404.19604/assets/images/Ablation/knee/knee_input.png

## 6 Analysis and Ablation Study

### 6.1 Volume Averaging

We study the effect of volume averaging at inference as detailed in Equation [4](#S3.E4). We note (from Figure [7](#S4.F7)) how the averaging volumes indeed increase the performance up to a certain point. The results 3D PSNR (dB) for the 31-slices X-Diffusion on $N=$ 1, 2, 3, 5, and 10 volumes are 35.48, 35.94, 36.17, 37.40, and 36.72 respectively. This is consistent with multi-view understanding literature [].

### 6.2 Why does X-Diffusion Work?

The Effect of Pretraining.
We hypothesize that the massive pretraining of our X-Diffusion based on Stable Diffusion weights [] played an important role. Another aspect is that the Zero-123 [] weights which are modified Stable Diffusion weights that understand viewpoints and fine-tuned on large 3D CAD dataset Objaverse [] can indeed be the reason why X-Diffusion generalizes well. The PSNR for 1-slice on BRATS dataset are (SD-pretraining): 21.52 dB, (Zero-123-pretraining): 23.13 dB, (no-pretraining): 17.14 dB. These results highlight the importance of pertaining to X-Diffusion.

Leveraging Context.
Since we train on a predominantly cancerous brain dataset, one question that might arise is whether X-Diffusion generated MRIs preserve tumour information when the given inputs do not intersect with any tumour. We perform experiments varying the input slice index used to generate the 3D brain MRIs and measure the performance for input slices with no intersection with the tumour (not a single pixel with tumor label in the input slice). We also measure performance when only input slices are selected from tumor range. The Dice Scores of the random slices, no-tumour, and only-tumour are 83.09, 79.23, and 83.68 respectively.
As can be seen here, the brain volumes generated from input slices with no tumour still preserve tumour information in reconstructed brain volumes despite a small drop in performance. This indicates that X-Diffusion is leveraging the context to preserve key information, such as tumor locations. This observation is consistent with how tumor segmentation models with global context [] perform better than local-based U-Nets. More details are provided in supplementary material.

**Table 2: Out-of-Domain Results. We evaluate 3D knee generation using 3D PSNR and mean SSIM on the test set of $n=109$ knees. We show a comparison between X-Diffusion and baselines that are all trained on brain MRIs from BRATS.**
|  | PSNR$\uparrow$ | SSIM$\uparrow$ |  |  |
| --- | --- | --- | --- | --- |
| Method |  | Axial | Coronal | Sagittal |
| ScoreMRI (multi) [] | 30.88 | 0.82 | 0.82 | 0.81 |
| TPDM (multi) [] | 33.76 | 0.87 | 0.86 | 0.87 |
| X-Diffusion (single) | 34.17 | 0.88 | 0.87 | 0.88 |
| X-Diffusion (multi) | 36.57 | 0.89 | 0.88 | 0.89 |

### 6.3 When does X-Diffusion Fail?

To see when and how X-Diffusion fails, we conducted an experiment on healthy brains (no tumour) using IXI dataset, by running an X-diffusion trained on BRATS brain tumor dataset. Our X-Diffusion achieved a PSNR of 35.86 dB on the IXI dataset despite being trained on the BRATS dataset. We then ran the tumour segmenter on the set of 582 healthy scans and corresponding generated MRIs. The segmenter predicted tumours in 9.9% of the real healthy brains and in 11.3% of the generated brain MRIs. Some of these tumor hallucination examples fron X-Diffusion generation are shown in Figure [10](#S6.F10).

Figure: Figure 10: Tumour Hallucination and Failure Cases in X-Diffusion Generation. We show two cases of failure (red arrow) of our model hallucinating tumour in healthy sample scans. These tumour hallucinations represent only 2% of the healthy sample test set. Also, we show a failure case for the out-of-domain knee generation with the reference ground truth MRI slice.
Refer to caption: /html/2404.19604/assets/x24.png

## 7 Conclusions and Future Work

X-Diffusion advances 3D MRI generation, achieving high precision with limited inputs, as confirmed by tests on BRATS and UK Biobank data. Future directions include extending its application to dynamic MRI types and exploring its utility in other domains like environmental sciences.

Limitations.
X-Diffusion occasionally exhibits minor artifacts in complex tissue interfaces, a known issue in generative models operating in input-sparse scenarios. An instance of this is discussed in Section [6.3](#S6.SS3) and supplementary material with additional examples.

Acknowledgement.
We thank Professor Andrew Zisserman for his insightful discussions. This work was supported by the Centre for Doctoral Training in Sustainable Approaches to Biomedical Science: Responsible and Reproducible Research (SABS: R3), University of Oxford (EP/S024093/1), and by the EPSRC Programme Grant Visual AI (EP/T025872/1). We are also grateful for the support from the Novartis-BDI Collaboration for AI in Medicine. Part of the support is also coming from KAUST Ibn Rushd Postdoc Fellowship program.

## Appendix A Detailed Setup

### A.1 Evaluation Metrics

To quantify the efficacy of X-Diffusion, we employed a suite of evaluation metrics, namely:

- •
Peak Signal-to-Noise Ratio (PSNR): Indicates the quality of the reconstructed MRI by assessing the fidelity of the generated MRI in relation to the original.
$PSNR(x,\hat{x})=10log_{10}(\frac{max(x)^{2}}{\frac{1}{n}\sum_{i,j,k}(x_{i,j,k}-\hat{x}_{i,j,k})^{2}})$
(5)
where x represents the ground truth volume, $\hat{x}$ is the predicted volume, and n is the total number of voxels in the ground truth volume.
- •
Structural Similarity Index (SSIM): Captures the perceived changes between the original and generated MRI images.
$SSIM(x,\hat{x})=\frac{(2\mu_{x}\mu_{\hat{x}}+C_{1})+(2\sigma_{x{\hat{x}}}+C_{2})}{(\mu_{x}^{2}+\mu_{\hat{x}}^{2}+C_{1})(\sigma_{x}^{2}+\sigma_{\hat{x}}^{2}+C_{2})}$
(6)
where $x$ denotes the ground truth slice, $\hat{x}$ is the predicted slice, $\mu{x}$ is the average of $x$, $\sigma_{x}^{2}$ is the variance of $x$, $\sigma_{x\hat{x}}$ is the covariance between $x$ and $\hat{x}$, $C_{1}$=$(k_{1}L)^{2}$, $C_{2}$=$(k_{2}L)^{2}$, L is the dynamic range of pixel values, and $k_{1}$=0.01 and $k_{2}$=0.03.
We measured the random PSNR on the whole test set for reference on the UKBiobank, BRATS, and knee fastMRI dataset. For the UKBiobank, two randomly sampled MRIs have a PSNR of 15.95 $\pm$ 0.36 dB. For BRATS, it is of 19.89 $\pm$ 1.59 dB, and for the knee fastMRI of 20.21 $\pm$ 2.58 dB.
On BRATS dataset only
- •
Dice Score: We use the average Dice score to evaluate the performance of our model at segmenting the brain tumours []: Dice Score = $\frac{2|Y\cap\hat{Y}|}{N(|Y|+|\hat{Y}|)}$, where $Y$ is the prediction, $\hat{Y}$ is the ground-truth label and $N$ the total number of slices in the set.
- •
Brain Volume:
We measure brain volume in $mm^{3}$ by counting the non-zero voxels in the volume multiplied by the volume in $mm^{3}$ of each voxel [].
$NonZeroVoxCount=\sum_{i}^{N}V(x_{i},y_{i},z_{i})>0$
(7)
$VoxVol(mm^{3})=v_{x}*v_{y}*v_{z}$
(8)
$BrainVol=NonZeroVoxCount*VoxVol$
On UK Biobank dataset only
- •
Ground-truth Correlation Index:
Pearson’s correlation coefficient $r$ measures the strength of a linear association between two variables. The formula in [9](#Pt0.A1.E9) returns a value between -1 and 1, where: 1 denotes a strong positive relationship; -1 denotes a strong negative relationship; and zero denotes no relationship [].
$r=\frac{\sum_{i=1}^{n}(x_{i}-\bar{x})(y_{i}-\bar{y})}{\sqrt{\sum_{i=1}^{n}(x_{i}-\bar{x})^{2}}\sqrt{\sum_{i=1}^{n}(y_{i}-\bar{y})^{2}}}$
(9)
- •
Spine Curvature
Let $\gamma(t)=(x(t),y(t))$ be the equation of a twice differentiable plane curve parametrized by $t\in[0,209]$. We measure the spine curvature $\kappa$ with the standard mathematical formula []:
$\kappa=(y^{{}^{\prime\prime}}x^{{}^{\prime}}-x^{{}^{\prime\prime}}y^{{}^{\prime}})/(x^{{}^{\prime}2}+y^{{}^{\prime}2})^{\frac{3}{2}}$.

### A.2 Implementation Details

We conducted our experiments on two primary datasets:
BRATS. 5,880 brain MRIs annotated by expert clinicians for three classes: Whole Tumour (WT), Tumour Core (TC), and Enhanced Tumour Core (ET) [, , ]. We split the 5,880 MRIs into train (n = 4704), validation (n = 588), and test (n = 588) sets.
UKBiobank. 48,384 full-body Dixon MRIs from more than 500,000 volunteers[]. We split the UKBiobank MRIs into train (n = 38 707), validation (n = 4839), and test (n = 4838) sets. These Dixon MRI patches could not be stitched seamlessly with our current pipeline. These artifacts appear on all scans of the UKBiobank that we stitch. Therefore, the X-Diffusion trained on this data will recreate these artifacts regardless of input. The same pattern is present on all samples in the dataset for a fixed depth, while different depth indices will have different fixed patterns.
We made sure there was a coherence split, such that each patient was in a unique set. We will publish the unique IDs used for train-validation-testing to confirm there is no leakage, nor retrieval of images.

We implement X-Diffusion based on the Stable Diffusion [] U-Net with additional controls and conditions. We detail some of the hyperparameters and design choices below.

For the first stage of autoencoder training,
the encoder downsamples the image $x\in R\textsuperscript{H\texttimes W\texttimes 3}$, where $H=W=256$ by a factor $8$ to allow the DPM to focus on the semantic features of the latent space in a computationally efficient manner.
KL regularization is added to mitigate high variance latent space. In the second stage, a DPM is trained on the learned lower-dimensional latent space. The configuration of the U-Net is as follows: 2 residual blocks, channels multiples: [ 1, 2, 4, 4 ], attention resolutions: [ 4, 2, 1 ], 8 heads, using spatial transformer with depth = 1. For the DDPM Latent Diffusion, we use a base learning rate of $1.0\textsuperscript{-06}$, timesteps $T=1000$, image size = 32, channels = 4, and hybrid conditioning (concatenation and cross attention). Sampling is performed with classifier free guidance (see Figure[I](#Pt0.A1.F1) for example of test time sampling).

We use image-conditioned stable diffusion $v2$ checkpoint from [Lambda Labs](https://huggingface.co/spaces/lambdalabs/stable-diffusion-image-variations). We follow the novel view synthesis training from [Zero-123](https://github.com/cvlab-columbia/zero123). X-Diffusion is trained on a single GPU a6000, 48GB of RAM for four days.

Figure: Figure I: Test Time Brain Generation at Different Sampling Steps. For the input slice (top left), we show the ground-truth slice (bottom) and corresponding brain slice generating at different sampling steps $t$ in the denoising diffusion process.
Refer to caption: /html/2404.19604/assets/images/Ablation/Iterations/Slide1.png

Figure: Figure II: Residual Error of Generated MRIs. For the input slice (left), we show a difference map(bottom) between generated MRI (top) and ground truth. Below the (bottom) row, we indicate the mean squared error between generated and ground-truth images. Brighter pixels indicate greater disparity.
Refer to caption: /html/2404.19604/assets/images/MultipleSlices/Brains/76_sample_input_76.png

Figure: Figure III: Generated MRIs with Segmentation Maps Overlaid. We show ground-truth segmentation maps(bottom) and generated MRI (top). Red is used for the non-enhancing and necrotic tumor core, green for the peritumoral edema, and blue for the enhancing tumor core. The 3D Dice Score for this example is 77.26.
Refer to caption: /html/2404.19604/assets/images/segmentation/Seg_Maps/Slide2.png

Figure: Figure IV: Probabilistic Output for Different Volumes Generated from Single Slice by X-Diffusion. For the same input slice (top left), we show 12 generated output slices ( at index $d=88$) using 12 different inputs Gaussian noise for X-Diffusion U-Net.
Refer to caption: /html/2404.19604/assets/images/Ablation/probabilistic/Slide1.png

Figure: Figure V: Extension of Figure 5 from Main Paper. We show input DXA 1st column, generated MRI 2nd column, ground-truth MRI 3rd column, and Difference Map 4th column with residual value $\epsilon$.
Refer to caption: /html/2404.19604/assets/x30.png

## Appendix B Additional Results

### B.1 Fat Validation

We ran further experiments to investigate whether the generated MRIs (see an example in Figure [V](#Pt0.A1.F5)) preserve fat information. We use an image-based regression network trained on the UKBiobank to estimate DXA metadata information from 2D compressed middle coronal and sagittal MRIs []. Pearson’s correlations using Equation [9](#Pt0.A1.E9) comparing reference values and generated values are reported in Figure [VIII](#Pt0.A2.F8) with most fields having high correlation $r>0.9$. We show that the generated MRIs preserve crucial internal information.

### B.2 Brain Volumes Preservation

The comparison of generated MRIs versus reference MRIs suggest a nearly perfect preservation of brain volume (in mm^3) with median volume of reference MRIs of $1.31e^{6}$ $mm^{3}$ versus generated MRIs $1.28e^{6}$ $mm^{3}$ (see an example of brain generation in Figure [II](#Pt0.A1.F2)).

### B.3 Tumour Information Preservation

On the test set with human ground-truth annotations ($n=333$), the brain volumes generated from single slice input preserve the volume of the different tumour components (paired t-test, $p-value<0.05$ for all 3 classes) (see Table [I](#Pt0.A2.T1)). The real MRI Dice scores are put for reference to our generated MRIs. X-Diffusion outperforms baselines TPDM [] and ScoreMRI [] in tumour preservation (see Table [I](#Pt0.A2.T1) and Figure [III](#Pt0.A1.F3)).
We ran experiments comparing the tumour segmentation Dice Score varying X-Diffusion configurations. The multi-slice input X-Diffusion achieves marginally better Dice Score than the single slice input model (83.47 $\rightarrow$ 83.09). We also ran experiments with slice input used for volume reconstruction intersecting or not with tumour. We observe on average a drop of 6% Dice Score (see Table [I](#Pt0.A2.T1)). Further away from the tumour the input slice for volume reconstruction is selected, we observe a linear decrease in tumour segmentation Dice Score with lowest value of 77.21 Dice Score (see Figure [VI](#Pt0.A2.F6)).

This shows how the generated MRIs indeed preserve the tumour information and can act as an affordable and informative pseudo-MRI, before conducting an actual costly MRI examination in hospitals. Given that our model has been trained on brain scans all with tumours, we expect to see hallucinations of tumours in healthy scans. We report two cases of failure of our model in Figure [VII](#Pt0.A2.F7). Hallucinations of tumours on healthy samples represent 2% of the test set.

**Table I: Dice Score for Brain Tumour Segmentation on Real *vs*. Reconstructed MRIs with our X-Diffusion. We show Dice Score for generated MRIs ($n=587$) by our X-Diffusion when input only intersection with tumour (only-tumour) and when input does not intersect with tumour (no-tumour) for single slice and multi-slice input (31 slices). Note how X-Diffusion predicts the correct 3D tumour locations even when the input 2D slice does not intersect the tumour in most cases (drop from 83.47 to 81.65 Dice Score).**
|  | Test Dice Score $\uparrow$ |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
| X-Diffusion Generated MRIs | ET | WT | TC | Average Dice | 3D PSNR(dB)$\uparrow$ |
| single slice | 75.48 | 89.24 | 84.57 | 83.09 | 35.81 |
| multi-slice | 75.82 | 89.56 | 85.04 | 83.47 | 36.13 |
| multi-slice (only-tumour) | 76.12 | 90.04 | 85.87 | 84.01 | 36.98 |
| multi-slice (no-tumour) | 70.14 | 84.29 | 81.65 | 78.69 | 33.24 |
| Real | 76.47 | 91.13 | 86.24 | 85.15 | N/A |

Figure: Figure VI: Dice Score versus Distance to Tumour. We show the decrease in Dice Score for slice selection at increasing distance from the centre ($0mm$) of the tumour up to $80mm$ (where slice index $\in[1,5]\cup[151,155]$, total number of slices is 155 per scan, $n=587$).
Refer to caption: /html/2404.19604/assets/x40.png

Figure: Figure VII: X-Diffusion Failure Cases. We show two cases of failure (red arrow) on BRATS generations.
Refer to caption: /html/2404.19604/assets/x41.png

### B.4 Preservation of Spine Curvature and Fat

For the spine segmentation on UK Biobank, we use a UNet++ model [] with Dice Loss. We use a model trained to predict curves on DXA on UK Biobank []. We show in Figure [IX](#Pt0.A2.F9) that generated MRIs preserve the spine curvature from normal to severe scoliosis cases. We also study the case when DXA is used to generate the MRIs and show in Figure [IX](#Pt0.A2.F9) how the correlation to real curvatures compares to the input MRI case. The curvatures of the MRI generated from the coronal plane match the DXA curvatures more than the curvatures generated from sagittal MRI. This is expected since the antero-posterior plane of DXA is equivalent to the coronal plane for MRIs. This also explains the greater Pearson’s correlation coefficient $r$ of the coronal MRI (0.89) and DXA-generated curvature (0.88) compared to sagittal-generated curvature (0.87) relative to the reference curvature on the coronal plane. We observe though that MRI generation using X-Diffusion from another plane than the conventional plane for scoliosis assessment is valid.

**Table II: Model Performance on Synthetic Cone Data. We report the test 3D PSNR on synthetic volume generation of our model X-Diffusion for varying input slice number in training.**
|  | Test 3D PSNR $\uparrow$ |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Input Slices | 1 slice | 2 slices | 3 | 5 | 10 | 31 | 60 | 120 |
| X-Diffusion | 22.30 | 23.50 | 24.63 | 25.77 | 26.79 | 25.55 | 24.44 | 24.24 |

Figure: Figure VIII: Evaluation of Fat Estimation for Reference versus Generated MRIs. We show correlation (left) and mean difference plot (right) for DxaArmsTotalMass (top) with $r$ < 0.95, and DxaTotalFreeMass (middle), and DxaTotalTissueFatPercentage (bottom) with $r$ > 0.95 from Table .
Refer to caption: /html/2404.19604/assets/x45.png

Figure: Figure IX: Curvature Preservation of Generated MRIs. We plot spine curvature measured on reconstructed MRIs where the input was either, (i) a single MRI coronal slice, (ii) a single sagittal slice, or (iii) from the paired DXA, against the curvature of reference real MRIs of the same samples. The correlation coefficients are 0.89 for the MRIs, 0.88 for the MRIs from sagittal plane generation, and 0.87 for the DXAs.
Refer to caption: /html/2404.19604/assets/x48.png

## Appendix C Additional Analysis

### C.1 Ablation Study

Repeated Input Single Slice in Multi-Slice Models.
We try to see whether the multi-slice models are better than single slice models by studying if we used repeated input single slice multiple times. The 3D PSNR results for multi-slice input with 1, 2, 3, 5, 10, 31, and 60 repeated slices are 23.1, 23.256, 23.638, 23.921, 24.379, 25.125, and 24.921 respectively.

The Effect of Pretraining.
We hypothesize that the massive pretraining of our X-Diffusion based on Stable Diffusion weights [] played an important role. Another aspect is that the Zero-123 [] weights which are modified Stable Diffusion weights that understand viewpoints and fine tuned on large 3D CAD dataset Objaverse [] can indeed be the reason why X-Diffusion generalizes well to out-of-domain dataset (see generalization to knee MRIs in Figure [XI](#Pt0.A3.F11)).

We show the results in the following Table [III](#Pt0.A3.T3).

**Table III: X-Diffusion with Pre-Training versus no Pre-Training. We show comparison of X-Diffusion with fine-tuning pre-trained Stable Diffusion weights versus no pre-training.**
|  | 3D PSNR$\uparrow$ |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Models | 1 slice | 2 slices | 3 slices | 5 slices | 10 slices | 31 slices | 60 slices |
| X-Diffusion (pre-training) | 23.13 | 25.25 | 29.43 | 31.25 | 33.27 | 35.48 | 33.18 |
| X-Diffusion (no-pretraining) | 21.52 | 23.42 | 25.16 | 27.06 | 29.32 | 27.86 | 27.43 |

Different Mechanisms for Multi-Volume Aggregation.
We used view-dependant volume averaging as described in the main paper in all of the main results in the work. We show probabilistic outputs in Figure [IV](#Pt0.A1.F4) for different brain volumes generated for a single slice. We show the results of varying the number of volumes in Figure [X](#Pt0.A3.F10). We see that as the number of volumes averaged increase, the performance increase up to a certain point before saturating (as noted in the multi-view literature []). We did try to use other ways to aggregate the view-dependant volumes (*e.g*. by max pooling the volumes) and show the results as follows in Table [VII](#Pt0.A3.T7).

Figure: Figure X: Effect of the Number of Volume Averaged. We plot the test PSNR *vs*. number of volume averaged for X-Diffusion and our baselines i.e. TPDM and ScoreMRI for 31-slice mode on BRATS dataset. The PSNR increases up to 5 volumes averaged before dropping and stabilising at 60 volumes averaged for X-Diffusion and baselines.
Refer to caption: /html/2404.19604/assets/x49.png

MRI Volumes Specificity.
One hypothesis that can justify why the X-Diffusion model works very well on MRIs is that MRI data is not ordinary volume data since it is obtained by actually running an inverse Fourier transform on different k-frequency components, which means that the 3D information is embedded in every slice of the MRI. Introducing this Fourier effect on our synthetic Cone volumes dataset by applying masks on the high frequencies and then inverse Fourier results in slight improvement of volume reconstruction of $+0.51$ PSNR (dB) higher than with no Fourier masking ($26.788\rightarrow 27.298$ dB). This indicates that the Fourier frequency effect is negligble and does not explain away the performance of X-Diffusion.

### C.2 Time and Memory Requirements

Lowering reconstruction speed is important for greater accessibility, MRI re-acquisition purposes, and to monitor surgery in the case of dynamic MRI. The number of model parameters should be kept low to enable implementation on machines with lower memory capacity. X-Diffusion is in par with other diffusion based baseline models, albeit higher in memory requirements than classical methods. However, X-Diffusion is the only 3D medical imaging diffusion model that shows the capacity to generalize beyond the training data, opening the potential for foundation models in 3D MRIs. We show the cost analysis in Table [V](#Pt0.A3.T5).

**Table IV: Effect of Volume Averaging on The Performance. We show best performing model on BRATS with number of volumes averaged from view-dependent 3D MRI generation. We see that the PSNR reaches a peak for 5 volumes averaged before stabilising at 10 volumes.**
|  | Test 3D PSNR $\uparrow$ |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models | 1 vol | 2 vol | 3 vol | 5 vol | 10 vol | 20 vol | 31 vol | 60 vol |
| ScoreMRI [] | 31.67 | 31.86 | 32.23 | 33.49 | 33.19 | 32.16 | 30.43 | 29.86 |
| TPDM [] | 32.57 | 32.76 | 33.25 | 34.68 | 33.43 | 32.93 | 32.85 | 31.75 |
| X-Diffusion (max-pool) | 35.48 | 35.48 | 35.52 | 35.48 | 35.31 | 35.46 | 35.19 | 35.33 |
| X-Diffusion (averaging) | 35.48 | 35.94 | 36.17 | 37.40 | 36.72 | 36.35 | 36.83 | 36.53 |

**Table V: Cost Analysis. We show compute cost and runtimes that are measured on a computer with a single GPU a6000, 48GB of RAM.**
| Models | #Params | Runtime (s) |
| --- | --- | --- |
| Score-MRI[] | 860M | 139.142 |
| TDPM[] | 1720M | 149.468 |
| X-Diffusion | 990M | 141.461 |

Figure: Figure XI: Out-of-Domain Knee Generations of X-Diffusion 1. We show two examples of knee 3D MRI generation using X-Diffusion from the single input slice on the left. We show (top): different slices of the generated 3D MRI, (bottom): ground truth slices of the same sample as reference. Mean PSNR for top example is of 36.84 dB and for bottom example of 35.17 dB.
Refer to caption: /html/2404.19604/assets/images/Ablation/knee2/Slide1.png

### C.3 Compressed Sensing Experiment

Some of the previous works on MRI reconstruction [, ] target the task of compressive sensing, where the goal is to increase the frequency resolution of the MRIs when the k-space is undersampled. While this is not the goal of X-Diffusion, we adapted X-Diffusion to this task and train X-Diffusion on the k-space of the MRIs. The performance for our model in the compressive sensing task for under-sampling factor $\alpha=2$ is $PSNR=35.17$ dB. Results are shown in Table [VI](#Pt0.A3.T6).

**Table VI: Compressive Sensing Experiment. We show test 3D PSNR for brenchmark models DiffusionMBIR[], TPDM[], and ScoreMRI[], and X-Diffusion for input downsampled by acceleration factor 2, 4, and 6.**
|  | Test 3D PSNR $\uparrow$ |  |  |  |
| --- | --- | --- | --- | --- |
| Acceleration Factor | 2 | 4 | 6 |  |
| X-Diffusion | 35.17 | 34.41 | 34.16 |  |
| DiffusionMBIR[] | 37.16 | 36.12 | 35.85 |  |
| TPDM[] | 36.48 | 35.52 | 35.18 |  |
| ScoreMRI[] | 34.18 | 33.88 | 33.57 |  |

### C.4 Multi-Slice Inputs

The multi- slice inputs are sampled from the same axis of rotation during training and testing. To reduce the memory requirement for running the pipeline, the reduction operation of the $K>1$ input slices $(x_{1},x_{2},...,x_{K})$ is similar to what is followed in TPDM [] in the conditioning volume, and it can be described as follows: $x=\frac{1}{K-1}\sum_{j=1}^{K}x_{j}\cdot x_{j+1}$.
The difference in performance between the simple dot product reduction and the learned reduction with additional MLP is shown in Table [VII](#Pt0.A3.T7). During training, the slices do not need to be consecutive. The diffusion model implicitly learns to handle the slice gap since it is trained on multiple slices with different gaps. For the evaluation of multi-slice benchmarks, fixed input slices are sampled uniformly from the test set and used for all the compared models.

**Table VII: Comparing Model Performance of Multi-Input Aggregation Procedure on Brain Data. We compare the MRI reconstruction for X-Diffusion model for varying aggregation procedure i.e. dot averaging and multi-layer-perceptron (MLP) reduction and for varying input slice numbers. We report the mean 3D test PSNR on BRATS brain dataset. The results show that our aggregation method with dot product averaging increases model performance by a margin compared to MLP reduction method for varying number of input slices.**
|  | Test 3D PSNR $\uparrow$ |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Models | 1 slice | 2 slices | 3 slices | 5 slices | 10 slices | 31 slices |
| X-Diffusion (Avg. Dot) | 23.1 | 25.2 | 29.43 | 31.25 | 33.27 | 35.48 |
| X-Diffusion (MLP) | 22.7 | 24.91 | 28.89 | 30.73 | 32.82 | 35.16 |

### C.5 Synthetic Volumes Generation

We applied our X-Diffusion model on a completely different volumetric data modality to see if the MRI volume generation is indeed a simple task for X-Diffusion (see Figure [XIII](#Pt0.A3.F13)). To do this, we trained X-Diffusion on a synthetic volumes dataset. We show an example of the generated volume used for training and the corresponding prediction in Figure [XII](#Pt0.A3.F12) and quantitative results in Table [VIII](#Pt0.A3.T8).

**Table VIII: Synthetic Cone Generation with varying input slices and number of volumes averaged.**
|  | Test 3D PSNR $\uparrow$ |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| X-Diffusion | 1 | 2 | 3 | 5 | 10 | 31 | 60 | 120 |
| Input Slices | 22.302 | 23.504 | 24.625 | 25.771 | 26.788 | 25.548 | 24.439 | 24.238 |
| Volumes Averaged | 26.788 | 26.836 | 26.944 | 26.967 | 26.872 | 26.841 | 26.784 | 26.763 |

Figure: Figure XII: Synthetic Cone Generation with X-Diffusion. Top left image is the input to the model. We show (top): the generated 3D cone and slices from the 3D volume, (bottom): the ground-truth 3D cone and corresponding slices from 3D volume.
Refer to caption: /html/2404.19604/assets/x50.png

Figure: Figure XIII: Synthetic Volume Dataset. We show some samples of our proposed Synthetic Volumes dataset. The dataset consists of cones with different sizes, orientations, and colours (constant and gradient colours).
Refer to caption: /html/2404.19604/assets/x59.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: It would be interesting to see how X-Diffusion works with other types of augmentations, like blurring, changing colorspace etc. Intuitively, this should boost ability to generalize OOD or at least improve algorithm compared to using only SO3.

Author’s do not discuss what are the benefits of using pretrained VAE from SD without any additional finetuning (i.e VAE is frozen) to specific MRI domain. Additional discussion would be great. 

Is it possible to integrate and check performance using existing text-to-3D models. Seems like utilizing directly 3D generative model provides more inductive bias, which should improve learning.
soundness: 3
strengths: The authors demonstrate superior performance in several benchmarking tasks, including brain tumor and full body MRIs. 

The authors claim that the proposed model is able to generalize to novel domains, not seen during the training. 

The problem, motivation and contributions are clearly stated. 

Thorough experiments were conducted, methodically supporting the claims in paper. The idea of conditioning MRI reconstruction on different cross-sectional viewpoints seems novel.
summary: The authors introduces a new model for generating detailed 3D-MRI volumes from sparsified spatial-domain inputs, resulting in accurate 2D-to-3D scans reconstruction (in contrast to existing methods that require full 3D scans). The main idea is to incorporate view-dependent cross-sections during training, which is different from the current approaches.
weaknesses: There are no measurements on how fast the X-diffusion is. It would be beneficial to include comparisons with previous approaches. E.g., Figure 3 shows that for good quality reconstruction, almost T=1000 steps are required. 

The limitations discussed in Sections 6 and 7 appear to be strong, potentially limiting the practical value of the work. 

The authors dismiss the study of whether it is possible to reduce the number T in diffusion without sacrificing the quality? And how would this reduction be specific to medical imaging domain? (can simply substitute ideas of speeding up from text-to-image models?)

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 4
questions: Clarification:

- I really appreciate the ablation for healthy patient inpainting, c.f. 6.4, but it would have been really important to know how many healthy slices were used for the inpainting experiment. How many tumors do we see if we only provide slices at the outer parts of the brain where tumors will not lie? 

- For Figure 3, please clarify how many input slices are utilized from the target image. Considering the leftmost image is considered as input, am I right to assume that single-slice conditioning was utilized? However, given the crisp reconstruction, I would be surprised if this is the case. Was there tumor information in the conditioning - or is all of it impainted?

- For Figure 6, which image(s) were used to condition X-Diffusion - was any of the tumor information of the presented slice present in the slice? If so, it is not surprising that the X-diffusion model is able to recover this 

- For Figure 9, Figure II/II Appendix (and more)  - how many slices were used to condition the model? 

Suggestions:

- Please add SSIM scores alongside to the PSNR in the reconstruction experiments to the main paper (from supplementary).
- Consider adding deep learning-based metrics for image reconstruction, such as medical LPIPS (https://docs.monai.io/en/stable/losses.html#perceptualloss) or LPIPS (https://github.com/richzhang/PerceptualSimilarity), as these often correlate better with human perception and are widely used in medical image reconstruction as well.
soundness: 3
strengths: Originality: While the proposed method is built upon recent advantages in computer vision, its combination of key components and area of application is novel. Especially the cross-sectional MRI synthesis approach, which enables the stacking of slices from arbitrary viewing directions, not just axial/coronal and sagittal planes, is of particular interest to the medical imaging community and constitutes - to the best of my knowledge - a unique contribution of the paper.

Quality: The paper is written in a very clear, concise, and comprehensive way. I believe the authors did a great job in setting their method into the context of related work, thoroughly and clearly describing it, and performing an interesting experimental section, including further experiments in the supplementary material. The ablations clarified and addressed questions that I previously had when reading the manuscript and addressed different important parameters of the pipeline, demonstrating the effectiveness of the key components, such as the stable diffusion encoder, etc. 

Clarity: The presented method, related work, and experiments are on point. Ablations serve a purpose, the mathematical foundation behind key concepts was explained quite well, and figures and tables provided a great way of deepening the understanding of the reader. The supplementary offered a lot of great insights and the possibility for reproduction of experiments. 

Significance: The method of synthesizing a 3D volume with a cross-sectional model from arbitrary viewing directions is very interesting to the medical imaging community. Also, X-Diffusion tackles a unique reconstruction scenario, which is very much unexplored.
summary: This paper introduces X-Diffusion, a novel MRI reconstruction algorithm that can generate 3D volumes by conditioning on single or few slice 2D MR input. In contrast to traditional MRI reconstruction, X-diffusion takes on a different challenge and tries to predict the entire 3D volume given the sparse (but uncorrupted/high-quality slice(s)). The authors conduct a variety of different experiments on multiple datasets and downstream tasks and test the meaningfulness of pathological findings in the reconstructed images. Given their experimental setting, they are able to challenge the state of the art.  In a variety of ablation studies, the authors assess the impact of individual key components, providing an explanation of the capability of the model.
weaknesses: 1) Clinical relevance of the presented method: 

While the clinicians in this paper argue that the model is clinically relevant - I have a conflicting opinion. While "scout" or "localizer" scans with a single or a few slices may be acquired and thus used for X-Diffusion, I believe it is dangerous to assume that these will substitute a real volumetric scan. While the authors argue that tumor information is preserved, I would rather say this information is correctly hallucinated. Especially small tumors and lesions in other neurodegenerative diseases, such as MS, often occur on a few or single slices. Given that single or few-slice acquisitions will not be able to capture this information- (at all) - I would argue only undersampled or lower-resolved images pose a clinical tradeoff between finding anomalies and reducing scan time. I feel deep learning-based reconstruction really introduces a well-motivated use case for super-resolution or reconstruction in these problem settings, especially when considering multiple (complementary contrasts), but - to me - single/few slices are far-fetched and thus X-Diffusion solves a rather unrealistic clinical setting? Consequently, I would dial this down in the manuscript and frame this more as an exploratory or proof of concept study, i.e., "confirmed the potential usefulness of the generated MRIs" feels a bit bold to me.

I would also be curious to see how the introduced method (key components would still have merit for these problem settings!) would perform on SR settings where we try to super-resolve anisotropic to isotropic scans. On another note, I acknowledge and really appreciate the study of k-space reconstruction. This poses a much more clinical use case, where X-diffusion is powerful but not on par (at least in scores) with SOTA (please correct me if I am wrong). 

2) Difficulty in assessing how much is owed to conditioning with the "right" slices and the capability of the model:

While in regular medical image reconstruction, images are generally downsampled (in k-space) (c.f. ScoreMRI) or images with lower resolution are upscaled (TPDM), X-Diffusion has access to high-quality 2D slices that constitute a relevant part of the GT. In other words - the conditioning slices are not altered (i.e., they are identical and part of the GT). This is very different from the traditional scenarios - and thus poses a unique challenge in assessing the quality of the reconstruction. While the authors compare against SOTA baselines, these baselines were not originally designed to work as X-Diffusion works. This should be highlighted and discussed more prominently in the manuscript. I understand it is hard to come up with appropriate baselines, given that X-Diffusion is the first to attempt a single 2D slice to 3D reconstruction. However, the challenges that come with evaluating this should be more prominently discussed. 

3) Downstream Evaluation: 

Evaluation of the brain volume: While the brain volume estimation may serve as a first good downstream analysis task, I believe it constitutes an oversimplification of the problem. It would have been more meaningful to perform brain white-/grey matter analysis using, e.g., freesurfer on ground truth and reconstructed volume and comparing the results of this. 

4) Please share the SSIM for all tables and experiments; these values are often more relevant to PSNR values for reconstruction. 

5) A1 and A3 are almost identical. Consider merging or rewriting them to have one cohesive section.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: - Please see weaknesses above.
- Please fix the hyperref links, many references in the Appendix do not point to the good figure, this makes the paper even more difficult to follow.
soundness: 2
strengths: - Both the approach and the application appear quite novel.
- Many experiments were performed.
summary: The authors propose a diffusion-based approach to reconstruct 3D MRI volumes from a single or multiple 2D slices. Their approach is conditioned by the target rotation and slice index and relies on heavy pre-training. It is applied to brain MRI (with and without tumours), whole-body MRI and knee MRI.
weaknesses: - The experiments are confusing. It is as if the authors tried all they could think of and let the readers do what they can with the results.
- A major point that is often unclear in the experiments/results is what slice(s) was/were used as input. If the authors show in Figure VI that is has a impact on the results, it is not specified in other sections, e.g. 5.1, 5.2, 5.3, 5.4. 
- The way the authors split the BraTS dataset is not clear. On the one hand, they state that the dataset includes 5,880 MRI scans from 1,470 patients (roughly four MRI sequences per patient) and that they split it into Train (n=4704), Validation (n=588), and Test (n=588) sets, without mentioning if it is done at the subject level or not, which could mean data leakage. On the other hand, they say that they just use the FLAIR sequence, so then the numbers do not add up. Please clarify.
- When difference maps are displayed there is no scale, on top of the colormap poorly chosen, so they are impossible to analyse.
- No estimate of variation is provided in the tables.
- The evaluation by medical experts, though valuable, confirms that images generated are realistic but not that they show what they are supposed to, i.e. an image could be realistic but not correspond to the reality of a patient.
- More philosophically speaking, I fail to understand how the authors see what they propose used in practice, which is what they state they aim for. In their scenario, patients would come to the imaging centre, have 2D slices acquired and pseudo 3D volumes reconstructed and then what? Would they go home and maybe be called to come back and have a normal exam later? If so in what case? This does not seem cost-efficient nor practical.
- The related work section should cover super-resolution approaches. Also the Full-Body MRI Analysis paragraph is very generic and pretty empty. See for instance Tunariu et al., British Journal of Radiology, 2020 on the use of whole-body MRI in clinical practice and Küstner et al., Radiology: Artificial Intelligence, 2020 for the automatic analysis methods related to these applications.
- Some references are too generic or do not correspond to the statement they are supposed to comfort, e.g. Tran et al., 2015; Sohl-Dickstein et al., 2015 or Kawar et al., 2022 in the second paragraph of the introduction.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: *	Since H = W = D in this study, what would happen when the data is non-isotropic, as is most common in MRI applications and protocols.
*	For the example provided in Figure 5, it is interesting to think about why the method should work in the first place? Given that this is an ill-posed problem, why should a model be able to predict a slice that is 20 slices away accurately?
*	The better test would be to pick a conditioning slice that does not have a tumor but from a patient with a tumor. Then, will the model correctly produce a tumor? That is the real test. I hypothesize that there is no reason a model should be able to predict that. A brief sentence or two is included in the paper but this analysis can be made more rigorous. 
*	Spine curvature: What is the distribution of spine curvature? One may think this is regression to the mean of healthy spines. showing a Bland-Altmann or regression pot where would be better.
soundness: 2
strengths: *	Multiple datasets are used for the evaluation
*	The diffusion based benchmarks are chosen well and represent some of the latest approaches in the field.
summary: The proposed manuscript explores the use of a conditional generative model which can produce a full 3D MRI volume conditioned only on 1 or MRI slices. The authors conduct experiments with three MRI datasets that span brain, whole-body, and out-of-distribution knee MRI scans. Overall, this study could be improved with a justification of why such an approach that has no test-time data consistency should be favored over conventional deep learning based variational reconstruction methods that are very common and quite robust.
weaknesses: *	Many of the approaches that are used for MRI reconstruction rely on the notion of data consistency which can ensure that there is consistency between the reconstructed imaging data and the extent of signals that are actually obtained. The proposed approach provides no sense of such confirmation.
*	It is mostly an impractical scenario where a single high-resolution slice will be obtained instead of obtaining multiple undersampled slices. So the overall premise of the work unfortunately seems more academic rather than something that would have a clinical impact.
*	In prior work, there is a very incomplete analysis of work in image reconstruction and super-resolution. It is unclear what “classical” means. based on the references, does classical mean anything without foundation models? This doesn’t seem accurate.
*	In figure 4, the acquisition is actually axial and study depicts the acquisition to be coronal. Since there is always a fully sampled dimension readout dimension in MRI, the reformatting of this 3D volume does not accurately convey the representational capacity of the model. Actual slice traversal needs to be in the *true* slice direction.
*	A better evaluation would be show a progression of PSNR values the further away a test slice is being computed from a conditioning slice.
*	All datasets had 155 slices. In the case of using 31 conditioning slices, it is not clear why one should use this approach instead of just undersampling the 3D data by 5x and then reconstructing with variational methods. These methods can easily provide accelerations of 8-10x with much higher metrics.
*	In figure 6, up until 10-13 slices, there are significant hallucinations in the structure and contrast of the grey and white matter. These results would cause radiologists quite a bit of anxiety! Like the previous point, at this point, why not just use a variational reconstruction method?
*	The results in Figure 6 and the massive impact in hallucinated details entirely reduces confidence in all 1-5 slice results. This further shows how poor a metric PSNR is.
*	Total brain volume is not the correct metric to assess fidelity. Cortical volumes and folding patters would be better. One would hypothesize that it is likely trivial to regress total brain volume from a single brain MRI slice anyway.
*	It is pretty shocking to me why the model should be able to work in the OOD setting to synthesize knee MRI scans? Can the authors please provide a rationale why this should work in the first place? The UKB dataset includes IDEAL fat-water sequences that are MUCH lower resolution than the 2D TSE knee MRI fastMRI sequences that have a proton density contrast. Why should the model know this distribution Section 6.3 attempts to answer this but simply claiming Stable Diffusion training and large non-domain specific datasets unfortunately does not seem convincing.

Meta-review: This manuscript proposes diffusion model based volumetric conditional generation (inpainting with extremely sparse inputs) intended for MRI (structural contrasts, presumably). They provide results on a variety of datasets, including UK Biobank (whole body MR), BRATS (brain tumor public challenge data), IXI (healthy human brain data), and the fastMRI knee data.

Overall, reviewer discussion is focused on experiments and clinical viability. There are only a few notes from reviewers about the method itself; it appears to be exactly conditional latent diffusion with an augmented training process. Reviewer opinion is mixed on viability; while reported results have high PSNR which is quite promising, several reviewers voice concerns that PSNR does not inform well upon important features (and their possible changes). Both reviewer `YwSq` and `F5BK` state as much, with `F5BK` suggesting LPIPS or other metrics, and `YwSq` highlighting hallucinated tissue/ventricle/tumor differences between volumes with different inputs. These are addressed in an additional Appendix (D, and subsections therein).

Further, reviewers also find the broader use case to be somewhat unrealistic. (paraphrasing: "when would you go in for only a partial scan?").

In summary, results appear quite astonishing in the applied domain, though metrics may hide large issues, the methodology itself is unremarkable given prior work, and the broader context is questionable by some. By scores, this is definitionally borderline for ICLR. I do not think ICLR is the correct venue to discuss the applied results as is necessary for understanding the authors' claim that "this work is a paradigm shift in MRI reconstruction"; the conference in general does not have the correct audience nor (relatedly) the depth of domain knowledge to make that assessment even if the paper is published in the proceedings. Without major methodological contribution there is less relevant contribution outside of those domain specific advancements, and thus I am choosing to recommend to reject this paper. I will defer to the SACs/PC Chair's final decision in either case.

Decision: Reject