PatchMorph: A Stochastic Deep Learning Approach for Unsupervised 3D Brain Image Registration with Small Patches

MRI and microscopy image stacks, $\mathbf{I}_{i}$, stored in formats such as Nifti and Tiff, are accompanied by affine 3D matrices $\mathbf{A}_{i}$ that map array coordinates to physical world coordinates. We denote the inverse of this matrix as $\textbf{T}_{i}:=(\mathbf{A}_{i})^{-1}$. Figure 1 illustrates how different arrays and orientations can represent the same image through such transformations, as expressed by $\mathbf{I}_{1}(\mathbf{T}_{1}\mathbf{r})=\mathbf{I}_{2}(\mathbf{T}_{2}\mathbf{r})=\mathbf{I}_{3}(\mathbf{T}_{3}\mathbf{r})$.

PatchMorph leverages this by directly sampling patch data from original arrays without the need for padding or resolution alteration, thus maintaining the intended resolution, orientation, and scale. It simplifies the use of fixed ($\textbf{I}_{\text{f}}$, $\textbf{T}_{\text{f}}$) and moving ($\textbf{I}_{\text{m}}$, $\textbf{T}_{\text{m}}$) image stacks. The affine transformation matrices, $\textbf{T}_{\text{i}}$, integrate rotation, scaling, and translation, providing 12 degrees of freedom.

Figure 5 provides a detailed representation of the PatchMorph-Block architecture. It accepts two input arguments: the world coordinates $\mathbf{X}_{\text{f}_{(t)}}$ for the current fixed image patch and, when $t>0$, a relative displacement patch $\mathbf{d}_{\text{in}_{(t)}}$ from the previous iteration. The primary output, which is utilized during both training and inference, is the updated relative displacement vector patch $\mathbf{d}_{\text{out}_{(t)}}$.

Given the inputs, PatchMorph determines the world coordinates of the current moving patch using $\mathbf{X}_{\text{m}_{(t)}}:=\mathbf{X}_{\text{f}_{(t)}}+\mathbf{d}_{\text{in}_{(t)}}$. For $t=0$, $\mathbf{X}_{\text{m}_{(t)}}=\mathbf{X}_{\text{f}_{(t)}}$. Based on both, the transformation matrices $\mathbf{T}_{\text{f}}$ and $\mathbf{T}_{\text{m}}$, we can map the world coordinate fields to array coordinate fields $\mathbf{x}_{\text{f}_{(t)}},\mathbf{x}_{\text{f}_{(m)}}$. Based on them, we extract the corresponding image patches $\mathbf{p}_{\text{f}_{(t)}}$ and $\mathbf{p}_{\text{m}_{(t)}}$ from the image stacks $\mathbf{I}_{\text{f}}$ and $\mathbf{I}_{\text{m}}$. Post-extraction, both patches are concatenated and fed into a CNN that estimates the relative voxel displacement field $\hat{\mathbf{d}}_{\text{local}_{(t)}}$ between the two image patches. Using the affine matrix $\mathbf{T}_{\text{w}_{(t)}}$, we translate the update of the displacement field back to the world coordinates, represented as $\hat{\mathbf{d}}_{(t)}$. The matrix $\mathbf{T}_{\text{w}_{(t)}}$ can be computed from the patch array size, the matrix $\mathbf{T}_{\text{ref}}$, and the matrices $\mathbf{T}_{\text{p}_{t=\{0,\cdots,t\}}}$. If available, we update these coordinates with $\mathbf{d}_{\text{in}_{(t)}}$, deriving the refreshed displacement field

During training, we cropped zero-boundaries of the images to minimize the array size, and we modify the translation vectors of $\mathbf{T}_{\text{f}}$ and $\mathbf{T}_{\text{m}}$ so that the world coordinate center of the images aligns with the origin. We also added the left-right mirrored copy of each image to the training set.

For each $t$, we update the world coordinates of the moving patch $\mathbf{X}_{\text{moved}_{(t)}}:=\mathbf{X}_{\text{f}_{(t)}}+\mathbf{d}_{\text{out}_{(t)}}$ to obtain an updated version of the moving image patch denoted as $\mathbf{p}_{\text{moved}_{(t)}}$. The training objective is to optimize the similarity between the outputs, $\mathbf{p}_{\text{moved}_{(t)}}$, and the fixed image patches, $\mathbf{p}_{\text{f}_{(t)}}$ across all scales in an unsupervised way using an image similarity loss such as the normalized cross correlation (NCC).

After the training process, equation $\mathbf{X}_{\text{moved}_{(t)}}:=\mathbf{X}_{\text{f}_{(t)}}+\mathbf{d}_{\text{out}_{(t)}}$ determines the world coordinates of the image patch in the moving image that optimally align with the fixed image patch situated at location $\mathbf{X}_{\text{f}_{(t)}}$ within the fixed image.

Figure 6 displays examples of fixed, moving, and moved image patches after training, extracted based on coordinate patches in nine different scales.

During inference, we adjust the translation vector of $\mathbf{T}_{\text{m}}$ to align the center of the moving image’s world coordinates with that of the fixed image, correcting this shift afterward in post-processing.

The goal is to construct the global DDF $\mathbf{D}$, satisfying the image registration equation (1), where $\mathbf{X}_{\text{m}}=\mathbf{X}_{\text{ref}}+\mathbf{D}$.

Unlike training, inference does not require a cascade of nested patches. To assemble the global $\mathbf{D}$, PatchMorph generates successive updates $\mathbf{D}_{(t)}$ at each scale, culminating in: $\mathbf{D}:=\sum_{t}\mathbf{D}_{(t)}$.

Initially, PatchMorph creates a low-resolution displacement field $\mathbf{D}_{(0)}$ for the first scale. As the process advances to subsequent iterations, the spatial resolution progressively increases, with each iteration yielding a global update $\mathbf{D}_{(t)}$, as visualized in Fig. 7. The array resolution of each $\mathbf{D}_{(t)}$ is enhanced to match the increased spatial resolution of the patches. To accommodate the increase in resolution from $\mathbf{D}_{(t)}$ to $\mathbf{D}_{(t+1)}$, linear interpolation is utilized to upscale the fields with lower array resolutions.

At each iteration $t$, PatchMorph samples a large number of coordinate patches $\mathbf{X}_{\text{f}_{t}}$ from the canvas $\mathbf{X}_{\text{ref}}$. For $t>0$, displacement updates $\mathbf{d}_{\text{in}_{(t)}}$ are interpolated from $\mathbf{D}{(t-1)}$, input into the PatchMorph-Block for the current scale, and new updates $\mathbf{d}_{\text{out}_{(t)}}$ are obtained. These updates are then integrated into $\mathbf{D}_{(t)}$, with overlapping patch results being averaged.

The required number of patches for completing a scale $t$ is determined by the ratio of an individual patch’s volume to that of the coordinate field, ensuring that each voxel in $\mathbf{D}_{t}$ is derived from an average of 10 patch predictions $\mathbf{d}_{\text{in}_{(t)}}$.

In our experiments, we tested PatchMorph based on two kinds of CNN architectures: an affine CNN for local affine alignment, inspired by GlobalNet [17], and a dense CNN that allows for single-voxel displacement, similar to VoxelMorph [5, 4] or LocalNet [17].

Our CNN network architectures are built on four major components: the 3D convolution, a convolution-block, a down-convolution block, and an up-convolution block. The convolution-block consists of a batch norm, followed by a leaky ReLU and a 3D convolution, if not stated otherwise, with a kernel size of 3. A Down-convolution block is identical to a Convolution-block, but with the stride parameter set two 2. Consequently, it halves the spatial resolution of a patch. The up-convolution block is a convolution-block, but with a nearest neighbor up-scaling operation with a scaling factor of 2, placed before the convolution. An instance normalization layer (not shown) is normalizing the input patches. Figure 8 shows a sketch of our CNN architectures along the input and output channels of each layer.

Affine CNN: Panel (a) in Figure 8 shows the simple encoder network that predicts 12 parameters of an affine transformation. This matrix is used to predict the DDF using a matrix exponential layer: Let $\mathbf{A}$ be the 12 parameters in the form of an affine matrix. Let $\mathbf{T}=e^{(0.25\mathbf{A})}$ be the estimated affine matrix, were $e^{\mathbf{A}}$ is the matrix exponential of a square matrix $\mathbf{A}$ that we use as activation function. We heuristically determined the constant $0.25$ so that the network generated small deformations with randomly initialized network parameters. Let $\mathbf{T}_{3\times 3}$ be the left upper $3\times 3$ matrix of $\mathbf{T}$, representing rotation and scale, and $\boldsymbol{\tau}$ be the 3D translation vector of $\mathbf{T}$. Further, let $\mathbf{r}$ be a vector from the patch coordinate field which was normalized to the interval $[0,1]$. Then the matrix exponential layer is defined by:

The result is then scaled by the patch size to obtain the DDF $\hat{\mathbf{d}}_{\text{local}}$ in (patch) voxel spacing. Note that for the zero tensor, the matrix exponential is just the identity transform. The translation factor 0.1 has been determined heuristically. In our implementation, we used the approximation $e^{\mathbf{A}}\approx(\mathbf{A}/k+\mathbf{I})^{k}$, with $k=10$ for performance reasons.

Dense CNN: We estimate voxel-displacement fields using a simplified U-Net in combination with an integration layer to estimate a DDF from a dense velocity field (DVF), similar to the diffeomorphic VoxelMorph. We used the DVF2DDF block from MONAI [18]. Panel (b) in Figure 8 illustrates the architecture. The last convolution has a kernel size of 1.

In our figures, we show results for nine distinct scales. The number of which was determined through a preliminary experiment (details provided later in the experiments). We determine the coarsest and finest resolutions of the patches based on training data, with intermediate resolutions linearly interpolated between these two extremes. All patches possess isotropic voxel resolutions.

Using a uniform array size for the patches, for instance, $32^{3}$, we cycle through all training images. During each cycle, we calculate the maximum edge length (in millimeters) across all canvases $\mathbf{I}_{\text{ref}}$, which determines the edge length for the coarsest patch scale. Thus, at the coarsest scale, a patch, when centered on an image, encompasses the entire image in low resolution.

The voxel resolution for the most detailed patch scale is set to match the smallest voxel resolution found along any axis in all training images. Consequently, this voxel resolution is always equal to or finer than that of a source image.

The MindBoggle dataset consists of 101 T1-weighted MRI scans [20]. The dataset provides labels for 62 cortical regions across both hemispheres. Each dataset contains an image in its native image space and a version aligned to a standard MNI template, both accompanied by an affine world transformation matrix.

For our PatchMorph experiments, we used the original, skull-stripped images. Following the partitioning used by [13], we randomly divided the dataset into 73 training, 5 validation, and 20 test images, excluding one due to missing data. The remaining 2 images were used as fixed images for evaluation during validation and testing. They were added to the training set as well. The images were taken from 53 men and 47 women, ranging in age from 19 to 61 years. The images varied in array sizes, with a rounded average dimension of $[212\times 250\times 219]$ voxels with a rounded standard deviation of $[41\times 14\times 46]$ voxels, and presented diverse affine transformations, often involving swapped or mirrored axes. While most images had isotropic voxel resolutions of $1~{}mm^{3}$, 21 images featured an anisotropic resolution of $1.2\times 1\times 1~{}mm^{3}$.

We utilized a collection of 52 3D marmoset brain images acquired via serial two-photon microscopy from [21] for image registration, focusing on the first color channel that captures auto-fluorescent tissue signals. The dataset includes a population average template used as a fixed image in our experiments, along with cortical and subcortical labels for 36 brain regions. These labels were manually annotated on the template and projected onto individual brains using the ANTs toolkit. For our study, we split the dataset evenly, with 27 images for training, including the template, and 26 for testing. We used the raw, unaligned images with a mean size of approximately $212\times 257\times 327$ voxels (standard deviation of $13\times 16\times 17$) and an isotropic voxel resolution of $(0.1~{}mm)^{3}$.

The dataset presents challenges such as variable background signal quality and corruption from anterograde tracer injections, along with tissue damage and missing coronal sections due to acquisition problems.

We assessed three metrics: average Dice score (Dice${}_{\text{avg}}$), minimum Dice score (Dice${}_{\text{min}}$), and the integrity of the deformation field ($|J_{\Phi}|\leq 0$). For the latter, we calculated the proportion of negative Jacobian determinants from the normalized coordinate field, aiming for a 0% rate as the ideal. It is important to note that, in contrast to [13], we excluded voxels outside the brain volume from this analysis, increasing the sensitivity to any anomalies in the deformation field.

We tested three, six, nine and twelve distinct scales. Preliminary experiments determined that nine scales yielded optimal performance on the MindBoggle validation set. Table 1 shows the results on the test set. We also tested several patch sizes on the validation set, with $32^{3}$ performing best (details discussed later in the text).

Subsequently, we evaluated two PatchMorph configurations with nine scales. The first, our default setup, utilizes a shared Affine CNN for the coarsest scales (0-2), a separate Affine CNN for the intermediate scales (3-5), and a Dense CNN for the finest scales (6-8). The second configuration, referred to as A-PatchMorph, employs an additional Affine CNN in place of the Dense CNN for the finest scales.

In each training iteration, we process two pairs of moving and fixed images (four images in total). From each pair, we randomly sample 10 patches for each scale, resulting in a total batch size of 20. The training incorporates two data augmentation steps: First, we apply random affine transformations to the initial crop matrices $\mathbf{T}_{\text{p}_{0}}$, adjusting for rotation (within $\pm 15^{\circ}$) and scale (between 0.9 and 1.1). Second, we apply another set of random affine transformations to the moving image’s transformation matrix $\mathbf{T}_{\text{m}}$, introducing variations in rotation (up to $\pm 25^{\circ}$), scale (ranging from 0.8 to 1.2), and translation (up to $\pm 2$ cm). Note that $\mathbf{T}_{\text{p}_{0}}$ effects both moving and fixed image. We used the image intensities themselves as a saliency map to sample patches with a higher likelihood from the foreground than from the background, maintaining a ratio of 2:1.

We implemented PatchMorph in PyTorch 2.0.1 (https://pytorch.org/). An AdamW optimizer [22] with an initial learning rate of 0.001, and gradient norm clipping [23] with a max norm of 2 is used for optimization. The training regimen begins with 10,000 iterations focusing solely on the coarsest scale $t=0$, after which the second scale is included for another 10,000 iterations. This incremental process continues until all 9 scales have been incorporated. Upon reaching 90,000 iterations, we extend the training with 40,000 additional iterations, employing a scheduler to gradually decrease the learning rate. To enhance efficiency, the loss is calculated only for the last three scales, reducing the frequency of patching operations.

Following the loss used by [13], for the MindBoggle dataset, a global NCC loss is combined with a local NCC loss using a kernel size of 9. For the Marmoset brain dataset, we implement a mutual information loss, which is detailed later in the text.

We benchmarked PatchMorph against the methods evaluated in our prior study on implicit neural representation networks (INRs) [13]. As some methods necessitate pre-aligned images, comparative analyses were performed on affine-aligned brain images. Our comparison included a standard Affine baseline, ANTs SyN registration, diffeomorphic VoxelMorph (D-VMorph), and SIREN, the highest-performing INR model from  [13].

The results for PatchMorph are shown in the upper section of Table 2, and the comparative analyses are shown in the lower section. Figure 9 shows qualitative results for one example.

We tested four patch sizes: $16^{3}$, $24^{3}$, $32^{3}$, and $48^{3}$ voxels, with $32^{3}$ yielding the highest Dice score (0.546). In further tests, we repeated certain PatchMorph scales during inference. Considering the vast number of combinations, here we report only the scenario where we repeated the finest scale ($t=8$) on a patch size of $32^{3}$ twice more. This approach resulted in the highest Dice score observed (0.583). While SIREN outperformed other methods (0.579), ANTs SyN and PatchMorph showed significantly better deformation field quality, with SIREN struggling in nearly 2% of voxels. With lower Dice scores (0.424) on the MindBoggle dataset than other methods, A-PatchMorph still significantly outperformed the baseline.

Figure 10 shows the dice score assigned to each single scale for PatchMorph $32^{3}$. It shows that the biggest gain is in the three first and last three scales.

To comprehensively assess the efficacy and efficiency of PatchMorph, we performed a series of performance analyzes that scrutinize various aspects of the algorithm, ranging from regularization impacts to computational resource utilization.

Regarding training requirements, PatchMorph $32^{3}$ on the MindBoggle dataset with six and nine scales necessitated approximately 5 and 9 gigabytes of GPU memory, respectively. The training duration for models with nine scales varied significantly depending on the system and the load of the system, with times ranging between 7.5 and 20 hours.