点击下方卡片，关注
「AI视界引擎」
公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
基于注意力的方法在模拟球形皮层表面的长距离依赖关系方面表现出卓越的性能，超越了传统的几何深度学习（GDL）模型。然而，它们大量的推理时间和高的内存需求为将它们应用于有限计算资源的大型数据集带来了挑战。
受到计算机视觉中状态空间模型的启发，作者引入了无注意力的视觉Mamba（Vim）到球形表面，提出了一种领域无关的架构，用于分析球形流形上的数据。
作者的方法通过将球形数据表示为从细分二十面体派生出的三角片序列来实现表面修补。
所提出的表面视觉Mamba（SiM）在多个神经发育表型回归任务上进行了评估，使用来自新生儿大脑的皮层表面指标。
实验结果表明，SiM在Ico-4网格分区下，比基于注意力和GDL的方法表现更优，推理速度提高了4.8倍，与表面视觉 Transformer （SiT）相比，内存消耗降低了91.7%。敏感性分析进一步强调了SiM识别细微认知发育模式的潜力。
代码： https://github.com/Rongzhao-He/surface-vision-mamba
许多方法已被开发用于传统的欧几里得空间数据，例如卷积神经网络（CNNs）和基于注意力的方法[5, 6]。CNNs通过滑动常规卷积核对输入数据进行处理，计算每个位置的加权总和，而基于注意力的方法将数据视为一系列的 Patch 。然而，针对非欧几里得空间数据的模型却相对较少，这类数据包括图、流形和双曲空间数据，它们具有更复杂的几何形状和距离度量。这些类型的数据在神经科学、社交网络分析和理论物理学等领域通常至关重要，它们的独特结构提供了丰富但未充分利用的信息。
现有处理非欧几里得数据的方法大致可分为基于注意力机制和基于几何深度学习（GDL）的方法。基于注意力机制的方法在捕捉长距离依赖关系方面非常有效，但由于注意力机制在序列长度方面的二次复杂性，在资源受限的情况下受到限制，导致更高的内存消耗和更慢的推理时间。相反，基于GDL的方法直接作用于非欧几里得数据，在处理复杂的几何拓扑结构和距离度量方面非常有效。
然而，它们无法提取全局模式，尤其是在应用于大规模且高度复杂的数据时，这导致了性能的下降。因此，处理非欧几里得数据的一个关键挑战在于在保持相对优秀性能的同时提高效率。
随着状态空间模型（SSMs）[9]的兴起，传统的序列建模方法得到了复兴，展示了其在高效表示学习方面的良好潜力。近期的一种变体，Mamba [10]，通过整合基于输入参数自适应的筛选机制，并使用硬件感知算法并行化扫描，从而减少了内存I/O，提高了推理效率，显著超越了传统的SSMs。受到ViT [11]和ViG [12]、[13]等模型的影响，[13]将Mamba应用于计算机视觉领域，引入了双向SSM结构以解决对方向敏感的挑战，称之为视觉Mamba（Vim）。
非欧几里得数据，尤其是球形皮层表面数据，以其高分辨率、丰富特征和复杂的几何形状为特点，因为皮层表面本质上是一个高维流形。虽然这些数据为神经发育提供了有价值的见解，但它们的有效表示却面临着巨大的挑战，通常需要在性能和计算效率之间取得平衡。受Vim效率的启发，作者将它的应用扩展到大脑皮层分析——一个重要但尚未充分探索的领域——通过提出表面视觉蟒蛇（SiM）。为了使SiM适应皮层表面数据的独特特性，作者采用各种表面 Patch 方法调整了输入序列长度，如图1所示。
本研究的主要贡献可以概括如下：
- 作者引入了SiM，它是Vim的一种改进版本，作为分析映射到零类曲面的数据的一般性 Backbone 网络。
- 利用Mamba在处理具有长序列和自回归特性的任务中的适用性[14]，作者探讨了输入序列长度变化对非欧几里得空间中表面数据的影响。此外，作者进一步实施了自回归预训练，以验证该方法的有效性。
在三项神经发育表型回归任务（包括预测月经后年龄（PMA）和长期语言及运动结果）的广泛实验中，作者的SiM模型相较于基于注意力和GDL的模型表现出了令人鼓舞的效果，其在Ico-4网格分区下的批量推理速度比SiT快4.8倍，并且节省了91.7%的GPU内存。
2.1 几何深度学习
几何深度学习（GDL）已成为分析不规则几何形状的强大工具。尽管传统的卷积神经网络（CNN）擅长处理欧几里得数据，如图像，但在处理不规则数据（例如皮质表面）方面效果较差。GDL模型扩展了CNN的功能，使其能够应用于非欧几里得领域，从而捕捉大脑皮层复杂的拓扑和几何属性。虽然这些模型在捕捉局部特征方面表现出色，但它们往往因为高计算成本或固有的架构限制而面临学习长距离依赖关系的挑战，这限制了它们建模更复杂关系的能力。在脑表型预测任务中，对各种GDL方法（如MoNet[15]和Spherical UNet[16]）的系统比较强调了这些挑战。
2.2 基于注意力的方法
自注意力机制，在Transformer[17]中提出，通过捕捉长距离依赖关系，彻底改变了自然语言处理（NLP）。这种架构成为BERT[18]和GPT[19]等模型的基础。研究者将自注意力扩展到视觉表示学习，通过将图像分割成称为 Patch 的部分，提出了视觉Transformer（ViT）[11]。Swin Transformer[20]通过使用偏移窗口进行分层图像合并，显著提高了目标检测和图像分割等任务的高效性和可扩展性。在医学图像任务中，提出了表面视觉Transformer（SiT）[5]，用于处理不规则几何形状，如大脑皮层。受大脑结构不对称发展的启发，半球关系推理网络（HRINet）[6]被设计出来，以提取两侧半球之间的潜在协变关系。然而，基于注意力的模型不仅因为与序列长度成二次时间复杂度而面临挑战，导致在NLP中建模密集、长序列以及在计算机视觉中处理高分辨率图像时产生显著的计算成本，而且由于需要存储先前序列中的所有键值对，其内存复杂度也受到二次限制。为了解决这些问题，许多研究工作集中在减少二次时间复杂度和内存成本上，如[21-29]所示，通过改变注意力计算的操作，但这也导致了性能的下降。
2.3. 状态空间模型（SSMs）
最近，序列状态空间模型（SSMs）被提出以解决循环神经网络（RNNs）的关键局限性[30]，特别是非并行训练的挑战以及随着序列长度增加而遗忘早期信息的趋势。序列（S4）[9]的结构化状态空间采用零阶保持技术进行离散化，并使用高阶多项式投影算子（HiPPO）[31]将上下文压缩到更小的状态。然而，S4受到线性时不变性（LTI）的限制，导致其根据不同输入进行自适应推理的能力有限。此外，由于对每个片段同等对待，S4未能优先考虑并关注最关键的部分。为了克服这些挑战，Mamba[10]引入了一种选择性扫描机制，根据输入有选择性地提取相关信息。Vim[13]扩展了Mamba在计算机视觉中的应用，设计了一个基于双向SSM的通用视觉 Backbone 网络，以解决对方向敏感的问题，类似于ViT。此外，视觉状态空间模型（VMamba）[32]提出了一种交叉扫描模块，以弥合一维数组扫描和二维平面遍历之间的差异，使得Mamba能够适应视觉数据，同时保持感受野的大小。
- 材料与方法
3.1 图像采集与数据集
本研究中使用的影像数据来自公开的发育人类连接组项目（dHCP）和甘肃省妇幼保健院（GPMCH）。作者利用T1加权（T1w）和T2加权（T2w）图像计算大脑皮层的形态计量学指标。
dHCP已获得英国健康研究伦理委员会的批准（参考编号：14/LO/1169）。此外，作者从GPMCH收集了10名婴儿的T1w和T2w图像（2020- GSFY-05）。这些图像以
的分辨率获取，重叠度为
，并重构至
的各向同性分辨率。
关于dHCP的数据，本研究纳入了516名婴儿，包括早产儿和足月儿，他们的孕周从24周到45周不等。这些婴儿在18个月大时，使用贝利婴儿发展量表进行的神经发育评估数据也可获得。作者采用了以下排除标准：对于孕周预测，(i)排除了那些进行过两次扫描的婴儿的后期扫描，以避免宫外环境因素的影响；(ii)排除了有局灶性异常（放射学评分>2）的足月新生儿，以建立正常脑部发育评估的规范性模型。剩余的婴儿随后被分为两个子集：子集1：34至45孕周出生和扫描的408名参与者；子集2：16名早产儿，他们在34孕周前出生，并在足月等效年龄（>37孕周）进行扫描，以评估早产对大脑发育的影响。对于语言和运动评分的预测，作者保留了那些在两次扫描中识别为子集3的、出生在23至43周妊娠周数（GA）的婴儿的接近40周的扫描：410名婴儿。作者还利用GPMCH的数据作为重复数据集，包括在34至40孕周出生和扫描的10名新生儿，以评估模型的泛化能力。人口统计学细节见表1。
四种皮层表面指标——曲率、沟回深度、皮层厚度和髓鞘化（T1w/T2w比例）被用作特征。每个特征通道使用Z分数进行标准化。将子集1和3在每个标签间隔内按8:1:1的比例分为训练集、验证集和测试集。将四个子集的数据配准到dHCP 40周球形模板，该模板将皮层表面表示为一个由三角形组成的近似球体，每个半球有32,492个顶点。作者使用重心插值将模板球体重新采样为规则的六阶二十面体（Ico-6）。
3.2. 前期准备
SSMs通常被视为线性时不变（LTI）系统，它们通过一个隐藏状态
将输入刺激
映射到输出响应
。隐藏状态随时间的演变受参数矩阵
、
和
的控制。该系统用以下线性常微分方程（ODEs）进行数学描述：
代表控制潜在状态
的状态矩阵，
表示控制矩阵，
是输出矩阵。方程（1）旨在基于观测数据预测系统的状态。由于输入通常是连续的，状态空间模型（SSM）的主要应用在于连续时间表示。然而，由于计算机在处理连续信号方面存在困难，而作者使用的实际数据通常是离散的而非连续的，因此标准的做法是利用零阶保持（ZOH）技术对方程（1）进行离散化。这种方法假设输入信号在采样间隔之间保持恒定，其表达式如下：
和
是离散化的参数矩阵，其中
是离散化步长。输出
然后使用全局卷积核
计算，
是输入序列的长度。核的定义如下：
表示序列索引。
Mamba通过选择性扫描机制和更快的硬件感知算法解决了序列建模中的关键挑战。选择性扫描机制优先提取序列上下文中最重要的信息，将其压缩成精炼状态，避免了同等对待所有元素的低效问题。硬件感知算法通过利用现代硬件能力优化计算效率，实现了低于二次方的时间复杂度。这些特性使Mamba在处理长序列时，既能保持准确性和计算效率，又具有高度的效力和可扩展性。
3.3. S表面视觉Mamba
考虑到大脑的互联性质，一个区域内的变化不可避免地会影响到其他区域。为了捕捉这些长距离依赖关系，作者提出了SiM模型，如图2所示。值得注意的是，输入域被划分为
个块，表示为
，其中
是一个块中的顶点数，C表示特征通道数。随后，这些块被展平为
。接下来，作者使用一个可训练的全连接层将
投影到
维向量。遵循ViT和BERT的设计，一个可学习的类别 Token
被连接在左右半球之间，以表示块序列。为了保留位置信息，标准的一维位置嵌入
被添加到块特征中。
SiM的初始输入
属于
，
和
分别代表左右半球的第一个块。SiM的实现结构遵循与Vim相同。具体来说，对于给定的层
，前一层的输入
的处理方式如下：
3.4 表面修补方法
表面补片方法的选择可以显著影响模型性能。在大多数基于表面的视觉任务中，第二阶二十面体（Ico-2）的每个面通常被用作补片，面内的所有数据点被视为顶点。这种方法将表面分割成320个非重叠补片，每个补片包含153个顶点，补片之间仅共享公共边。由于Mamba在涉及长序列的任务上表现良好[14]，作者通过逐步细分二十面体到更精细的离散 Level 来扩展序列长度，并评估包括一阶到三阶二十面体在内的不同表面补片方法，如表2所示。二十面体细分过程包括三个步骤：（i）在上一细分 Level 的边的中点插入新顶点；（ii）在同一面的相邻新顶点之间生成新边；（iii）将新添加的顶点投影到二十面体的外接球上。不同的表面补片方法在图1中进行了视觉表示。为了解决补片方法的影响，作者使用平均绝对误差（MAE）作为性能指标，使用乘累加操作（MACs）来估计计算成本。
3.5. 训练方法
医学影像数据集通常比自然影像数据集要小，这主要是因为在图像获取方面存在的伦理、隐私和法律限制，成像设备和参数设置的多样性，以及标注大型数据集的挑战。为了克服这些限制，预训练方法对于学习鲁棒特征、提高下游任务性能至关重要。在本研究中，作者探讨了三种训练策略：(i)从零开始训练模型；(ii)微调来自ImageNet（如Vision Mamba发布）的预训练权重；(iii)进行自我监督预训练以学习视觉表示。鉴于Mamba架构适合自回归建模[33]，作者采用了自回归方法进行自我监督预训练，即模型根据先前的信息预测下一个 Token 。三种训练策略的性能均使用均方误差（MSE）作为损失函数进行评估。
- 结果与讨论
4.1. 模型变体
所提出的SiM配置基于三种Vim变体：Vim-Tiny、Vim-Small和Vim Base。表3总结了SiM的架构细节，使用简洁的符号表示模型大小和表面修补方法。例如，SiM-B/3指的是“Base”变体，输入尺寸为2560×180，采用Ico-3网格。在Ico-6恒定分辨率下，高阶二十面体球体产生更细粒度的 Patch ，从而实现更精确的分析。这一功能在医学成像中尤为宝贵，因为更细粒度的洞察对于捕捉与疾病相关的微妙变化至关重要。
请注意： Token 为
的案例代表未指定的值。（*OOM'）表示内存不足。粗体字分别表示MAE和MSE的最佳性能。详细设置见A.1。
4.2. 婴儿大脑年龄预测
如表4所示，SiM模型在PMA预测中与基准GDL模型[7, 8]和基于注意力的方法[5, 6]在三个训练策略（参见第3.5节）下的比较结果。值得注意的是，在采用ImageNet预训练权重进行微调时，三种SiM变体在Ico-3网格上优于所有GDL模型。在模型参数和MACs相当的情况下，SiM-S/3实现了
的性能，超过了SiT-S/3（
）。然而，当从头开始训练时，所有SiM变体的性能明显下降，这可能是由于缺少强先验权重而导致对小型数据集过度拟合的趋势，从而引起性能下降。尽管之前的研究表明自监督预训练可以有效提升模型性能，但这一优势在作者的结果中并不明显。实际上，只有SiM-T/3和SiM-S/3与从头开始训练相比有所改进。这可能归因于过拟合，这阻碍了泛化，或者是因为样本量有限，限制了模型捕捉足够特征的能力。
此外，作者还进行了进一步实验：
在自回归预训练中的解码器设计消融研究中，当解码器的深度为1且宽度为256时，取得了最佳性能。
早产对大脑发育的影响研究表明，在子集2中，早产儿的预测大脑年龄显著低于实际年龄，平均绝对误差（MAE）为0.89±0.87，均方误差（MSE）为1.56±2.71，这表明早产可能会在足月年龄时延缓大脑的发育。
在所有模型中，SiM在预测18个月的语言（SiM-Base/2）和运动（SiM-Small/1）结果方面取得了最佳性能，从零开始训练时，其均方误差（M AE）分别为2.82±2.39和1.55±1.17，均方根误差（M SE）分别为15.15±1.70和2.32±1.76。
总体泛化验证实验。作者使用Replication数据集来验证模型的泛化性能。结果显示，尽管所有结果都经历了一定程度的下降，但作者在使用ImageNet预训练权重进行训练时，SiM-Small/3模型展现出最佳的泛化性能，其均方误差（MSE）为1.17±0.95，平均绝对误差（MAE）为2.89±2.17。
上述实验内容详见附录B，所有实验的参数设置总结于附录A。
4.3 长序列与效率分析
图3a展示了不同表面修补方法下，微型SiM模型在性能和效率方面的表现。就平均绝对误差（MAE）而言，SiM在所有其他修补方法中略优于SiT，除了使用Ico-2网格时，且随着二十面体球阶数的增加，两个模型均显示出MAE的下降趋势。在每秒帧数（FPS）方面，当二十面体球阶数低于3时，SiM比SiT略慢，但随着阶数的增加，SiM的性能超过了SiT（图3b）。对于GPU内存使用情况，当二十面体球阶数提升时，SiM的效率优于SiT。值得注意的是，使用Ico-4网格时，SiM的速度比SiT快4.8倍，同时GPU内存消耗比SiT减少91.7%（图3c）。所有关于效率分析的实验都是在40G A100设备上进行的。这些结果突显了SiM在细粒度任务中的适用性及其在临床应用中的潜力。
4.4 对年龄预测有显著贡献的皮层区域
作者对子集1的测试数据集进行了敏感性分析[34]，以评估单个顶点在皮质表面对于大脑年龄预测的贡献，如图4所示。对于大脑表面的每个顶点，作者评估了四个形态学特征（即曲率、沟回深度、皮质厚度和髓鞘化）。"全部"取消所有特征，而其他结果则每次仅取消每个顶点的一个特征。贡献通过测量取消前后的性能变化来量化，变化越大表示影响越大。
每个顶点/特征的平均性能变化被计算出来，并使用z分数进行归一化，然后在皮层表面上进行可视化（图4）。热色的强度表示每个顶点/特征的影响。在右半球，关键区域包括颞叶、前中央回和前额叶及旁中央皮质。当所有特征或仅曲率被遮盖时，左半球显示出对前额叶皮质、感觉皮质、语言区域和顶叶皮质的类似关注。沟深度通过强调颞叶、中央沟区域和额顶叶区域来补充这些发现。皮质厚度和髓鞘化突出了额叶区域，特别是前岛叶、前额回和眶额叶皮质。
- 结论
在这项研究中，作者介绍了表面视觉曼巴（SiM），一种针对零类曲面设计的、具有次平方时间复杂度的创新视觉 Backbone 网络。作者验证了SiM在从皮质表面数据预测神经发育表型这一挑战性任务中，相较于SiT是一种更稳健和高效的替代方案。
利用曼巴在处理长序列和自回归建模方面的优势，作者通过各种表面 Patch 方法扩展了序列长度，并进行了自回归预训练。虽然SiM对序列长度敏感，但自回归预训练的益处有限，这可能是由于样本量小的限制。
使用更长的序列有助于进行更细粒度的划分，增强了识别潜在病理特征的能力，这些特征在临床应用中至关重要。此外，SiM提供更快的推理速度和更低的GPU内存消耗，使其既高效又实用。
敏感性分析还强调了SiM的可解释性，突显了其在医学研究和应用中的潜在价值。
参考
[1]. Surface Vision Mamba: Leveraging Bidirectional State Space Model for Efficient Spherical Manifold Representation .
点击上方卡片，关注
「AI视界引擎」
公众号