SHMoAReg: MoE-Driven Deformable Registration

• SHMoAReg is a deep learning-based deformable image registration method that incorporates Mixture of Experts in both encoder and decoder to specialize feature extraction and directional prediction.
• The encoder uses Mixture of Attention Heads within modified Swin Transformer blocks to dynamically select optimal attention combinations for each image token.
• The decoder applies Spatial Heterogeneous Mixture of Experts with varied convolution kernels in three Cartesian directions, yielding improved Dice scores and ASSD.

Searching arXiv for the target paper and closely related DIR context. arxiv_search(query="SHMoAReg Spark Deformable Image Registration via Spatial Heterogeneous Mixture of Experts and Attention Heads", max_results=5, sort_by="relevance") arxiv_search(query="TransMorph deformable registration Swin Transformer", max_results=10, sort_by="relevance") arxiv_search(query="VoxelMorph deformable image registration", max_results=10, sort_by="relevance") arxiv_search(query="SHMoAReg", max_results=10, sort_by="relevance") SHMoAReg is a deep learning-based deformable image registration (DIR) network that introduces a Mixture of Experts (MoE) mechanism into both the encoder and decoder of an encoder–decoder registration architecture. The method combines Mixture of Attention heads (MoA) in modified Swin-Transformer encoders with Spatial Heterogeneous Mixture of Experts (SHMoE) in the decoder, with the stated aim of specializing feature extraction for registration and predicting deformation heterogeneously in the three Cartesian directions rather than jointly and homogeneously. It operates on a moving image $I_m$ and a fixed image $I_f$, produces a deformation field $\phi:\mathbb R^3\to\mathbb R^3$, and was evaluated on OASIS and BTCV, where it reported consistent improvements in Dice and ASSD together with expert-level interpretability analyses (Zheng et al., 24 Sep 2025).

1. Architectural definition and data flow

SHMoAReg follows a classical two-stream encoder / pyramid-decoder design. The input consists of a moving image $I_m$ and a fixed image $I_f$, each of size $C\times D\times H\times W$. The encoder is implemented as two parameter-shared Swin-Transformer backbones, modified so that every Window/Shifted-Window MSA layer is replaced by a Mixture of Attention-head layer. This yields multi-resolution feature maps

$\{F_m^{(r)},F_f^{(r)}\}_{r\in\{1,1/2,1/4,1/8\}},$

where $r$ denotes the down-sampling factor (Zheng et al., 24 Sep 2025).

The pyramid decoder uses these features in a coarse-to-fine registration process. At each resolution level, features are first warped by the deformation predicted at coarser levels and are then passed to convolutional sub-networks that predict a residual deformation field $\Delta\phi^{(r)}$. At the two finest scales, $r=1$ and $I_f$0, these convolutional sub-networks are replaced by SHMoE layers. After the last scale, the full-resolution deformation is obtained by additive composition of all residuals. The data flow is summarized in the paper as

$I_f$1

This architecture is motivated by two limitations identified in existing encoder–decoder DIR systems: the lack of specialized extraction of features useful for registration and the use of homogeneous deformation prediction across all three spatial directions. A plausible implication is that SHMoAReg attempts to localize specialization at two separate stages: token-wise feature extraction in the encoder and voxel-wise directional deformation prediction in the decoder.

2. Mixture of Attention heads in the encoder

The encoder-side innovation is the Mixture of Attention heads mechanism. For a given Swin block, the method defines token-wise queries $I_f$2 together with shared keys $I_f$3 and values $I_f$4. With $I_f$5 attention experts $I_f$6, expert $I_f$7 computes

$I_f$8

followed by

$I_f$9

The key and value projections $\phi:\mathbb R^3\to\mathbb R^3$0 are shared across experts, while the query and output projections $\phi:\mathbb R^3\to\mathbb R^3$1 are expert-specific (Zheng et al., 24 Sep 2025).

Routing is handled by a small network $\phi:\mathbb R^3\to\mathbb R^3$2 producing unnormalized scores

$\phi:\mathbb R^3\to\mathbb R^3$3

after which the top-$\phi:\mathbb R^3\to\mathbb R^3$4 experts $\phi:\mathbb R^3\to\mathbb R^3$5 are selected and re-normalized: $\phi:\mathbb R^3\to\mathbb R^3$6 The MoA output is then

$\phi:\mathbb R^3\to\mathbb R^3$7

In the reported implementation, the encoder uses $\phi:\mathbb R^3\to\mathbb R^3$8 experts and $\phi:\mathbb R^3\to\mathbb R^3$9. The stated functional role of this mechanism is to enhance specialization of feature extraction by dynamically selecting the optimal combination of attention heads for each image token. This makes the encoder’s specialization explicitly token-dependent rather than fixed at model initialization.

3. Spatial Heterogeneous Mixture of Experts in the decoder

The decoder-side innovation is the Spatial Heterogeneous Mixture of Experts layer, used at the two fine-scale levels $I_m$0 and $I_m$1. At these scales, SHMoAReg predicts the residual field $I_m$2 by running three parallel experts in each Cartesian direction. The experts are 3D convolutions over an incoming feature map $I_m$3 with kernel sizes

$I_m$4

Each expert returns a one-channel map for residual displacement in the corresponding direction (Zheng et al., 24 Sep 2025).

Routing is performed by a learnable tensor

$I_m$5

which has the same spatial size as the feature volume and is normalized per voxel such that exactly one expert is “on” at each voxel, since $I_m$6. Writing $I_m$7, the per-direction residual at voxel $I_m$8 is

$I_m$9

The local deformation update is

$I_f$0

and stacking all levels yields the final deformation field.

This construction is explicitly heterogeneous in two senses. First, different voxels may activate different experts. Second, the three directions $I_f$1 may select distinct experts at the same voxel. The paper interprets this as per-direction adaptation in local deformation magnitude and reports that kernel-size differentiation is especially relevant under anisotropic voxels (Zheng et al., 24 Sep 2025).

4. Objective function, regularization, and implementation

Training uses a similarity term, a smoothness regularizer, and a routing-classification loss. Let $I_f$2 denote the warped image. The similarity term is mean squared error: $I_f$3 The deformation smoothness term is a standard diffusion penalty: $I_f$4 The routing-classification loss is defined on $I_f$5 relative to a per-voxel label map $I_f$6, constructed as in Dryden & Hoefler (2022) by thresholding the gradient-magnitude of the similarity loss: $I_f$7 The total objective is

$I_f$8

These definitions are given explicitly in the paper (Zheng et al., 24 Sep 2025).

The reported implementation details are as follows:

Component	Setting	Value
Framework and hardware	Training environment	PyTorch on a single NVIDIA L40 GPU
Optimizer	Optimization	Adam, fixed learning rate $I_f$9
Training schedule	Batch size and iterations	Batch size $C\times D\times H\times W$0, training for $C\times D\times H\times W$1 iterations
MoA	Experts and routing	$C\times D\times H\times W$2, $C\times D\times H\times W$3
SHMoE	Experts and routing	$C\times D\times H\times W$4 experts, kernel sizes $C\times D\times H\times W$5, $C\times D\times H\times W$6, routing quantile $C\times D\times H\times W$7
Loss weights	Regularization	$C\times D\times H\times W$8; $C\times D\times H\times W$9 on OASIS and $\{F_m^{(r)},F_f^{(r)}\}_{r\in\{1,1/2,1/4,1/8\}},$0 on BTCV

The encoder Swin settings are stated to be identical to TransMorph-small. This situates SHMoAReg within a Swin-based DIR design space while changing the attention and fine-scale decoder mechanisms (Zheng et al., 24 Sep 2025).

5. Experimental results and reported performance

The paper reports experiments on two publicly available datasets: OASIS for brain MRI and BTCV for abdominal CT without pre-affine. On BTCV, SHMoAReg achieves Dice $\{F_m^{(r)},F_f^{(r)}\}_{r\in\{1,1/2,1/4,1/8\}},$1 compared with $\{F_m^{(r)},F_f^{(r)}\}_{r\in\{1,1/2,1/4,1/8\}},$2 for the second-best cascade method, corresponding to a $\{F_m^{(r)},F_f^{(r)}\}_{r\in\{1,1/2,1/4,1/8\}},$3 absolute gain, and ASSD $\{F_m^{(r)},F_f^{(r)}\}_{r\in\{1,1/2,1/4,1/8\}},$4 mm compared with $\{F_m^{(r)},F_f^{(r)}\}_{r\in\{1,1/2,1/4,1/8\}},$5 mm. On OASIS, it reaches Dice $\{F_m^{(r)},F_f^{(r)}\}_{r\in\{1,1/2,1/4,1/8\}},$6 compared with $\{F_m^{(r)},F_f^{(r)}\}_{r\in\{1,1/2,1/4,1/8\}},$7 for the second best, and ASSD $\{F_m^{(r)},F_f^{(r)}\}_{r\in\{1,1/2,1/4,1/8\}},$8 mm compared with $\{F_m^{(r)},F_f^{(r)}\}_{r\in\{1,1/2,1/4,1/8\}},$9 mm (Zheng et al., 24 Sep 2025).

The paper also reports a diffeomorphic variant enforcing invertibility, which still outperforms prior art by $r$0 Dice on BTCV. Inference time is approximately $r$1 s per image pair, and the model size is approximately $r$2 M parameters. Qualitative visualizations are said to show improved alignment of complex structures, including kidneys, and include per-voxel expert-ID maps.

These results are presented as evidence that the gain is not limited to a single modality or anatomical region. A plausible implication is that the method’s gains derive from the interaction between encoder specialization and decoder heterogeneity rather than from a single isolated architectural modification.

6. Interpretability, novelty claim, and common points of confusion

A notable aspect of SHMoAReg is that it frames interpretability in terms of expert-load statistics across and within resolution levels. Across resolutions, the paper reports that certain encoder heads specialize at particular layers; for example, expert 4 dominates at $r$3 and expert 6 at $r$4. In the decoder, raw load appears uniform, but kernel-size differentiation is described as crucial for spatially varying displacement scales. Within a single resolution, different spatial tokens select different subsets of attention heads in the encoder, and at an organ boundary voxel the $r$5 directions may select distinct decoder experts, revealing per-direction adaptation (Zheng et al., 24 Sep 2025).

One common misunderstanding would be to read SHMoAReg as merely replacing standard attention with a larger attention block. The paper’s actual claim is more specific: MoA dynamically selects the optimal combination of attention heads for each image token, while SHMoE predicts deformation fields heterogeneously in three directions for each voxel using experts with varying kernel sizes. Another possible misunderstanding is to interpret decoder expert utility solely through global load balance. The reported analyses argue that even when raw load appears uniform, kernel-size differentiation remains functionally important.

The authors further state that, to the best of their knowledge, SHMoAReg is the first introduction of an MoE mechanism into DIR tasks. In the paper’s own summary, the method “tightly weaves together” Mixture of Attention heads in the encoder and Spatial Heterogeneous Mixture of Experts in the decoder, yielding both strong registration accuracy and fine-grained interpretability at the token, voxel, direction, and scale levels (Zheng et al., 24 Sep 2025).