Source-linked AI summary
Beyond Self-Attention: Deformable Large Kernel Attention for Medical Image Segmentation
Reza Azad, Leon Niggemeier, Michael Huttemann, Amirhossein Kazerouni, Ehsan Khodapanah Aghdam, Yury Velichko, Ulas Bagci, Dorit Merhof
TL;DR
Medical image segmentation needs efficient modeling of global and inter-slice context without the computational burden of self-attention. The paper introduces 2D and 3D Deformable Large Kernel Attention within D-LKA Net, achieving state-of-the-art results across several public segmentation datasets while retaining a reported limitation on noisy skin-lesion annotations.
Problem
Self-attention provides global context but has high computational demands, while pseudo-3D methods miss inter-slice information in volumetric images.
Method
D-LKA Net uses deformable large-kernel attention with adaptive sampling, including a 3D version for cross-slice feature extraction.
Results
D-LKA Net achieves state-of-the-art results on several public medical segmentation datasets, including 0.27% higher DSC than UNETR++ on Synapse.
Takeaways & Limitations
Large-kernel attention combines global-information aggregation with deformable sampling for medical image segmentation across 2D and 3D settings.
Takeaways & Limitations
On ISIC 2018, noisy or primitive segmentations in the training data may hinder accurate predictions.
Abstract
from arXiv · showhide
Medical image segmentation has seen significant improvements with transformer models, which excel in grasping far-reaching contexts and global contextual information. However, the increasing computational demands of these models, proportional to the squared token count, limit their depth and resolution capabilities. Most current methods process D volumetric image data slice-by-slice (called pseudo 3D), missing crucial inter-slice information and thus reducing the model's overall performance. To address these challenges, we introduce the concept of \textbf{Deformable Large Kernel Attention (D-LKA Attention)}, a streamlined attention mechanism employing large convolution kernels to fully appreciate volumetric context. This mechanism operates within a receptive field akin to self-attention while sidestepping the computational overhead. Additionally, our proposed attention mechanism benefits from deformable convolutions to flexibly warp the sampling grid, enabling the model to adapt appropriately to diverse data patterns. We designed both 2D and 3D adaptations of the D-LKA Attention, with the latter excelling in cross-depth data understanding. Together, these components shape our novel hierarchical Vision Transformer architecture, the \textit{D-LKA Net}. Evaluations of our model against leading methods on popular medical segmentation datasets (Synapse, NIH Pancreas, and Skin lesion) demonstrate its superior performance. Our code implementation is publicly available at the: https://github.com/mindflow-institue/deformableLKA
1. Introduction
Medical image segmentation seeks rapid, precise automation, but CNNs and Transformers each have complementary limitations in local and global context modeling. The paper motivates a hybrid Deformable-LKA architecture to adapt receptive fields while integrating contextual information.
- Medical image segmentation can reduce manual workload and support faster, more accurate diagnosis and image-guided surgery.These benefits motivate the development of rapid and precise segmentation algorithms.
- Object-scale variation creates a receptive-field trade-off: fields that are too small can under-segment, while overly large fields amplify background influence.Multi-scale kernels, pyramid pooling, and dilated convolutions address this challenge but can increase computational requirements or training complexity.
- CNNs capture local details effectively but can overlook global context, motivating vision-transformer approaches for long-range information.Vision Transformers aggregate information across the input sequence beyond CNN receptive-field limits.
- Vision Transformers model long-range dependencies but may represent local textures less effectively than convolutional operations.This limitation matters for diagnostic and prognostic tasks that depend on local texture information.
- Existing hybrid and pure-Transformer segmentation methods combine CNN and Transformer components or redesign attention to improve local and global feature representation.Examples include Swin-Unet, MISSFormer, D-Former, TransUNet, HiFormer, UNETR, CoTr, and TransBTS.
- The proposed Deformable-LKA module adapts receptive fields to data while preserving local descriptors, forming 2D and 3D architectures for semantic segmentation.The 3D version is designed to exchange information across volumetric slices.
2. Method
The method combines large-kernel attention with deformable convolutions to model broad contextual information while adapting sampling to irregular medical-image structures. It instantiates hierarchical 2D and 3D encoder-decoder networks with residual D-LKA blocks for segmentation.
- Large Kernel Attention: Large Kernel Attention approximates self-attention’s receptive field using depthwise, depthwise dilated, and 1 × 1 convolutions with lower computational cost.The construction forms a K × K receptive field from smaller convolutional components.
- Large Kernel Attention: FLOPs grow linearly with input-image size, while parameter growth depends quadratically on channels and kernel size and is reported as nonrestrictive in practice.The paper states that both parameter and computation quantities are usually small enough not to restrict the approach.
- Deformable Large Kernel Attention: Deformable convolutions learn offset fields from feature maps, dynamically warping the sampling grid to adapt kernel shapes to lesions, organ deformations, and object boundaries.The 2D formulation substitutes convolution layers with deformable convolutions because medical-image objects often have irregular shapes and sizes.
- Deformable Large Kernel Attention: The D-LKA module applies convolutional projections, deformable large-kernel operations, element-wise attention weighting, and residual fusion without requiring sigmoid or Softmax normalization.The input feature F is projected to F′ with GELU activation, and the attention map is multiplied element-wise with F′ before output projection and residual addition.
- 2D D-LKA Net: The 2D network uses a MaxViT encoder or deformable LKA layers, hierarchical downsampling and upsampling, and decoder stages containing D-LKA blocks with residual connections.The encoder produces four hierarchical representations, while the decoder uses four D-LKA stages and patch-expanding layers.
- 3D D-LKA Net: The 3D network uses a hierarchical encoder-decoder with three D-LKA stages, a bottleneck, symmetric decoding, and convolutional output layers for volumetric segmentation.Each encoder stage contains three D-LKA blocks; downsampling halves spatial resolution and doubles channels, while the bottleneck contains two D-LKA blocks.
3. Experiments
Experiments evaluate 2D and 3D D-LKA Net models across Synapse, skin-lesion, and Pancreas datasets, alongside robustness and ablation studies. Results show strong segmentation performance, efficient 3D computation, and measurable benefits from deformable convolutions and skip connections.
- Experimental setup: The experiments use 2D and 3D models, SGD training, dataset-specific augmentation, and cross-entropy plus Dice loss for 2D training.The 2D model trains for 400 epochs, while the 3D model uses 128 × 128 × 64 patches for 1000 epochs.
- 2D results: D-LKA Net outperforms prior methods on Synapse 2D DSC, including ScaleFormer by 1.41% and DAEFormer by 1.64%.The largest reported organ-level gain is a 2.04% improvement for the pancreas over the second-best method.
- 3D results: D-LKA Net improves Synapse 3D DSC by 0.27% over UNETR++ and 0.92% over nnFormer while using 42.35 M parameters and 66.96 G FLOPs.The model uses about 17% and 31% of the computations required by Swin-UNETR and nnFormer, respectively, while achieving better performance.
- 3D results: On the NIH Pancreas dataset, D-LKA Net achieves the best performance across all four reported metrics and uses 62.07 M parameters.Relative to UNETR++, the reported changes are +0.63% DSC, +0.82% Jaccard, −1.04 HD95, and −0.26 ASD.
- Robustness: Five-run Synapse 2D evaluation shows stable gains for most organs, while spleen performance is slightly worse and improvements are significant for gallbladder, pancreas, and stomach.The evaluation visualizes performance variation across repeated training runs.
- Ablation studies: Ablations report a 0.99% DSC gain from an additional 3D convolution and a further 0.63% gain when it is replaced by a deformable convolution.Removing skip connections reduces performance, while adding the highest-level skip connection improves DSC by 0.42%.
4. Conclusion
D-LKA Net is a hierarchical hybrid Vision Transformer-CNN architecture that uses deformable large-kernel attention to combine adaptive sampling with broad contextual aggregation. Its 3D version adds cross-slice feature extraction, and the models achieve state-of-the-art results on several public segmentation datasets.
- D-LKA Net combines a hierarchical Vision Transformer-CNN architecture with Deformable Large Kernel Attention.The proposed attention mechanism learns a deformation grid and aggregates global information within large convolutional kernels.
- The 3D network performs cross-slice feature extraction for stronger volumetric representation.
- The proposed models achieve state-of-the-art results on several publicly available segmentation datasets.
Supplementary Material
The supplementary material extends the empirical analysis with ablations, visualizations, and discussions focused on the robustness, efficacy, and role of the D-LKA module.
- The supplementary material includes an extended ablation study of the method for semantic segmentation.
- Supplementary visualizations and discussions examine the impact of the proposed D-LKA module.
A. Computationl complexity of the D-LKA
The analysis shows that decomposed and deformable decomposed convolutions reduce parameter growth relative to standard convolution, while deformable LKA adds computational overhead without slowing inference overall. For batch sizes above one, inference can even become faster.
- Decomposed convolution has fewer parameters and slower parameter growth than standard convolution as channel counts increase.Deformable decomposed convolution adds parameters over decomposed convolution but remains substantially smaller than standard convolution.
- Deformable LKA increases model parameters and FLOPS but does not reduce overall inference speed.
- For batch sizes > 1, deformable LKA can reduce inference time.The authors specifically report this effect in experiments with batch size 16.
B. Performance vs Efficiency
D-LKA Net is evaluated against performance, parameter, and memory considerations on the Synapse 2D dataset. It uses approximately 101M parameters, remains below ScaleFormer’s parameter count, and outperforms the lighter DAEFormer despite its larger model size.
- Approximately 101M parameters are used by D-LKA Net on the Synapse 2D comparison.This is fewer than ScaleFormer’s 111.6M parameters.
- Figure 7 reports single-image inference times in milliseconds for 2D methods using input size 3×224×224.
- Figure 8 compares normalized DSC, HD, and memory values across D-LKA Net and several state-of-the-art approaches.The chart includes ScaleFormer, MISSFormer, SwinUnet, and DAEFormer on the Synapse dataset.
- D-LKA Net achieves better performance than the lighter DAEFormer despite requiring more parameters.The comparison is presented as a performance-versus-parameter tradeoff.
C. Qualitative results on the Synapse dataset
The 3D visualization reveals inconsistencies among 2D segmentation slices because the 2D network lacks neighboring-slice information exchange. The 3D network mitigates these limitations.
- 2D segmentation produces slice-to-slice inconsistencies in the Synapse organ visualizations.These discrepancies are attributed to missing information exchange between neighboring slices.
- The 3D network mitigates the limitations observed in the 2D approach.
D. Limitations on the Skin dataset
On ISIC 2018, the proposed approach fails on samples with noisy or primitive segmentation patterns, which are also present in the training data.
- The approach fails on some ISIC 2018 samples with noisy or primitive segmentation patterns.Because these patterns also occur in the training data, they could hinder accurate segmentation learning.
E. Robustness Visualization
The robustness visualization combines repeated Synapse evaluations with qualitative comparisons across Synapse and ISIC 2018. It reports five-run statistical analysis and highlights reduced noise in the 3D model relative to the 2D model.
- Each Synapse 2D model was evaluated five times, with statistical significance reported in Figure 13.The analysis compares methods visually on the Synapse 2D dataset.
- Additional qualitative Synapse results omit the liver and stomach to improve visibility of smaller occluded organs.Figures 9 and 10 provide the additional Synapse visualizations.
- The 3D D-LKA Net creates less noise than the 2D D-LKA-Former in qualitative Synapse visualizations.The comparison is shown after visualizing the 2D results in 3D.
- Figure 12 provides additional qualitative results for the 2D D-LKA-Former on ISIC 2018.
- Figure 13 compares single-organ Synapse performance among state-of-the-art methods and the proposed 2D D-LKA Net using DSC.