Modeling the Neonatal Brain Development Using Implicit Neural Representations

TL;DR

This work uses implicit neural representations (INRs) to model healthy neonatal brain development across PMA by predicting subject-specific brain images from sparse longitudinal MRI data. It introduces two disentanglement strategies, Subject Specific Latent Vectors (SSL) and Stochastic Global Latent Augmentation (SGLA), to separate age information from individual identity within the INR latent space, enabling consistent cross-time predictions. Compared to an age-conditioned denoising diffusion baseline, the INR with SSL and SGLA yields superior PSNR, SSIM, MAE, and, in 3D, more accurate head-circumference correlations, while remaining memory-efficient via pixel-wise training and micro-batching. The approach demonstrates potential for generating age-conditional brain atlases and lays groundwork for extensions to segmentation and atlas construction, with limitations in cortical folding fidelity and dependency on dataset size.

Abstract

The human brain undergoes rapid development during the third trimester of pregnancy. In this work, we model the neonatal development of the infant brain in this age range. As a basis, we use MR images of preterm- and term-birth neonates from the developing human connectome project (dHCP). We propose a neural network, specifically an implicit neural representation (INR), to predict 2D- and 3D images of varying time points. In order to model a subject-specific development process, it is necessary to disentangle the age from the subjects' identity in the latent space of the INR. We propose two methods, Subject Specific Latent Vectors (SSL) and Stochastic Global Latent Augmentation (SGLA), enabling this disentanglement. We perform an analysis of the results and compare our proposed model to an age-conditioned denoising diffusion model as a baseline. We also show that our method can be applied in a memory-efficient way, which is especially important for 3D data.

Figures (7)

• Figure 1: Overview: The input to the network $f_\vartheta$ consists of the spatial coordinates $\hbox{\boldmath$x$}$ of the desired output pixel $\hbox{\boldmath$\hat{I}$}_{\hbox{\boldmath$x$}}$, the desired PMA $t \in T$, and a latent vector $\hbox{\boldmath$l$} \in \Lambda$, encoding the subject identity. The switch with probability $p$ between $\hbox{\boldmath$l$}$ and the global latent vector $\hbox{\boldmath$l$}_G$ represents the SGLA.
• Figure 2: Ablation of our model w.r.t. the HC. The blue circles show the HC ground truth, while the red x-es show the HC of our models' predictions on the test set.
• Figure 3: Six examples from the test set, along with the PMA $t_1$ of the input, and the PMA $t_2$ of the target ground truth image.
• Figure 4: Comparison of our "average" brain development with the IMAGINE atlas and the atlas of Schuh et al. in terms of HC growth curves \ref{['fig:imagine:hc']}, along with the quantiles reported in fenton2003new. Furthermore, we show the corresponding axial slices in \ref{['fig:imagine:imgs']}.
• Figure 5: PSNR and SSIM on the test set as a function of the number of noising and denoising steps $L$ and the gradient scale $c$ for the DDM+GG baseline wolleb2022swiss.