Trustworthy Longitudinal Brain MRI Completion: A Deformation-Based Approach with KAN-Enhanced Diffusion Model

TL;DR

Longitudinal brain MRI analyses suffer from missing data due to participant attrition. The authors propose DF-DiffCom, a deformation-field–based diffusion framework powered by Kolmogorov-Arnold Networks (KAN) that warps a source image to target ages, producing modality-agnostic completions. Key contributions include deformation-field generation via conditional diffusion, DiffKAN as a high-capacity backbone, and the Flexible Temporal Information Enhancement (F-TIE) module for variable guidance, along with a dual loss incorporating deformation fidelity and age-consistency. Empirical results on the OASIS-3 dataset show state-of-the-art image quality (PSNR $= 25.52 \pm 1.21$, SSIM $= 0.84 \pm 0.03$) and robust ablations confirming the value of DiffKAN and F-TIE; qualitative analyses demonstrate realistic aging trajectories and cross-modality applicability. Overall, the method offers a trustworthy, flexible, and modality-agnostic solution for longitudinal brain MRI completion with potential to improve downstream clinical analyses.

Abstract

Longitudinal brain MRI is essential for lifespan study, yet high attrition rates often lead to missing data, complicating analysis. Deep generative models have been explored, but most rely solely on image intensity, leading to two key limitations: 1) the fidelity or trustworthiness of the generated brain images are limited, making downstream studies questionable; 2) the usage flexibility is restricted due to fixed guidance rooted in the model structure, restricting full ability to versatile application scenarios. To address these challenges, we introduce DF-DiffCom, a Kolmogorov-Arnold Networks (KAN)-enhanced diffusion model that smartly leverages deformation fields for trustworthy longitudinal brain image completion. Trained on OASIS-3, DF-DiffCom outperforms state-of-the-art methods, improving PSNR by 5.6% and SSIM by 0.12. More importantly, its modality-agnostic nature allows smooth extension to varied MRI modalities, even to attribute maps such as brain tissue segmentation results.

Figures (2)

• Figure 1: The schematic pipeline of our proposed DF-DiffCom. (a) During training, we use a source image, target age, and optional additional images to guide a KAN-enhanced diffusion model (DiffKAN) to predict a deformation field. A dual loss is computed on both the deformation field and the warped image. (b) During inference, the trained model generates a deformation field, which is applied to the source image via a Spatial Transformer Network (STN) to produce the completed image. (c) The Flexible Temporal Information Enhancement (F-TIE) module encodes a variable number of additional images into a fixed-length guidance vector. (d) The DiffKAN backbone integrates KAN blocks for superior non-linear feature extraction.
• Figure 2: (a) Our method generates a deformation field (DF) to warp a source image (age 60) and its segmentation map to a target age (70). The generated images show plausible aging effects like ventricular enlargement. (b) Examples of longitudinal generation with increasing age gaps, demonstrating the model's ability to capture progressive and anatomically consistent changes, such as the gradual expansion of ventricles over time.