White Matter Geometry-Guided Score-Based Diffusion Model for Tissue Microstructure Imputation in Tractography Imaging

TL;DR

This work addresses missing tissue microstructure data in diffusion MRI tractography caused by incomplete white matter parcellation. It introduces WMG-Diff, a geometry-guided conditional score-based diffusion model that uses white matter atlas-based geometric distances to guide the reverse diffusion denoising for imputing missing FA values in fiber clusters. On a large harmonized ABCD dataset ($n=9{,}342$), WMG-Diff outperforms standard imputation baselines (e.g., mean/median/zero, GAIN, MICE) and a baseline diffusion model in RMSE, and it yields competitive, near-full-data accuracy in a downstream sex-prediction task ($ACC$ approaching $0.779$). The approach demonstrates that incorporating anatomical geometry into diffusion-based imputation enhances microstructure recovery, with practical implications for large-scale tractography analyses and non-imaging phenotype prediction; the code is publicly available at the linked repository.

Abstract

Parcellation of white matter tractography provides anatomical features for disease prediction, anatomical tract segmentation, surgical brain mapping, and non-imaging phenotype classifications. However, parcellation does not always reach 100\% accuracy due to various factors, including inter-individual anatomical variability and the quality of neuroimaging scan data. The failure to identify parcels causes a problem of missing microstructure data values, which is especially challenging for downstream tasks that analyze large brain datasets. In this work, we propose a novel deep-learning model to impute tissue microstructure: the White Matter Geometry-guided Diffusion (WMG-Diff) model. Specifically, we first propose a deep score-based guided diffusion model to impute tissue microstructure for diffusion magnetic resonance imaging (dMRI) tractography fiber clusters. Second, we propose a white matter atlas geometric relationship-guided denoising function to guide the reverse denoising process at the subject-specific level. Third, we train and evaluate our model on a large dataset with 9342 subjects. Comprehensive experiments for tissue microstructure imputation and a downstream non-imaging phenotype prediction task demonstrate that our proposed WMG-Diff outperforms the compared state-of-the-art methods in both error and accuracy metrics. Our code will be available at: https://github.com/SlicerDMRI/WMG-Diff.

Figures (3)

• Figure 1: Tractography parcellation into fiber clusters (top row) allows the calculation of FA values, represented in tabular format (bottom right), with missing values in yellow. The proposed WMG-Diff pipeline (bottom row) imputes missing FA data for input to a downstream sex prediction testbed task. Sub: Subject, cls: Cluster.
• Figure 2: An illustration of the pairwise distances (white arrows) between example fiber clusters (left) and the resulting matrix of distances between all of the 953 white matter fiber clusters (right). D: Distance.
• Figure 3: Overview of the architecture of our WMG-Diff guided process. In the input (upper left), actual missing data (white) and non-missing observable data (blue) are shown. Then, the non-missing data is split into observable values (green) and imputation targets (red) using information about cluster geometric relationships. At timestep $t$, the imputation target undergoes the forward process to add noise until it becomes a fully noisy target. The reverse diffuse process leverages the guided observable data at each timestep to denoise the noisy target and minimize the difference between the noise and the output.