[Paper Review] Learning Conditional Deformable Templates with Convolutional Networks

Adrian V. Dalca, Marianne Rakic+2 more|arXiv (Cornell University)|Aug 7, 2019

Advanced Numerical Analysis TechniquesEngineering65 references59 citations

TL;DR

The paper presents a probabilistic, CNN-based framework to learn deformable templates that can be conditioned on attributes, enabling on-demand templates and fast image alignment, demonstrated on MNIST/QuickDraw and large-scale brain MRI data.

ABSTRACT

We develop a learning framework for building deformable templates, which play a fundamental role in many image analysis and computational anatomy tasks. Conventional methods for template creation and image alignment to the template have undergone decades of rich technical development. In these frameworks, templates are constructed using an iterative process of template estimation and alignment, which is often computationally very expensive. Due in part to this shortcoming, most methods compute a single template for the entire population of images, or a few templates for specific sub-groups of the data. In this work, we present a probabilistic model and efficient learning strategy that yields either universal or conditional templates, jointly with a neural network that provides efficient alignment of the images to these templates. We demonstrate the usefulness of this method on a variety of domains, with a special focus on neuroimaging. This is particularly useful for clinical applications where a pre-existing template does not exist, or creating a new one with traditional methods can be prohibitively expensive. Our code and atlases are available online as part of the VoxelMorph library at http://voxelmorph.csail.mit.edu.

Motivation & Objective

• Motivate deformable templates for analyzing anatomical and geometric variability.
• Propose a probabilistic model using diffeomorphic deformations and a learnable conditional template.
• Develop an end-to-end neural network that jointly learns the template and deformation fields.
• Enable on-demand conditional templates as a function of attributes (e.g., age, sex).
• Demonstrate efficiency and usefulness on benchmark and neuroimaging datasets.

Proposed method

• Formulate images as diffeomorphic deformations of a conditional template t = f_theta_t(a).
• Model deformations v_i as stationary velocity fields and integrate via scaling and squaring.
• Specify a prior p(V) to encourage smooth, unbiased templates and topology-preserving deformations.
• Define data likelihood p(x_i|v_i,a_i) with a warp x_i ≈ t ∘ φ_v_i, using a Gaussian or robust alternative.
• Use a two-network architecture: g_t,theta_t(a_i) outputs the template; g_v,theta_v(t, x_i) outputs the velocity field; train end-to-end via maximum likelihood (negative log-likelihood) loss.
• Train with stochastic gradient descent to jointly optimize template and deformation networks; at test time, generate on-demand templates and deformations quickly.

Experimental results

Research questions

• RQ1 Can a single network learn a conditional deformable template across an entire dataset?
• RQ2 Does conditioning templates on attributes (e.g., age, sex) produce unbiased, central templates with smaller deformations?
• RQ3 How does the proposed method compare to exemplar or decoder-based baselines in terms of template centrality and deformation regularity?
• RQ4 Is the conditional template function capable of generating anatomically meaningful templates for unseen attribute configurations?
• RQ5 Can the framework handle 2D benchmarks and 3D neuroimaging data efficiently?

Key findings

• Conditional templates are more central and require smaller deformations than baselines on MNIST variants.
• The method yields smooth deformation fields with comparable or better MSE than baselines and exemplars.
• On 3D brain MRI, unconditional and age/sex conditional templates achieve higher Dice scores than a baseline atlas (0.800±0.110 vs 0.795±0.116 vs 0.731±0.153).
• Conditional templates enable rapid synthesis of on-demand templates (training ~1 day; test-time conditioning in under a second).
• The framework captures age-related anatomical variability, e.g., ventricles grow and hippocampus shrinks with age, demonstrated by volume trends.
• Latent attribute analyses show the model can recover meaningful template representations even when attributes are partially missing or unknown.