[论文解读] Surface Vision Transformers: Attention-Based Modelling applied to Cortical Analysis

Simon Dahan, Abdulah Fawaz+7 more|arXiv (Cornell University)|Mar 30, 2022

3D Shape Modeling and Analysis被引用 6

一句话总结

本文提出表面视觉变换器（SiT），一种基于注意力机制的新型深度学习框架，通过将大脑皮层表面投影到球面流形上，并将其处理为三角形补丁序列，将视觉变换器架构适配于皮层表面数据。该方法在神经发育表型回归任务中优于表面卷积神经网络（CNN），并揭示了与围产期皮层发育相关的生物意义明确的注意力模式。

ABSTRACT

The extension of convolutional neural networks (CNNs) to non-Euclidean geometries has led to multiple frameworks for studying manifolds. Many of those methods have shown design limitations resulting in poor modelling of long-range associations, as the generalisation of convolutions to irregular surfaces is non-trivial. Motivated by the success of attention-modelling in computer vision, we translate convolution-free vision transformer approaches to surface data, to introduce a domain-agnostic architecture to study any surface data projected onto a spherical manifold. Here, surface patching is achieved by representing spherical data as a sequence of triangular patches, extracted from a subdivided icosphere. A transformer model encodes the sequence of patches via successive multi-head self-attention layers while preserving the sequence resolution. We validate the performance of the proposed Surface Vision Transformer (SiT) on the task of phenotype regression from cortical surface metrics derived from the Developing Human Connectome Project (dHCP). Experiments show that the SiT generally outperforms surface CNNs, while performing comparably on registered and unregistered data. Analysis of transformer attention maps offers strong potential to characterise subtle cognitive developmental patterns.

研究动机与目标

• 解决卷积神经网络（CNN）在建模非欧几里得、不规则皮层表面长程关联方面的局限性。
• 将原本为图像设计的基于注意力的视觉变换器架构——适配于通用亏格为零的表面数据（如皮层表面）。
• 在发育人类连接组计划（dHCP）提供的皮层度量数据上，评估所提出的表面视觉变换器（SiT）在发育表型回归任务中的性能。
• 通过解释注意力图，揭示围产期皮层发育的时空模式。
• 研究数据去混杂（如考虑扫描时的孕龄）对新生儿出生年龄预测模型性能的影响。

提出的方法

• 使用细分的二十面体球面将皮层表面数据投影到球面流形上，以实现规则的补丁化处理。
• 将球面表面表示为源自二十面体镶嵌的三角形补丁序列，从而支持基于序列的建模。
• 应用带有多头自注意力（MSA）层的视觉变换器架构，以在保留序列分辨率的同时建模补丁之间的依赖关系。
• 集成可学习的位置嵌入，以在补丁序列中编码空间上下文信息，类似于图像上的视觉变换器。
• 通过将扫描时的孕龄（PMA）归一化并投影到可学习嵌入中，实现去混杂策略，该嵌入被添加到补丁嵌入中，以减少出生年龄预测中的混杂因素。
• 从零开始训练SiT模型，并采用自监督预训练，使用原始和模板对齐的皮层表面数据。

实验结果

研究问题

• RQ1 视觉变换器架构能否有效适配于非欧几里得、亏格为零的流形上的皮层表面数据？
• RQ2 所提出的表面视觉变换器（SiT）在预测神经发育表型方面，与现有几何深度学习框架（如表面CNN）相比表现如何？
• RQ3 SiT模型中的注意力图是否揭示了围产期皮层发育的生物上合理的模式，特别是足月和早产新生儿？
• RQ4 对扫描时孕龄（PMA）进行去混杂在多大程度上提高了新生儿皮层表面数据中出生年龄预测的准确性？
• RQ5 自监督预训练和网络架构设计选择在多大程度上影响注意力图模式以及模型在未注册和注册数据上的泛化能力？

主要发现

• SiT模型在预测神经发育表型方面优于表面CNN，尤其在原始（未注册）数据设置中表现更优，经去混杂和自监督训练后，出生年龄（GA）预测的平均误差为1.61周。
• 经自监督训练的SiT模型注意力图将注意力从内侧壁（表面处理的伪影）转移至外侧皮层区域，特别是联合区，与已知的发育轨迹一致。
• 不同注意力头的注意力模式具有互补性，并在与围产期从初级皮层到联合皮层发育序列相关的区域表现出高度激活。
• 该模型在注册和未注册数据上均保持强大性能，表明其对解剖变异具有鲁棒性。
• 在表现最佳的SiT-tiny模型中，对扫描时孕龄（PMA）进行去混杂使预测误差降低了0.24周（从1.85降至1.61），证明了在发育表型分析中考虑扫描时间的重要性。
• 即使经过自监督训练，GA预测任务的注意力图对内侧壁仍保持敏感，表明该区域可能携带预测胎龄的信号，尽管其模式一致性低于PMA任务。