# A Scalable Distributed Framework for Multimodal GigaVoxel Image Registration

arXiv: https://arxiv.org/abs/2509.25044 · HTML: https://mlsanity.com/paper/2509.25044/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 123 references that are in the map.
- First posted on arXiv 29 Sep 2025.

## Builds on (most cited first, 30 of 41)

- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems](https://mlsanity.com/paper/1603.04467/) (2016)
- [VoxelMorph: A Learning Framework for Deformable Medical Image Registration](https://mlsanity.com/paper/1809.05231/) (2018)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [SynthSeg: Segmentation of brain MRI scans of any contrast and resolution without retraining](https://mlsanity.com/paper/2107.09559/) (2021)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [A Deep Learning Framework for Unsupervised Affine and Deformable Image Registration](https://mlsanity.com/paper/1809.06130/) (2018)
- [SynthMorph: learning contrast-invariant registration without acquired images](https://mlsanity.com/paper/2004.10282/) (2020)
- [Large Deformation Diffeomorphic Image Registration with Laplacian Pyramid Networks](https://mlsanity.com/paper/2006.16148/) (2020)
- [Sequence Parallelism: Long Sequence Training from System Perspective](https://mlsanity.com/paper/2105.13120/) (2021)
- [Learn2Reg: comprehensive multi-task medical image registration challenge, dataset and evaluation in the era of deep learning](https://mlsanity.com/paper/2112.04489/) (2021)
- [Quicksilver: Fast Predictive Image Registration - a Deep Learning Approach](https://mlsanity.com/paper/1703.10908/) (2017)
- [Recursive Cascaded Networks for Unsupervised Medical Image Registration](https://mlsanity.com/paper/1907.12353/) (2019)
- [Unsupervised 3D End-to-End Medical Image Registration with Volume Tweening Network](https://mlsanity.com/paper/1902.05020/) (2019)
- [uniGradICON: A Foundation Model for Medical Image Registration](https://mlsanity.com/paper/2403.05780/) (2024)
- [Region Mutual Information Loss for Semantic Segmentation](https://mlsanity.com/paper/1910.12037/) (2019)
- [U-Net vs Transformer: Is U-Net Outdated in Medical Image Registration?](https://mlsanity.com/paper/2208.04939/) (2022)
- [Zero Bubble Pipeline Parallelism](https://mlsanity.com/paper/2401.10241/) (2023)
- [Affine Medical Image Registration with Coarse-to-Fine Vision Transformer](https://mlsanity.com/paper/2203.15216/) (2022)
- [DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training](https://mlsanity.com/paper/2310.03294/) (2023)
- [Learning General-Purpose Biomedical Volume Representations using Randomized Synthesis](https://mlsanity.com/paper/2411.02372/) (2024)
- [NODEO: A Neural Ordinary Differential Equation Based Optimization Framework for Deformable Image Registration](https://mlsanity.com/paper/2108.03443/) (2021)
- [Deep Learning in Medical Image Registration: Magic or Mirage?](https://mlsanity.com/paper/2408.05839/) (2024)

## Built on by (most cited first, 1 of 1)

- [Registration is a Representation Problem: Solving Whole-Body PET/CT Registration](https://mlsanity.com/paper/2610.18804/) (2026)

## Abstract (arXiv)

In this work, we propose FFDP, a set of IO-aware non-GEMM fused kernels supplemented with a distributed framework for image registration at unprecedented scales. Image registration is an inverse problem fundamental to biomedical and life sciences, but algorithms have not scaled in tandem with image acquisition capabilities. Our framework complements existing model parallelism techniques proposed for large-scale transformer training by optimizing non-GEMM bottlenecks and enabling convolution-aware tensor sharding. We demonstrate unprecedented capabilities by performing multimodal registration of a 100 micron ex-vivo human brain MRI volume at native resolution - an inverse problem more than 570x larger than a standard clinical datum in about a minute using only 8 A6000 GPUs. FFDP accelerates existing state-of-the-art optimization and deep learning registration pipelines by upto 6 - 7x while reducing peak memory consumption by 20 - 59%. Comparative analysis on a 250 micron dataset shows that FFDP can fit upto 64x larger problems than existing SOTA on a single GPU, and highlights both the performance and efficiency gains of FFDP compared to SOTA image registration methods.