Beyond Self-Attention: Deformable Large Kernel Attention for Medical Image Segmentation

In plain words

Segmentation means labelling every pixel of an image, here an organ, a skin lesion or a pancreas in a medical scan. Transformer networks see the whole image at once, which helps, but their cost grows quickly with image size. Many medical methods also treat a 3D scan as a stack of 2D slices, so they lose information between neighbouring slices.

The paper builds a block that gets a wide view cheaply. It uses a few large, cheap convolutions to mix information over a big area, and lets the sampling pattern of those convolutions bend to follow the shape of the structure. A 2D version works on single slices; a 3D version lets neighbouring slices inform each other. The authors place these blocks in a hierarchical network. On the datasets they test, they report higher average overlap scores (Dice) than the networks they compare against.

TL;DR

D-LKA Net is a hybrid transformer-CNN for medical segmentation. Its attention replaces self-attention with stacked large depthwise convolutions and deformable sampling. The 2D version uses a MaxViT encoder; the 3D version is meant to share information across slices. Author-reported: average DSC 84.27 on Synapse 2D, 87.49 on Synapse 3D and 81.22 on the NIH Pancreas set, each the highest average in its table. The released code's defaults differ from the paper in several places (see Paper vs code).

Problem

• Self-attention cost grows with the square of the token count, which limits network depth and input resolution [abstract].
• Most methods process 3D volumes slice by slice and so miss information between slices §1.
• CNNs capture local texture but miss global context; ViTs capture global context but model local detail less well §1.
• Lesions and organs deform in shape, so a learned operator should handle deformation while staying cheap enough for 3D volumes §1.

The paper states no assumptions about data beyond the datasets in §3.2.

Method

Overview. The core block is D-LKA attention. It starts from Large Kernel Attention (LKA) §2.1 and makes its convolutions deformable §2.2. Two networks use it: a 2D network with a MaxViT encoder and D-LKA decoder §2.3, and a 3D encoder-decoder §2.4.

Step 1. LKA as large-kernel, low-cost attention §2.1. A K×K kernel is built from a depthwise convolution with kernel (2d−1)×(2d−1), a depthwise dilated convolution with kernel ⌈K/d⌉×⌈K/d⌉ at dilation d, and a 1×1 convolution eq. 1 eq. 2. For C channels the parameter count is P(K,d) = C(⌈K/d⌉² + (2d−1)² + 3 + C) eq. 3. FLOPs are F(K,d) = P·H·W eq. 4. The 3D analogues are eq. 6 and eq. 7. The paper says a dilation of about 3.37 minimises parameters for K = 21 §2.1 (see Pitfalls for the sign of the derivative in eq. 5).

Step 2. Deformable depthwise convolutions §2.2. Each deformable convolution learns offsets from the feature map with an extra convolution. The offset convolution follows the kernel size and dilation of its layer. Off-grid samples use bilinear interpolation. The 2D D-LKA attention is:

• F′ = GELU(Conv(F)) eq. 8;
• Attention = Conv1×1(DDW-D-Conv(DDW-Conv(F′))) eq. 8, where DDW is a deformable depthwise convolution and DDW-D its dilated form;
• Output = Conv1×1(Attention ⊗ F′) + F eq. 8, with ⊗ the element-wise product.

No sigmoid or softmax normalises the attention map. The paper argues that such normalisers suppress high-frequency information §2.2.

# Deformable LKA attention, 2D  [eq. 8]
# [code: 2D/deformable_LKA/deformable_LKA.py#L93-L94]
# Input: features (B, C, H, W); output: (B, C, H, W)

def deformable_lka_attention(features):
    # F' = GELU(Conv1x1(F))  [eq. 8]
    gated_input = gelu(conv1x1(features))          # (B, C, H, W)
    # Deformable depthwise 5x5, then 7x7 with dilation 3  [§2.2]
    spatial = deform_dw_5x5(gated_input)           # (B, C, H, W)
    spatial = deform_dw_7x7_dilated(spatial)       # (B, C, H, W)
    # Attention map: 1x1 conv, no sigmoid or softmax  [eq. 8]
    attention = conv1x1(spatial)                   # (B, C, H, W)
    # Element-wise gating, 1x1 projection, residual  [eq. 8]
    mixed = conv1x1(attention * gated_input)       # (B, C, H, W)
    return mixed + features

Step 3. Deformable LKA in 2D §2.2 §2.3. Both convolutions in the attention are deformable. The 2D block is eq. 9 to eq. 11: LayerNorm, the deformable LKA attention with a residual, then a Conv1×1 → depthwise conv → GELU → Conv1×1 MLP with a residual. The kernel size of the depthwise conv in the MLP is not stated.

# 2D D-LKA block  [eq. 9-11]
# Input: x_in (B, C, H, W); output: (B, C, H, W)

def d_lka_block_2d(x_in):
    # Attention branch with residual  [eq. 9]
    x_1 = deformable_lka_attention(layer_norm(x_in)) + x_in
    # MLP branch with residual  [eq. 10]
    x_out = mlp_2d(layer_norm(x_1)) + x_1
    return x_out


def mlp_2d(x):
    # Conv1x1, depthwise conv, GELU, Conv1x1  [eq. 11]
    x = conv1x1(x)
    x = depthwise_conv(x)       # kernel size not stated
    x = gelu(x)
    return conv1x1(x)

Step 4. 2D network §2.3 fig. 1 fig. 2. A convolutional stem reduces the input to H/4 × W/4 × C. Four MaxViT stages follow, each with downsampling §2.3. The decoder has four stages; each has two D-LKA blocks, then patch-expanding layers that upsample and reduce channels. A final linear layer produces the output §2.3.

Step 5. Deformable LKA in 3D §2.2 §2.4. Offset generation in 3D would multiply the channel count (the paper says output channels grow to 3·k×k×k), so the 3D version uses one deformable convolution after the depthwise convolutions instead of two §2.2. The 3D block is eq. 12 and eq. 13: LayerNorm, the 3D attention with a residual, then Conv1×1 applied to a feed-forward network Conv3 (two convolutions and activations, per the text §2.4) with a residual.

# 3D D-LKA block  [eq. 12-13]
# Input: x_in (B, C, H, W, D); output: (B, C, H, W, D)

def d_lka_block_3d(x_in):
    # LayerNorm, 3D deformable LKA attention, residual  [eq. 12]
    x_1 = d_lka_attention_3d(layer_norm(x_in)) + x_in
    # Feed-forward Conv3 (two convs, activations), then Conv1x1  [eq. 13]
    x_out = conv1x1x1(feed_forward_3d(x_1)) + x_1
    return x_out

Step 6. 3D network §2.4 fig. 1. A patch embedding maps H×W×D to H/4 × W/4 × D/2 §2.4. The encoder has three stages of three D-LKA blocks; each stage ends with a downsampling step that halves resolution and doubles channels. The bottleneck has two D-LKA blocks. The decoder mirrors the encoder with transposed convolutions, three blocks per stage §2.4. Skip connections add features from other stages; an extra convolutional skip links the input to the output §2.4. The output is a 3×3×3 convolution followed by a 1×1×1 convolution §2.4. Ablation shows the number of skip connections matters tab. 6.

Step 7. Loss §3.1 eq. 14. For the 2D method, the paper trains on 0.6 × Dice + 0.4 × cross-entropy eq. 14. The text attaches this to the 2D setup only; the 3D loss is not stated.

# Training loss, 2D  [eq. 14]
# [code: 2D/trainer_MaxViT_deform_LKA.py#L137-L139]
# Input: logits (B, K, H, W), labels (B, H, W); output: scalar

def training_loss_2d(logits, labels):
    # Cross-entropy term  [eq. 14]
    loss_ce = cross_entropy(logits, labels)
    # Soft Dice term over the softmax  [eq. 14]
    loss_dice = dice_loss(softmax(logits), labels)
    # Weighted sum as stated in eq. 14
    return 0.6 * loss_dice + 0.4 * loss_ce

Hyperparameters

Values are as stated in the paper unless marked "code". "not stated" means the paper is silent. Code values are from the release at code_commit; they are not the paper's claim.

Input and output details:

• Synapse: 30 cases, 3779 axial CT images in total; 85–198 slices per volume, 512×512 each; voxel spacing 0.54–0.54 × 0.98–0.98 × 2.5–5.0 mm §3.2. The split follows TransUNet and UNETR++ §3.2.
• ISIC 2017: 2000 train, 150 validation, 600 test §3.2. ISIC 2018: the division follows a scheme "described in previous literature", which the card does not reproduce §3.2. PH2: 200 images (160 nevi, 40 melanomas); split not stated §3.2.
• NIH Pancreas: 82 volumes; 62 for training and the rest for testing, so 20 by arithmetic §3.2.

Results

Author-reported. Setup: single RTX 3090 §3.1; each dataset as above §3.2. The paper marks best and second-best with colour, which is not reproduced here.

Synapse, 2D tab. 1. Params in millions, FLOPs in billions, DSC and HD95 as averages over organs.

Skin lesions tab. 2. DSC, SE (sensitivity), SP (specificity), ACC for ISIC 2017, ISIC 2018 and PH2.

Synapse, 3D tab. 3.

NIH Pancreas, 3D tab. 4.

Ablations (Synapse, 3D) tab. 5 tab. 6.

Parameter counts of the deformable layers tab. 7 app. A. Kernels are 21×21 nominal. Values for C = 32, 64, 128, 256, 512 channels.

Claims in the text.

• Synapse 2D: DSC 1.41 points above ScaleFormer and 1.64 above DAEFormer §3.3.1. Pancreas is said to be 2.04 points above the second-best method §3.3.1; see Pitfalls.
• Synapse 3D: DSC 0.27 above UNETR++ and 0.92 above nnFormer §3.3.2. HD95 is second-best §3.3.2 tab. 3.
• NIH Pancreas: best on all four metrics; DSC +0.63, Jaccard +0.82, HD95 −1.04, ASD −0.26 relative to UNETR++ §3.3.2 tab. 4.
• Parameter counts: 42.35 M is the lowest 3D parameter count in tab. 3; the text says only UNETR++ has fewer FLOPs §3.3.2 (see Pitfalls).
• Robustness: five training runs per method on Synapse 2D; the spleen is slightly worse; gallbladder, pancreas and stomach improve most §3.4 fig. 13. No run-to-run numbers are given in the text.
• Inference time: with a batch of 16 on 224×224 inputs, 8.01 ms with the deformable layer versus 17.38 ms without app. A fig. 7. The paper claims a speed-up only for batch sizes above 1 app. A.
• 2D parameters: 101.64 M, below ScaleFormer's 111.6 M and above most other 2D rows tab. 1 app. B fig. 8.

Paper vs code

Code: https://github.com/xmindflow/deformableLKA at d2d299267a765b07ca8956b18d94876e729e79f2 (cloned, depth 1). The paper links mindflow-institue/deformableLKA, which points to this repository; its owner is the xmindflow organisation, and an author's account appears in the commit history. The code has 2D/ and 3D/ subfolders, each with a MIT LICENSE; there is no root LICENSE. The two MIT files name different copyright holders (2D/LICENSE: Amirhossein Kazerouni, 2022; 3D/LICENSE: jsguo, 2021), so keep attribution to both when redistributing code. Code is cited, not quoted.

1. Synapse 3D script uses the wrong block. The shipped Synapse script passes --trans_block TransformerBlock run_training_synapse.sh. That class is the UNETR++ EPA attention block with no LKA transformerblock.py. The D-LKA block is TransformerBlock_3D_single_deform_LKA transformerblock.py. The pancreas script uses the D-LKA block train_pancreas.py. The paper does not say which block produced tab. 3, so the Synapse 3D number may need the flag changed to reproduce.
2. 2D batch size. The paper says 20 §3.1; the code default is 24 train_MaxViT_deform_LKA.py. The script scales the base learning rate by batch/24 when the batch is not 24 and is a multiple of 5 train_MaxViT_deform_LKA.py. At batch 20 with base LR 0.05, the code gives 0.05 × 20/24 ≈ 0.0417 (our arithmetic).
3. 2D learning-rate schedule. Not stated in the paper; the code decays the LR polynomially with power 0.9 per iteration over max_epochs × iterations per epoch trainer_MaxViT_deform_LKA.py trainer_MaxViT_deform_LKA.py.
4. 3D optimiser details. Paper: SGD, LR 0.01, weight decay 3e-5 §3.1; code matches these d_lka_former_trainer_synapse.py Trainer_synapse.py. The code also sets Nesterov momentum 0.99 d_lka_former_trainer_synapse.py and decays the LR with power 0.9 per epoch d_lka_former_trainer_synapse.py. If the mean foreground Dice is zero at epoch 100, the trainer lowers momentum to 0.95 d_lka_former_trainer_synapse.py. None of these is in the paper.
5. 3D loss. The paper's loss (eq. 14) is written in the 2D paragraph and the 3D loss is not stated. The code uses Dice + cross-entropy with weights 1 and 1, sum aggregation, smooth 1e-5, background excluded Trainer_synapse.py dice_loss.py. Deep supervision weights halve per level and are normalised d_lka_former_trainer_synapse.py. Neither is in the paper.
6. 3D patch size. Paper: 128 × 128 × 64 §3.1. The trainer sets crop_size = [64, 128, 128] d_lka_former_trainer_synapse.py, but the training patch size is read from the plans file Trainer_synapse.py. The plans file is not in the repository, so the effective patch size and axis order cannot be confirmed.
7. 3D batch size. Paper: 2 §3.1. The code reads it from the plans file Trainer_synapse.py, which is not in the repository.
8. 3D block differs from eq. 12 and 13. The code's block scales the attention branch by a learnable per-channel gamma, initialised to 1e-6 transformerblock.py transformerblock.py. The convolution after attention is a residual UNet block with batch normalisation transformerblock.py, and the output projection is preceded by dropout (rate 0.1) transformerblock.py. The paper's equations show none of these.
9. 3D deformable layer. Paper: "a sole deformable convolution layer" after the depthwise convolutions §2.2; kernel not stated. Code: standard depthwise 5×5×5 and 7×7×7 (dilation 3) convolutions, then one deformable 3×3×3 convolution with in = out = dim, which is not depthwise transformerblock.py. The 3D deformable convolution needs a custom build README.md (dcn/make.sh).
10. NIH Pancreas script. The paper gives no training hyperparameters §3.2. The released script uses SGD with LR 0.01, momentum 0.9, weight decay 1e-4, batch 8 with one labelled sample, 6000 iterations, 96³ patches, and LR ×0.1 every 2500 iterations train_pancreas.py train_pancreas.py train_pancreas.py train_pancreas.py. The loss is cross-entropy + Dice on the labelled part only train_pancreas.py. A consistency-weight function is defined train_pancreas.py but I did not find it added to loss_total. Whether this script produced tab. 4 is not stated.
11. Skin scripts. The paper gives no skin training hyperparameters. The skin scripts default to batch 4, LR 0.01, SGD momentum 0.9, weight decay 1e-4 train_skin.py train_skin.py train_skin.py. They carry their own model copies under 2D/skin_code/model/; I did not compare these with the main 2D network.
12. Kernel span. The code's two depthwise kernels (5×5, then 7×7 at dilation 3) cover a combined span of 5 + 3·6 = 23 pixels (our arithmetic). The paper's nominal kernel is 21×21 app. A deformable_LKA.py.
13. Paper's 2D code path matches the design. Four MaxViT encoder stages, four decoder stages with two D-LKA blocks each, and the 2D deformable offsets are all implemented MaxViT_deform_LKA.py MaxViT_deform_LKA.py.
14. Paths and data. The Synapse script hard-codes /home/leon/repos/deformableLKA/3D/DATASET run_training_synapse.sh, and the README asks users to rename folders containing unetr_pp to d_lka_former README.md.

No discrepancies were found for the 2D loss weights (0.6 Dice, 0.4 CE), SGD momentum 0.9, 2D weight decay, 2D epochs (400), 3D epochs (1000), 3D learning rate (0.01), 3D weight decay (3e-5), or 250 patches per epoch.

Pitfalls

• Printed derivative has a sign error. Eq. 5 is written as C(8d̂ + 2K²/d̂³ − 4), which has no positive root for K = 21. With the sign of the 2K²/d̂³ term flipped to minus, d ≈ 3.37 comes out (our check). The derivative also ignores the ceiling in eq. 3.
• Table 7 and eq. 3 disagree by a term. The "decomposed" column equals C(⌈K/d⌉² + (2d−1)² + C) for every row (C = 32: 32·(49 + 25 + 32) = 3,392). Eq. 3 adds "+3" that the table does not include (our arithmetic). The offset columns are full convolutions, not depthwise, which is why offsets dominate the deformable parameter cost app. A.
• Text disagrees with Table 3 on FLOPs. The text says only UNETR++ has fewer FLOPs than D-LKA Net §3.3.2. In tab. 3 the LKA Baseline (48.79 G) and UNETR++ (47.98 G) both have fewer; D-LKA Net (66.96 G) is third lowest.
• Text disagrees with Table 1 on the pancreas margin. The text says pancreas DSC is 2.04 points above the second-best method §3.3.1. In tab. 1 the second-best pancreas value is 66.76 (LKA Baseline), a margin of 0.95 points. The 2.04 figure is measured against MISSFormer (65.67).
• Text overstates the skin results. The paper says D-LKA Net "consistently outperforms its competitors across various evaluation metrics" §3.3.1. In tab. 2 it has the best DSC on all three sets, but its specificity (SP) is not the best on any of them (DeepLabv3+ leads on ISIC 2017 and PH2; Att-UNet on ISIC 2018).
• 2D HD95 is not the best. D-LKA Net's 2D HD95 is 20.04 tab. 1; several methods are lower (HiFormer-B 14.70, DAEFormer 16.39, ScaleFormer 16.81, LeViT-UNet 16.84). The text does not discuss HD95 for 2D.
• Pancreas FLOPs are the highest in Table 4. D-LKA Net has 166.63 G versus 63.53 G (UNETR) and 102.44 G (UNETR++) tab. 4. The text reports only parameters.
• Robustness has no numbers in the text. The "statistical significance" claim §3.4 rests on five runs shown only in fig. 13. The text names no test and gives no variance figures.
• Baseline numbers may be copied. A commented-out caption line in the source says the baseline rows were "aggregated from" earlier papers tab. 1 source. The visible text does not say whether they were re-run.
• Split and augmentation details are missing. The ISIC 2018 split and the PH2 split are not stated §3.2. The augmentations are only named by reference to TransUNet, nnFormer and UNETR++ §3.1. Test-time procedure and the HD95 definition are not given.
• Speed claim is narrow. Timing uses random 224×224 inputs, 1000 runs after 50 warm-up iterations on an RTX 3090 app. A. The paper claims a speed-up only for batch sizes above 1 app. A.
• 3D needs a custom build. The 3D deformable convolution is a CUDA op built from 3D/dcn/make.sh README.md; the 2D path uses torchvision.ops.DeformConv2d deformable_LKA.py.
• Encoder initialisation not checked. The paper does not say whether the MaxViT encoder is pretrained. The 2D training script has resume and checkpoint options but no pretrained-weights argument train_MaxViT_deform_LKA.py; I did not check the model constructor.
• Decoder and MLP kernels. The paper's MLP depthwise kernel and the 3D Conv3 kernels are not stated §2.3 §2.4.
• Layer count in 3D. The release defaults to depths 3 run_training_synapse.sh but the bottleneck depth is not confirmed in the release.

Lineage

• Builds on: Large Kernel Attention, Guo et al. (arXiv:2202.09741), which D-LKA revisits and extends §2.1. Deformable convolutions, Dai et al. §2.2. The MaxViT encoder, Tu et al. §2.3. The 3D training recipe follows nnFormer and UNETR++ §3.1; the nnFormer card is nnFormer.
• Compared against (2D, Synapse): TransUNet, Swin-UNet, LeViT-UNet-384, MISSFormer, ScaleFormer, HiFormer-B, DAEFormer, TransDeepLab, PVT-CASCADE, LKA Baseline tab. 1. Skin: U-Net, Att-UNet, DAGAN, TransUNet, MCGU-Net, MedT, FAT-Net, TMU-Net, Swin-Unet, DeepLabv3+, HiFormer-B, LKA Baseline tab. 2. 3D Synapse: UNETR, Swin-UNETR, nnFormer, UNETR++, LKA Baseline tab. 3. Pancreas: UNETR, UNETR++ tab. 4.
• Superseded by: not known from content in this wiki.