AtlasMorph: Learning conditional deformable templates for brain MRI

TL;DR

AtlasMorph presents a neural framework to learn conditional, deformable brain MRI templates and corresponding segmentation maps conditioned on subject attributes. It combines a template-decoder and a registration network to produce on-demand templates $t=f_{ heta_t}(a)$ and velocity fields $v_i$ that warp subjects to the template, optimized via a probabilistic model with a centrality loss grounded in kernel density estimation. The method yields sharper, population-representative templates, improves registration accuracy, and captures known aging and disease-related trends, while enabling segmentation without post-hoc label propagation. It outperforms baseline unconditional and subpopulation templates (e.g., ANTs, Aladdin) and offers a scalable route to rapid, attribute-conditioned template construction across large neuroimaging cohorts.

Abstract

Deformable templates, or atlases, are images that represent a prototypical anatomy for a population, and are often enhanced with probabilistic anatomical label maps. They are commonly used in medical image analysis for population studies and computational anatomy tasks such as registration and segmentation. Because developing a template is a computationally expensive process, relatively few templates are available. As a result, analysis is often conducted with sub-optimal templates that are not truly representative of the study population, especially when there are large variations within this population. We propose a machine learning framework that uses convolutional registration neural networks to efficiently learn a function that outputs templates conditioned on subject-specific attributes, such as age and sex. We also leverage segmentations, when available, to produce anatomical segmentation maps for the resulting templates. The learned network can also be used to register subject images to the templates. We demonstrate our method on a compilation of 3D brain MRI datasets, and show that it can learn high-quality templates that are representative of populations. We find that annotated conditional templates enable better registration than their unlabeled unconditional counterparts, and outperform other templates construction methods.

Figures (23)

• Figure 1: Conditional Templates developed with AtlasMorph as a function of age.
• Figure 2: AtlasMorph Architecture. We employ two main networks, a decoder which creates templates (3D volumes and corresponding segmentation maps) given an attribute, and a registration network that aligns subjects to templates. We learn these jointly using a loss that is a function of the subject scan, the warped template and the deformation field. The decoder and the UNet contain learnable weights. We estimate both a voxel template and the corresponding label maps, shown with the shaded optional box.
• Figure 3: Statistics for subjects in the training set. Left: Distribution of subject age. Right: Distribution of subject sex.
• Figure 4: Template construction by averaging the first 100 subjects in our training data. Top: Intensity template obtained by averaging. Bottom: Corresponding label map to the average intensity template. The templates constructed by averaging is notably less sharp than the AtlasMorph templates.
• Figure 5: Intensity Templates obtained with AtlasMorph. Left: Unconditional Intensity Template. Right: Conditional Intensity Templates sampled from our learned template function from age 10 to 90, left to right. The conditional templates capture well-known signs of age-associated atrophy as analyzed in Figure \ref{['fig:3D_vent_vol_kernel']}.