QID$^2$: An Image-Conditioned Diffusion Model for Q-space Up-sampling of DWI Data
Zijian Chen Jueqi Wang Archana Venkataraman




Diffusion models for Q-space Up-sampling of DWI





Z. Chen et al.



Department of Electrical and Computer Engineering, Boston University 
 {zijianc,jueqiw,archanav}@bu.edu

: An Image-Conditioned Diffusion Model for Q-space Up-sampling of DWI Data
    Zijian ChenEqual Contribution, Jueqi Wang^⋆ and Archana Venkataraman
    
==========================================================================



§ ABSTRACT 
We propose an image-conditioned diffusion model to estimate high angular resolution diffusion weighted imaging (DWI) from a low angular resolution acquisition. Our model, which we call , takes as input a set of low angular resolution DWI data and uses this information to estimate the DWI data associated with a target gradient direction. We leverage a U-Net architecture with cross-attention to preserve the positional information of the reference images, further guiding the target image generation. We train and evaluate on single-shell DWI samples curated from the Human Connectome Project (HCP) dataset. Specifically, we sub-sample the HCP gradient directions to produce low angular resolution DWI data and train to reconstruct the missing high angular resolution samples. We compare with two state-of-the-art GAN models. Our results demonstrate that not only achieves higher-quality generated images, but it consistently outperforms the GAN models in downstream tensor estimation across multiple metrics. Taken together, this study highlights the potential of diffusion models, and in particular, for q-space up-sampling, thus offering a promising toolkit for clinical and research applications.








§ INTRODUCTION




Diffusion weighted imaging (DWI) is a non-invasive technique that capitalizes on the directional diffusivity of water to probe the tissue microstructure of the brain <cit.>. A typical DWI acquisition applies multiple magnetic gradients, with the field strength controlled by the b-value and the gradient directions given by the b-vectors. Mathematically, these gradients can be represented by a set of coordinates on the sphere, where the magnitude and direction of each coordinate is related to the corresponding b-value and b-vector, respectively. The domain of all such coordinates is called the q-space <cit.>. In general, a denser sampling of directions in the q-space, also known as the angular resolution, leads to higher quality DWI. For example, higher angular resolution acquisitions can improve the tensor estimation <cit.> and facilitates the progression from single-tensor models <cit.> to constrained spherical deconvolution models <cit.> that estimate an orientation distribution function, which captures more complex fiber configurations. However, increasing the angular resolution also prolongs the acquisition time, which can be impractical in clinical settings. Not only are longer acquisitions more expensive, but they are also difficult for some patients to tolerate, which in turn increases the risk of artifacts due to subject motion <cit.>. Given these challenges, it is necessary to explore computational methods that can achieve high-quality DWI with a minimal number of initial scan directions.


Several studies have applied generative deep learning to DWI data. For example, the work of <cit.> uses a spherical U-Net to directly estimate the ODF using DWI acquired with only 60 gradient directions. More recently, generative adversarial networks (GANs) have also been used to estimate DWI volumes. Specifically, the work of <cit.> generates DWI for a user-specified gradient direction based on a combination of T1 and T2 images <cit.>. Similarly, the authors of <cit.> use the Pix2Pix model introduced by <cit.> to synthesize DWI with 6 gradient directions from data originally captured with only 3 gradient directions <cit.>. Further variants of the GAN, such as CycleGAN and DC^2Anet, have been applied to simulate a high b-value image from a low b-value one <cit.>. Autoencoders have also been used to adjust the apparent b-value <cit.>. While these works are seminal contributions to the field, none of them consider the clinically relevant problem of up-sampling a low angular resolution DWI acquisition.

Diffusion models (DMs) have emerged as powerful tool for image generation. At a high level, DMs work by successively adding Gaussian noise to the input and then learning to reverse this noising process <cit.>. DMs have been employed in several medical imaging tasks <cit.>, including image translation between modalities <cit.>, super-resolution and artifact removal <cit.>, registration <cit.>, and segmentation <cit.>. We will leverage DMs to up-sample the DWI gradient directions, as task which to our knowledge, has not been explored in prior work. 

In this paper, we propose an image-conditioned DM, which we call , that can estimate high angular resolution DWI from a low angular resolution acquisition. One highlight is that automatically identifies several closest available directions and uses the corresponding images as prior knowledge for generating images from any target direction not included in the initial scan. This target image generation process, carried out using a U-Net based structure conditioned on this prior information, can be seen as an extrapolation based on the identified directions and images. By focusing on the most relevant data, solicits more targeted prior information and is more computationally efficient. We train and evaluate on DWI curated from the Human Connectome Project (HCP) dataset <cit.>. Our model demonstrates superior performance over GAN-based approaches, particularly when the available low angular resolution images are sparsely distributed across the sphere.



§ METHODS


Fig. <ref> provides an overview of our framework. For any user-specified target gradient 𝐛_g, our model will find and take as input the R closest reference b-vectors 𝐛̅ = (𝐛_1, …, 𝐛_R) available in the low angular resolution scan and the corresponding DWI slices 𝐗̅ = (𝐗_1, …, 𝐗_R). will then output the estimated target image 𝐗_𝐛_g. We can obtain a high angular resolution DWI by sweeping the target gradient directions across the sphere and aggregating the generated images with the original low angular resolution scan. 





 §.§ A Diffusion Model for Q-space Up-sampling of DWI



Inspired by recently-introduced image-conditioned Denoising Diffusion Probabilistic Models (DDPMs) <cit.>, we design a position-aware diffusion model that leverages “neighboring" DWI data to estimate the image associated with a target gradient direction. Similar to traditional diffusion models <cit.>, is comprised of both a forward noising process and a reverse denoising process. 

In the forward process, Gaussian noises are added successively at each time step t∈{0,1,…, T} to the generated image. This corruption process is

    q (𝐗_𝐛_g^(t)|𝐗_𝐛_g^(t-1)) = 𝒩(𝐗_𝐛_g^(t) ;  √(1-β_t)𝐗_𝐛_g^(t-1)  ,  β_t𝐈),    t≥ 1,

where {β_t} are the forward process variances, and 𝐗_𝐛_g^(t) is the noisy image at time t. By repeatedly applying Eq. (<ref>) to the starting image 𝐗_𝐛_g^(0), we have


    q (𝐗_𝐛_g^(t)|𝐗_𝐛_g^(0)) = 𝒩(𝐗_𝐛_g^(t)   √(α_t)𝐗_𝐛_g^(0)  ,   (1-α_t)𝐈)

where α_t=1-β_t and α_t=∏_s=1^t  α_s. Therefore, at step t, the generated image 𝐗_𝐛_g^(t) can be represented as a function of the initialization 𝐗_𝐛_g^(0):

    𝐗_𝐛_g^(t) = √(α_t) 𝐗_𝐛_g^(0) + √(1-α_t) ϵ,    ϵ∼𝒩(0,𝐈).


While the forward noising process operates solely on 𝐗_𝐛_g^(0), the reference DWI slices {𝐗_1, …, 𝐗_R } will be used to guide the subsequent denoising process. Rather than constructing a separate network to encode the reference images, which greatly increases the number of parameters and may introduce information loss, we opt to simply concatenate these slices with the target image being generated (i.e., denoised) at each time t as 𝐗̅_𝐛_g^(t) = 𝐗_𝐛_g^(t)⊕_i=1^R 𝐗_𝐛_i.


Starting from the fully corrupted image 𝐗̅_𝐛_g^(T), the reverse process aims to gradually recover the original image 𝐗̅_𝐛_g^(0). We denote this process as p_θ(·), where θ denotes the learnable parameters of the underlying neural network. By restricting the denoising to be Gaussian, the process p_θ(·) can be written:

    p_θ(𝐗̅_𝐛_g^(t-1)|𝐗̅_𝐛_g^(t)  ;  {𝐛_g,𝐛̅}) = 𝒩(𝐗̅_𝐛_g^(t-1)  ;  μ_θ(𝐗̅_𝐛_g^(t),{𝐛_g,𝐛̅})   ,  σ^2_t 𝐈),

where the variances σ_t^2 are hyperparameters of the model. We note that the denoising process relies on the references DWI data {𝐗_1, …, 𝐗_R } and the corresponding gradient directions 𝐛̅ = {𝐛_1,…,𝐛_R}, and the target direction 𝐛_g. This combination of inputs allows to be position-aware. 

To reverse the forward noising process, we train by minimizing the KL-divergence between p_θ(·) and q(·) at each time step t. As shown in <cit.> this loss minimization is equivalent to matching the mean functions, i.e.,

    ℒ = 𝔼_t,q[ 1/√(α_t)( 𝐗̅_𝐛_g^(t) - β_t/√(1-α̅_t)ϵ) - μ_θ(𝐗̅_𝐛_g^(t),{𝐛_g,𝐛̅}) ^2 ] .


The mean function μ_θ(·) is generated with a U-Net architecture <cit.> with the cross-attention mechanism based on the concatenated gradient vectors 𝐛=[ 𝐛_g 𝐛_1   ⋯ 𝐛_R ]. Specifically, the encoding block is computed as follows:

    𝐇_1 = FF(𝐇_0)+𝐇_0,   𝐇_2 = Attn(𝐇_1, 𝐛) + 𝐇_1,

where FF(·) denotes a feed-forward network, 𝐇_0 denotes the block input, and 

    Attn(𝐇_1, 𝐛) = Softmax((W_Q 𝐇_1)(W_K 𝐛)^⊤/√(d_k))W_V𝐛,

with W_Q, W_K, W_V being the learned weights and d_k being the dimension of 𝐛. The decoding block follows a similar expression but includes skip connections from the corresponding encoding block. This design ensures that image features are effectively attended to and integrated with positional information.

Once is trained, we can generate DWI for arbitrary gradient directions by sampling from the standard normal distribution and applying the reverse process in Eq. (<ref>) recursively with the corresponding reference images, namely:

    𝐗̅_𝐛_g^(t-1) = μ_θ(𝐗̅_𝐛_g^(t),{𝐛_g,𝐛̅}) + σ_t  ϵ,     ϵ∼𝒩(0,𝐈).




 §.§ Baseline Comparison Methods

We compare with two state-of-the-art GAN models. 
The first model is a conditional GAN (cGAN) for image generation proposed by <cit.>. We use the same cross-attention U-Net architecture for the generator as used in . We use a PatchGAN discriminator <cit.> and inject the gradient direction information {𝐛_g,𝐛̅} into the discriminator with cross-attention mechanism. We train the generator to minimize the minimax GAN objective plus a regularization term that encourages voxel-level similarity of the generated and ground-truth images:

    G^* = min_Gmax_Dλ_G  [𝔼_𝐗[log D(𝐗,𝐛)] + 𝔼_𝐗[log (1 - D(𝐗, 𝐛))] ] + λ_V ℒ_1(G)

where 𝐗 = [𝐗_𝐛_g, 𝐗_1,…,𝐗_R] is the concatenated real sample with 𝐗_𝐛_g drawn from the (high resolution) training data and 𝐗 = [G(𝐗_1:R, 𝐛), 𝐗_1,…,𝐗_R] represents the synthesized data of generated DWI and real reference slices. Finally, λ_G and λ_V balance the adversarial and similarity L_1 losses, respectively.

The second model is the Q-space conditional GAN (qGAN) proposed by <cit.>. Unlike and the cGAN baseline, qGAN incorporates the gradient directions and reference DWI data using a feature-wise linear modulation scheme <cit.>. The qGAN discriminator is also a conditional U-Net <cit.> and combines the gradient directions and reference DWI data via an inner product. 


Finally, as a sanity check, we compare the deep learning models to a simple interpolation scheme (Interp), in which we express the target gradient direction as a linear combination of the reference gradients and then use the linear coefficients to interpolate between the reference DWI slices to obtain the target.



 §.§ Implementation Details


For , we use a linear noise schedule of 1000 time steps. The U-Net employs [128, 128, 256] channels across three levels with one residual block per level. We use the Adam optimizer with a learning rate of 2.5 × 10 ^-5, β_1=0.5 and β_2=0.999. These hyperparameters are selected based on a relevant study <cit.> and not fine-tuned. We use the same U-Net architecture for the cGAN generator with the same set of hyperparameters. We fix λ_G = 1 and λ_V=100 in the loss function weights for both GAN methods. The discriminator is updated once for every two updates of the generator during training <cit.>. For qGAN, we use a learning rate of 5 × 10^-5 with the Adam optimizer. For comparison, we evaluate all models with both R=3 and R=6 reference DWI data.

To avoid memory issues <cit.>, we train the deep learning models to generate 2D axial slices, which we stack into 3D DWI volumes. Each 2D image has a size of (145, 174). During training, we independently normalize each slice from its original intensity to a range of [0,1]. Data augmentation is employed to enhance model training. Specifically, we use rotations by random angles in [- 15^∘, 15^∘] and random spatial scaling factors in [0.9, 1.1]. The final output is rescaled voxel-wise back to the original intensity and masked by the subject 𝐗_𝐛_g image.



§ EXPERIMENTAL RESULTS
 


Dataset Curation and Preprocessing:
We curate a total of 720 subjects from the HCP S1200 release <cit.>. The remaining HCP subjects are excluded due to an inconsistent number of gradient directions at b=1000  s/mm^2. The DWI is acquired on a Siemens 3T Connectome scanner at 3 shells (b=1000,2000 and 3000 s/mm^2). Each shell has exactly 90 gradient directions sampled uniformly on the sphere. The voxel size is 1.25 × 1.25 × 1.25 mm^3. The data is preprocessed with distortion/motion removal and registration to the 1.25mm structural space.

Clinical diffusion imaging typically uses lower b-values with approximately 30 gradient directions <cit.>. To better accommodate this situation, we focus our evaluation on the b=1000 s/mm^2 shell. From here, we construct low angular resolution DWI data by subsampling the 90 gradient directions to 30 evenly spaced ones that preserve the uniformity of the sphere <cit.>. The data for the remaining 60 directions serve as the targets for model training and evaluation. 




Each volume is broken down into 145 axial slices. The deep learning models are trained to predict the image slices for each target direction. Based on this scheme, we create 60 samples for each slice. Each sample consists of one 2D slice for the target gradient direction and R reference slices corresponding to the closest low resolution gradient directions. The distance between gradients is defined by the geodesic distance on the sphere: d(𝐛_1,𝐛_2) = arccos(𝐛_1𝐛_2^⊤).  


Finally, we use 576 HCP subjects for training, 72 for validation, and 30 for testing. The original DWI scans are treated as the gold standard for evaluation.

Comparing Reconstructed Image Quality:
Fig. <ref> presents qualitative results that compare the ground-truth DWIs to those generated by and the baseline methods. The GAN models and Interp fail to preserve high-frequency details in the synthesized DWI data, while succeeds in capturing the finer details more accurately, as highlighted in the zoomed-in blue boxes.

Table <ref> (left) reports the Fréchet inception distance (FID) <cit.> and the structural similarity index measure (SSIM) <cit.> of the synthesized DWI data. We observe that achieves nearly a two-fold improvement (i.e., decrease) in FID than the GAN models for both R=3 and R=6, which indicates that the DWI data generated by diffusion possess higher quality and greater diversity. Although the GAN models achieves a slightly higher SSIM than , the difference is not statistically significant using a two-sample (paired) t-test. Interestingly, the simple interpolation technique achieves better FID than when R=3. This is likely because the interpolation tracks the closest reference image, which is more akin to the original DWI distribution.
However, the improved FID does not generalize to better tensor estimation, as seen in the next section.







Impact on Tensor Estimation:

We estimate the fractional anisotropy (FA) using the standard single-tensor model <cit.>. Fig. <ref> shows the fiber direction and FA value maps among the ground-truth, and baseline methods for R=3 and R=6. Similarly to the finding in the reconstructed image, we observe that the qGAN and cGAN methods capture the general FA trends but fail to capture the high frequency features. Conversely, the diffusion-generated image by more closely resembles the ground-truth data by capturing finer details more accurately. This shows that the visual differences in the reconstructed images in Fig. <ref> are important when estimating tensors. The Interp method fails to generate realistic FA maps for R=3. Empirically, we also observe quality issues with Interp for R=6 even though they are less evident in the figure.

Table <ref> (right) reports the mean absolute error and SSIM, as compared to the FA computed from the ground-truth high angular resolution DWI. As seen, consistently outperforms the GAN-based model and the Interp method for both R=3 and R=6. Specifically for R=3, the error in FA is roughly three times lower for than for the GANs. also achieves significantly higher SSIM values. These trends persist when the number of reference images increases to R=6, i.e., even when more prior information is provided. However, the relative performance gain over the GANs shrink. Additionally, although the image-based metrics are better for the interpolation-generated (Interp) images, outperforms this baseline by a large margin when estimating FA. Taken together, these results suggest that is particularly effective in scenarios where the images are scarce and distributed sparsely, i.e., smaller values of R.






§ CONCLUSION

We introduce an image-conditioned diffusion model () that can generate high angular resolution DWI from low angular resolution data, effectively estimating high-quality imaging with limited initial scan directions. Our approach takes advantage of similar DWI data as prior information to predict the data for any user-specified gradient direction. The results demonstrate that diffusion-generated DWIs by achieve superior quality and significantly outperform those generated by GAN models in downstream tensor modeling tasks. Although our method currently exhibits longer training times due to the denoising characteristics of DDPMs, this limitation could be mitigated by employing more efficient sampling techniques <cit.> in future work. We will also focus on on clinical datasets to evaluate its robustness in real-world clinical scenarios. 



 §.§ Acknowledgments


This work was supported by the National Institutes of Health R01 HD108790 (PI Venkataraman), the National Institutes of Health R01 EB029977 (PI Caffo), the National Institutes of Health R21 CA263804 (PI Venkataraman).








splncs04



