MAE再次升级，FocusMAE开源 | 源于MAE又高于MAE，有更高质量的表征&全新的架构设计
MAE再次升级，FocusMAE开源 | 源于MAE又高于MAE，有更高质量的表征&全新的架构设计
近年来，自动胆管癌（GBC）检测引起了研究者的关注。目前基于超声波成像（US）图像的领先（SOTA）方法显示出有限的泛化能力，这强调了需要变革性的方法。作者观察到单个US帧可能缺乏捕捉疾病表现的信息。 本研究提倡向基于视频的GBC检测范式的转变，利用时空表示的内在优势。采用 Mask 自动编码器（MAE）进行表示学习，作者解决了传统基于图像方法的不足。作者提出了一种新颖的设计，称为FocusMAE，以系统地偏向于从高信息区域选择 Mask 标记，从而促进对恶性的更精细的表示。此外，作者为GBC检测贡献了最广泛的US视频数据集。作者还注意到，这是首次关于基于US视频的GBC检测的研究。 作者在精选的数据集上验证了所提出的方法，并为GBC检测问题报告了96.4%的新SOTA准确率，相较于当前基于图像的SOTA - GBCNet和RadFormer的84%，以及基于视频的SOTA - AdaMAE的94.7%。作者进一步在公开的基于CT的新冠检测数据集上展示了所提出的FocusMAE的泛化能力，相较于当前 Baseline 提高了3.3%的准确率。 源代码和预训练模型：https://github.com/sbasuz76/FocusMAE
1 Introduction
胆管癌（GBC）。近来，自动检测GBC的技术引起了研究者的极大兴趣。GBC在早期很难被检测出来，而且当疾病被发现时大多数患者已经处于晚期，手术切除变得不可行。因此，这种疾病的生存统计数据非常惨淡。晚期GBC患者的5年生存率仅为5%，平均生存时间仅为六个月。因此，早期发现GBC对于及时干预和提高生存率至关重要。
用于胆管癌（GBC）检测的超声（US）技术。 由于其成本低廉、易于获取和无离子辐射的特点，超声已成为首选的非侵入性诊断成像方式。在资源匮乏的国家，通常它是针对腹部疾病患者进行的唯一成像检查。然而，与结石或息肉等良性病症不同，放射科医生从常规超声中识别恶性肿瘤的迹象颇具挑战性。如果胆管癌未能被检测出来，它可能会默默进展。因此，从早期阶段就开始通过超声识别胆管癌至关重要。
自动检测胆囊癌（GBC）。使用深度神经网络（DNNs）从超声（US）图像中检测GBC是具有挑战性的。由于传感器问题，US图像通常质量较低，这会在DNNs中造成偏差，使得难以准确确定胆囊（GB）区域。探头的便携性质也意味着视图不对齐，这增加了难度。与具有清晰解剖结构的非恶性情况不同，恶性情况的检测由于缺乏明显的GB边界或形状以及肿块的存在而变得困难。
尽管有最近的努力来克服US在准确检测GBC方面的挑战，这些技术主要是基于图像的。由于前面讨论的挑战，单一图像可能缺乏明确用于恶行病变检测的特征。作者在实验中还观察到，以图像为中心的方法并不能很好地泛化到未见过的数据集。因此，作者主张转变范式，支持从US视频进行GBC检测。值得注意的是，在文献中尚未尝试过基于US视频的GBC检测。
Mask 自编码器（MAEs）。近年来， Mask 自编码器已成为一种有前景的视觉任务表征学习方法。MAE背后的思想是Mask输入的某些部分（也称为标记），然后尝试从可见部分重建被Mask的部分作为预训练任务。通常，一个视觉 Transformer （ViT）生成标记的嵌入。Mask采样策略在有效使用MAEs进行学习方面发挥着重要作用。目前，大多数MAE方法采用了随机Mask策略。对于视频中的随机Mask，常用的有块Mask，帧Mask，或基于管状（tube-based）的Mask（在几连续帧中相同空间位置丢弃标记）。
基于管状的Mask策略被认为在防止时间维度冗余引起的信息泄露方面更有效。然而，研究表明，由于视频中的场景多样性、获取条件以及高低时空信息区域的差异，单一的Mask策略可能并不适合所有数据集。例如，VideoMAE在SSv2数据集上通过随机管状Mask取得了最佳的动作分类效果。对于Kinetics-400数据集，MAE-ST通过随机块Mask达到了最佳性能。图1a展示了不同Mask策略的例子。
在美国的视频中，GB（假定是一种特定的区域）和恶性区域通常只占很小的部分。注意到这些区域相对于帧中的非GB部分含有更多信息，后者是信息量较低的区域。因此，随机 Mask （均匀分布）不利于学习恶性的有效表达。一些最近的方法建议使用自适应 Mask 采样策略以获得更有意义的语义表示。MGMAE建议使用物体运动来指导 Mask 采样。AdaMAE利用策略梯度优化策略，通过最大化期望的标记重建误差来提高属于物体的标记的采样概率。由于在US视频或CT体积中器官大多是静止的，因此运动指导策略不适用于作者的案例。
另一方面，作者的实验表明，AdaMAE并没有显著优于VideoMAE。仅仅关注重建误差，模型可能会低估数据中的关键特征或模式。相比之下，作者采用了一种简单的策略，即FocusMAE，在MAE Pipeline 中采样有效的 Mask 。作者识别出候选的高信息区域，并通过这些区域先验来偏置采样策略，从而从这些Focus的候选区域采样 Mask 标记。通过对高信息区域进行更强的 Mask ，并重建这些标记，FocusMAE学习到了更精细的表达。
本工作的主要贡献是：
- 作者提出，现有在超声图像中进行肝包虫病（GBC）检测的最先进（SOTA）技术显示出次优的准确性和泛化性能。因此，作者主张向基于视频的超声肝包虫病检测范式转变。另外，利用机器学习进行以超声视频为中心的肝包虫病检测问题在文献中尚未被尝试过。作者提供了这一问题的首个解决方案，并展示了一个强有力的 Baseline 。
- 尽管基于视频的GBC分类在准确度、特异性和敏感性方面相较于基于图像的方法有所改进，但作者观察到在MAE中的随机 Mask 提供了进一步改进的机会。特别是，表明恶性的时空区域通常只占视频的一小部分。随机选择 Mask 的标记引入了冗余的背景信息，这需要一种更系统化的方法。为了解决这个问题，作者提出了一种新颖的设计，即FocusMAE，以系统地将 Mask 标记的选择偏向于具有语义意义的候选区域。因此，网络在重建 Mask 标记的同时，被迫学习更精细的GB恶性表现。作者采用这种方法获得了96.4%的准确度，而当前的GBCNet和Radformer的SOTA准确度为84%。
- 作者的Focus Mask 理念是通用的，并且作者通过将这种方法应用于一个公开的基于CT的新冠识别任务[1]来验证其通用性。作者报告了作者的方法相对于现有技术水平（SOTA）[47]的准确度提升了2.2%。
- 同时，作者精心策划了目前可用于GBC检测的最全面的美国视频数据集。作者通过将27个展现GBC的美国视频样本添加到公开可用的GBUSV数据集中来建立该数据集。该数据集将向社区开放使用。
2 Proposed Method
Object Priors in MAE
视觉数据通常表现为由前景物体主导的、在语义上有意义的稀疏信息分布。当前的MAE技术主要使用随机 Mask ，这可能导致次优结果，因为信息可能不是均匀分布的。对于美国（US）视频，GBC在帧中通常只占很小的部分。随机 Mask 往往使网络偏向于学习包含其他器官或腹膜腔的冗余背景的表征。
为了缓解这个问题，作者提倡利用高信息密度的物体位置先验来增强MAE中的表征学习。作者在图1(b)中展示了通过物体定位先验提升 Mask 标记概率的潜在优势的初步证据。作者选择了一个随机的验证分割，包含约20%的GB US视频数据集。作者使用数据集提供的恶性ROI框来确定物体位置。作者手动增加了边界框区域内数据样本的 Mask 概率，并用于自监督预训练。作者变化了概率提升值，用
表示，与背景中的 Patch 相比，边界框内 Patch 的增加概率。
作者的实验揭示了，与随机 Mask 相比，增加边界框内 Patch 的 Mask 概率会导致结果明显提升。然而，过分提高边界框内 Patch 的 Mask 概率可能会损害预训练任务的完整性，并导致性能下降。这些发现强调了认识到不同的图像 Patch 对视觉表征学习贡献不同的重要性。此外，平衡地强调重建前景物体对于最佳性能至关重要。
FocusMAE Architecture
视频子采样。 视频数据中包含时间上的冗余，因为连续的帧在内容上有很高的重叠。作者通过子采样视频来减少时间冗余。假设一个视频包含
帧的话，作者首先以步长为
的方式抽取
帧。
尽管在US帧中的视角可能变化得非常快，但根据作者对数据的观察，帧与帧之间，当距离相当于步长为
的时候，其变化是不显著的。每一帧的大小为
，
和
分别代表帧的高度和宽度，每个帧有三个通道（RGB）。作者将这些子采样帧进一步划分为视频片段——每个片段包含 16 帧。然后，在预训练阶段，作者随机抽取四个片段进行使用。在传递到预训练流程之前，这些帧被调整为
尺寸。
Token 生成。 作者首先将大小为
的视频
划分为不重叠的立方体 Token ，大小为
。
是帧数（时间维度），
和
是帧的高度和宽度。每一帧都有RGB通道。作者使用尺寸为
的3D卷积核，步长为
，以及
个输出通道。
使用这个3D卷积层，作者生成了总共
个 Token ，每个 Token 的维度为
（在作者的设计中
），对应于每个视频。接下来，作者使用在[48]中引入的固定的3D周期性位置编码方案，向 Token 中添加位置信息。
生成目标定位先验。 作者使用深度目标检测网络作为区域 Proposal 网络（RPN），以识别帧内潜在的GB区域。预测的边界框被用作可能包含目标（恶性）的潜在候选区域。作者使用了公开的GBCU 数据集来训练目标检测器。GBCU数据集提供了带有用边界框标记的兴趣区域的超声图像。
训练重点关注两个类别：背景和GB区域。作者降低预测框的置信度阈值以生成多个候选区域。这些区域被用作 Mask Token 采样器中的先验，以提高 Token 的 Mask 概率。如果一个 Token 的空间中心点落在一个先验区域内，那么它的 Mask 概率就会增加。为了定义整个剪辑的一个候选区域，作者取该剪辑内每帧候选区域的并集。
基于区域先验的 Mask 标记采样。 为了生成标记的 Mask 概率，作者遵循[4]的方法，使用一个辅助网络，该网络包括多头注意力（MHA）以及其后的线性层和一个Softmax (
)层。对于给定的嵌入标记
，所有标记的概率分数
生成方式如下：
区域先验接着按照以下方式提升概率得分：
如果第
个标记在空间上位于候选区域内，那么作者会将这个标记的 Mask 概率膨胀为
，其中
是小于
的小分数。然后作者以
的概率为第
个标记无替换地选择一组可见标记索引集合
。
Mask 标记索引的集合由
给出。采样的可见标记数量
是根据预定义的 Mask 比例
计算出来的，等于
。
编码器。 为了计算效率，只有可见（未 Mask ）的标记被传递给编码器。可见标记的数量是
。作者采用了具有空间时间注意力的标准ViT架构。ViT编码器有12层的深度，每层有6个头。嵌入维数是384。
解码器。 在传递给解码器之前，将编码的可见标记与Mask标记拼接在一起。Mask块是可学习的标记，解码器通过这些标记值与其重建之间的均方误差损失来学习重建它们。通常，在MAE中的解码器是一个浅层且窄的ViT。
然而，作者的实验表明，增加解码器的深度有助于性能提升。在为最佳深度进行网格搜索后，作者将解码器深度保持在10。解码器从编码和Mask的标记中重建原始视频立方体，其大小为
。
Training
Mask 重建损失。 作者使用了 Mask 标记的预测值与真实值之间的 均方误差 （MSE）作为预训练MAE的目标函数。损失函数如下给出：
在这里，
和
分别表示预测的标记和标记的标准化真实RGB值。
指的是 Mask 标记的数量。
标记采样损失。 作者使用一种标记采样损失，
，来训练生成采样概率的采样网络。作者采用了AdaMAE提出的采样损失，并通过最大化平均重构误差来定义损失。
这种形式的制定受到强化学习中的REINFORCE算法期望奖励最大化的启发。在这里，可见标记的采样过程是 _动作_，MAE是环境，而 Mask 标记的重构误差是回报。与低信息背景区域相比，高信息区域的重构误差较高。因此，最大化期望重构误差将导致网络预测高信息区域的概率得分更高。损失公式如下：
AdaMAE中的一个关键区别在于，在作者的公式中， Token 概率被区域先验增强，而AdaMAE使用的是针对整个图像的分布的 Token 概率。因此，作者得到了自适应 Token 采样的更精细版本。对数概率处理了下溢和浮点错误。在采样网络中的梯度流保持与主MAE的ViT编码器和解码器独立。
4 Dataset
Curated US Video Dataset for GBC Detection
视频数据收集与整理。 作者使用了公共胆囊超声视频数据集（GBUSV）[8]以及由作者放射科团队收集的额外一组超声视频。GBUSV数据集包含了64个胆囊超声视频，其中32个被标记为良性，另外32个被标记为恶性。为了增强作者基于视频的gbc检测任务的数据集，作者纳入了27个额外描绘胆囊恶性的超声视频。
作者从被转介到印度旁遮普邦昌迪加尔的PGIMER医院接受腹部超声检查的疑似胆囊病变患者那里获得了视频样本。每位患者在招募期间都提供了知情书面同意，并且作者通过完全匿名化数据来确保患者隐私。患者在至少空腹6小时以确保胆囊充分膨胀。
作者的放射科团队在扫描过程中使用了1-5 MHz弯曲阵列换能器（C-1-5D，Logiq S8，通用电气医疗保健）。扫描协议涵盖了整个胆囊（包括底、体和颈）以及任何相关的病变或病理。作者从视频中心裁剪帧以保护患者隐私和标注。处理后的帧大小为360x480像素。图3展示了数据集中的样本序列。
标注。 GBUSV中的视频标签已提供。对于作者的其他视频，作者依赖于活检报告进行标注。此外，作者还咨询了两名分别具有2年和10年腹部超声（US）专业经验的放射科医生，在视频帧中绘制边界框，覆盖整个GB及相邻的肝脏实质。
数据集统计。 该数据集包括59个恶性视频和32个非恶性视频，分别从41个恶性患者和32个良性患者那里收集。总的来说，该数据集包含21,955帧，其中18,406帧归因于被标记为恶性的视频。
数据集划分。 作者为关键实验报告了整个数据集的5折交叉验证指标。交叉验证的划分是以患者为单位的，确保在交叉验证期间，特定患者的所有视频要么仅在训练划分中出现，要么仅在验证划分中出现。
Public CT Dataset for Covid Detection
作者使用了公开可获得的COVID-CT-MD数据集[1]来评估作者提出的方法在不同模态和疾病上的泛化能力。COVID-CT-MD数据集包含了169名（108名男性和61名女性）确诊为阳性COVID-19病例的肺部CT扫描，76名（40名男性和36名女性）正常病例以及60名（35名男性和25名女性）社区获得性肺炎病例的肺部CT扫描。所有样本都由三名不同的放射科医生在患者、肺叶和切片层面上进行了标注。作者使用西门子SOMATOM Scope扫描仪获取扫描，重建图像的输出尺寸设置为
像素。
此外，数据集还包含了临床数据，包括患者的年龄、性别、体重、症状、手术史、随访和RT-PCR测试报告。然而，在作者的实验中，作者没有使用临床数据。作者使用分层随机80:20的比例来获取训练和验证数据。
5 Implementation and Evaluation
预训练。 作者使用PyTorch实施了实验。在MAE权重初始化中，作者使用了Kinetics-400的预训练权重。尽管自然图像数据与医学图像数据之间存在领域差距，但研究表明，在自然图像数据上进行预训练可以提高网络在医学成像任务上的表现。作者采用了第3.2节讨论的视频子采样方案。作为预训练的数据增强部分，作者应用了随机调整大小的裁剪、随机水平翻转和随机缩放。
作者选择ViT-S作为基础网络架构。作者使用的 Patch 大小为
，对于尺寸为
的输入视频，生成的token数量为
。预训练阶段使用AdamW优化器进行训练，学习率LR为
，层衰减为
，权重衰减为
，以300个周期的MSE损失最小化为目标。批处理大小为2。前3个周期使用LR
进行预热。
微调。 在微调过程中对视频进行子采样时，作者使用了更密集的采样率3。作者用16帧组成一个片段。从每个视频中，作者均匀地抽取5个片段。在推理过程中，作者预测每个片段的标签。如果任何一个片段被预测为恶性，那么整个视频就会被标记为恶性。作者使用AdamW优化器，学习率LR
，层衰减
，权重衰减
，最小化软目标交叉熵损失，共训练30个周期。作者使用的批处理大小为4。
作者在实验中使用了一台配备有英特尔至强金牌5218@2.30GHz双核处理器和8块Nvidia特斯拉V100 32GB GPU的机器。
评估指标。 作者使用视频级准确度、特异性（真阴性率）和敏感性（真正阳性率/召回率）来评估基于视频的GBC识别。
6 Experiments and Results
Efficacy of FocusMAE over SOTA Baselines
作者探讨了五种先进的视频分类方法在美 国视频上的GBC分类性能，这些方法分别是Video-Swin，TimeSformer，VidTr，VideoMAEv2，以及AdaMAE。
此外，作者还探索了三种针对超声图像上GBC检测的SOTA技术。除了这些专用模型之外，作者分析了基于图像的CNN分类器和检测器的的性能。作者还研究了三种流行的基于Transformer的分类器 - ViT, DEiT 和 PvT，用于GBC检测。
采用基于图像的方法进行视频分类。 作者使用了与FocusMAE的微调阶段（参见第5节）中相同的视频子采样方案来获取帧和片段。然后，作者使用以图像为中心的方法来预测每个片段中每帧的标签。如果一个片段中大多数帧被预测为恶性，那么该片段就被预测为恶性。如果视频中任何片段被预测为恶性，那么整个视频就被归类为恶性。这些基于图像的方法是在公开的GBCU 数据集上预训练的。
定量分析。作者在表1中展示了基于准确度、特异性和灵敏度的5折交叉验证性能，对比了基准方法和作者提出的FocusMAE。显然，基于视频的技术超越了以图像为中心的GBC检测的最新状态（SOTA）方法，这支持了作者对于该问题向基于视频分类范式的转变的推荐。此外，作者也证明了FocusMAE在检测GBC中的有效性。
定性分析。 作者在图4中展示了定性分析。VideoMAE的随机 Mask 并没有充分 Mask 高信息恶变区域。相比之下，由区域先验引导的FocusMAE通过偏向于恶变定位区域进行 Mask ，生成更强的 Mask 以学习恶变表征。作者在下游任务期间可视化了注意力展开。显然，与VideoMAE相比，FocusMAE的注意力区域突显出语义上更有意义的区域，如胆囊边界和解剖结构。
Generality of the Proposed Method
作者探讨了所提出的FocusMAE方法在公开可用的CT数据集上检测新冠病毒任务的泛用性。表2显示了FocusMAE在准确度、特异性和灵敏度方面都取得了更好的效果，这表明了FocusMAE在学习疾病表征方面的优越性。FocusMAE在两项不同任务上的适用性——1)从超声视频检测胆管癌（GBC），以及2)从CT检测新冠病毒——证实了该方法在两种诊断方式和不同疾病间的泛用性。
Ablation Study
作者对使用了ViT Backbone 网络的FocusMAE进行了消融研究，在US视频数据集上进行了实验。作者展示了在5折交叉验证中的平均分数。
Mask 比例。图4(a)展示了 Mask 比例
如何影响FocusMAE的性能。对于FocusMAE，95%的 Mask 比例达到了最佳的96.4%准确度。VideoMAE使用了90%的 Mask 比例，并采用基于随机管状 Mask 。区域优先引导的方法帮助FocusMAE比VideoMAE采样更具信息性的标记，同时降低了冗余度。
重建损失。 在作者的研究中，作者探讨了改变重建损失函数的影响。作者尝试了两种变体：平均绝对误差或L1损失，以及均方误差（MSE）。图4(b)显示的结果表明，在预训练期间使用MSE损失在准确性方面能获得更好的性能。与使用L1损失训练的模型相比，使用MSE损失训练的模型的平均准确率高出2.2%。
编码器主干网。 图4(c)展示了ViT变体在标记编码任务上的效果。作者实验了ViT-S和ViT-B。作者观察到，较大的主干网在作者的数据上表现不佳，这表明可能存在过拟合现象。
解码器深度。 作者实验了不同数量的解码器块，并将结果展示在图4(d)中。作者发现，当解码器深度从4变化到6时，性能有所提升。然而，当解码器深度增加到8时，性能出现了下降。这一观察与[4, 46]中的结果一致。有趣的是，当作者进一步增加深度时，准确度有所提高，这表明解码器可以从增加深度中受益，并不一定需要是浅层网络。
Analysis on Candidate Region Selection
图6展示了RPN的样本目标区域定位。作者采用了基于FasterRCNN的RPN生成候选区域，以作为FocusMAE中的先验。该RPN实现了0.712的mIoU和0.994的召回率。
7 Conclusion
本研究针对当前基于超声图像的胆囊癌（GBC）检测技术的局限性，强调需要向基于超声视频的方法转变范式。作者名为FocusMAE的新颖设计，战略性地偏向于从高信息区域选择 Mask 标记，并学习GB恶性的高质量表征。FocusMAE在基于超声视频的GBC检测上取得了最先进的结果。作者期望作者的工作对社区有价值，激发对从超声视频中检测GBC这一挑战性问题的兴趣。
此外，作者通过成功将FocusMAE应用于公共肺CT基础上的新冠病毒检测任务，展示了其泛用性，证明了它在两种模态和疾病上的适用性。这表明FocusMAE在未来可能会找到更广泛的应用场景，使其成为多功能诊断解决方案的有希望的步骤。
Appendix A Region Selection Network
作者尝试了多种深度检测器来定位候选区域。在训练候选区域网络时，作者没有使用视频中边界框的标注。相反，作者使用了公开的GBCU数据集来预训练检测器以定位恶性区域。然后，作者降低了阈值，为FocusMAE实验中使用的视频帧生成了多个候选区域。
为了计算在GB定位阶段中的精确度和召回率，遵循[42]的建议，如果预测区域的中心位于真实区域边界框内，作者就判定为一个真正例。反之，如果中心位于边界框外，作者将这种预测归类为由于定位错误导致的假正例。表S1展示了不同候选区域检测器的mIoU和召回率。
图S1展示了RPN的样本目标区域定位。作者采用了基于FasterRCNN的RPN来生成候选区域，以便作为FocusMAE中的先验，因为该检测器实现了最佳的召回率。
Appendix B Region Selection Network Implementation
作者采用了Faster-RCNN模型进行候选区域的选择。使用了冻结的Resnet50特征金字塔作为主干网络。输入尺寸为
。作者使用了SGD优化器，学习率（LR）= 0.005，动量= 0.9，权重衰减= 0.0005。作者设置了批次大小为16，并在GBCU数据集上训练了60个周期。
Appendix C Visualization
图S2和图S3展示了所提出的FocusMAE方法在附加数据样本上的注意力可视化。显然，FocusMAE能够关注到疾病检测的显著区域。
参考
[1].FocusMAE: Gallbladder Cancer Detection from Ultrasound Videos with.