A Stage-Aware Mixture of Experts Framework for Neurodegenerative Disease Progression Modelling

Abstract

The long-term progression of neurodegenerative diseases is commonly conceptualized as a spatiotemporal diffusion process that consists of a graph diffusion process across the structural brain connectome and a localized reaction process within brain regions. However, modeling this progression remains challenging due to 1) the scarcity of longitudinal data obtained through irregular and infrequent subject visits and 2) the complex interplay of pathological mechanisms across brain regions and disease stages, where traditional models assume fixed mechanisms throughout disease progression. To address these limitations, we propose a novel stage-aware Mixture of Experts (MoE) framework that explicitly models how different contributing mechanisms dominate at different disease stages through time-dependent expert weighting. This architecture is a key innovation designed to maximize the utility of small datasets and provide interpretable insights into disease etiology. Data-wise, we utilize an iterative dual optimization method to properly estimate the temporal position of individual observations, constructing a cohort-level progression trajectory from irregular snapshots. Model-wise, we enhance the spatial component with an inhomogeneous graph neural diffusion model (IGND) that allows diffusivity to vary based on node states and time, providing more flexible representations of brain networks. We also introduce a localized neural reaction module to capture complex dynamics beyond standard processes.The resulting IGND-MoE model dynamically integrates these components across temporal states, offering a principled way to understand how stage-specific pathological mechanisms contribute to progression. When used to model tau pathology propagation in human brains, IGND-MoE outperforms purely pathophysiological and purely neural baselines in long-term prediction accuracy. Moreover, its stage-wise weights yield novel clinical insights that align with literature, suggesting that graph-related processes are more influential at early stages, while other unknown physical processes become dominant later on. Our findings highlight the necessity of designing hybrid and expert-constrained models that account for the evolving nature of neurodegenerative processes.