# SHMoAReg: Spark Deformable Image Registration via Spatial Heterogeneous Mixture of Experts and Attention Heads

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Vision mixture-of-experts, one of 12 directions in Mixture-of-experts.
- Builds on 33 of its 50 references that are in the map.
- First posted on arXiv 24 Sep 2025.

## Placed in

- Mixture-of-experts / [Vision mixture-of-experts](https://mlsanity.com/moe/vision-mixture-of-experts/)

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://mlsanity.com/paper/2405.04434/) (2024)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [VoxelMorph: A Learning Framework for Deformable Medical Image Registration](https://mlsanity.com/paper/1809.05231/) (2018)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [OpenMoE: An Early Effort on Open Mixture-of-Experts Language Models](https://mlsanity.com/paper/2402.01739/) (2024)
- [Large Deformation Diffeomorphic Image Registration with Laplacian Pyramid Networks](https://mlsanity.com/paper/2006.16148/) (2020)
- [ViT-V-Net: Vision Transformer for Unsupervised Volumetric Medical Image Registration](https://mlsanity.com/paper/2104.06468/) (2021)
- [Hard Mixtures of Experts for Large Scale Weakly Supervised Vision](https://mlsanity.com/paper/1704.06363/) (2017)
- [Mixture of Attention Heads: Selecting Attention Heads Per Token](https://mlsanity.com/paper/2210.05144/) (2022)
- [Go Wider Instead of Deeper](https://mlsanity.com/paper/2107.11817/) (2021)
- [JetMoE: Reaching Llama2 Performance with 0.1M Dollars](https://mlsanity.com/paper/2404.07413/) (2024)
- [Network of Experts for Large-Scale Image Categorization](https://mlsanity.com/paper/1604.06119/) (2016)
- [Dual-Stream Pyramid Registration Network](https://mlsanity.com/paper/1909.11966/) (2019)
- [XMorpher: Full Transformer for Deformable Medical Image Registration via Cross Attention](https://mlsanity.com/paper/2206.07349/) (2022)
- [Dense Training, Sparse Inference: Rethinking Training of Mixture-of-Experts Language Models](https://mlsanity.com/paper/2404.05567/) (2024)
- [Patch-level Routing in Mixture-of-Experts is Provably Sample-efficient for Convolutional Neural Networks](https://mlsanity.com/paper/2306.04073/) (2023)
- [How Does Attention Work in Vision Transformers? A Visual Analytics Attempt](https://mlsanity.com/paper/2303.13731/) (2023)
- [Robust Mixture-of-Expert Training for Convolutional Neural Networks](https://mlsanity.com/paper/2308.10110/) (2023)
- [M4oE: A Foundation Model for Medical Multimodal Image Segmentation with Mixture of Experts](https://mlsanity.com/paper/2405.09446/) (2024)
- [ModeT: Learning Deformable Image Registration via Motion Decomposition Transformer](https://mlsanity.com/paper/2306.05688/) (2023)
- [Joint Progressive and Coarse-to-fine Registration of Brain MRI via Deformation Field Integration and Non-Rigid Feature Fusion](https://mlsanity.com/paper/2109.12384/) (2021)
- [Non-iterative Coarse-to-fine Transformer Networks for Joint Affine and Deformable Image Registration](https://mlsanity.com/paper/2307.03421/) (2023)
- [Biased Mixtures Of Experts: Enabling Computer Vision Inference Under Data Transfer Limitations](https://mlsanity.com/paper/2008.09662/) (2020)
- [Non-iterative Coarse-to-fine Registration based on Single-pass Deep Cumulative Learning](https://mlsanity.com/paper/2206.12596/) (2022)
- [Low-Rank Mixture-of-Experts for Continual Medical Image Segmentation](https://mlsanity.com/paper/2406.13583/) (2024)

## Abstract (arXiv)

Encoder-Decoder architectures are widely used in deep learning-based Deformable Image Registration (DIR), where the encoder extracts multi-scale features and the decoder predicts deformation fields by recovering spatial locations. However, current methods lack specialized extraction of features (that are useful for registration) and predict deformation jointly and homogeneously in all three directions. In this paper, we propose a novel expert-guided DIR network with Mixture of Experts (MoE) mechanism applied in both encoder and decoder, named SHMoAReg. Specifically, we incorporate Mixture of Attention heads (MoA) into encoder layers, while Spatial Heterogeneous Mixture of Experts (SHMoE) into the decoder layers. The MoA enhances the specialization of feature extraction by dynamically selecting the optimal combination of attention heads for each image token. Meanwhile, the SHMoE predicts deformation fields heterogeneously in three directions for each voxel using experts with varying kernel sizes. Extensive experiments conducted on two publicly available datasets show consistent improvements over various methods, with a notable increase from 60.58% to 65.58% in Dice score for the abdominal CT dataset. Furthermore, SHMoAReg enhances model interpretability by differentiating experts' utilities across/within different resolution layers. To the best of our knowledge, we are the first to introduce MoE mechanism into DIR tasks.