Memory efficient training for 3D brain image registration networks using PatchMorph.

Area of Science:

• Medical Imaging
• Computational Neuroscience
• Machine Learning

Background:

• Deep learning models for 3D brain image registration often require substantial memory, limiting their application to large or high-resolution datasets.
• Existing methods struggle with global registration tasks, varying voxel resolutions, and large image dimensions.
• Unsupervised convolutional and transformer-based networks are powerful but memory-intensive.

Purpose of the Study:

• To introduce PatchMorph, a novel stochastic framework to decrease memory demand in 3D brain image registration.
• To enable the use of existing deep learning architectures (CNN, Transformer) for large-scale and heterogeneous brain image datasets.
• To handle global registration tasks and varying voxel resolutions efficiently.

Main Methods:

• PatchMorph employs a patch-based, coarse-to-fine strategy operating in world coordinate space.
• It decouples spatial logic from network architecture, matching patches across multiple scales.
• The framework supports integration with VoxelMorph-like architectures (CNN or Transformer).

Main Results:

• Reduced training memory for a transformer-based network on high-resolution images from ~40 GB to <10 GB.
• Maintained state-of-the-art registration performance on human T1 MRI and marmoset brain images.
• Demonstrated efficient handling of large image sizes, differing array dimensions, and varying voxel resolutions.

Conclusions:

• PatchMorph effectively removes memory bottlenecks in deep learning-based 3D brain image registration.
• It facilitates the deployment of sophisticated registration networks on challenging, real-world datasets.
• The method offers a significant reduction in memory footprint while preserving registration accuracy.