即插即用系列 | WACV 2024 CSAM：面向各向异性医学图像分割的 2.5D 跨切片注意力模块

论文标题：CSAM: A 2.5D Cross-Slice Attention Module for Anisotropic Volumetric Medical Image Segmentation

论文原文 (Paper) ：https://openaccess.thecvf.com/content/WACV2024/html/Hung_CSAM_A_2.5D_Cross-Slice_Attention_Module_for_Anisotropic_Volumetric_Medical_WACV_2024_paper.html
代码 (code) ：https://github.com/aL3x-O-o-Hung/CSAM

GitHub 仓库链接（包含论文解读及即插即用代码） ：https://github.com/AITricks/AITricks
哔哩哔哩视频讲解 ：https://space.bilibili.com/57394501?spm_id_from=333.337.0.0

- - [1. 核心思想](#1. 核心思想)
  - [2. 背景与动机](#2. 背景与动机)
  - - [2.1 文本背景总结](#2.1 文本背景总结)
    - [2.2 动机图解分析](#2.2 动机图解分析)
  - [3. 主要创新点](#3. 主要创新点)
  - [4. 方法细节](#4. 方法细节)
  - - [4.1 整体网络架构](#4.1 整体网络架构)
    - [4.2 核心创新模块详解：CSAM (Cross-Slice Attention Module)](#4.2 核心创新模块详解：CSAM (Cross-Slice Attention Module))
    - [4.3 理念与机制总结](#4.3 理念与机制总结)
    - [4.4 图解总结](#4.4 图解总结)
  - [5. 即插即用模块的作用](#5. 即插即用模块的作用)
  - [6. 实验部分简单分析](#6. 实验部分简单分析)
  - [7. 获取即插即用代码关注【AI即插即用】](#7. 获取即插即用代码关注【AI即插即用】)

1. 核心思想

本文针对医学图像（如 MRI）中普遍存在的**各向异性（Anisotropic）问题（即层间分辨率远低于平面内分辨率），提出了一种名为 CSAM (Cross-Slice Attention Module) 的轻量级即插即用模块。该模块旨在弥补纯 2D 方法缺乏体积信息和纯 3D 方法在各向异性数据上参数冗余且效果不佳的缺陷。其核心是通过语义（通道）、位置（空间）和切片（深度）**三个维度的注意力解耦，以极少的参数量捕获跨切片的上下文依赖，实现高效的 2.5D 分割。

2. 背景与动机

2.1 文本背景总结

在临床医学影像中，为了平衡扫描时间和信噪比，MRI 图像通常是各向异性的，即切片厚度（Through-plane）远大于像素间距（In-plane）。

3D 方法的局限：直接使用 3D CNN（如 3D U-Net）假设数据是各向同性的，这导致在 Z 轴上卷积核学习困难，且参数量巨大，容易在小样本医疗数据上过拟合。
2D 方法的局限：逐层分割完全忽略了人体器官的 3D 连续性，导致分割结果在 Z 轴上不连贯。
2.5D 方法的现状：现有的 2.5D 方法（如输入相邻的多张切片）通常只是简单堆叠通道，缺乏深度的特征交互机制。

2.2 动机图解分析

看图说话与痛点分析：

左侧 (2D Approach)：模型只看单张切片，完全丢失了上下层的信息（Volumetric Information Loss），容易导致器官边缘判定错误。
中间 (3D Approach)：模型试图构建一个立体的感受野，但由于数据在 Z 轴被拉伸（各向异性），标准的 3D 卷积核难以适应这种形变，且计算开销巨大（Inefficient）。
右侧 (Ours - CSAM) ：本文提出的方法。输入是一个切片堆栈（Stack），但不是像 3D 那样暴力卷积，而是通过一个专门的跨切片注意力模块 来提取层间关系。它精准地定位了问题的核心：如何在保留 2D 高分辨率特征的同时，引入 Z 轴的关联性？

3. 主要创新点

CSAM 模块：设计了一个轻量级的跨切片注意力模块，能够从语义、位置和切片三个维度聚合特征，有效解决了各向异性数据的分割难题。
极低的参数量：相比于现有的 3D 注意力机制或 RNN 融合模块，CSAM 引入的参数量几乎可以忽略不计（minimal trainable parameters），极大地降低了模型过拟合的风险。
多尺度集成：该模块可以灵活地集成到 CNN 的深层特征提取阶段，利用多尺度特征图来捕获不同层级的上下文信息。
SOTA 性能：在三个公共医学数据集（前列腺、海马体等）上，验证了该方法优于现有的 2D、3D 以及其他 2.5D SOTA 方法。

4. 方法细节

4.1 整体网络架构

数据流详解 ：

模型整体采用经典的 Encoder-Decoder (U-Net Like) 架构，但在输入和中间层做了特殊设计：

输入 (Input) ：输入不再是单张图像，而是 S S S 张相邻切片组成的堆栈（例如 S = 3 S=3 S=3 或 5 5 5）。
编码器 (Encoder) ：使用标准的 2D CNN（如 ResNet 或 EfficientNet）作为骨干。注意，这里的卷积是逐切片共享权重的，或者将 S S S 视为通道维度的扩展。
核心组件：CSAM 插入 ：在编码器的瓶颈处（Bottleneck）或者解码器的不同尺度上，插入 CSAM 模块 。这是特征发生 3D 交互的关键场所。
- 特征图维度变化： B × C × S × H × W B \times C \times S \times H \times W B×C×S×H×W。CSAM 接收这个 5D 张量（在 Batch 维度隐含 S 或显式 S），并在内部进行跨切片融合。
解码器 (Decoder) ：经过 CSAM 增强后的特征被送入解码器，最终输出针对中心切片的分割预测图。

4.2 核心创新模块详解：CSAM (Cross-Slice Attention Module)

CSAM 的设计灵感来源于将复杂的 3D 注意力分解为三个正交的轻量级注意力分支：

分支 A：语义注意力 (Semantic Attention)
- 流动：输入特征 F F F。在空间 ( H , W ) (H, W) (H,W) 和切片 ( S ) (S) (S) 维度上进行 Global Average Pooling。
- 目的：生成一个 1 × C × 1 × 1 1 \times C \times 1 \times 1 1×C×1×1 的权重向量。它关注**"什么特征"**（What）在当前的一组切片中是最重要的（例如，增强肿瘤组织的通道响应）。
分支 B：位置注意力 (Positional Attention)
- 流动：在通道 ( C ) (C) (C) 和切片 ( S ) (S) (S) 维度上进行 Max Pooling 或 Average Pooling。
- 目的：生成一个 1 × 1 × H × W 1 \times 1 \times H \times W 1×1×H×W 的空间热力图。它关注**"哪里"**（Where）是目标区域，利用相邻切片的信息来抑制当前切片的背景噪声（例如，确定的器官边界）。
分支 C：切片注意力 (Slice Attention) ------ 核心创新
- 流动：在空间 ( H , W ) (H, W) (H,W) 和通道 ( C ) (C) (C) 维度上进行压缩。
- 目的：生成一个 1 × 1 × S × 1 1 \times 1 \times S \times 1 1×1×S×1 的深度权重向量。
- 机制：它显式地计算相邻切片对中心切片的重要性。例如，如果相邻切片包含清晰的边界，它的权重就会增加，从而辅助中心切片的分割。这是解决各向异性问题的关键一环。
特征融合：三个分支生成的权重分别通过 Sigmoid 激活后，与原始特征进行逐元素相乘（Element-wise Multiplication），实现全方位的特征校准。

4.3 理念与机制总结

CSAM 的理念是 "维度解耦 (Dimension Decomposition)"。

机制：与其计算昂贵的 ( H × W × S ) × ( H × W × S ) (H \times W \times S) \times (H \times W \times S) (H×W×S)×(H×W×S) 全局注意力矩阵，不如将其拆解为 Channel、Spatial、Depth 三个独立的 1D/2D 注意力。
公式解读 ： F o u t = F i n ⊗ M s e m a n t i c ⊗ M p o s i t i o n a l ⊗ M s l i c e F_{out} = F_{in} \otimes M_{semantic} \otimes M_{positional} \otimes M_{slice} Fout=Fin⊗Msemantic⊗Mpositional⊗Mslice。
这种设计既保留了 3D 上下文，又将参数量控制在与 2D 模型几乎持平的水平。

4.4 图解总结

回到"动机图解"的核心问题：

Slice Attention 解决了 2D 模型"看不见邻居"的问题。
Positional Attention 解决了各向异性导致的空间对齐难题（利用 2D 强先验）。
整体轻量化设计解决了 3D 模型"参数爆炸"的问题。

5. 即插即用模块的作用

CSAM 是一个高度模块化的单元，适用于多种医疗影像任务：

适用场景 ：
- 各向异性数据分割：前列腺 MRI、脊柱 MRI、海马体分割等，层厚较厚的数据。
- 序列数据分析：甚至可以推广到视频分割任务中，处理时序依赖。
具体应用 ：
- 改进 U-Net：可以将 CSAM 插入到 nnU-Net 的 2D 版本中，在不增加推理显存压力的前提下，提升 Z 轴的分割连续性。
- 多模态融合：在处理 MRI 不同序列（T1, T2）时，CSAM 也可以作为一种模态间的特征筛选器。

6. 实验部分简单分析

论文在 MSD Prostate (前列腺) , Promise12 , 和 MSD Hippocampus (海马体) 三个数据集上进行了验证。

定量对比 ：
- 在 Prostate 数据集上，CSAM 取得了 Dice Score 的显著提升，超越了标准的 2D U-Net 和 3D U-Net。
- 特别是在 HD95 (Hausdorff Distance) 指标上表现优异，这说明 CSAM 生成的边界更加平滑、准确，有效减少了层间"断层"现象。

消融实验：
- 作者分别移除了 Semantic、Positional 和 Slice 注意力。结果显示，Slice Attention 的贡献最大，直接证明了在各向异性数据中，显式建模层间关系是提升性能的关键。
可视化分析：
- 从分割结果的 3D 重建图可以看出，CSAM 的结果在 Z 轴上非常连续，而 2D 方法则呈现出明显的"阶梯状"伪影。

到此，所有的内容就基本讲完了。如果觉得这篇文章对你有用，记得点赞、收藏并分享给你的小伙伴们哦😄。