A Scalable Distributed Framework for Multimodal GigaVoxel Image Registration
Authors: Rohit Jena
Vedant Zope
Pratik Chaudhari
James Gee
Venue: ICLR 2026 Oral (2026)
Keywords: image registration
distributed optimization
CUDA kernels
neuroanatomy
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: we propose non-GEMM CUDA kernels and distributed primitives to scale multimodal image registration to arbitrary image sizes

Abstract: 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μm ex-vivo human brain MRI volume at native resolution – an inverse problem more than 570× 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 − 7× while reducing peak memory consumption by 20 − 59%. Comparative analysis on a 250μm dataset shows that FFDP can fit upto 64× 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.

Published as a conference paper at ICLR 2026
A SCALABLE DISTRIBUTED FRAMEWORK FOR MULTIMODAL GIGA VOXEL IMAGE REGISTRATION
Rohit Jena*.t
Vedant Zope!
Pratik Chaudhari
James C. Gee
University of Pennsylvania Philadelphia, PA 19104, USA
{rjena, zope, pratikacJ8sear upenn. edu, geelupenn.
ABSTRACT
In this work, we propose FFDP, a set of IO-aware non-GEMM fused kernels supplemented with distributed framework for image registration unprecedented scales. Image registration is 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-awan tensor sharding. We demonstrate unprecedented capabilities by performing multimodal registration of a 100pm 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 of-the-art optimization and deep learning registration pipelines by upto 6 - 7x while reducing peak memory consumption by 20 59%. Comparative analysis on 250 dataset shows that FFDP can fit upto 64x larger problems than existing SOTA on single GPU, and highlights both the performance and efficiency gains of FFDP compared to SOTA image registration methods.
INTRODUCTION
Image registration (also called 'image alignment' or 'image matching') is non-linear inverse problem ubiquitous in biomedical and life sciences. Given d-dimensional images F S2 R and M : Ω2 8d defined on domain Q (usually compact subset of Rd), image registration seeks to find a coordinate transform 8 S + Zt that deforms the moving image M to look similar to the fixed image F. Mathematically, we minimize the following objective (Fig. 1):
= arg min L(x) = C(F, + R(())
(1)
FEG
where C is cost or dissimilarity function, and o is the interpolation operator, i.e. (Tog)(x) T(g(a))
for all E Ω. Popular choices of yo are affine and deformable transforms, i.e. <p(x) Ar + 1 and (xx)) = + u(x). Modern registration pipelines (Hoffmaan et al., 2021; Jena et al., 2024a) consider an affine matching followed by a deformable matching step. resulting in a composite transform (xxx: Ar t + u(x). is called the displacement field, modeled
grid of per-voxel vectors u(x) E IRd For an image of size N, the displacement field a tensor of size dN. We use [x] A + t, and ju], to denote the identity grid, grid of affine transformed coordinates, and deformation grid defined on Ω respectively. Common choices of C are mean squared error, Localized Normalized Cross Correlation (Avants al., 2008a), and Mattes Mutual Information (Mattes et al., 2001). Common choices of R include Sobolev norm of the gradient or warp fields (Beg et al., 2005; Mang al., 2019; Avants et al., 2008b), total variation, and inverse-consistency (Christensen & Johnson, 2001). To optimize Eq. (1), iterative methods optimize y" directly using gradient descent, and deep learning methods learn a deep neural network & fo M). Image registration establishes common coordinate system, aligning scans across individuals and atlases (Hering et al., 2022;
"Corresponding author.
"hhee authors contributed equally.

---

Published as a conference paper at ICLR 2026
[a]
(b)
Moving image
image
Warp fleld
image
Movine
Warp
campler
Los, MI)
M
Mo ip
Figure 1: Image Registration Problem. (a): The task is to find coordinate transform that warps the moving image M to the fixed image F. Individual corresponding points are shown violet arrows; the per-pixel coordinate transform is shown as a warp field ye, and the transformed image yo. (b): A typical registration pipeline the grid sampler warps the moving image, that is then compared to the fixed image using loss function. Green denotes the optimizable warp, red denotes the primary bottlenecks that we optimize in this paper.
Marcus et al., 2007; Murphy al., 2011). This alignment is prerequisite for multimodal data fusion, cross-subject comparison, morphometric analysis (Das et al., 2009), and construction of large-scale atlases (Wang et al., 2020b). Establishing such voxelwise correspondence is fundamental for studying anatomical variability, detecting pathological signatures (Ravikumar et al., 20,1), and advancing precision medicine (Börner et al., 2022; Jonsson et al., 2022). The saliency and centrality of the task across various biomedical and life science applications has spurred numerous methodological advances the field, spanning more than three decades of research (Gee et al., 1993; Tian al., 2024).
Over the past decade, advances MRI, CT, PET, STPT, and microscopy have enabled ultra-high- resolution imaging, often more than three orders of magnitude larger than macroscopic biomedical domains (Balchandani
Naidich, 2015; Esquivel et al., 2022; Badawi et al., 2019; Gambarotto et al., 2019; Wassie et al., 2019; Kleven al., 2023; Wang et al., 2020b; Mansour et al., 2025; Kleinfeld et al., 2011). While a typical clinical registration problem involves ~200 parameters, high-resolution ex-vivo human brain scans can require solving up 11B parameters, far beyond the ~-50M-parameter scale at which current registration methods remain reliable. As a result, state-of-the-art deformable image alignment struggles scale the resolutions demanded in modern neuroimaging, computational pathology, developmental biology, and connectomics, creating a substantial performance gap. In parallel, innovations large-scale transformer training such as 10-aware fused operations (Dao et al., 2022; Dao, 2023; Spector et al., 2025) and 5D parallelism for distributing surgerthhan memory workloads (Shoeybi et al., 2019; Li al., 2023; Jacobs al., 2024; Li et al., 2024; Zhao et al., 2023; Ansel et al., 2024) optimize GEMM-like workflows. However, the fundamental concepts utilized by these methods (lO-awareness, recomputing and aggregating intermediates shared memory to minimize high bandwidth memory (HBM) storage, identifying partial aggregates across hosts to minimize communication overheads for distributed optimization) broadly applicable to wide class of problems of the non-GEMM nature.
than-memory
In this paper, we apply these concepts to scale image registration algorithms to match parity with the developments in both inereasing resolution of image acquisition and compute capabilities. To that end, our contributions are twofold. First, we identify key compute and memory bottlenecks in image registration algorithms, and propose novel components that fit problems upto 64x larger than existing algorithms on a single GPU. Second, we propose Flash Fused Distributed Primitives (FFDP), distributed framework to scale registration an arbitrary number of GPUs, thereby scaling ultra nigh-resolution problems. We present first-of-its-
demonstration: aligning a 250μmm in-vivo MRI (Lüseerink et al., 2017) to a 100gmm ex-vivo human brain FLASH volume (Edlow al., 2019) multimodal registration problem more than 570x larger than standard clinical datum (Mareus et al., 2007), with over 11.8B transform parameters - completed one minute using only 8 A6000 GPUs FFDP accelerates existing traditional registration pipelines by upto 7.48x while reducing memory consumption by upto 59%, and deep learning pipelines by upto 6. 111 while consuming upto 24% less memory. We highlight the necessity of performing high-resolutie comparing our method with various SOTA optimization and deep learning baselines on a 250pmm T1-weighted MRI dataset, showing unprecedented performance and gains in efficiency.
registration by
2

---

Published as a conference paper at ICLR 2026
2
RELATED WORK
2.1 MEMORY EFFICIENT AND LARGE SCALE OPTIMIZATION
Recent years have also witnessed tremendous innovations in large-scale transformer model training. 10-aware implementations typically include individual fused kernels (Dao et al., 2022; Dao, 2023) and nomain--peccifi languages (Spector et al., 2025; PyTorch, 2025) to minimize launch latency and large memory overheads. To distribute larger-than-memrry model training workloads across multiple GPUs, 5D parallelism techniques (Shoeybi et al., 2019; Li et al., 2023; Jacobs et al., 2024; Li et al., 2024; Zhao et al., 2023; Ansel et al., 2024) have been proposed. Many of these techniques leverage a divide-and conquue approach to break down a larger GEMM-like operation like matrix multiplication attention into smaller sub- pol-proobles that can be executed on multiple GPUs and synchronized to compute the final result. To our knowledge, most of these techniques are tailored to trasformer-specifie architectures and GEMM-like operations (self attention, FeedForward, LayerNorm, etc.) only, and a Model Parallel variant for convolution-aware tensor sharding and synchronization is not available.
2.2 LARGE SCALE REGISTRATION IN LIFE SCIENCES AND BIOMEDICAL IMAGING
Ex-vivo neuroimaging, Neuroanatomical studies often integrate high-resolution ex-vivo MRI, blockface imaging, and histology to bridge the gap between in-vivo imaging and microscopic "gold standard" pathology (Casamitjana et al., 2025; Ravikumar et al., 2024). While large-scale consortia like SEA-AD and HMBA, along with submillimeter whole-brain datasets (Edlow et al., 2019; Lüsebrink et al., 2017), aim map cellular and molecular organization across species, computational costs often limit analysis to local effects. Specifically, existing tools cannot register these datasets at native resolution due to excessive memory requirements. We overcome this limitation, demonstrating native-resolution registration of whole-brain datasets in one minute using eight A6000 GPUs (see Section 5.2), thereby preserving fine anatomical details typically lost to downsampling.
Large-scale registration in model organisms. Over the past decade, imaging across the life sciences and biomedical domains has progressed from mesoscale surveys to organ- and organism-win acquisitions at cellular or even subcellular resolution. These span transparent organisms and small animal models (e.g., C. elegans, zebrafish, adult Drosophila) ( (Varol et al., 2020; Venkatachalam et al., 2016; Marquart et al., 2017; Gupta et al., 2018; Peng et al., 201 Brezovec et al., 2024), adult mouse and brains imaged sub-micron resolutions (Gong et al., 2016; Wang al., 2020a; Kleven al., 2023) using Light Sheet Fluorescence Microscopy (LSFM) and Serial Two-Photon Microscopy (STPT) imaging. Such modalities routinely generate gigavoxel to teravoxel volumes (Kutten et al., 2016; Nazib et al., 2018). Their scientific utility. however, hinges the ability to perform registration at the native resolution of acquisition, i.e. aligning specimens (or modalities) in a common coordinate system without sacrificing the fine-scale morphologies including cell bodies, layers, axon bundles, synaptic neighborhoods, etc. that motivate high-resolution acquisition in the first place (Nazib et al., 2018: Goubran et al., 2013).
Across these diverse domains, the unifying requirement demands access to scalable multimodal registration algorithms- a challenge we address in this work. We provide an extended discussion of more related work and the necessity of our approach in Section A.
3 FUSED KERNELS FOR MEMORY EFFICIENT REGISTRATION ON A SINGLE GPU
Bottlenecks of a deformable image registration pipeline Our primary objective is to identify compute and memory bottlenecks in lurge-scale image matching tasks. In identifying these bottlenecks, training-free optimization methods are better suited than deep networks since the latter has much larger activation memory footprint, which forms the primary memory bottleneck (Tazi et al., 2024). For instance, for a 250 image pair, a standard deep learning method (Hoffmann al., 2021) generates an activation map of size 27GB only after the first layer. Extrapolating memory usage for clinical data, existing deep networks will require upto 1.2TB of GPU memory at inference to process these image volumes native resolution. In contrast, training ffree optimizer can fit this problem in less than 45GB of GPU memory. We use FireANTs (Jena et al., 2024a) as our base framework to identify compute and memory bottlenecks in typical image registration problem. We analyze
3

---

Published as a conference paper at ICLR 2026
Grid Sempler
Cross Correetton/ / Mutuel Informet
Fo
Mp
bilg
Alxip +1
MI((X))
[F. M. MM. FM....
OiFg by.
13,215
Drjy Day
Mad E1.
Apilavab
om Bay
F MiLLa
SRAM: 10OTR/S
+ +
M
Reg
SRAM: 20T0/4
Read amage
MF
PM
Losa
Wetne statogreet wee sicertyy
THE
Fused General and Grid Sample'
Fused Informetion
Figure 2: Left: FFDP uses fused kernels to eliminate intermediate HBBM memory usage (in dark red) for memory-bound workhorse operations (grid.sampler, LNCC, MI) for large-scale image registration. For grid_sample and LNCC, additional intermediate per-pixel variables (warp coordinates, patchwise statistics) computed per pexep in registers (bluc). For MI, the Parzen Windowing and histogram aggregation is performed using shared memory (green), avoiding large HBM overheads. Right: Pie charts show breakdown of memory overheads for storing the image, grid, optimizer state, and intermediate variables for MI and LNCC losses.
the flamegraph of a typical climical MRI registration task from the OASIS brain dataset (Marcus al., 2007) in Fig. 20. We identify three key memory bottlenecks image matching pipelines (I) deformable interpolation and warp composition (2) cross-correlation loss, and (3) mutual information loss (see Fig. 2rright)). We first propose efficient designs fit larger problems on single GPU, and then extend the framework to distributed registration.
3.1 COMPOSITE IMPLICIT GRID SAMPLER
fundamental operation used in image registration is the grid sampler. This operator allows us to warp an image M using deformation field yo Ω - Ω and computes the image M' M² (p)z)). Virtually every image registration pipeline uses this operation to warp the moving image using an affine, deformable, initializes a regular grid [x] a grid of size 3N. The affine grid A A'x] + t is another grid of size 3.. If a deformable grid (ulg is optimized, then third grid A + t + [ul, is materialized, costing total of 9N overhead for an image of size N. To consolidate these memory overheads, we propose a composite implicit grid sampler. This is a fused CUDA kernel that performs the following operation:
composite transform. For affine and composite transforms, the operator
fused_grice www.griddssameeeeee.... A, [u], S, Thounds
(z)) (x) [(Ax + Su(x))
where A, S' E GL(d, R) are affine matrices, t is a translation vector, [u] is the deformation grid, and Thounds are the bounds of the (implicit) identity grid [x] There are three benefits of this approach. First, the kernel avoids materializing any additional grids in HBM, reducing the memory overhead of the kernel from (n) to O(1) with no loss in runtime or accuracy. Second, when the warp (u)p is sharded across hosts in a distributed setting, the identity grid [x]o needs to be sharded correctly too. Since the identity grid implicitly defined by its bounds Thounds = (xmin- #max) E 2 implementation can be easily used in distributed optimization setting without instantiating partial shards xl],,- Finally, the matrix is used to rescale the deformation field to sample from the coordinates of the sharded images In which lie the grid Qi instead of Ω (see Section I.2) without initializing additional memory. The backward pass is very similar to the existing Py PyTorch implementation, with the exception of the gradient of the affine matrix. We discuss the derivation and pseudocode of the forward and backward pass in the Section H.
3.2
IMPLICIT PARZEN WINDOWING MUTUAL. INFORMATION
Mattes Mutual Information (MI) is one of the most commonly used loss functions for multimodal image matching (Chen et al., 2022; Avants et al., 2009; Mattes et al., 2001). For random variables X and Y. MI the KL divergence between the joint distribution P(X, Y) and product of marginal
AA GPU's memory hierarchy spans multiple tiers: registers (per-thread, single-cycle), shared memory/L1 cache (on-chip, tens of KB, low latency within block). L2 cache (MBs, shared across SMs, moderate latency), and global memory (HBM). Our work focuses reducing HBM usage for key non-GEMM operations used in image registration, by maximizing register and shared memory usage while minimizing global memory traffic.
4

---

Published as a conference paper at ICLR 2026
spAnce
Sin.
(M-+it
Sebotty
Ieln,
usage (in GB)
10
Mins
Ikln,
Iclo,
CL ARE
10
SynthMorph
unigadiCON
VFA
Icin
Ieloa
10"
10
10
10"
10
Voxel Size (x109)
updade
Figure 3: Left: Overview of our distributed framework. GridParallel (GP) shards the fixed and moving images (F, M) and the warp field [u] across multiple GPUs. Yellow blocks and arrows denote synchronized halo boundaries between GPUs, enabling smoothing on images and warp fields without an allgather. The ring sampler (violet) computes interpolated image shards on the fly, avoiding materialization of the full moving image. We then compute losses (MSE, LNCC, MI), compute gradients w.r.t. each warp shard, apply Sobolev regularization with GP, and update shards by gradient descent. Right: Scaling efficiency compared to deep methods and CLAIRE (Mang et al., 2019). a distributed registration method. Most SOTA deep learning baselines require orders-of-mant of-mannitude more memory for the same problem size and scalability is limited to a single GPU (dotted line). Our framework scales to arbitrarily large prohlem sizes while using about 5x less memory than CLAIRE.
distributions P(X) P(Y) of the intensities of the two images. For image matching, X and Y are the pixel intensities for the images I, J. The distributions are estimated using kernel density estimator:
Pv( N Σ
- n(U-TA)) Pu.J)(u, w) N k 1 Tt)r(w - Jk))
1 Σ
where & is a kernel function of choice. Common choices of x are the Gaussian (Guo, 2019) and 3rd order B-Spline kernels (Thévenaz & Unser, 2000). To empirically compute the KL divergence, the distributions Eq. (2) are discretized over B equally spaced bins on the domain of u E L, v E J. However, to compute the joint histogram of size B this method requires materializing the entire Parzen Block Ve (j. k)) = x(b, I&) of size 2kp BN, where kp is a kernel constant. Since N >>>>> B (B is typically chosen to be 32), this operation becomes significant memory bottleneck for large N. For instance, a typical clinical image volume (N x 30MB) with 32 bins will consume 7.5GB of HBM a significantly huge cost that grows much faster for larger problems.
Our efficient implementation leverages the fact that B is small to avoid materializing the tensors ,, v E BxX altogether and use high-throughput shared memory to compute and accumulate the histogram entries and partial gradients for each image pixel. We provide the detailed derivation in Section G. This leads to an efficient implementation that consumes ((1) additional HBM instead of O(N) (holding B constant). This leads to upto 98% lesser HBM usage for images considered in experiments, and an asymptotic 100% reduction in HBM usage for large images (Fig. 7tto-righht).
3.3 EFFICIENT IMPLICIT FUSED CROSS-CORRELATION
Local Normalized Cooss-oorellaion (LNCC) is used ubiquitously in signal and image processing as a similarity metric. In deformable image registration, it is used as robust similarity function to compare anatomical similarities (Chen et al., 2022; Hoffmann et al., 2021; Avants et al., 2008b; Wu et al., 2024). Most LNCC implementations are memory- bound due to the large number of intermediate variables. Our analysis in Section F shows that the computational graph adds 16x HBM overhead, and upto another 16x HBM overhead for computing gradients with respect to all intermediates.
To avoid these huge memory overheads, we fuse all the intermediate computation in fused kernel. Our fused forward pass requires only 5x memory for storing all intermediates (I, J, 12, J2 IJ convolved with matrix w). In Section F we analytically derive the gradient and show that the input gradients can be computed by modifying the saved intermediates in-place. This leads upto 76.5% reduction in memory (see Table 3) and outperforms even torch. compi le implementations.
EXTENDING IMAGE REGISTRATION TO MULTIPLE GPUs
Our composite implicit grid sampler and improved loss functions allows optimizing problems with image sizes that are upto two magnitudes larger than other baselines on single A6000 GPU (Fig.
5

---

Published as a conference paper at ICLR 2026
Image
shards across GPU
M
Sharded warp
Agur egoted moved Image
Shard
Shards
Shard
Shard H
(a)
(b)
I(((()))
(c)
Figure 4: (a) Neighboring coordinates in the warp field may refer to pixel locations on arbitrary image shards due to the deformable nature of the warp field, making distributed interpolation non-trivial. (b) Ring Sampler interleaves fetching of image shards and aggregating the partial sums of interpolated values, avoiding a eeemory-expensive allgather. (c) Bilinear Interpolation is decomposed into partial sums over image shards, which are accumulated with a ring topology communication, similar to Liu et al. (2024b).
However, many applications using mesoscopic and microscopic data require registration of images that do not fit on a single GPU. Inspired by distributed frameworks for LLM training (Shoeybi et al., 2019; Rajbhandari et al., 2020) and initial work on distributed image registration (Mang et al., 2019), we propose distributed framework that allows sharding large images across multiple GPUs to efficiently scale to arbitrarily large problem sizes with any similarity loss function.
Distributed Setting. For distributed registration with H hosts or GPUs, we partition the domain P(2)) = {12,, S2 12H} ... } such that [94] = N/H. Ω nΩ,
Vi * j and UgΩ, Ω. We use [x]a, . A[x]a, + t, and [ to denote the sharded tensors defined on domain 22p
4.1 GRID PARALLEL FOR BOUNDARY-SYNCHRONTPH IMAGE SHARDING
Parallel have been tremendously successful in distributed optimization by sharding large models and sequences across multiple GPUs (Shoeybi et al., 2019; Li et al., 2023; Liu et al., 2024b;a). However, these techniques work for transformer- like architectures and input sequences where the model parameters and activations do not require boundary synchronization. In contrast, image registration contains operations that require boundary synchronization between image and grid shards to perform mathematically correct convolutions. Examples of such operations include convolutions for calculating LNCC, total variation loss, Sobolev norm of the gradient and warp fields (Mang et al., 2019; Avants et al., 2008b; Beg et al., 2005).
Techniques like Tensor/Squencee/xprr/CCoooe
To enable these functionalities and complement existing parallelism techniques, we propose 'Grid abstraction on tensor. GP shards a tensor across hosts, stores the sharded
Parallel' (GP) dimension and bounds as metadata, and provides synchronization operations to augment the tensor with sufficient boundary padding from neighboring shards prior to performing convolution operation.
GP allows us to partition the fixed images, [u]. and the optimizer state [m,], [m2] essentially
sharding the entire problem across H/ hosts while allowing the user apply convolutional operations seamlessly. We compare the performance of GP with naive DTensorr sharding in Section D.
4.2 DISTRIBUTED RING SAMPLER
Despite the sharding in GP, the moving image M cannot be sharded across GPUs due to the random-access nature of the grid_sample operation applied on M. In general, the warp vector (p()) residing on GPU can point to coordinates that reside on the sharded image on GPU j for any j * í. Even for neighboring coordinates Zz. Fu [x], the coordinates sp(x,) and ((xx) can point to different shards jı # j2 # i. This is illustrated in Fig. 4(a). Keeping the entire moving image in memory limits the maximum problem size N < V. where V is the memory per GPU, regardless of the number of hosts H. However, we want the maximum problem size to scale with H. Therefore, we propose distributed grid_s d_s ampler that allows us to correctly interpolate the moving image with sharded images scattered across multiple hosts without performing an al allgather operation on the moving image.
Our approach leverages the key observation that (bi/tri)linear interpolation can be decomposed as an aggregate of partial sums of interpolated values on individual image shards. Fig. 4(b) illustrates this example. These individual image shards are sent across hosts in a ring topology, similar to Liu et al. (2024b), and the partial sum aggregated inplace. This operation only incurs an additional N/ H HBM overhead for fetching the sharded image from other hosts, scaling efficiently to arbitrary large

---

Published as a conference paper at ICLR 2026
(a) Performance comparison across methods and resolutions.
SynthMooph
fiddeeg
Aratoma
FROANS
VIA
CON-.d
TramNerph
CLARR
Resslution
Method
Avgilice Score
InvDire Score
AvgiiDe0 (mm)
Baseline
0.033
0.14:
1,587
Dice vs GPU Cost Tradeoff
Aratomis
0.790 0.035
0.39
01388
0.008 0.137
0.5
CLAIRE
0.776 1: 0.044
0.34 0.120
0.554 à 0.150
FiRANTS
0.822 a: 0.002
0,430
0.147
63300
0.6
ITK dreg
0,002 0.005
-
1.002 0.277
SymthMorph
0,801 1: 0.022
0,378 00133
0.405 0.008
0.7
TrunsMrphh
0,851 0.001
0.468 0.161
0:310 a: 0.064
UniGradICON
0.826 0,022
0,091
0.384
UniGadlCOON
0.815 0.020
0
0:419
Resoktion
VPA Oux
0.851 = 0.022
0.491 00102
0.321 2: 0.006
0,838 + 0.028
0.156 100-000 0.108
0.311 0.100
2.4
250
500
Buseline
0,580 4: 0.055
0 0.143
1,357
CLAIRE
Atotomis
0.738
0.779 2:
0 0.210
0
0.570 0 0:211
0.619
10a7
10-7
10-1
(CBB-brum)))
10"
10
10'
10
500gmm
ITKdiieg
PRRANT
0,841 1 0.031
0.066
0.489 0.163
0:340 2: 0.127
0 0.254
(b) Accuracy vs. GPU Compute Cost.
SynthMoophh
0.771 0.035
0.887 0.81
0.557 0.144
TraneMopph'
0.759+-0.02
0.300 0.175
(.021 0.127
Dice vs Wall- Hours Tadeoff
UniGradICON"
0.630:
0.183 0.122
1.251 0.202
0.9
UniGradiOON (jO)"
0.005-00.017
0.805
0.119 0.181
0.102 0.103
0.0
Onrs
0.822
0.26 0.180
0,258 0.000
Boscline
0.680 3 0.055
0.136 0.141
1.400 0.322
Dioa Score
Aratomis
CLAIRE
0.80
0.378 01133
0.570 0.11
200g
FiRANNI
ITK-drg
0.777 0.006
0.758 0.008
0.200 0.125
0341 0.100
0.013 0.101
0.029 a 0.29
Synthhlorph"
0.000
0260-0116
0.882
0.4
'InneMoeph'
0.680 ::
0.90 0.132
0.973 0 0:255
UnIGradICOON (BO)"
0 1: 0.002
0.86 0.07L
3,491 à: 3.195
10-3
30
UniGradiOON'
0,330 0.044
0.045 0.056
0.670
Wol-Cllock Hours
Ouss
VEA
0,855 4 0.029
0.724 0.000
0.281 0.216
0.697 = 0204
0.216 2 0.098
0.821
(c) Accuracy vs. Wall-clock Time.
Figure 5: Registration performance on Faux-OASIS dataset 1 mm, 500 pm, and 250 pm (native 250 pm); mean + std over pairs. higher is better; lower is better. HD90 values are reported using our cumulative definition (see Sec. K.2). (Green)/ (Yellow) = best/second; "= patch- patch-based
problem sizes for sufficiently large H. The detailed derivation and correctness of this operation is shown in Section I.
4.3 DISTRIBUTED Loss FUNCTIONS
Since the moved image and fixed image are sharded cross H hosts, the loss function must take this into account compute the loss function correctly.
Mean Squared Error (MSE). Since MSE is a per-pixel loss, we compute the individual MSE on host h and perform a reduce operation.
Localized Normalized Cross Correlation (LNCC). The LNCC computes per-pixel patch similarities
for each pixel, using a convolution over its neighbors. For sharded images, the patch statisties the boundary requires a boundary synchronization with its neighboring shards which is provided by our GP implementation. After computing the LNCC for all pixels in each shard, we perform another allreduce to compute the LNCC over the entire image.
Mutual Information (MI). The MI loss computes the joint histograms P(.,)) (I. y) and marginals (r()), (v)). However, these distributions are partial aggregates from the sharded images on each GPU, Eq. (2) can be rewritten as pr(v) = Σn N No ExER, x(v - -a)) (1 w)
Σn No computation. Performing an all llluuecce to compute the weighted average of these histograms (with weights /2/N) results This also leads to only a B² + 2B communication overhead regardless of N, making a distributed implementation highly practical.
Nr EtEo, <(v - [x)(w= Jb)) where the red terms correspond to the per-host histogram
valid and correct joint and marginal distributions over all hosts.
7

---

Published as a conference paper at ICLR 2026
Fixed (100pm FLASH)
Moved (250jm - 100pmm)
Moving (250 T1)
Figure 6: Qualitative comparison registration of 100pmm ex-vivo brain MRI (T1 + FLASH) image. Fine details like cerehellar white matter are visible macroscopic scales, but aligned 100pm. Fixed image size 1760 x 1760 x 1278. Best viewed zoomed in. More results Fig. 11.
5 EXPERIMENTS
Our primary goals are to (a) accelerate both optimization and neural network based registration workflows, and (b) solve significantly larger image registration problems. We show the efficacy of our method by accelerating existing registration workflows on standard clinical data. This is followed by optimizing multimodal registration task with more than 11.8B optimizable parameters, an unprecedented result in large-scale registration. We compare the performance and computational efficiency of our method with various stateoof-the-ae baselines simulated 250μmm ex-vivo brain MRI dataset, followed by ablations on various components of our framework.
Baselines. To accelcrate existing registration workflows, we compare against TransMorph (Chen et al., 2022) and FireANTs (Jena et al,,,, 2024a), which are state-of-the- the-art deep learning and optimization based registration frameworks respectively. methods explicitly designed for large-scale registration: ITKK-DReg (itk) (CPU-based) and CLAIRE (Mang et al., 2019) (multi-GPU), and several SOTA learning-basse approaches for clinical data SynthMorph (Hoffmann et al., 2021), Vector-Field Attention (Liu et al., 2024c), unigradICON (Tian et al., 2024) (with/without instance optimization), anatomix*Convex (Dey et al., 2025).
addition, we perform comparative evaluation with two
5.1 ACCELERATING EXISTING REGISTRATION WORKFLOWS AND ABLATIONS
For deep networks, we train TransMorph large under three loss configurations: (a) LNCC+Dice, (b) MI+Dice, and (c) shown in Table 1, we either use the vanilla Py Torch implementation (Baseline) or our kernels (Ours). For classical optimization, we benchmark runtime and memory against multiple LNCC backends (FireANTs, SahelMorpphTransMorog
aeassalling-nddsuuuriiin (Ashburner, 2007) +Dice. For each configuration
Fast LNCC, torch. compile, and Ours) and MI backends (PyTorch and Ours with and without torch. compile) Tables 1 and 4 and Fig. 12 show that during network training our kernels converge 6. faster with LNCC while using 16.5 less memory, and reduce MI memory usage by 24.7%. Despite being designed for very large images, the runtime and memory benefits are significant for clinical-scale data (i.e., 30MB for OASIS). Optimization frameworks see larger gains: FireANTs achieves up to 95.2% memory savings and 2.6 x speedup with MI, and a 7.5x speedup over FastLNCC (Jia et al., 2025) (and 2.9x over FireANTs' LNCC backend which applies separable convolutions Fast NCCC with 44-59% lower memory usage overall.
5.2 REGISTRATION A 100 MICRON EX-VIVO BRAIN MRI VOLUME
To showcase the efficacy of our method on real large scale images, we register 250μm in-vivo MRI image (Lüsebrink et al., 2017) to a 100 pmn ex-vivo FLASH human brain volume (Edlow et al., 2019). This represents an inverse problem with more than 1 11.2B optimizable parameters (compared to ~00 for clinical datasets), or 44.8GB of GPU memory. The entire problem does not fit on most GPUs, necessitating distributed multimodal registration. We optimize composite transform affine followed by a diffeomorphic mapping: details can be found in Section E. Multimodal deformable registration took ~58 seconds on 8 NVIDIA A6000 GPUs, which is unprecedented at this resolution. Fig. 6 shows qualitative results, highlighting the ability to register highly detailed structures such cerebellar white matter; these structures are not visible at macroscopic scales. The resultant advantages of performing registration at this scale can allow researchers to characterize the
8

---

Published as a conference paper at ICLR 2026
Abtision on Sackends soarneyy
Ablation on Mutual Information Backands Fylees
tackccnppa
QuB
tectccmppe
Der
Occhrenn Thee In
Veceee 11:0
Vie See Oinem Abeation on Affice Baciends dut
per Dimerake
Sit
per
Abletion on SKS Beckends
Ous Noo
Jhage
Six Deno
Diriciie
State
Numbo Mtan
Mtnder Sign
Figure 7: Ablations on key workhorse operations: LNCC, MI, gr dedddpsmpppss er, and scaling-and-squaring operations. Our fused kernels consume significantly less HBM and runtime.
neuroanatomy at microscopic resolutions and allow morphometric analysis of cortical layers and subcortical nuclei among other structures.
Registration accuracy in these studies measured using privately annotated fiducial markers, hindering reproducibility and comparability of methodological advances. Due to lack of scalable frameworks, most high-resolution studies simply run ANTs significantly downsampled resolution (Kleven et al., 2023; Mansour et al., 2025; Wang et al., 2020b; Kronman et al., 2024; Bogovic et al., 2020; Edlow et al., 2019) and upsample the warp field to the native resolution.
5.3 COMPARATIVE ANALYSIS ON A SIMULATED EX-VIVO BRAIN MRI DATASET
The faux-OASIE dataset To compare regis- tration performance at high resolutions and leverage existing methods baselines, we synthesize the faux-OASE dataset, which mimics the anatomical distribution of an MRI dataset at 250μm isotropic resolution (more details in Section K). At 250 the deformation field has 1.32B degrees of free- dom per image pair, compared to ~00 for OASIS.
Table 1: Accelerating Torph (Top) and FireANTs (Bottom) training with various computation backends.
Variant
Loss
Diffeumorphic
Training Time
GPU Mem (GB)
Val DSC
Baseline
20011
Ous
LNCC
27.84
16.95
87.23
Baseline
LNCC
171.42
21.28
36.35
LNCC
27.93
17.34
$7.09
Baseline
22.34
36.74
24.04
1680
Rackend
Dice
Runtins
Memory
LNCC
FIRANNTS
78.51 1 3.87
1.44 1 0.08
10M15 0.0
LNUC
FaSILNCU
76,96 +: 3.00
3.76 :: 0.16
1026.3 = 0.0
LNCC
VXM/TM
76,96 1. 3,60
$7.08 2.45
1418.5
LNCC
Luzch, cumgil
4.09
0.82 0.04
540.7
LNUC
Ours
78.67 :: 3.04
0.50 a: 0.01
577.5 22 0.0
Baselines and evaluation. All methods (including CLAIRE and FireANT's without FFDP) run out of memory at 250 res- olution. We proposed two modifications to deep learning based methods to enable them to work on this dataset: (a) inspired by several
MI
PyTrech
25.88 + 3.45
7.11 0037
1200.3 U:0
Fozah.compi..
MI
Ours
75.88 1: 3.44
2.90 J 0.16
577.5 0.0
MI
torch. comp compilc-Ous
75.90 1 3.47
2.95 1 0.166
667.3 0.0
high-resolution studies (Wang et al., 2020b; Mansour al., 2025; Edlow et al., 2019), we register the images at downsampled resolution, and then upsample the deformation field (b) inspired by several histology registration methods (Wodzinski et al., 2024; Lotz al., 2015; Liang et al., 2021), we perform patchwise registration and mosaicing of the final deformation. We compare the methods at three resolutions: 1mm, 500μm,, and 250pm. At Imm, the full image fits within patch, providing a baseline reference comparable to reported OASIS performance. At higher resolutions, patches are defined by each method's default input size with stride equal to 50% of the patch size. FireANTss augmented with FFDP is denoted as Ours. We report Dice, inverse-weightes Dice (InvDice; Mang et al. (2019)), and average Haussdorf distance capped at 90 percentile (AvgHD90). To compare efficiency, we measure both wall-clock time and GPU-hours.
Results. Fig. 5a summarizes performance metrics. At 1mm, most methods achieve performance consistent with their reported performance on OASIS, including VFA and TransMorph which were trained on the OASIS dataset with label supervision. At higher resolutions, nearly all methods degrade, especially for InvDice and HD90, which emphasize alignment of fine structures. In contrast, our method improves in accuracy: at 250jm,, we improve Dice by 18.1 points, InvDice by 31.6 points, and reduce AvgHD990 by 62. 1%. The correlation between resolution and performance is also

---

Published as a conference paper at ICLR 2026
observed in (Mang et al., 2019; Mang & Ruthotto, 2017; that patch-based methods degrade in performance at higher resolutions.
Nazib et al., 2018); in addition we verify
This degradation among patchwise methods is expected; histology-sti pipelines typically register consecutive slides with small deformations after affine alignment. At high res- olution, patching reduces anatomical context and the patches become progressively more out-of-distribution (see Fig. 19). Patchwise or downsampling strategies are therefore insuffi- cient for ultra-high resolution large-scale registration, and ex- isting deep methods cannot be repurposed to work at higher olutions efficiently. Accuracy-efficinny tradeoffs in Figs. 5b and 5c show that our method is Pareto-efficient compared to all other methods (CPU, deep learning, and distributed GPU methods), requiring up to 500 fewer GPU-hours compared to alternatives at 250pm..
Figure 8: Scaling and GP ablations.
Weak Scaling
Ring Sampler
Ring Sampler
Efticency
22.4%
50
Memory Consumption per
Ring Sampler -
With Ring Sampler
Memory per GPU (GB
30
25
5.4 ABLATION STUDIES
We ablate on the efficiency of various workhorse opera- tions used in image registration in Fig. 7 and Table 3. We compare our implementations to community-standant Torch implementation (Jia al., 2025; Chen al., 2022) and torch. compile le versions. For grid sampler and MI kernels, our kernels have ((1) extra HBM overhead in- stead of (NN) in the PyTorch implementation. For LNCC, implementation achieves
20
Number of GPUs
Py-
(a) Weak scaling and Per-GPU memory consumption of FFDP.
average speedup in the
forward pass by 5.22 x and 56.98x x in the backward pass. Our prridssapple also leads an efficient scaling-and- squaring operation, commonly used in deep learning registra- tion pipelines (Chen al., 2022), with a memory reduction of 50% compared to the baseline implementation. Scalabil- ity Analysis. We test the weak scaling of our distributed framework by registering synthetic images with increasing voxel sizes. For И GPUs, we instantiate an image pair size 700 700 700/ and shard the images, warp, and optimizer state across H GPUs. Fig. 8a shows weak scaling of FFDP with and without ring sampler. Without the ring sampler, the grid d.samppleeeee operation requires storing the moving image of size 700 700 700H each GPU, leading to peak HBM memory increasing linearly with H. This implies the framework would not scale to arbitrarily large problem sizes, regardless of cluster size H. Peak Memory consumption is independent of with the Ring Sampler, and scaling efficiency only minimally affected. Ablation on GP. We ablate the effect of GP by replacing it with DTensor sharding (no boundary sync). Figs. 8b, and 10 show that incorrect boundary synchronization leads to undesirable artifacts in the moved images, and reduces labelmap overlap.
(b) Qualitative ablation of GP synchro- nization in FFDP the IMOST mouse brain dataset (Tustison al., 2024). Red arrows highlight regions affected by incor- rect boundary effects due to no GP. See Fig. 10 for more examples.
6 CONCLUSION
We propose novel distributed framework for arbitrarily large image registration problems. Our work identifies and proposes 10-aware and distributed-friendil implementations of workhorse operations in image registration algorithms, enabling registration of images at arbitrarily large resolutions GPU. Our fused primitives demonstrate compelling results in both improving existing registration pipelines and scaling to arbitrarily large, multimodal problems pertinent in modern life science applications, that were previously infeasible without approximations. FFDP shows unprecedented registration capabilities that will enable researchers to leverage and effectively work with large-scale image volumes and unearth new insights leveraging the large resolution images.
single
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was supported by the National Institutes of Health (NIH) under grants EB031722 and NS135568.
REFERENCES
Allen brain atlas. URL htt.ht:://tttiiiie httt::://At....bbbr.mm.....o..
Itk-dreg: A framework for distributed, large-scale image registration. URL httpt:t//ttkkkeeeeee
readthedocs hitee//io///eee/eeeeeeee het//aattsss.tto
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, et Tensorflow: Large-scale machine learning heterogeneous distributed systems. arXiv preprint arXiv: 1103.0000055 2016.
Jesper LR Andersson, Mark Jenkinson, Stephen Smith, et al. Non-linear registration, aka spatial normalisation fmrib technical report tr07ja2. FMRIB Analysis Group of the University of Oxford, (1):1-22, 2007.
Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, al. Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International Conference on Architectural P Programming Languages and Operating Systems, Volume 2, pp. 929-947, 2024.
for
ANTXX. Antsx:
Advanced normalization tools (ants). URL ANTs. GitHub repository.
ANTSSX/
John Ashburner. A fast diffeomorphic image registration algorithm. Neuroimage, 38(1):95-113, 2007.
John Ashburner. Spm: history. Neuroimage, 62(2):791-800, 2012.
B. B. Avants, C. L. Epstein, M. Grossman, Jnd C. Gee. Symmetric diffeomorphic image registration with cross-correlation: evaluating automated labeling of elderly and neurodegenerative brain. Medical Image Analysis, 12(1):26-41, February 2008a. ISSN 1361-8423, doi: 10. 10.1016/j.media. 3007.06.00.
B. B. Avants, C. L. Epstein, M. Grossman, and J. C. Gee. Symmetric diffeomorphic image registration with cross-correlation: Evaluating automated labeling of elderly and neurodegenerative brain. Medical Image Analysis. 12(1):26-41, February 2008b. ISSN 1361-8415, doi: 10.1016/ j.media.2001 2007 06.04. URL https::////. wwww. sciencedirect com science/article pii/s1361 41517000666
Brian B. Avants, P. Thomas Schoenemann, and James C. Gee. Lagrangian frame diffeomorphic image registration: Morphometric comparison of human and chimpanzee cortex. Medical Image Analysis, 10(3):397-412, June 2006. ISSN 13618415. doi: 10.10166j.ecdaa2000.0033
URL
https: htte:://iinninggnge elsevier wtw/:coto/////eeeee
Brian B Avants, Nick Tustison, Gang Song, et al. Advanced normalization tools (ants). Insight j
(365):1-35, 2009.
Ramsey D Badawi, Hongcheng Shi, Pengcheng Hu, Shuguang Chen, Tianyi Xu, Patricia M Price, Yu Ding, Benjamin A Spencer, Lorenzo Nardo, Weiping Liu, et al. First human imaging studies with the explorer total-body pet scanner. Journal of Nuclear Medicine, 60(3):299-303, 2019.
Guha Balakrishnan, Amy Zhao, Mert R. Sabuncu, John Guttag, and Adrian V. Dalca. VoxeIMorph: A
Learning Framework for Deformable Medical Image Registration. IEEE Transactions on Medical Imaging, 1788) August 2019. ISSN 0278-0062, 1558-254X. doi: 10. 10.1109/TMM.200 2897538. URL http: / arxiv. orgoorgggao.... abs / 1809. 9. 05231. 10X.02809.00231 [cs].
Balchandani and TP Naidich. Ultra-high-fiel mr neuroimaging. American Journal of Neuroradiol- ogy,
2015.
11

---

Published as a conference paper at ICLR 2026
Erin S Beck, Pascal Sati, Varun Sethi, Tobias Kober, Blake Dewey, Pavan Bhargava, Govind Nair, Irene Cortese, and Daniel Salo Reich. Improved visualization of cortical lesions in multiple sclerosis using 7t mp2rage. American Journal of Neuroradiology, 39(3):459 466, 2018.
M Faisal Beg, Michael Miller, Alain Trouvé, and Laurent Younes. Computing large deformation metric mappings via geodesic flows of diffeomorphisms. International journal of computer vision, 61:139-157, 2005.
Ganesh Bikshandi and Jay Shah. case study in cuda kernel fusion: Implementing flashattention-2 on nvidia hopper architecture using the cutlass library. arXiv preprini arXiv:2312. 2023,
Benjamin Billot, Douglas N Greve, Oula Puonti, Axel Thielscher, Koen Van Leemput, Bruce Fischl, Adrian Dalca, Juan Eugenio Iglesias, et al. Synthseg: Segmentation brain mri scans of any contrast and resolution without retraining. Medical image analysis, 86: 06:102789, 2023.
John Bogovic, Hideo Otsuna, Larissa Heinrich. Masayoshi Ito, Jennifer Jeter. Geoffrey Meissner, Aljoscha Nern, Jennifer Colonell, Oz Malkesman, Kei Ito, et al. An unbiased template of the drosophila brain and ventral nerve cord. Plos one, (5(12):e0336495, 2020.
Katy Börner, Andreas Bueckle, Bruce W Herr, Leonard E Cross, Ellen M Quardokus, Elizabeth G Record, Yingnan Ju, Jonathan C Silverstein, Kristen M Browne, Sanjay Jain, et al. Tissue registration and exploration user interfaces in support of a human reference atlas. Communications Biology, 5(1):1369, 2022.
Bella E Brezovec, Andrew B Berger, Yukun A Hao, Feng Chen, Shaul Druckmann, and Thomas R Clandinin. Mapping the neural dynamics of locomotion across the drosophila brain. Current Biology, (4(4):710-724 2024.
Xiaohuan Cao, Jianhua Yang, Jun Zhang, Dong Nie, Minjeong Kim, Qian Wang, and Dinggang Shen. Deformable image registration based on similarity-steered cnn regression. In Medical Image Computing and Computer Assisted Intervention- MICCAI 2017: 20th International Conference, Quebec City, QC, Canada, September 11-13, 2017, Proceedings, Part I 20, pp. 300-308. Springer, 2017.
Adrià Casamitjana, Matteo Mancini, Fleanor Robinson, Loic Peter, Roberto Annunziata, Juri Althonayan, Shauna Crampsie, Emily Blackburn, Benjamin Billot, Alessia Atzeni, et al. probabilistic histological atlas of the human brain for mri segmentation. Nature, pp. 1-8, 2025.
Junyu Chen, Eric C. Frey, Yufan He, William P. Segars, Li, and Yong Du. TransMorph: Transformer for unsupervised medical image registration. Medical Image Analysis, 82:102615, November 2022. ISSN 13618415. doi: 10.1016/j.mm www.mmaa..22.212222222222222222. 10480. arXiv:2111.10488 [cs, eess].
URL Stthhtp:/// /ttp/varivvvv//v//////////////aa org/abs /2111.
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. TVM: An automated End-to-End optimizing compiler for deep learning. In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), pp. 578-594, Carlsbad, CA, October 2018. USENIX Association. ISBN 978-1-9391333008-3 URL https:/ hattps://wwert //www uswnix. org/ con.ference/ httpaa/eppeeeeaaaeeeeeeeee chen.
Gary Christensen and Hans J Johnson. Consistent image registration. IEEE transactions on medical
imaging, 20(7):568-582. 2001.
Gilberto Corso, Gabriel MF Ferreira, and Thomas M Lewinsohn. Mutual information general measure of structure in interaction networks. Entropy, 22(5):528, 2020.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint wtxXiv2307.08691 2023.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in neural information processing systems, 35: 16344-16359, 2022.
12

---

Published as a conference paper at ICLR 2026
Sandhitsu R Das, Brian B Avants, Murray Grossman, and James C Gee. Registration based cortical thickness measurement. Neuroimage, 45(3):867-879, 2009.
Chris Davis and A Murat Maga. Image registration and template based annotation of great ape skulls. In American Journal of Physical Anthropology, volume 165, pp. 60-61. WILEY 111 RIVER ST, HOBOKEN 07030-5774, NJ USA, 2018.
Bob D De Vos, Floris Berendsen, Max A Viergever, Hessam Sokooti, Marius Staring, and Ivana lguum A deep learning framework for unsupervised affine and deformable image registration. Medical image analysis, 52:128 111128884,, 2019.
Neel Dey, Benjamin Billot, Hallee E. Wong, Clinton Wang, Mengwei Ren, Ellen Grant, Adrian V Dalca, and Polina Golland. Learning general-purpose biomedical volume representations using randomized synthesis. In The Thirteenth International Conference on Learning Representations, 2025. URL attpp::/openzeweee wwweeviwwwwweeeeeeeeeeeeeeeeeeee www.nnee/ooouua htt.//ormmmicc ieeee-idddxxcse 2im.eiidddiiiiiieee VuN.
Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He. Flex attention: A programming model for generating optimized attention kernels. arXiv preprint arXiv:2412.0549. 2024.
Carmen Echávarri, P Aalten, Harry BM Uylings, HILL Jacobs, Pieter Jelle Visser, EHBM Gronenschild, FRJ Verhey, and S Burgmans. Atrophy in the parahippocampal gyrus as an early biomarker of alzheimer's disease. Brain Structure and Function, 215(3):265-27) 2011.
Brian L Edlow, Azma Mareyam, Andreas Horn, Jonathan R Polimeni, Thomas Witzel, M Dylan Tisdall, Jean C Augustinack, Jason P Stockmann, Bram R Diamond, Allison Stevens, et al. 7 tesla mri of the ex vivo human brain 100 micron resolution. Scientific data, 6(1):244, 2019.
Andrea Esquivel. Andrea Ferrero, Achille Mileto. Francis Baffour, Kelly Horst, Prabhakar Shantha Rajiah, Akitoshi Inoue, Shuai Leng, Cynthia McCollough, and Joel G Fletcher. Photon-counting detector et: key points radiologists should know. Korean journal of radiology. 23(9):854, 2022.
Fenja Falta, Christoph Großbröhmer, Alessa Hering, Alexander Bigalke, and Mattias P Heinrich. Lung250m-4b: A combined 3d dataset for CT- and point louedbased intra-patient lung registration. In Thirty-sevnnth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023. https:/ httpt::ttpp/eeeeeeeeetee openreview net/ fo derumii==FOddsguul
Miriam Friedel, Matthijs C van Fede, Jon Pipitone, M Mallar Chakravarty, and Jason P PLerch. Pydpiper: flexible toolkit for constructing novel registration pipelines. Frontiers neuroinformatics, 8:67, 2014.
Sarah F Frisken. Surfacenets for multi-label segmentations with preservation of sharp boundaries. The Journal of computer graphics techniquex, 1 1(1):34, 2022.
Brett M Frye, Suzanne Craft, Thomas C Register, Jeongchul Kim, Christopher T Whitlow, Richard Barcus, Samuel N Lockhart, Kiran Kumar Solingapuram Sai, and Carol A Shively. Early alzheimer's disease-like reductions gray matter and cognitive function with aging in nonhuman primates. Alzheimer's s & Dementia: Translational Research & Clinical Interventions, 8(1):e122844 2022.
Davide Gambarotto, Fabian U Zwettler, Maeva Le Guennec, Marketa Schmidt-Cernohorske Denis Fortun, Susanne Borgers, Jörn Heine, Jan-Gero Schloetel, Matthias Reuss, Michael Unser, et al. Imaging cellular ultrastructures using expansion microscopy (u-exm). Nature methods, 16(1): 71-74, 2019.
James C Gee, Martin Reivich, and Ruzena Bajosy, Elastically deforming a three-dimensional atlas to match anatomical brain images. 1993.
Hui Gong, Dongli Xu, Jing Yuan, Xiangning Li, Congdi Guo, Jie Peng. Yuxin Li, Lindsay A Schwarz, Anan Li, Bihe Hu, et al. High-throughpe dual-colour precision imaging for brain- connectome with cytoarchitectonic landmarks at the cellular level. Nature communications, 711):2214,, 2016.
Maged Goubran, Cathie Crukley, Sandrine De Ribaupierre, Terence M Peters, and Ali R Khan. Image registration of vivo mri to sparsely sectioned histology of hippocampal and ncocortical temporal lobe specimens. Neuroimage, 83:770-781, 2013.
13

---

Published as a conference paper at ICLR 2026
Courtney K Guo. Multi-modal image registration with unsupervised deep learning. PhD thesis, Massachusetts Institute of Technology. 2019.
Tripti Gupta, Gregory D Marquart, Eric J Horstick, Kathryn M Tabor, Sinisa Pajevic, and Harold A Burgess. Morphometric analysis and neuroanatomical mapping of the zebrafish brain. Methods, 150:49-62, 2018.
Alessa Hering, Keelin Murphy, and Bram van Ginneken. Learn2reg challenge: Ct lung registration- training data, 2020. URL https: /ooi. wtw.ooo.o.....ooooo org/10. w2211/2eooo....... 3835682.
Alessa Hering, Lasse Hansen, Tony CW Mok, Albert CS Chung, Hanna Siebert, Stephanie Häger, Annkristin Lange, Sven Kuckertz, Stefan Heldmann, Wei Shao, al. Learn2reg: comprehensive multi-task medical image registration challenge, dataset and evaluation in the era of deep learning. IEEE Transactions on Medical Imaging. 42(3):697-712, 2022.
Malte Hoffmann, Benjamin Billot, Douglas N Greve, Juan Eugenio Iglesias, Bruce Fischl, and Adrian V Dalca. Synthmorph: learning contrast-invariant registration without acquired images. IEEE transactions on medical imaging, 41(3):543 558, 2021.
Junhao Hu, Weijie Gan, Zhixin Sun, Hongyu An, and Ulughek S. Kamilov. Plug-and-Ping Image Registration Network, March 2024. URL arXiv:2310.04297 [eess].
Phillip Isola, Daniel Zoran, Dilip Krishnan, and Edward H Adelson. Crisp boundary detection using pointwise mutual information. In European conference on computer vision, pp. 799-814. Springer, 2014.
Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Reza Yazdani Aminadabi, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. System optimizations for enabling training of extreme long sequence transformer models. In Proceedings of the 43rd ACM Symposium on Principles of Distributed Computing. PODC '24, pp. 121-130, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400706684, doi: 0411145/3662158.6666666
URL wttptttps///cc / /toi. or1/1...g111...00 1145/115550 3 36631866688. 3 3662806.
Rohit Jena, Pratik Chaudhari, and James C Gee. Fireants: Adaptive riemannian optimization for multi-scale diffcomorphic registration. arXiv preprint arXiv:2444.01249 2024a.
Rohit Jena, Deeksha Sethi, Pratik Chaudhari, and James C. Gee. Deep learning in med- ical image registration: Magic or mirage? grave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang (eds.), Advances in Neural In- formation Processing Systems, volume 37, pp. 108331-108353 Curran Associates, Inc., 2024b. URL https::tp//c htttts::/pproceediing wtwp:eeeediinn....n.eeeeeeeeeeee neurips. #ttle//cfe22000eccc1118055ee3663
In A. Globerson, L. Mackey, D. Bel-
wwwerrps.cc/pappeepeee cc/ wttp//pppr_iiieeee hacperr_filleeeeeeeeeeeeeeeessee www.fi/pppsppee//ee/../ Bett223335--PPpere- Conference .df.
Rohit Jena, Pratik Chaudhari, and James C. Gee. Deep implicit optimization enables robust learnable features for deformable image registration. Medical Image Analysis, 103: 2025, ISSN 1361-8415. doi: https://doi.org 10. http://0...1....emia..0000000 htt.fff.mmmeaa222200033777777771 URL https: // //www. www.ecedieeect...... com/
Xi Jia, Joseph Bartlett, Tianyang Zhang, Wenqi Lu, Zhaowen Qiu, and Jinming Duan. U-net vs transformer: Is u-net outdated in medical image registration? arXiv preprint arXiv:2208.04933 2022.
Xi Jia et al. A naive trick to accelerate training of Ince-based deep image registration models. Preprints, February 2025. doi: 10.29944/peprrrnn222222222200200
Bailiang Jian, Jiazhen Pan, Morteza Ghahremani, Daniel Rueckert, Christian Wachinger, and Benedikt Wiestler. Mamba? catch the hype or rethink what really helps for image registration. arXiv preprint 2rXiv:2407.19274 2024.
Hanna Jonsson, Simon Ekstrom, Robin Strand, Mette A Pedersen, Daniel Molin, Hakan Ahlstrom, and Joel Kullberg. An image registration method for voxel-wise analysis of whole-body oncological pet-ct. Scientific Reports, 12(1):18768, 2022.
14

---

Published as a conference paper at ICLR 2026
Michael I Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2):181-214, 1994.
Ankita Joshi and Yi Hong. Diffeomorphic Image Registration using Lipschitz Continuous Residual Networks. pp. 13.
Justin W Kenney, Patrick Steadman, Olivia Young, Meng Ting Shi, Maris Polanco, Saba Dubaishi, Kristopher Covert, Thomas Mueller, and Paul W Frankland. A 3d adult zebrafish brain atlas (azba) for the digital age. Elife, 10re69988, 2021.
Stefan Klein, Marius Staring, Keelin Murphy, Max A Viergever, and Josien PW Pluim. Elastix: a toolbox for intensity- based medical image registration. IEEE transactions medical imaging, 29 (1):196-205, 2009.
David Kleinfeld, Arjun Bharioke, Pablo Blinder, Davi D Bock, Kevin L Briggman, Dmitri B Chklovski,, Winfried Denk, Moritz Helmstaedter, John P Kaufhold, Wei-Chung Allen Lee, et al. Large-scale automated histology in the pursuit of connectomes. Journal of Neuroscience, 31(45): 16125-16138, 2011.
Heidi Kleven, Ingvild Bjerke, Francisco Clascá, Henk Groenewegen, Jan G Bjaalie, and Trygve Leergaard. Waxholm space atlas of the rat brain: 3d atlas supporting data analysis and integration. Nature methods, 20( I): 1822-1829, 2023.
Julian Krebs, Tommaso Mansi, Hervé Delingette, Li Zhang, Florin C Ghesu, Shun Miao, Andreas K Maier, Nicholas Ayache, Rui Liao, and Ali Kamen. Robust non-rigid registration through agent-based action learning. In Medical Image Computing and Computer Assisted Intervention- MICCAI 2017: 20th International Conference, Quebec City, QC, Canada, September 11-13, 2017, Proceedings, Part I 20, pp. 344 352. Springer, 2017.
Fae N Kronman, Josephine K Liwang. Rebecca Betty, Daniel J Vanselow, Yuan-Ting Wu, Nicholas J Tustison, Ashwin Bhandiwad, Steffy B Manjila, Jennifer A Minteer, Donghui Shin, et al. De- velopmental mouse brain common coordinate framework. Nature communications, 15(1):9072, 2024.
Kwame S. Kutten, Joshua T. Vogelstein, Nicolas Charon, Li Ye, Karl Deisseroth M.D., and Michael I. Miller. Deformably registering and annotating whole CLARITY brains an atlas via masked LDDMM. In Peter Schelkens, Touradj Ebrahimi, Gabriel Cristóbal, Frédéric Truchetet, and Pasi Saarikko (eds.), Optics, Photonics and Digital Technologies for Imaging Applications IV, volume 9896, pp. 989616. International Society for Optics and Photonics, SPIE, 2016. doi: 1,21177/2.2227444, URL https: wttp///1///....11......111 2227444,
Joel Lamy-Poirier. Breadth-first pipeline parallelism. Proceedings of Machine Learning and Systems, 5:48-67,2022
Leo Lebrat, Rodrigo Santa Cruz, Frederic de Gournay, Darren Fu, Pierrick Bourgeat, Jurgen Fripp, Clinton Fookes, and Olivier Salvado. CorticalFlow: A Diffeomorphic Mesh Transformer Network for Cortical Surface Reconstruction. In Advances in Neural Information Processing Systems, vol- ume 34, pp. 29491-29505. Curran Associates, Inc., 2021. URL https: //apttrss///ppp.....p/ nips.cc. cc/ html.
paper/20215
Dacheng Li, Rulin Shao, Anze Xie, Eric P. Xing, Xuezhe Ma, Ion Stoica, Joseph Gonzalez, and Hao Zhang. DISTFLASHATTN; Distributed memory-efficient attention for long-context LLMs training. In First Conference on Language Modeling, 2024. URL net/Eorunt wrwuunndddUUDEXyPD...
Shenggui Li, Fuzhao Xue, Chaitanya Baranwal, Yongbin Li, and Yang You. Sequence parallelism: Long sequence training from system perspective. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics ( (Volume 1: Long Papers). Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 00.18653/v1/2000 acl-
long 134. URL https: / ac anthology org/
2023 3.23.. lo...oo..../.... 134/.
15

---

Published as a conference paper at ICLR 2026
Cher-Wei Liang. Ruey-Feng Chang, Pei-Wei Fang, and Chiao-Min Chen. Improving algorithm for the alignment consecutive, whole-slide, immunohistochemier section images. Journal of Pathology Informatics, 12(1):29, 2021.
Jiayong Liang, Xiaoping Liu, Kangning Huang, Xia Li, Dagang Wang, and Xianwei Wang. Automatic registration of multisensor images using an integrated spatial and mutual information (smi) metric. IEEE transactions on geoscience and remote sensing, 52(1):603- 615, 2013.
Aixin Liu, Bei Feng. Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint 29X4:22412.19437, 2024a.
Hao Liu, Matei Zaharia, and Pieter Abbeel. Ringattention with blockwise transformers for near-infinite context. In The Twelfth International Conference Learning Representations, 2024b. URL
https: http:///ppeeeeieeeeeeeeee net forum?id--ssREpiie
Yihao Liu, Junyu Chen, Lianrui Zuo, Aaron Carass, and Jerry L Prince. Vector field attention for deformable image registration. Journal of Medical Imaging. 116):064011-04400) 2024c.
Josephine K Liwang, Hannah C Bennett, Hyun-Jae Pi, and Yongsoo Kim. Protocol for using serial two-photon tomography to map cell types and cerebrovasculature at single-cell resolution in the whole adult mouse brain. STAR protocols, 1(2)410 2023.
Josephine Liwang, Fae N Kronman, Hyun-Jae Pi, Yuan-Ting Wu, Daniel Vanselow, Steffy B Manjila, Deniz Parmaksiz, Donghui Shin, Yoav Ben-Simon, Michael Taormina, al. epdevatlas: mapping gabaergic cells and microglia in the early postnatal mouse brain. Nature Communications, 16(1):9538, 2025.
William Lorensen and Harvey E Cline. Marching cubes: A high resolution surface construction algorithm. In Seminal graphics: pioneering efforts that shaped the field, pp. 347-353. 1998.
Johannes Lotz, Janine Olesch, Benedikt Müller, Thomas Polzin, P Galuschka, JM Lotz, Stefan Heldmann, Hendrik Laue, Margarita Sanzzlez-Vallinaae Arne Warth, et al. Patch-based nonlinear image registration for gigapixel whole slide images. IEEE Transactions on Biomedical Engineering, 6319):1812-1811 2015.
Xin Luo, Zhigang Liu, Mingsheng Shang, Jungang Lou, and MengChu Zhou. Highly-accurate community detection via pointwise mutual Coormation-incorooni symmetric non-negative matrix factorization. IEEE Transactions Network Science and Engineering, 6(1):463-476 2021. doi: 10. N0.1109//NSE.2200030000000
Falk Lüsebrink, Alessandro Sciarra, Hendrik Mattern, Renat Yakupov, and Oliver Speck. TI-weighted in vivo human whole brain mri dataset with ultrahigh isotropic resolution of 250 p.m Scientific data, 4(1):1-12, 2017.
Mads AJ Madsen, Vanessa Wiggermann, Stephan Bramow, Jeppe Romme Christensen, Finn Sellebjerg, and Hartwig Siebner. Imaging cortical multiple sclerosis lesions with ultra-high field mri. Neurolmage: Clinical, 32:102847, 2021.
Andreas Mang and Lars Ruthotto. A lagrangian eauss-newton-rrylor solver for mass-and intensity- preserving diffeomorphic image registration. SIAM Journal on Scientific Computing. 39(5): B860-B885, 2017.
Andreas Mang, Amir Gholami, Christos Davatzikos, and George Biros. CLAIRE: A distributed- memory solver for constrained large deformation diffeomorphic image registration. SIAM Journal on Scientific Computing, 41(5):C548-C584 January 2019. ISSN 1064-8275. 1095-7197. doi: 00.113718812078111 URL http: wwwpiaxxvvvvoooaaaaaaa org /abs//b////// 11011100880 04487. arXiv:1808.000 1808.04487 [cs, math].
Harrison Mansour, Ryan Azrak, James J Cook, Kathryn Hornburg. Yi Qi, Yuqi Tian, Robert W Williams, Fang-Cheng Yeh, Leonard E White, and G Allan Johnson. The duke mouse brain atlas: Mri and light sheet microscopy stereotaxic atlas of the mouse brain. Science Advances, 1 11(18): cadq8089, 2025.
16

---

Published as a conference paper at ICLR 2026
Daniel S Marcus, Tracy H Wang, Jamie Parker, John G Csernansky, John C Morris, and Randy L Buckner. Open access series of imaging studies (oasis): cross-sectionss mri data in young. middle aged, nondemented, and demented older adults. Journal of cognitive neuroscience, 19(9):
1498-1507, 2007.
Gregory D Marquart, Kathryn M Tabor, Eric Horstick, Mary Brown, Alexandra K Geoca, Nicholas F Polys, Damian Dalle Nogare, and Harold A Burgess. High-precisten registration between zebrafish brain atlases using symmetric diffeomorphic normalization. GigaScience, 6(8):gix056, 2017.
David Mattes, David R Haynor, Hubert Vesselle, Thomas K. Lewellyn, and William Eubank. Nonrigid multimodality image registration. In Medical imaging 2001: image processing, volume 4322, pp. 1629-1620. Spie, 2001.
Michael P Milham, Lei Ai, Bonhwang Koo, Ting Xu, Céline Amiez, Fabien Balezeau, Mark G Baxter, Erwin LA Blezer, Thomas Brochier, Aihua Chen, et al. An open resource for non-human primate imaging. Neuron, 100(1):61-74, 2018.
Tony C. W. Mok and Albert C. Chung. Large Deformation Diffeomorphic Image Registration with Laplacian Pyramid Networks, June 2020. URL http::/awxw...rgg.///// 2006. arXiv:2006. 6188 [cs, eess].
2ttp:////.bbb/0..61116.... 16148.
Tony CW Mok and Albert Chung. Affine medical image registration with coarse-to-fine vision transformer. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 08835-20844, 2022.
Keelin Murphy, Bram Van Ginneken, Joseph M Reinhardt, Sven Kabus, Kai Ding, Xiang Deng, Kunlin Cao, Kaifang Du Gary E Christensen, Vincent Garcia, et al. Evaluation of registration methods on thoracic ct: the empire challenge. IEEE transactions medical imaging, 30(11): 1901-1920, 2011.
Abdullah Nazib, James Galloway, Clinton Fookes, and Dimitri Perrin. Performance of registration tools on high-resolution 3d brain images. In 2018 40th Annual International Conference of the IEEE Engineering in Medicine and Biology Society (EMBC), pp. 566-569, 2018. doi: 10.1100/MMBBC..218883322
OpenAl. Triton. https::////.. httpp::/ppenaa..cm///ceeoc/e/// ctm hptppaaa..om//eeeee/tttttoo/oooo 2021.
Hanchuan Peng, Phuong Chung, Fuhui Long, Lei Qu, Arnim Jenett, Andrew M Seeds, Eugene W Myers, and Julie H Simpson. Brainaligner: 3d registration atlases of drosophila brains. Nature methods, 8(6):493 498, 2011.
Zhen Peng, Minnan Luo, Wenbing Huang, Jundong Li, Qinghua Zheng, Fuchun Sun, and Junzhou Huang. Learning representations by graphical mutual information estimation and maximization. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(1):722-737, 2023. doi: 10.1 1009/TPAM..2022.31144
PyTorch. Fusing convolution and batch norm using custom function. https://docs http://ccccp.pttcchchcc txq/tutors oittorrllss als entermediaa cust om function _en_coooooooyyoooee
conv _o__bnn_____ tustuuoorall. html, 2023. Created July 22, 2021; Last updated April 18, 2023; Last verified November 5, 2024.
PyTorch. Helion. https:t////t httpss::/ttttthht
2025.
Penghui Qi, Xinyi Wan, Guangxing Huang, and Min Lin. Zero bubble pipeline parallelism. arXiv preprint ar 102241,40, 2023.
Huaqi Qiu, Chen Qin. Andreas Schuh, Kerstin Hammernik. and Daniel Rueckert. Learning diffeomorphic and modality-invarial registration using b-splines. 2021.
Quan Quan, Qingsong Yao, Heqin Zhu, a Kevin Zhou. Igu-aug: Information-guider unsupervised augmentation and pixel-wise contrastive learning for medical image analysis. IEEE Transactions on Medical Imaging, 2024.
17

---

Published as a conference paper at ICLR 2026
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Compuring, Networking. Storage and Analysis, pp. 1-16, 2020. doi: 10.1 10.1109/SC41405.2000 www.:0...20000022202
Sadhana Ravikumar, Laura E. M. Wisse, Sydney Lim, Ranjit Itteerhh, Long Xie, Madigan L. Bedard, Sandhitsu R. Das. Edward B. Lee, M. Dylan Tisdall, Karthik Prabhakaran, Jacqueline Lane, John A. Detre, Gabor Mizsei, John Q. Trojanowski, John L. Robinson, Theresa Schuck, Murray Grossman, Emilio Sttacho-Pérula, Maria Mercedes Iñiguez de Onzoño Martin, María del Mar Arroyo Jiménez, Monica Muñoz, Francisco Javier Molina Romero, Maria del Pilar Marcos Rabal, Sandra Cebada Sánchez, José Carlos Delgado González, Carlos de la Rosa Prieto, Marta Córcoles Parada, David J. Irwin, David A. Wolk, Ricardo Insausti, and Paul A. Yushkevich, Ex vivo mri atlas of the human medial temporal lobe: characterizing neurodegeneration due to tau pathology. Acta Neuropathologica Communications, 9(1):173, 2021. ISSN 2051-5960. doi: 10.1186/ URL https://ddoooo org/ 11/.//1....0 / s40478 10444404720001100 021-00275-
Sadhana Ravikumar, Amanda E Denning, Sydney Lim, Eunice Chung, Niyousha Sadeghpour, Ranjit Ittyrrah, Laura EM Wisse, Sandhitsu R Das, Long Xie, John L Robinson, et al. Postmortem imaging reveals patterns of medial temporal lobe vulnerability tau pathology in alzheimer's disease. Nature Communications, 45:033, 2024.
Marc-Michel Rohé, Manasi Datar, Tobias Heimann, Maxime Sermesant, and Xavier Pennec. Svf-net: learning deformable image registration using shape matching. In Medical Image Computing and Computer Assisted Intervention- MICCAI 2017: 20th International Conference, Quebec City, QC, Canada, September 11-13, 2017, Proceedings, Part I 20, pp. 266-274. Springer, 2017.
Suman Sarkar and Biswajit Pandey. A study on the statistical significance of mutual information between morphology of galaxy and its large-scale environment. Monthly Notices of the Royal Astronomical Society, 497(4):4077-40090 2020.
Will Schroeder and Spiros Tsalikis. Really fast isocontouring. https: //ww: k www.ww.ree............ com/ June 13 2023. Kitware blog / announcement.
Jay Shah, Ganesh Bikshandi, Ying Zhang. Vijay Thakkar, Pradeep Ramani, and Dao. Flashattention-
3: Fast and accurate attention with asynchrony and low-precision. Advances in Neural Information Processing Systems, 77:68658-88633 2024.
Noam Shazeer, Azalia Mirhoseini, *Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- garsely-gated mixture-of-experts layer. In International Conference on Learning Representations, 2017. URL https: /oppnrevie.. net/ forrmm? i 1d=31 ckMDqlg.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training 2019.
multiingguulibbiiiene multi-billion parameter language models using model parallelism.
arXiv preprint
Henrik Skibbe, Muhammad Febrian Rachmadi, Ken Nakae, Carlos Enrique Gutierrez, Junichi Hata, Hiromichi Tsukada, Charissa Poon, Matthias Schlachter, Kenji Doya, Piotr Majka, et al. The brain/minds marmoset connectivity resource: An open-access platform for cellular-level tracing and tractography in the primate brain. PLoS biology. 21(6))e3002118 2023.
Hessam Sokooti, Bob De Vos, Floris Berendsen, Boudewijn PF Lelieveldt, Ivana lsgum, and Marius Staring. Nonrigid image registration using multi-scale 3d convolutional neural networks. In Medical Image Computing and Computer Assisted Intervention- MICCAI 2017: 20th International Conference. Quebec City, QC, Canada, September 11-13, 2017, Proceedings. Part I 20, pp. 232-239. Springer, 2017.
Benjamin Frederick Spector, Simran Arora, Aaryan Singhal, Arjun Parthasarathy, Daniel Y Fu, and Christopher Re. Thunderkittens: Simple, fast, and $ texttit Adorable) kernels. In The Thirteenth International Conference on Learning Representations, 2025. URL https: /pperrevie.. net/ bet/forume Eorum? www.iidoo0ooosss.
18

---

Published as a conference paper at ICLR 2026
Bane Sullivan and Alexander Kaszynski. PyVista: 3D plotting and mesh analysis through streamlined interface for the Visualization Toolkit (VTK). Journal of Open Source Software, 4(37):1450, May 2019. doi: 10.21 10.21105/joss.0444 01450. URL https: /ttpp//////o/o/////////ooooooooo / 10. http:////.......//....... 211.1/oossssso... 01450,
Kathryn M Tabor, Gregory D Marquart, Christopher Hurt, Trevor S Smith, Alexandra K Geoca, Ashwin A Bhandiwad, Abhignya Subedi, Jennifer L Sinclair, Hannah M Rose, Nicholas Polys, et al. Brain-wide cellular resolution imaging of ere transgenic zebrafish lines for functional Aarcuit-mapping Elife. 8:e42687, 2019.
Nouamane Tazi, Ferdinand Mom, Haojun Zhao, Phuc Nguyen, Mohamed Mekkouri, Leandro Werra, and Thomas Wolf. The ultra-scale playbook: Training LLMs GPU clusters, 2024. URL https: //buggingffae. co /oo/ll////// ceee-ultrasseee www.erraaaallllllllaaaaee paaaaleeellllobooooooo HuggingFace Blog.
Philippe Thévenaz and Michael Unser. Optimization of mutual information for multiresolution image registration. IEEE transactions image processing, (112):2833-209) 2000.
Lin Tian, Hastings Greer, Roland Kwitt, François Xavier Vialard, Raúl San José Estépar, Sylvain Bouix, Richard Rushmore, and Marc Niethammer. unigradicon: foundation model for medical image registration. In International Conference on Medicul Image Computing and Computer-Assisted Intervention, pp. 749-760. Springer, 2024.
Lazaros C Triarhou. Dopamine and parkinson's disease. In Madame curie bioscience database [internet]. Landes Bioscience, 2013.
Nicholas J Tustison and Brian B Avants. Explicit b-spline regularization in diffcomorphic image registration. Frontiers in neuroinformatiis, 7:39, 2013.
Nicholas J Tustison, Min Chen, Fae N Kronman, Jeffrey T Duda, Clare Gamlin, Mia G Tustison, Michael Kunst, Rachel Dalley, Staci Sorenson, Quanxin Wang, et al. The antsx ecosystem for mapping the mouse brain. bioRriv, pp. 2024-05, 2024.
Erdem Varol, Amin Nejatbakhsh, Ruoxi Sun, Gonzalo Mena, Eviatar Yemini, Oliver Hobert, and Liam Paninski. Statistical atlas of c. elegans neurons. In Medical Image Computing and Compuler Assisted Intervention-MICCAA 2020: 23rd International Conference, Lima, Peru, October 4-8, 2020, Proceedings, Part V 23, pp. 119-129. Springer, 2020.
Vivek Venkatachalam, Ni Ji, Xian Wang, Christopher Clark, James Kameron Mitchell, Mason Klein, Christopher J Tabone, Jeremy Florman, Hongfei Ji, Joel Greenwood, et al. Pan-neuronal imaging in roaming caenorhabditis elegans. Proceedings of the National Academy of Sciences, 113(8): E1082-E1088. 2016.
Guoxia Wang, Jinle Zeng, Xiyuan Xiao, Siming Wu, Jiabin Yang, Lujing Zheng, Zeyu Chen, Jiang Bian, Dianhai Yu, and Haifeng Wang. Flashmask: Efficient and rich mask extension of flashattention. arXiv preprint arXiv:2410.01359, 2024.
Quanxin Wang, Song-Lin Ding, Yang Li, Josh Royall, David Feng, Phil Lesnar, Nile Graddis, Maitham Naeemi, Benjamin Facer, Anh Ho, Tim Dolbeare, Brandon Blanchard, Nick Dee, Wayne Wakeman, Karla E. Hirokawa, Aaron Szafer, Susan M. Sunkin, Seung Wook Oh, Amy Bernard, John W. Phillips, Michael Hawrylycz. Christof Koch, Hongkui Zeng, Julie A. Harris, and Lydia Ng. The Allen Mouse Brain Common Coordinate Framework: A 3D Reference Atlas. Cell, 181(4):936-953.e2) May 2020a. ISSN 00928674. doi: 10. 2..00166,cel11220000000000 https:
URL elsevier. com/ http:///eeeeeeeee////////22222// http:///p11000000772000330
Quanxin Wang, Song-Lin Ding, Yang Li, Josh Royall, David Feng, Phil Lesnar, Nile Graddis, Maitham Naeemi, Benjamin Facer, Anh Ho, et al. The allen mouse brain common coordinate framework: a 3d reference atlas. Cell, 181(4)9336-953, 2020b.
Asmamaw T Wassie, Yongxin Zhao, and Edward S Boyden. Expansion microscopy: principles and uses in biological research. Nature methods, 16(1 16(1):33-44 41,3-41,, 2019.
Thomas Welton, Septian Hartono, Yao-Chia Shih, Stefan T Schwarz, Yue Xing, Eng-King Tan, Dorothee P Auer, Noam Harel, and Ling-Ling Chan. Ultra-high-fee 7t mri in parkinson's 's disease: ready for clinical use?-a a narrative review. Quantitative Imaging in Medicine and Surgery, 13(11): 7607, 2023.
19

---

Published as a conference paper at ICLR 2026
Marek Wodzinski, Niccolo Marini, Manfredo Atzori, and Henning Müller. Deeperhistreg: robust whole slide images registration framework. arXiv preprint arXiv:2404.1443, 2024.
Yifan Wu, Tom Z. Jiahao, Jiancong Wang, Paul A. Yushkevich, M. Ani Hsieh, and James C. Gee. NODEO: A Neural Ordinary Differential Equation Based Optimization Framework for Deformable Image Registration. arXiv: 2xXiv:2000 108. 03443 [cs], February 2022. URL http: waa.iiiiiiii//.ia//aaaaaaaaaa..a org org//n abs/ arXiv: 2108.03443.
2108. 03443.
Yifan Wu, Mengjin Dong. Rohit Jena, Chen Qin, and James C Gee. Neural ordinary differen- tial equation based sequential image registration for dynamic characterization. arXiv preprint arXiv:2404.02/1 2024.
Xiao Yang, Roland Kwitt, Martin Styner, and Mare Niethammer. Quicksilver: Fast predictive image registration- deep learning approach. Neurolmage, 158:378-399 -97788--6,,10 2017.
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. arXiv preprint UrXiv:2502./1089 2025.
Liutong Zhang, Lei Zhou, Ruiyang Li, Xianyu Wang, Boxuan Han, and Hongen Liao. Cascaded feature warping network for unsupervised medical image registration. In 2021 IEEE 18th International Symposium on Biomedical Imaging (ISBI), pp. 913-916. IEEE, 2021.
Shengyu Zhao, Yue Dong, Eric I-Chao Chang, and Yan Xu. Recursive cascaded networks for unsupervised medical image registration. In Proceedings the on Computer Vision (ICCV), October 2019a.
the IEEF/CVF International Conference
Shengyu Zhao, Tingfung Lau, Ji Luo, I Eric, Chao Chang, and Yan Xu. Unsupervised 3d end-to-ent medical image registration with volume tweening network. IEEE journal of biomedical and health informatics, 24(5):1394-140) 2019b.
Shuai Zhao, Yang Wang, Zheng Yang, and Deng Cai. Region mutual information loss for semantic segmentation. Advances in Neural Information Processing Systems, 32, 2019c.
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, al. Pytorch fsdp: experiences scaling fully sharded data 2023.
parallel. arXiv preprint
20


Reviews:
[{"review_id": "SbPFIivEAG", "reviewer": "Reviewer_oLeD", "summary": "This paper makes large 3D medical image registration run fast and fit in GPU memory. \nKey idea the reviewer finds: real bottlenecks at giga-voxel scale are not matrix multiplies but ops like grid sampling, LNCC, MI. They write fused, IO-aware CUDA kernels plus a ring sampler so GPUs don’t need full image copies. Result: 6–7× speed, up to ~59% memory saved, and a demo of 11.8B-parameter multimodal brain registration in about a minute on 8×A6000. This work is quire impressive.\n\nThe reviewer is familiar with registration but not with optimization work. But this work appears impressive and practical. Therefore, the reviewer will assign a low confidence score but high rating.", "questions": "See weakness.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- Writing: Problem importance is well motivated: when the image is 100–1000× larger, registration algorithms did not keep up. It is an important gap.\n- Distributed story is novel for registration.\n- Empirical section is broad: classical (ANTs-like), learning (TransMorph, SynthMorph, VFA, UniGradICON), and purpose-built (CLAIRE, ITK-DReg). The 250µm setting, where every other method degrades while FFDP improves, is persuasive.\n- The accelerations for existing code (TransMorph training 6× faster, big memory drops, and also for FireANTs) mean this work is really a critical progress.", "weaknesses": "- This work enables faster/larger registration, but it lacks a clear section on how to apply this. It does not specify which registration methods are directly supported, how to integrate the kernels, or which methods are not supported and why.\n- Only mentioned related work on GPU/system acceleration. The author does not survey other GPU acceleration outside registration.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cvp3ZOzdV5", "reviewer": "Reviewer_TnBN", "summary": "In this paper, we propose the FFDP framework, which aims to solve the problem of insufficient algorithmic scalability in ultra-high resolution multimodal image alignment.FFDP designs three I/O-aware non-GEMM fusion kernels (composite implicit grid sampler, implicit Parzen window mutual information estimation, and efficient fusion inter-correlation) as well as distributed architectures (grid parallelism, ring sampler, and distributed loss computation) to realize the fast processing of tasks at the level of hundreds of billions of voxels. Experiments show that FFDP outperforms existing methods on both synthetic and real datasets, increasing the runtime speed by 6-7 times, reducing peak memory by 20-59%, and significantly reducing GPU computational overhead while maintaining higher alignment accuracy (e.g., Dice coefficient of 89.5% at 250µm resolution).", "questions": "Please refer to weaknesses", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Motivation is clear: It precisely targets the core contradiction between high-resolution acquisition capabilities and the insufficient scalability of existing algorithms in medical imaging, with real scenarios like 100µm ex-vivo human brain MRI clearly illustrating the necessity.\n- Experiments are rich and comprehensively prove the effectiveness of its method: It covers multiple types of datasets and various dimensions of experiments, from component-level ablation to system-level comparison with mainstream baselines, fully demonstrating the method's validity.\n- Performance improvement is significant", "weaknesses": "While this paper targets at improving efficiency, it would be better provide experimental evaluations on accuracy against counterparts.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tNf06oQYUa", "reviewer": "Reviewer_bohM", "summary": "The paper proposes a set of IO-aware fused kernels and a distributed framework enabling deformable image registration at extremely large scale. It significantly reduces the memory usage for interpolation, and 2 very common loss functions. Experimental results show significant acceleration for both classical iterative optimization and deep learning–based registration pipelines, and demonstrate the first native-resolution multimodal registration at extremely high spatial resolutions.", "questions": "Please see the section above.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. Enables giga-voxel registration in its native resolution that previously was infeasible.\n2. Identifies and proposes solutions for major memory bottlenecks speeding up both iterative and learning based registration.\n3. Demonstrates performance on multimodal registration of 250μm to 100μm images completing in 1 minute running on 8 gpus.\n4. Ablates each components contribution.\n5. Benchmarks the performance of registration algorithms on clinical scale and very high resolution brain images.\n6. The language of the paper is clear. \n7. I really enjoyed reading the paper.", "weaknesses": "**Weaknesses and questions**\n\n1. I was wondering why the affine transformations is denoted as φ(x)=Ax+t. Isn’t it the case normally that the translation is incorporated in A along with rotation and scaling? \n2. Although the scalability demonstrated on ex-vivo MRI is impressive, it is unclear why the method was not evaluated on multi-gigapixel histopathology datasets, which seem like an ideal target application given their fine-scale structural alignment requirements. Such experiments would directly validate the paper’s stated motivation and better highlight FFDP’s advantages over existing patch-based histology workflows.\n3. While the paper emphasizes improvements for multimodal registration, LNCC is widely used for monomodal scenarios too, so the proposed optimizations appear broadly applicable. I am therefore unsure why the authors chose to position their contribution primarily around multimodal registration and not a general purpose registration with more speedup in the multi-modal case.\n4. Not sure if I missed it, but are all methods (learning-based and iterative optimization ones) trained with the same objectives?\n5. Multi-resolution strategies are commonly used in high-resolution registration to manage memory and computation, and while they suffer from similar limitations as those discussed in the paper, acknowledging them would improve completeness and situate the work more clearly within existing practice.\n6. Regarding the learning-based baselines, their encoders are not designed for extremely high-resolution inputs, and while patchification is clearly suboptimal, it seems that alternative architectural designs (e.g., improved downsampling kernels in the feature space, larger capacity models, more spatially aware context handling) might yield stronger performance even if at higher memory cost. It would be helpful for the authors to comment on whether such adaptations were considered and how FFDP compares in those scenarios. I would be interested in the authors’ perspective on whether the limitations observed are inherent to deep approaches at scale or simply due to the baselines being used outside their intended operating regime.\n7. No codebase or link to code is available.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6U2AXFuyxo", "reviewer": "Reviewer_hXkU", "summary": "This paper introduces FFDP (Flash Fused Distributed Primitives) as a scalable and distributed framework for large-scale image registration. The main idea is to fuse non-GEMM operations using IO-aware kernels and distributed tensor sharding, enabling registration of multi-billion-voxel images. The proposed framework includes fused kernels for single-GPU efficiency and GridParallel plus RingSampler modules for distributed multi-GPU scalability without full-image allgather operations. Experimental results show that FFDP significantly accelerates registration and reduces memory consumption for both iterative optimization-based and deep learning-based methods.", "questions": "1. Could you provide evidence or references that shows the necessity of full-resolution registration?\n\n2. In Table 1, what do the three “Baseline” rows represent, what do “Top” and “Bottom” refer to, and which dataset was used to generate Table 1? Besides, Table 1 is not referenced in the main paper.\n\n3. Regarding the faux-OASIS dataset experiment, you register the images at a downsampled resolution and perform patchwise registration followed by mosaicing of the final deformation for deep learning methods. This setup may adversely affect their performance, as many hyperparameters proposed in their paper(such as the weighting of the smoothness term) are designed for whole-image registration rather than patchwise processing. Moreover, you did not report the performance of these methods on the original OASIS dataset. If the deep learning methods achieve better performance on the original dataset, it would suggest that your experimental setup is flawed.\n\n4. In Figure 5, you denote “fire-ants with the proposed method” as Ours and compare it with other methods. However, this comparison appears unfair since you did not include the performance of fire-ants itself. It is highly likely that fire-ants alone could outperform many of the other baselines. You should include the performance of the original fire-ants for a fair comparison. If memory limitations are a concern, you may apply the same two modifications used for the deep learning–based methods.", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "1. System contribution: \nFFDP elegantly adapts concepts like ring topology and sharded synchronization to a domain with boundary dependencies.\n\n2. Significant scalability results on iterative methods: \nDemonstrates unprecedented registration (11 B + parameters) within realistic compute budgets.", "weaknesses": "1. The introduction part does not demonstrate the motivation for full-resolution registration:\nThe introduction does not sufficiently justify the necessity of performing registration at full resolution. While image acquisition capabilities have improved, in many biomedical and clinical applications registration on downsampled images is adequate. The authors should provide stronger evidence or references showing that downsampling significantly degrades registration accuracy or downstream analysis.\n\n2. Limited benefit for deep learning pipelines:\nThe proposed framework provides modest gains (16.5% and 24.7% memory reduction) for deep learning networks comparing with it in iterative methods. It would be helpful if the authors could analyze whether further optimizations could narrow this gap.\n\n3. Combined Introduction and Related Work:\nThe paper merges the introduction and related work into a single section, which reduces readability and makes it harder to follow the motivation versus prior context. \n\n4. Poorly organized experiments:\nThe experimental section is difficult to follow and lacks clear structure.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)