Invertible Modeling of Bidirectional Relationships in Neuroimaging With Normalizing Flows: Application to Brain Aging

Many machine learning tasks in neuroimaging aim at modeling complex relationships between a brain's morphology as seen in structural MR images and clinical scores and variables of interest. A frequently modeled process is healthy brain aging for which many image-based brain age estimation or age-conditioned brain morphology template generation approaches exist. While age estimation is a regression task, template generation is related to generative modeling. Both tasks can be seen as inverse directions of the same relationship between brain morphology and age. However, this view is rarely exploited and most existing approaches train separate models for each direction. In this paper, we propose a novel bidirectional approach that unifies score regression and generative morphology modeling and we use it to build a bidirectional brain aging model. We achieve this by defining an invertible normalizing flow architecture that learns a probability distribution of 3D brain morphology conditioned on age. The use of full 3D brain data is achieved by deriving a manifold-constrained formulation that models morphology variations within a low-dimensional subspace of diffeomorphic transformations. This modeling idea is evaluated on a database of MR scans of more than 5000 subjects. The evaluation results show that our bidirectional brain aging model (1) accurately estimates brain age, (2) is able to visually explain its decisions through attribution maps and counterfactuals, (3) generates realistic age-specific brain morphology templates, (4) supports the analysis of morphological variations, and (5) can be utilized for subject-specific brain aging simulation.