MICCAI 2022 | 用于脑MRI分割的注意力对称自动编码器
-
论文标题：Attentive Symmetric Autoencoder for Brain MRI Segmentation
论文地址：https://arxiv.org/pdf/2209.08887.pdf
代码地址：暂无
摘要
- 背景介绍： 基于图像块重构的自监督学习方法在训练自动编码器方面取得了巨大成功，其预训练的权重可以被迁移到图像理解的其他下游任务中。
- 现存问题： 现有方法在应用于3D医学图像时很少研究重建patch块的各种重要性和解剖结构的对称性。
- 作者提出了一种新的基于视觉transformer（ViT）的注意力对称自动编码器（ASA），用于3D脑MRI分割任务。 与恢复平滑的图像块相比，强制自动编码器恢复信息图像区域可以获得更多的区分性表示。然后采用基于梯度的度量来估计每个图像块的重要性。在预训练阶段，所提出的自动编码器更加注重根据梯度度量重构信息块。此外，利用大脑结构的先验性，开发了对称位置编码（SPE）方法，以更好地利用长距离但空间对称区域之间的相关性，从而获得有效的特征。
- 实验结果： 对称自动编码器在三个脑MRI分割基准上优于最先进的自监督学习方法和医学图像分割模型。
算法
注意力对称自动编码器
提出了一种新的注意力对称自动编码器（ASA），可以对其进行训练以获得适用于脑MRI分割任务的通用模型权重。如图1所示，
所提出的ASA由一对编码器和解码器组成，具有对称位置编码（SPE）和注意力的重建损失。
在ASA的自监督训练过程中，输入的3D图像被分成规则的非重叠图像块（大小为s×s×s）。这些图像块的P%被随机遮蔽，并且只有未遮蔽的块可见。在线性投影之后，每个可见patch被嵌入到特征向量中，该特征向量与对称位置编码（SPE）相加，以产生编码器输入。编码器输出与其输入相同数量的vector。掩码tokens是添加了不同SPE的可学习vector。每个掩码token对应于掩码图像patch。编码器输出与掩码tokens连接以形成解码器输入。解码器重建所有图像块，并且仅使用遮蔽的图像块来计算所提出的损失。
注意力重建损失
考虑到学习恢复flatten区域对于鼓励模型获得有区别的表达没有太大帮助。作者开发了一种基于注意力的重建损失函数，它强调大脑MRI的信息区域。
为了估计图像块的信息，对3D图像采用基于梯度的度量。受3D VHOG的启发，通过应用滤波器掩码[-1，0，1]计算每个体素的梯度向量g=（gx，gy，gz）。在球坐标中，使用两个标量θ和φ来表示体素的方向。θ和φ可计算为：
对于每个图像块，构建了一个2D直方图G，并且bins的数量是b×b。为了计算G值，遍历图像块中的每个体素。设θ，φ表示当前体素的方向。首先确定体素的bin索引为r=⌊θ/(π/b)⌋,c=⌊ϕ/(π/b)⌋。然后，我们将||g||（当前体素的梯度大小）累加到2D直方图g的相应bin（r，c）。在处理图像块中的所有体素之后，对直方图G执行L2范数。对于每个图像块。将相对重要性pi表征为。
损失函数采用均方误差（MSE）来测量恢复图像区域与原始图像区域之间的像素级差异，并使用基于梯度的权重pi来更加关注信息脑区域。总损失可以公式化为：
对称位置编码
观察了大脑结构的左右对称性，并提出了对称位置编码（SPE）方法。
所提出的方法缩小了两个对称图像位置的编码差异，并可以鼓励模型从这两个相关区域中获得更好的特征。对于同一水平面上的斑块（图2（a）），左上角的vanilla位置编码与右上角的位置编码有很大不同（图2b），尽管这些区域的内容相似。然而，使用SPE，最左边和最右边的位置（在同一行中）可以共享相同的编码。设T×H×W，（t，h，w）表示图像的块数和图像块的坐标。对称位置编码被计算为
网络架构
所提出的ASA模型为下游任务提供预先训练的模型权重。
ASA的编码器和解码器基于视觉transformer（ViT）。标准ViT使用vanilla自注意力（SA），这导致了高计算成本，尤其是在处理3D图像时。为了提高效率，作者开发了基于线性窗口的多头自注意力（LW-MSA）和基于移位线性窗口的多头自注意力(SLW-MSA)。受SwinT的启发，将3Dpatches flatten为patches嵌入序列，并将序列拆分为大小为S的窗口。LW-MSA计算每个1D窗口内的自注意。SLW-MSA在计算LW-MSA模块之前将序列移位S/2，并在LW-MSA模块之后将序列反向移位S/2。LW-MSA和SLW-MSA是在块级别上计算的，因为在一开始就通过块嵌入层将每个图像块转换为特征向量。LW-MSA和SLW-MSA交替堆叠以提取交叉窗口特征，并为ASA模型构建移位窗口ViT（SW-ViT）。对于脑MRI分割，构建了一个以ASA编码器为主干的U形网络，如图3所示。
实验
