Identity Preserving Diffusion Model for Brain Aging Modeling
Poster No:
1159
Submission Type:
Abstract Submission
Authors:
Gexin Huang1, Mengwei Ren2, Guido Gerig2, Xiaoxiao Li1
Institutions:
1The University of British Columbia, Vancouver, BC, 2New York University, Brooklyn, NY
First Author:
Co-Author(s):
Introduction:
Brain aging is a natural process that affects brain structure and function, leading to cognitive decline and increased risk of neurodegenerative diseases. Aging brain image generation is crucial for understanding aging, diagnosing disorders, predicting risk, and planning treatments. Deep-learning-based image-to-image generation methods could effectively study structural changes associated with aging. Keeping identity consistent in this process is crucial for understanding the underlying mechanisms and enabling personalized brain aging modeling. However, preserving identity consistency in longitudinal brain images has not been adequately addressed by previous works [1,2,3]. Thus, we propose a novel Longitudinal Transformation Diffusion model with an identity consistency module (LT-Diff) that achieves accurate brain age transformations while preserving identity using the OASIS Brain Dataset.
Methods:
Aging brain generation is achieved via an image-2-image model that generates the target image based on the source image and the target age. We utilized the OASIS-3 brain dataset [4], consisting of longitudinal MRI imaging data collected over 15 years from ongoing studies. It includes data from 1098 individuals, 493 of whom were at various stages of cognitive decline ranging in age from 42 to 95 years. The model takes the middle slice of each 3D MRI volume as its input. The dataset covers an age range from 40 to 100, divided into four groups: 40-55, 56-70, 71-85, and 86-100. Each MRI image is labeled with the corresponding age group. The stratified sampling strategy [5] is used for train-validation-test data split.
Model Architecture: The LT-Diff model (Fig. 1a) consists of three components: the Identity Preservation Module (Fig. 1b), the Age Encoder, and the DDIM Decoder (Fig. 1c) based on the DDIM backbone [4]. Inspired by the Diffusion Autoencoder [6], the Identity Preservation Module decouples identity information using a ResNet-18 encoder, preserving identity consistency. The age encoder utilizes a four-layer MLP to generate a shared latent age feature. The DDIM decoder reconstructs the difference map between the target and source image, conditioned on concatenated identity and age features. The predicted target image is synthesized by fusing the difference map with the source image.
Objective functions and Optimization: Two objectives are used for training (Fig. 1): Mean Square Error loss for difference map reconstruction and Identity Consistency Loss for preserving subject identity. The Identity Consistency Loss combines triplet loss, cosine similarity, and collapse regularization terms. The triplet loss maximizes inter-identity distance while minimizing intra-identity distance. The cosine similarity term promotes similarity within the same identity, and the collapse regularization term prevents identical features and avoids local minima during optimization [9].
Model Architecture: The LT-Diff model (Fig. 1a) consists of three components: the Identity Preservation Module (Fig. 1b), the Age Encoder, and the DDIM Decoder (Fig. 1c) based on the DDIM backbone [4]. Inspired by the Diffusion Autoencoder [6], the Identity Preservation Module decouples identity information using a ResNet-18 encoder, preserving identity consistency. The age encoder utilizes a four-layer MLP to generate a shared latent age feature. The DDIM decoder reconstructs the difference map between the target and source image, conditioned on concatenated identity and age features. The predicted target image is synthesized by fusing the difference map with the source image.
Objective functions and Optimization: Two objectives are used for training (Fig. 1): Mean Square Error loss for difference map reconstruction and Identity Consistency Loss for preserving subject identity. The Identity Consistency Loss combines triplet loss, cosine similarity, and collapse regularization terms. The triplet loss maximizes inter-identity distance while minimizing intra-identity distance. The cosine similarity term promotes similarity within the same identity, and the collapse regularization term prevents identical features and avoids local minima during optimization [9].
Results:
LT-Diff outperforms the state-of-the-art method (lifespan GAN) [9] with a better FID score of 23.59 and KID score of 11.59e-4 (Fig. 2a). The brain age transformation (Fig. 2b) shows ventricles growing larger with increasing brain aging while maintaining subtle geometric variations for each subject, consistent with the transformation of brain aging in prior studies [10,11]. The t-SNE visualization (Fig. 2c) confirms LT-Diff's accurate reconstruction of brain images while preserving identity consistency.
Conclusions:
We propose LT-Diff, a method that preserves identity consistency in longitudinal brain imaging by leveraging decoupled identity features. By incorporating these features and age information as a condition for the diffusion decoder, LT-Diff successfully generates brain images with the desired target age. Experimental results on the Oasis-3 dataset show that LT-Diff outperforms the SOTA architecture in brain aging generation performance w.r.t aging modeling and identity preservation.
Lifespan Development:
Aging 1
Lifespan Development Other
Modeling and Analysis Methods:
Methods Development 2
Keywords:
Aging
Machine Learning
STRUCTURAL MRI
Other - Diffusion Model, Identity Consistency
1|2Indicates the priority used for review
Provide references using author date format
[2] Kanizo, Y., Hay, D., & Keslassy, I. (2013, April). Palette: Distributing tables in software-defined networks. In 2013 Proceedings IEEE INFOCOM (pp. 545-549). IEEE.
[3] Sasaki, H., Willcocks, C. G., & Breckon, T. P. (2021). Unit-ddpm: Unpaired image translation with denoising diffusion probabilistic models. arXiv preprint arXiv:2104.05358.
[4] LaMontagne, P. J., Benzinger, T. L., Morris, J. C., Keefe, S., Hornbeck, R., Xiong, C., ... & Marcus, D. (2019). OASIS-3: longitudinal neuroimaging, clinical, and cognitive dataset for normal aging and Alzheimer's disease. MedRxiv, 2019-12.
[5] Tong, C. (2006). Refinement strategies for stratified sampling methods. Reliability Engineering & System Safety, 91(10-11), 1257-1265.
[6] Preechakul, K., Chatthee, N., Wizadwongsa, S., & Suwajanakorn, S. (2022). Diffusion autoencoders: Toward a meaningful and decodable representation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 10619-10629).
[7] Dey, N., Ren, M., Dalca, A. V., & Gerig, G. (2021). Generative adversarial registration for improved conditional deformable templates. In Proceedings of the IEEE/CVF international conference on computer vision (pp. 3929-3941).
[8] Song, J., Meng, C., & Ermon, S. (2020). Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502.
[9] Or-El, R., Sengupta, S., Fried, O., Shechtman, E., & Kemelmacher-Shlizerman, I. (2020). Lifespan age transformation synthesis. In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part VI 16 (pp. 739-755). Springer International Publishing.
[10] Peng, L., Wang, N., Xu, J., Zhu, X., & Li, X. (2022). GATE: graph CCA for temporal SElf-supervised learning for label-efficient fMRI analysis. IEEE Transactions on Medical Imaging, 42(2), 391-402.
[11] Blinkouskaya, Y., Caçoilo, A., Gollamudi, T., Jalalian, S., & Weickenmeier, J. (2021). Brain aging mechanisms with mechanical manifestations. Mechanisms of Aging and Development, 200, 111575. doi:10.1016/j.mad.2021.111575
[12] Blinkouskaya, Y., & Weickenmeier, J. (2021). Brain shape changes associated with cerebral atrophy in healthy aging and Alzheimer’s disease. Frontiers in Mechanical Engineering, 7, 705653.