[Paper Review] Large-kernel Attention for Efficient and Robust Brain Lesion Segmentation

Liam Chalcroft, Ruben Lourenço Pereira+6 more|arXiv (Cornell University)|Aug 14, 2023

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Large-kernel Attention (LKA), a hybrid convolutional-attention module that enhances U-Net for 3D brain lesion segmentation by combining the global receptive field of transformers with the parameter efficiency and inductive bias for local patterns of CNNs. LKA outperforms state-of-the-art models in robustness to domain shifts and achieves competitive performance on BraTS and ATLAS datasets, particularly in shape-based generalization and out-of-distribution generalization.

ABSTRACT

Vision transformers are effective deep learning models for vision tasks, including medical image segmentation. However, they lack efficiency and translational invariance, unlike convolutional neural networks (CNNs). To model long-range interactions in 3D brain lesion segmentation, we propose an all-convolutional transformer block variant of the U-Net architecture. We demonstrate that our model provides the greatest compromise in three factors: performance competitive with the state-of-the-art; parameter efficiency of a CNN; and the favourable inductive biases of a transformer. Our public implementation is available at https://github.com/liamchalcroft/MDUNet .

Motivation & Objective

• To address the limitations of standard CNNs in capturing long-range spatial dependencies critical for brain lesion segmentation.
• To overcome the computational inefficiency and lack of local inductive bias in standard vision transformers for medical imaging.
• To develop a model that balances performance, parameter efficiency, and robustness to domain shifts in MRI contrast sequences.
• To investigate whether a hybrid architecture with large-kernel attention can achieve the inductive bias of transformers while maintaining CNN-like efficiency.

Proposed method

• Proposes a novel all-convolutional transformer block based on large-kernel attention (LKA), decomposing a large kernel into depth-wise, dilated depth-wise, and pointwise convolutions.
• Integrates the LKA module into a U-Net architecture, replacing standard self-attention or convolutional blocks to enable global context modeling with reduced parameters.
• Employs a factorized attention mechanism: X₁ = dDWConv(DWConv(X)), X₂ = PConv(X₁), enabling efficient spatial and channel mixing.
• Trains the model end-to-end on multi-modal MRI data using cross-entropy and Dice loss, with data augmentation including Gaussian blurring and intensity jittering.
• Uses test-time augmentation and half-precision inference to improve inference stability and speed on BraTS.
• Evaluates generalization via zero-shot transfer to unseen MRI sequences (T1, T2, DWI, FLAIR) from ISLES 2015.

(a) LKA

Experimental results

Research questions

• RQ1 Can a large-kernel attention mechanism effectively model long-range dependencies in 3D brain lesion segmentation while maintaining computational efficiency?

Key findings

• On the BraTS 2021 public test set, the LKA-E model achieved a Dice score of 0.939 for whole tumor (WT), outperforming both the CNN-based nnUNet (0.902) and the Swin-based model (0.938), with a statistically significant improvement in HD95 (2.449 vs. 3.606, p < 0.0001).
• On the ATLAS dataset, the LKA model maintained competitive performance across all metrics, with a Dice score of 0.888 for enhancing tumor (ET), and significantly outperformed the CNN baseline in terms of average surface distance (AVD).
• In out-of-distribution (OOD) evaluation on ISLES 2015, the LKA model showed superior robustness to MRI contrast shifts, maintaining high performance on T1, T2, DWI, and FLAIR sequences, while the nnUNet baseline degraded significantly.
• The LKA model demonstrated greater invariance to image blurring (σ = 0.5 to 2.0), indicating a stronger shape bias—closer to the Swin transformer—compared to the texture-biased CNN (nnUNet), as shown by stable Dice and HD95 metrics under Gaussian smoothing.
• The LKA model achieved a 2.000 median HD95 for tumor core (TC), significantly better than the Swin-based model (2.236) and the CNN (2.236), indicating improved surface accuracy.
• The ablation study confirmed that varying kernel size and dilation in LKA allows fine-tuning of inductive bias, enabling a spectrum between CNN and transformer behavior for task-specific adaptation.

(b) Attention