Normative Diffusion Autoencoders: Application to Amyotrophic Lateral Sclerosis

TL;DR

The paper tackles ALS survival prediction under data scarcity by uniting normative modelling with diffusion autoencoders through a hierarchical diffusion autoencoder conditioned on age. It learns a healthy brain semantic latent, computes a cosine-based similarity to ALS patient latents, and uses this metric in Cox proportional hazards modeling to predict survival. The approach outperforms both generative and non-generative normative benchmarks, with a unit increase in cosine similarity linked to a $HR=0.73$ hazard reduction and clear separation in survival via KS tests. This yields an interpretable, powerful biomarker for ALS prognostication and outlines paths toward fully 3D MRI integration and unified training-survival analysis.

Abstract

Predicting survival in Amyotrophic Lateral Sclerosis (ALS) is a challenging task. Magnetic resonance imaging (MRI) data provide in vivo insight into brain health, but the low prevalence of the condition and resultant data scarcity limit training set sizes for prediction models. Survival models are further hindered by the subtle and often highly localised profile of ALS-related neurodegeneration. Normative models present a solution as they increase statistical power by leveraging large healthy cohorts. Separately, diffusion models excel in capturing the semantics embedded within images including subtle signs of accelerated brain ageing, which may help predict survival in ALS. Here, we combine the benefits of generative and normative modelling by introducing the normative diffusion autoencoder framework. To our knowledge, this is the first use of normative modelling within a diffusion autoencoder, as well as the first application of normative modelling to ALS. Our approach outperforms generative and non-generative normative modelling benchmarks in ALS prognostication, demonstrating enhanced predictive accuracy in the context of ALS survival prediction and normative modelling in general.

Figures (7)

• Figure 1: The normative diffusion autoencoder. We first train our hierarchical diffusion autoencoder, $p_{\theta}(\mathbf{x}_{0:T}| \mathbf{z})$ on a healthy cohort. Then we use $s_{\phi}(\mathbf{x}_{0}, \alpha)$ to compute the latent representations, of all MRIs from our healthy cohort (N=4120), conditional on their ages $\alpha$. Next, we average these representations to produce $\boldsymbol{\mu}_{HC}$. After which we compute the latent representation of MRIs from ALS patients to produce $\mathbf{z}_{ALS}$. We quantify the similarity to the healthy cohort by computing the cosine similarity between $\boldsymbol{\mu}_{HC}$ and $\mathbf{z}_{ALS}$. Finally, we fit a Cox model using this similarity metric to model its effect on survival time in days.
• Figure 2: The graphical model induced by our hierarchical diffusion autoencoder. The conditional variable $\alpha$ condition's the semantic latent $\mathbf{z}$ alongside the original image $\mathbf{x}_{0}$. The reverse process then iteratively denoises starting from our final stochastic latent $\mathbf{x}_{T}$ conditional on $\mathbf{z}$ to recover $\mathbf{x}_{0}$.
• Figure 3: Our dataset. On the left is a description of our ALS cohort, and on the right are the names of the 7 datasets that comprised our healthy cohort.
• Figure 4: Our pre-processing pipeline. All pre-processing was conducted in the Python programming language using the HD bet and Simple ITK packages.
• Figure 5: The architure of our noise predictor model, $\epsilon_{\theta}^{(t)}(\mathbf{x}_{t},\mathbf{z})$. Each set of bars are convolutional layers, and the numbers above denote the output channels.