[论文解读] 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 Applications被引用 4

一句话总结

该论文提出了一种大感受野注意力（LKA）模块，一种混合卷积-注意力结构，通过结合Transformer的全局感受野与CNN的参数效率和局部模式归纳偏差，增强了U-Net在3D脑部病灶分割中的性能。LKA在应对领域分布偏移方面表现出更强的鲁棒性，并在BraTS和ATLAS数据集上取得了具有竞争力的性能，尤其在基于形状的泛化能力和分布外泛化方面表现优异。

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 .

研究动机与目标

• 为解决标准CNN在捕捉对脑部病灶分割至关重要的长程空间依赖关系方面的局限性。
• 克服标准视觉Transformer在医学影像中计算效率低下且缺乏局部归纳偏差的问题。
• 开发一种在MRI对比序列中平衡性能、参数效率与领域分布偏移鲁棒性的模型。
• 探究是否可通过结合大感受野注意力的混合架构，在保持CNN式效率的同时实现Transformer的归纳偏差。

提出的方法

• 提出一种基于大感受野注意力（LKA）的新型全卷积Transformer模块，将大感受野分解为逐通道、空洞逐通道和逐点卷积。
• 将LKA模块集成至U-Net架构中，替代标准的自注意力或卷积模块，实现在减少参数量的同时实现全局上下文建模。
• 采用因子化注意力机制：X₁ = dDWConv(DWConv(X))，X₂ = PConv(X₁)，以实现高效的通道与空间混合。
• 在多模态MRI数据上端到端训练模型，损失函数采用交叉熵与Dice损失，数据增强包括高斯模糊和强度抖动。
• 在推理阶段使用测试时增强和半精度推理，以提升BraTS上的推理稳定性和速度。
• 通过零样本迁移至ISLES 2015中未见的MRI序列（T1、T2、DWI、FLAIR）评估泛化能力。

(a) LKA

实验结果

研究问题

• RQ1 大感受野注意力机制是否能在保持计算效率的同时，有效建模3D脑部病灶分割中的长程依赖关系？

主要发现

• 在BraTS 2021公开测试集上，LKA-E模型在全肿瘤（WT）上的Dice分数达到0.939，优于基于CNN的nnUNet（0.902）和Swin基模型（0.938），HD95指标显著提升（2.449 vs. 3.606，p < 0.0001）。
• 在ATLAS数据集上，LKA模型在所有指标上均保持了具有竞争力的性能，增强肿瘤（ET）的Dice分数为0.888，且在平均表面距离（AVD）上显著优于CNN基线。
• 在ISLES 2015的分布外（OOD）评估中，LKA模型对MRI对比度变化表现出更强的鲁棒性，在T1、T2、DWI和FLAIR序列上均保持高性能，而nnUNet基线性能显著下降。
• LKA模型对图像模糊（σ = 0.5至2.0）表现出更强的不变性，表明其具有更强的形状偏好（更接近Swin Transformer），而CNN（nnUNet）则表现出纹理偏好，如在高斯平滑下Dice与HD95指标保持稳定。
• LKA模型在肿瘤核心（TC）上的中位HD95为2.000，显著优于Swin基模型（2.236）和CNN（2.236），表明表面精度得到提升。
• 消融实验表明，通过调节LKA中的感受野大小与空洞率，可精细调控归纳偏差，实现从CNN到Transformer行为的连续谱，以适应特定任务需求。

(b) Attention