点击下方卡片，关注「AI视界引擎」公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
基于内窥镜视频的任务，如视觉导航和手术阶段识别，在微创手术中起到了关键作用，它们能够提供实时辅助。尽管近期的视频基础模型显示出前景，但其应用受到计算效率低下(1)以及由于内窥镜预训练数据有限导致的性能不足(2)的问题阻碍。
为了应对这些问题，作者提出了EndoMamba，这是一种为实时推理而设计的基础模型，同时学习泛化的时空表示。首先，为了缓解计算效率低下的问题，作者提出了一种优化用于实时推理的EndoMamba主干。
该主干借鉴了近期状态空间模型的发展成果，通过在单个帧内部集成双向的Mamba块来进行空间建模，并通过在时间域内使用基本的Mamba块进行从前向后的推理解释，从而实现了强大的时空建模和高效的在线视频流推理。
其次，作者提出了一种自监督层次化预训练架构，利用内窥镜视频增强EndoMamba的表示学习，并整合一般视频领域知识。
具体而言，作者的方法结合了去噪重构与辅助监督，利用低级重构来捕获时空结构，并通过高级对齐来从预训练的一般视频领域基础模型中转移更广泛的知识。在四个下游任务——分类、分割、手术阶段识别和定位上进行的广泛实验表明，EndoMamba优于现有的基础模型及特定任务方法，同时保持实时推理速度。源代码将在接受后发布。
unsetunset1 Introductionunsetunset
基于内窥镜视频的任务，如导航[24]和手术阶段识别[14]，已成为医学图像分析中的一个重要研究领域。这些任务通过提供实时自动化辅助，对微创手术产生了重要贡献。内窥镜视频分析的核心挑战在于高效估计时空信息[11]。为了提高性能，大多数研究都集中在设计专门模块来学习特定任务领域的知识。然而，特定任务的方法在泛化能力和适应新数据方面往往表现不佳，限制了其在实际临床环境中的应用[35]。
近年来，基础模型在通用视频理解领域展现了出色的表现[29,30,12]。这些模型已成为各种视频下游应用的主干，相比于专门针对特定任务的方法，提供了更好的泛化能力和扩展性[31]。鉴于其在内窥镜领域的潜在应用价值，EndoFM [32] 引入了一种基于 Transformer 的基础模型来进行内窥镜视频分析，并在下游任务上取得了令人鼓舞的结果。
然而，仍然存在两个主要挑战未得到解决：
（1）在线推理过程中的计算效率低下问题；
（2）有效的预训练数据量有限且质量不足[32]。推理速度的限制主要源自于主干架构。现有的视频基础模型，例如基于 Transformer 的模型[29,30]和双向状态空间模型（SSMs）[11]，在推理过程中，每当有新帧到来时，都需要重复计算所有历史视频帧，这种低效性严重影响了实时应用的效果。至于数据限制，主要的挑战在于相对较小的数据集规模以及内窥镜领域缺乏成对的视觉语言数据。这些限制制约了大规模学习[30]和对比学习[22,12]的发展，而这些都是提高模型性能的关键。
在本研究中，作者提出了EndoMamba，这是一种通过分层预训练构建的高效内窥镜视频基础模型。为了解决计算 Bottleneck 问题，作者设计了EndoMamba Backbone 网络。受近期SSMs [8,4] 在长期记忆和快速推理方面的进展启发，EndoMamba 结合使用双向Mamba（Bi-Mamba）进行单个帧内的空间建模，以及普通的Mamba用于时间域上的因果推理。该架构能够实现强大的时空推理并支持在线视频流中的高效推理。
为了缓解数据限制问题，作者提出了一种分层的自监督预训练方案，使EndoMamba能够在捕获时空结构的同时，从预训练的一般领域模型中利用更广泛的先验知识。
这一方案包括两个主要组成部分：低层次视频重建和高层次特征对齐。对于低层次重建部分，受VideoMAE [29] 的启发，作者采用 Mask 自编码器策略进行数据高效的预训练，从而允许EndoMamba从高度遮挡的视觉线索中重建视频片段。此外，作者还提出了一种辅助的预训练模块来实现高层次特征对齐。 该模块利用在大规模视觉-文本数据上训练的一般领域视频基础模型来进行特征的有效对齐。作者在四个下游任务上进行了广泛的实验。实验结果表明，与EndoFM相比，EndoMamba在分割Dice分数上提升了11.5%，在手术阶段识别精度上提高了21.3%，而推理速度则从每秒9.2帧提升到了每秒46.7帧，内存长度为32帧。
unsetunset2 Methodsunsetunset
在本节中，作者首先简要介绍时空记忆网络（SSMs）（2.1）。接着，作者详细介绍了EndoMamba架构（2.2），该架构能够有效地捕捉泛化的时空表示，同时保证高效的推理过程。最后，为了缓解内镜数据带来的挑战，作者提出了EndoMamba的预训练策略，以增强其表示学习能力（2.3）。
2.1 Preliminaries
深度状态空间模型（SSMs），例如Mamba [8]，由于能够建模长距离依赖关系同时保持计算效率，最近在自然语言处理中取得了成功。这些模型基于经典的SSMs，后者将一维函数或序列
映射到输出
，并通过隐藏状态
实现。系统动力学由演化矩阵
描述，其中包含投影参数
和
。连续系统的工作原理如下：
为了高效训练离散数据，在连续系统中使用零阶保持（ZOH）方法将其离散化，定义如下：
其中离散的对应物
定义为：
基于SSMs，Mamba进一步提出根据当前Token选择性地沿序列传播或遗忘信息，从而实现序列长度的高效线性可扩展性和强大的长距离依赖建模能力。
然而，由于vanilla Mamba以过去到现在的顺序处理输入Token，它不适用于视觉任务。为了克服这一限制，VisionMamba [36] 使用了Bi-Mamba块，该块采用双向扫描方式从前后两个方向处理序列，从而进行图像分析。
2.2 EndoMamba Architecture
受到近期时间序列建模（SSMs）进展的启发，作者提出了一种基于Mamba的 Backbone 网络EndoMamba，旨在提升视频基础模型的计算效率。图1(a)概述了所提出的EndoMamba。由于Mamba块以1-D序列作为输入，作者首先将输入视频
转换为展平的空间-时间路径
，其中
这里的
表示核尺寸，
是批量大小，
表示帧数，而
表示帧的空间分辨率。
为了保留空间-temporal位置信息，作者添加了一个可学习的空间位置嵌入
和一个固定的正弦时间位置嵌入
，以捕获时间位置信息并允许灵活的输入序列长度。嵌入后的输入为：
为了在保持递归推理的同时进行空间-时间扫描，EndoMamba通过在每个图像块内执行双向扫描，并沿时间轴执行因果扫描来处理嵌入输入
. 具体来说，嵌入输入
首先被重新排列为
，其中
且
。然后，
通过
个并行堆叠的Bi-Mamba块进行双向扫描，以在每个图像帧内同时执行。经过处理后，tokens被重新排列回
并传递给
个堆叠的Mamba块，这些块沿时间轴执行因果扫描。最后，处理过的tokens被送入多层感知机（MLP）头以生成预测。对于帧级任务，MLP头输出
。
在训练过程中，EndoMamba以视频片段作为输入，并同时为每一帧生成预测，利用当前帧及其过去帧的信息。对于实时应用，EndoMamba逐帧处理实时视频流，并将过去的状态作为记忆传递。模型基于当前帧本身及其历史上下文生成最终的预测。
2.3 Pretraining Diagram
为了解决内窥镜数据的限制，作者提出了一种分层预训练图用于EndoMamba，如图1(b)所示。该预训练结合了低级视频重构与高级特征对齐，以增强表示学习。对于低级重构部分，作者随机掩盖输入视频的大部区域，促使EndoMamba推理缺失的部分。这一过程通过利用空间时间相关性[29]，自然地鼓励EndoMamba捕捉上下文依赖关系。形式上，重构损失定义如下：
其中，
表示被遮掩的视频区域集合，
是被遮掩的 Token 索引，而
是由 EndoMamba 重建的视频片段。
为了进一步增强表示学习，作者提出将EndoMamba的特征与一个在大规模视频数据上预训练的一般领域视频模型VideoMamba对齐，该视频模型作为教师模型。对齐通过余弦相似性损失实现，定义如下：
其中，
和
分别是 EndoMamba 和教师模型的输出特征；
是未被 Mask 的视频区域集合，
是未被 Mask 的 Token 索引。
该特征对齐功能使EndoMamba能够从更广泛的领域继承知识，但在预训练过程中引入了额外的挑战。与教师模型不同，EndoMamba在显著稀疏的输入上进行了训练。在空间上，它操作的是高度遮掩的视频片段，而VideoMamba则利用完整的输入。在时间维度上，VideoMamba利用了Bi-Mamba，能够在时间轴上进行双向扫描，而EndoMamba则遵循因果扫描。因此，对这些特征进行对齐需要在时间和空间两个维度上完成一个稳健的一般域表示。
最后，EndoMamba 的预训练使用以下损失函数：
其中，
是特征对齐损失的权重。作者的预训练框架使模型能够捕获上下文依赖关系，并从更广泛的领域继承知识，从而增强泛化能力。
6 Tian et al.
unsetunset3 Experimentsunsetunset
3.1 Implementation Details
为了增加数据的多样性和数量，EndoMamba 在 MIX12 上进行了预训练，MIX12 是 12 个内窥镜数据集的组合，包括 MIX7（与 EndoFM 共享）、四个来自公共数据集的专业精选集以及一个来自作者猪支气管镜检查数据集 (PB) 的数据集，详情参见表1 和图2。从 MIX7 到 MIX12，预训练数据量从约 5 百万增加到 11 百万帧。模型使用 Adam W 优化器进行训练，基础学习率为 1e-8，在 500 个周期中采用余弦学习率调度，并且批量大小为 48，前五个周期用于线性预热。对于 EndoMamba 架构，作者设置
。每段视频帧被调整为 224×224 的空间尺寸和 16 帧作为预训练输入。对于特征对齐，预训练的 VideoMamba-S 作为教师模型，根据初步实验将特征对齐损失权重设置为
。
为了评估EndoMamba的泛化能力，作者将其应用于一系列下游任务，包括1）PolypDiag分类[28]和2）CVC-12K结肠息肉分割[2]，采用与EndoFM相同的微调设置。此外，作者将应用范围扩展到其他两个任务：3）在Autolaparo数据集上的手术阶段识别[33]和4）患者支气管镜检查数据集中的分支级定位。与分类和分割任务不同，手术阶段识别和定位任务更侧重于长期记忆和实时推理。对于后两个任务，作者在微调时使用了尺寸为
的32帧输入片段，并按照现有方法进行数据划分[14,26]。手术阶段识别任务的训练轮数设置为50轮，而定位任务的训练轮数设置为20轮。没有使用特定任务的模块，而是仅在EndoMamba的主干网络上添加一个MLP头以产生最终输出。
3.2 Comparison with State of the Art
作者在四个下游任务中将EndoMamba与最近的最优方法（SOTA）进行了比较。1) 分类：与现有的预训练方法和基础模型[5,20,21,19,32,30,11]相比，EndoMamba在F1分数上提高了4.3%，排名第一（表2）。2) 分割：在分割任务中，EndoMamba继续保持着领先的性能（表2），相比于最佳的SOTA模型，Dice分数提高了6.9%。3) 手术阶段识别：如表3所示，EndoMamba优于所有评估的基础模型[30,11,32]，在视频级准确性上比VideoMamba提高了2.7%。此外，它还与最顶尖的任务特定方法[6,7,13,14]（例如SKiT，采用多阶段框架）表现相当，而EndoMamba仅使用其主干网络和MLP头就实现了端到端的估计，达到了可竞争的结果。4) 定位：EndoMamba在评估的基础模型[26,30,11,32]和任务特定模型[24,27]中取得了显著的优势（表4）。支气管镜介入手术通常面临着严重的视觉遮挡问题，如液体和气泡，这使得精确定位依赖于强大的时空推理能力。EndoMamba能够有效地在遮挡后恢复位置，从视频级准确性变异的减少可以看出其在不同病例中的稳定表现。这些结果表明，EndoMamba能够捕捉到足够鲁棒的迁移表示，能够在复杂的下游任务中表现出色，而无需进行任务特定的架构修改。
3.3Ablation Studies
作者探索了EndoMamba在预训练设置和实时性能方面的各个方面。1）层次化预训练：作者将从零开始训练与层次化预训练进行比较，通过在分类和分割任务上的下游性能评估来实现。如图3所示，“MIX12 w. teacher”代表了作者完整的预训练策略，显著优于“Scratch”，后者表示随机初始化。这一改进突显了作者预训练的有效性，其中包括低级视频重建以学习时空结构以及高级对齐以提升性能。2）教师指导：如图3所示，“MIX7 w. teacher”代表了使用教师模型指导的MIX7数据集预训练，在两个任务上的表现提高了约2%，而“MIX7 w/o teacher”则是使用相同的MIX7数据集但无教师指导。这一改进来自于高级特征对齐，它能够将更广泛的知识从一个通用领域的预训练基础模型中转移出来，从而生成更具泛化的表示。3）数据集扩展：为了探索由预训练数据量扩展带来的改进，作者将预训练与MIX7进行比较，MIX7与现有的内窥镜视频基础模型EndoFM对齐，同时还将预训练与作者扩展的MIX12进行比较。结果分别在图3的“MIX7 w. teacher”和“MIX12 w. teacher”中展示。数据集的扩展分别在更具挑战性的分割任务上带来了3%的改进，在分类任务上则带来了适度的改进。4）速度分析：为了评估EndoMamba的实际应用能力，作者将它的速度与其他现有的基础模型进行了比较，包括基于Transformer的EndoFM、VideoMAE v2以及基于Mamba的VideoMamba。所有模型的推理速度、计算复杂度和参数数量结果展示在表5中。为了公平比较，推理速度是在具有32、64和128帧内存长度的NVIDIA A800 GPU上以每秒帧数（FPS）测量的。结果显示，在评估的内存长度下，EndoMamba的推理速度在所有测试模型中都是最高的，显著优于基于Transformer和Mamba架构的模型。值得注意的是，随着内存大小的增加，其推理效率保持稳定，得益于它可以不进行冗余计算就传播过去的记忆的能力。相比之下，基于Transformer的模型受到了二次复杂度注意力机制的影响，而Bi-Mamba块则需要重新计算，这两者都妨碍了效率的提升。通过最小的计算开销保持时间连续性，EndoMamba特别适合实时内窥镜应用程序。
unsetunset4 Conclusionunsetunset
在本研究中，作者提出了EndoMamba，这是一种专门为实时内窥镜视频分析设计的基础模型。提出的EndoMamba Backbone采用空间双向扫描和时间因果扫描，实现了强大的空时建模和高效的推理。
为了应对数据限制问题，作者提出了一种层次化的预训练方案，该方案结合了低层次的视频重建用于时空表示学习，以及高层次与预训练通用领域基础模型的对齐，利用其更广泛的视觉知识来增强表示学习。
实验结果表明，在4个不同的下游任务上，EndoMamba在保持实时推理效率的同时，优于现有基础模型和最新的特定任务方法。
unsetunset参考unsetunset
[0]. EndoMamba: An Efficient Foundation Model for Endoscopic Videos via Hierarchical Pre-training .
点击上方卡片，关注「AI视界引擎」公众号