点击下方卡片，关注 「AI视界引擎」 公众号
全幅病理图像（WSIs）对于现代病理诊断至关重要，但其千兆像素 Level 的分辨率和稀疏的有信息区域提出了重大的计算挑战。由于数据量巨大以及非信息区域的冗余处理，传统在计算机视觉和自然语言处理中广泛使用的密集注意力机制对于WSI分析是不切实际的。
为了解决这些挑战，作者提出了记忆效率优化的稀疏金字塔注意力网络与移位窗口（SPAN），从其他领域中先进的稀疏注意力技术中汲取灵感。SPAN引入了一种稀疏金字塔注意力架构，它分层次地关注WSI内的有信息区域，旨在减少内存开销同时保留关键特征。
此外，引入移位窗口使模型能够捕捉到对准确分类至关重要的长距离上下文依赖。作者在多个公开的WSI数据集上评估了SPAN，观察到其具有竞争力的性能。
与现有方法不同，由于内存限制，这些方法通常难以建模空间和上下文信息，作者的方法能够准确建模这些关键特征。作者的研究还强调了注意力机制设计中的一些关键元素的重要性，如移位窗口方案和分层结构，这些元素对SPAN在WSI分析中的有效性贡献很大。
因此，SPAN在内存高效和有效的WSI数据分析方面的潜力得到了展示。
1 Introduction
全幅病理切片图像（WSIs）已成为现代数字病理学中不可或缺的工具，它使得组织病理学切片的数字化成为可能，并促进了计算机辅助诊断[6; 1]的发展。然而，WSIs的吉像素分辨率对自动化分析提出了重大的计算挑战，其数据量远远超过了传统图像分析技术的处理能力，这些技术是为自然图像设计的。
近年来，深度学习在各个领域都取得了显著的进展，改变了作者接近和解决复杂问题的方式。这一进展主要是由能够从海量数据中学习丰富、分层表示的强大架构的发展推动的。特别是自然语言处理（NLP）领域，在引入基于 Transformer （transformer）的模型[10; 5; 23]后，取得了重大突破。这些模型通过有效地捕捉文本数据中的长距离依赖和上下文信息，革新了语言理解、生成和翻译等任务。同样，计算机视觉（CV）领域也经历了快速的发展，这主要归功于卷积神经网络（CNNs）[14; 25]的成功，以及最近 Vision Transformers（ViTs）[12; 9; 11]的出现。这些最先进的架构在图像分类、目标检测和语义分割等各项任务中，通过学习从视觉数据中提取有意义的特征和表示，取得了卓越的性能。
尽管这些进展革命化了深度学习领域，并引入了能够有效捕捉长距离依赖和关注相关信息的注意力机制，但它们在可扩展性和效率方面也提出了挑战。密集注意力的二次复杂度在处理更长的序列或更大的数据量时带来了重大挑战。为解决这一计算瓶颈，提出了各种技术。Sparse Transformers[40; 3]选择性地关注一部分标记（tokens），将计算复杂度从二次降低到次二次。另一方面，Linear Transformers[38; 17]近似自注意力机制以实现线性计算复杂度，从而能够处理更长的序列。此外，还有许多其他在一般领域的发展，例如位置编码技术[32; 29]。这些进展的成功表明，将类似技术应用于WSIs的分析具有潜力，因为它们可能有助于解决由WSIs的大尺寸和空间复杂性带来的挑战。
在WSI分析中，占主导地位的模式是采用基于两阶段 Patch 的框架。这种方法首先将WSI分割成较小的、不重叠的 Patch ，并去除背景。然后，每个 Patch 由一个固定的特征提取器处理，以生成高维特征表示。这些特征通过多实例学习（MIL）模型，例如基于注意力的MIL（ABMIL）[16]，进行聚合以预测切片 Level 的结果。尽管许多WSI分析方法试图通过加入额外的损失或训练策略来扩展ABMIL，但它们将 Patch 视为独立且同分布（i.i.d.）的实体（图1，底部），忽略了WSI固有的丰富空间结构和长距离依赖。WSI的吉像素特性和大量非信息区域的存在，使得在依赖建模中直接应用通用计算机视觉（CV）和自然语言处理（NLP）的进展面临挑战。为了弥合通用深度学习领域与WSI分析之间的差距，作者提出了内存高效稀疏金字塔注意力网络（SPAN）。SPAN引入了一种新颖的框架，该框架高效地利用了WSI的层次性质和长距离上下文信息，同时保持了计算效率。
SPAN的关键组成部分旨在解决当前基于 Patch 的方法的局限性。稀疏金字塔注意力架构在层次上关注WSI内的信息区域，在减少计算开销的同时保留关键的诊断特征。通过采用金字塔结构，SPAN在多个尺度上高效处理WSI，捕捉局部和全局上下文。稀疏注意力机制有选择地关注信息区域，减轻由大型非信息区域引起的计算负担。此外，SPAN融合了移位窗口和全局 Token ，以增强模型捕获长距离上下文依赖和全局信息的能力。此外，SPAN与各种通用技术兼容，允许无缝整合并适应WSI数据的特定属性。这种灵活性为未来探索和精炼SPAN框架提供了机会。
本文的主要贡献如下：
- 作者提出了SPAN，一个新颖的框架，它将稀疏金字塔注意力与移位窗口相结合，专门为高效有效的WSI分析而设计。
- 作者引入了一种稀疏金字塔注意力架构，它分层关注信息丰富的区域，在降低计算复杂性的同时保留关键的诊断特征。
- 作者将移位窗口和全局信息载体标记结合起来，以增强模型捕捉长距离上下文依赖的能力，这对于准确的疾病分类至关重要。
- 作者在多个公开的WSI数据集上评估了SPAN，与现有最先进的方法相比，在下游分类任务中展示了其卓越的性能。
作者的方法能够精确建模空间和上下文信息，这对于由于内存限制而往往对现有方法构成挑战。
2 Related Works
Attention Mechanisms
注意力机制，特别是自注意力，已经改变了包括自然语言处理（NLP）和计算机视觉（CV）在内的各个领域。基于Transformer的模型，如BERT [10]和GPT [5]的引入，标志着从传统的循环神经网络在语言建模方面的范式转变。通过利用自注意力捕捉文本中的长距离依赖关系，Transformers在一系列任务上取得了最先进的表现，并在NLP中确立了主导地位。然而，自注意力的二次计算复杂度对于处理长序列可能是禁止的。为了解决这个问题，提出了稀疏注意力机制，如Longformer [3]和BigBird [40]，将注意力计算限制在固定的窗口内，显著降低了计算复杂度，同时仍然捕捉重要的长距离依赖关系。
Vision Transformer (ViT) [11]挑战了卷积神经网络（CNN）长期以来的主导地位，证明了自注意力在学习视觉表示方面的有效性。为了进一步提高ViT的性能和效率，几个变体，如Swin Transformer [24]和FasterViT [12]，引入了窗口注意力机制。与NLP中主要用于降低计算复杂度的窗口注意力不同，CV中的窗口注意力主要目的是引入层次结构并融入归纳偏置，从而在各种计算机视觉任务上取得最先进的表现。
位置编码是注意力机制的另一个关键方面，它允许模型融入输入标记的位置信息。在NLP中，绝对位置编码[34; 10; 22]和相对位置编码[31; 8]已经被广泛研究。同样，对ViT中位置编码的研究也成为一个高度活跃的领域，从最初的绝对位置嵌入（APE）[11]到最近的相对位置偏置（RPB）[24]。最近的研究还积极地将大型语言模型（LLMs）中的旋转位置编码技术引入到CV模型中，以增强下游分类、分割任务和高分辨率图像生成的性能[27; 12; 15]。自注意力机制与位置编码的结合极大地提高了模型捕捉长距离依赖和关系的能力，并在广泛任务上提升了性能。
Pyramid Structures in Computer Vision
多尺度特征提取和表示的概念一直是计算机视觉数十年的基本方面。与自然语言处理中通常对数据统一处理不同，视觉数据本质上是分层的，各种尺度上都存在信息。这种层次性质的早期认识可以追溯到如SIFT描述子[26]的开创性工作，该工作采用了尺度空间金字塔来提取尺度不变特征。
深度学习和卷积神经网络（CNN）的出现进一步加深了层次处理在计算机视觉中的重要性。从开创性的AlexNet[18]到更先进的架构如ResNet[14]和ConvNeXt[25]，CNN本质上是以层次化的方式处理视觉数据。特征图的逐步下采样和通道深度的增加使这些网络能够捕捉到多尺度的特征，浅层提取细节信息，深层捕捉更抽象的语义信息。在这个隐含的层次结构基础上，作者提出了显式的金字塔架构，以进一步增强CNN的多尺度能力。SPP-Net[13]引入了空间金字塔池化，以在多个尺度上聚集上下文，激发了一系列多尺度CNN设计。FPN[20]提出了一种自上而下的架构，通过横向连接在所有尺度上构建高级语义特征图。HRNet[35]采取了不同的方法，通过网络中的并行多分辨率卷积和重复的多尺度融合来保持高分辨率表示。
金字塔结构在ViT中也得到了关键认可。虽然原始的ViT[11]使用等向性结构统一处理图像块，但许多后续研究探索了将金字塔结构与有效的注意力机制结合，以提升ViT的性能和效率。PVT[37]将金字塔结构整合到 Transformer 架构中，逐步降低空间分辨率并增加通道维度以创建分层表示。Swin Transformer[24]结合了层次化设计与移位窗口机制，以实现更好的跨窗口信息交换。Focal Transformer[39]提出了一种焦点自注意力机制，在细粒度和粗粒度层面上操作，创建多级层次。FasterViT[12]结合了CNN和ViT的载体标记，以促进不同尺度局部窗口之间的全局信息交换。这些并行的发展加强了对多尺度金字塔表示学习在计算机视觉中的基本重要性的认识。
Whole Slide Image Analysis: Characteristics and Challenges
基于 Transformer 的模型的进步以及金字塔结构在捕捉多尺度信息方面的有效性，有望显著提高NLP和CV领域的性能。然而，由于预处理后的稀疏性和吉像素级的大小，将这些进步直接应用于WSIs是具有挑战性的。
原始的WSIs通常以金字塔格式存储，具有多个放大 Level 供病理学家在不同尺度下检查组织。然而，金字塔结构强调了分层信息对于人类分析WSIs的重要性。
然而，主流的WSI分析方法通常以各向同性的方式操作，统一对待输入数据，而未考虑WSI金字塔结构的内在多尺度特性。这些方法建立在ABMIL [16]之上。例如，CLAM [28]使用额外的网络从ABMIL中预测注意力得分高的块，并将它们与相应的WSIs分组到同一类别中。MHIM 采用孪生ABMIL网络的注意力输出随机丢弃块。DTFD [41]将包划分为子包，并使用这些子包进行ABMIL训练。然而，这些方法未能捕捉到准确分析WSI所必需的关键空间和分层特征。另一种方法，由TransMIL 提出，将块展平成序列，然后 Reshape 成方形以保留一些空间信息。然而，这种方法扭曲了真实的空间关系，在某些情况下可能表现不佳，也忽略了分层信息。
从计算机视觉中多尺度和分层结构的成功以及位置编码技术的进步来看，有效地将空间信息和多尺度特征融入WSI分析显然是至关重要的。作者提出的框架SPAN，采用了一种有效的编码策略，可以在合理的内存使用和速度下精确分析WSIs，使得可以使用与其他活跃研究领域相同的建模技术。
3 Method
Overview
SPAN是一种为高效有效进行WSI分析而设计的稀疏金字塔注意力架构。SPAN的主要组成部分包括一个稀疏卷积块、一个窗口生成块和一个稀疏注意力块。给定输入特征矩阵和坐标矩阵，SPAN首先对输入进行索引。该架构在参数化卷积层和参数化稀疏注意力层之间交替。稀疏注意力层捕获窗口内的局部依赖关系以及使用全局注意力的长距离依赖关系。这个过程聚焦于当前尺度下的信息性区域和交互作用。卷积层逐渐降低空间分辨率以捕获空间和层次特征。SPAN架构的流程在图3中展示。最后，分类头汇总所学习的特征以进行切片 Level 的预测。
Window Generation Block
给定特征输入 ，其中 是非空块的数量， 是特征维度，通常在一般领域使用的传统窗口生成方法[24; 3; 40]在作者的稀疏矩阵情况下可能在效率上不是最优的。这些方法通常直接在密集特征矩阵上操作，通过在矩阵内存中跨越一定数量的元素来获得同一密集矩阵的不同视图。然而，由于矩阵位置的稀疏性，应用相同的处理方法需要先将特征矩阵和坐标矩阵填充为密集形式。由于 通常很大，这种方法将导致内存消耗和计算开销显著增加，这是由于包含了许多不必要的填充操作。
为了解决窗口生成中的这个问题，作者提出了一个利用索引进行高效窗口生成和注意力机制的模块。通过对索引矩阵进行填充，并使用索引来指定后续的窗口注意力计算，作者避免了高维零向量的填充和特征矩阵的复制。如图2所示的过程是一个可并行化的高速模块，只涉及索引操作和执行。
作者将输入的WSI表示为一个稀疏张量 ，以及坐标输入 。此外，作者引入了一个索引矩阵 ，它编码了非空块原始空间位置，建立了一个索引-特征-位置映射。这个映射使作者能够在1D索引矩阵上执行填充操作，而不是直接填充高维特征向量，从而大幅度减少内存开销。窗口生成模块的步骤如下：
算法1 窗口生成模块
Parameterized Feature Extraction Block
作者提出的模型融合了参数化的特征提取块，这些块能够有效地捕获稀疏WSIs中的层次特征和长距离依赖。该架构由两种类型的层组成：卷积层和 Transformer 层。
卷积层鉴于输入特征及其对应位置的稀疏性，作者采用稀疏卷积[7; 21]进行下采样和特征编码。稀疏卷积直接在输入的非零元素上操作，与密集卷积相比，它们在计算上更高效，对内存更友好。
在第一个特征提取块中，作者对输入特征应用卷积，以避免直接下采样并保留初始空间分辨率。这有助于维持输入数据的细粒度细节。在随后的层中，使用尺寸为2、步长为2的稀疏卷积逐步下采样空间形状。此下采样操作将块的数量减少大约4倍，从而形成一个层次编码结构。块数量的减少也加速了后续的注意力计算，并提高了计算和内存效率。
Transformer 层在卷积层之后，作者利用由窗口生成块产生的计算图来利用稀疏窗口注意力机制。它涉及定制图操作，以有效地管理稀疏数据结构并优化注意力机制[36]。此块生成非重叠和移位的窗口， Transformer 层利用这些窗口在局部语境中执行注意力计算。通过使用从这些非重叠和移位窗口计算出的索引，作者避免了大量样本的复制，并可以直接在原始特征向量上执行 Transformer 操作。
尽管这种方法扩展了感受野并在移位窗口内捕获依赖关系，但它可能仍然不足以捕获超出窗口的长距离依赖。为了解决这一限制，作者引入了可学习的全局信息载体标记。这些标记作为可以被所有块标记访问的全局上下文，无论它们所在的局部窗口如何。通过关注这些全局标记，每个块标记都可以将其表示中融入全局信息。同样，全局标记关注所有块标记，使它们能够从整个输入序列中收集信息。这种全局标记与块标记之间的双向交互使模型能够捕获跨越多个窗口的长距离依赖，增强了其在WSI内建模复杂关系的能力。
作者初始化了大小为的可学习相对位置偏差，以在每个窗口内编码位置信息，增强了模型考虑标记相对空间排列的能力。在所有模型中，窗口大小被设定为默认值6。
Classification Head
在堆叠三个模块之后，作者获得了一个凝练且分层的WSI表示。为了执行最终的分类任务，作者引入了一个额外的注意力池化层，以汇总所获得的特征图进行分类。
与通常采用全局平均池化或最大池化进行最终特征聚合的传统卷积神经网络不同，作者使用了一个注意力池化层。这一选择是由WSI数据的独特特性所驱动的。即使在两次下采样步骤之后，样本数量（即， Patch ）仍然相对较大，并且在不同WSI之间差异显著。简单的池化方法可能无法有效处理这种变异性，可能导致重要信息的丢失。注意力池化层根据每个 Patch 对最终分类任务的贡献动态地加权其重要性。
4 Experiments
Experimental Setup
为了评估作者提出的SPAN架构的性能，作者在两个公开的全切片图像（WSI）数据集上进行了实验：CAMELYON-16 [2] 和 BRACS [4]。作者遵循了数据集提供的官方训练、验证和测试划分。如果官方划分不可用，作者使用了以下协议：
- 测试集划分：如果没有提供官方测试集，作者使用种子42随机选择了三分之一的样本作为测试集。
- 验证集划分：如果没有提供官方验证集，作者使用种子42随机选择了训练集的15%作为验证集。
本研究采用的预处理流程与CLAM [28] 对所有数据集的处理几乎相同，增加了一个关键步骤，即与网格大小的细胞对齐图像块。这一对齐步骤对于准确保持图像块之间的空间关系至关重要。通过将图像块边界扩展到最近的图像块大小的倍数（例如，224），作者确保了图像块的尺寸与模型的输入网格一致。这种方法使得图像块能够无缝地映射到整数坐标，保持了图像块之间的真实空间关系。相比之下，如果没有这一对齐步骤，图像块的坐标将是浮点数值，需要四舍五入到最近的整数。这种四舍五入过程可能会潜在地扭曲图像块之间的空间关系。这个预处理步骤可能会导致比原始CLAM预处理流程稍多的图像块数量。与WSI分析的常见做法一致，在作者的实验中，作者采用了ResNet50编码器作为特征提取器。具体来说，作者使用了ResNet50编码器倒数第二层的输出。
为了确保实验的可重复性和一致性，作者使用了固定的随机种子。每个 Baseline 模型都进行了五次运行，以不同的随机初始化来考虑训练中的变异性。所有模型遵循相同的超参数设置，学习率为1e-4，使用AdamW优化器和权重衰减5e-5。
Main Results
表2：在BRACS数据集上，与SPAN方法比较的各方法及其统计显著性（p值）。
表1和表2显示，作者提出的SPAN方法在CAMELYON-16和BRACS数据集上，在准确性和AUC方面显著优于现有 Baseline 。值得注意的是，TransMIL在CAMELYON-16上表现出显著的不稳定性，有一次运行完全未能学习到任何有效的特征，导致性能大幅下降。尽管SPAN的计算成本更高，作者认为在考虑到医疗诊断中准确性的重要性时，这种显著的性能提升是合理的。结果表明，SPAN的架构设计用于捕捉WSIs中的分层结构和长距离依赖，这有助于其性能表现，突显了其提高计算病理学工作流程准确性的潜力。
Ablation Studies
为了评估作者模型中各个组件对性能的贡献，作者使用CAMELYON-16数据集进行了消融研究。Baseline 模型包括可学习的相对位置偏差、注意力池化层、全局 Token 、带有下采样卷积层的金字塔结构、窗口大小为6，以及移位窗口机制。作者修改了位置编码、聚合方法、全局 Token 的有无、金字塔结构、移位窗口机制和窗口大小等方面，以探索不同的配置及其对模型性能的影响。
消融研究结果表明（表3），作者的模型具有鲁棒性和灵活性。即使没有位置编码，作者的模型也能表现良好，这可能是由于移位窗口注意力和卷积层捕获的内在位置信息。此外，作者的模型兼容来自其他领域各种先进的位置编码，如Alibi和RoPE，这突显了其在集成未来位置编码技术方面的可扩展性。尽管这些先进编码目前并没有超过带有可学习相对位置偏差的 Baseline ，但未来的工作可能会揭示更适合WSI特性的频率或变体，从而进一步提高模型性能。至关重要的是，作者的实验强调了在WSI分析中移位窗口机制和通过卷积层的分层下采样的重要性。与具有1D序列结构的文本数据不同，WSI具有2D空间结构。这需要仔细考虑相邻区域。移位窗口机制确保相邻窗口之间的有效通信，捕捉关键的空间关系。如消融研究中性能下降所示，如果没有此机制，非重叠窗口将导致某些相邻视觉 Token 无法计算任何注意力交互。此外，去除金字塔结构，即下采样卷积层，导致性能下降，这突显了它在捕捉多尺度特征方面的重要性。此外，作者观察到，超过一定点后增加窗口大小并不一定能提高性能。尽管这使用了更高的计算资源（图4）。这一现象可能是由于在捕获长距离依赖上的收益递减以及学习过程复杂性的增加，这可能导致窗口大小变得过大时，模型从训练数据中泛化的效率受阻。作者建议未来在WSI分析的研究中考虑融入这些技术以提高性能，并仔细平衡窗口大小和计算效率。最后，作者对全局 Token 的实验表明，它们在携带全局信息方面是有效的。尽管直接使用全局 Token 表示进行分类没有超过附加的注意力池化层，但它仍然产生了具有竞争力的结果。这一发现表明，全局 Token 可以成为捕捉WSI分析中全局上下文的有价值的工具。
5 Conclusion and Limitations
作者引入了SPAN，这是一种为分析吉像素级全切片图像（WSIs）而设计的内存高效的稀疏金字塔注意力网络。在实验中，SPAN在下游的WSI分类任务中展示了具有竞争力的性能。
然而，作者也认识到作者方法的几个局限性。尽管SPAN兼容各种位置编码技术，但在作者的测试中，直接应用现代编码方法并没有带来性能的提升。
未来的研究可以探索可学习的位置编码频率或特定于WSI的频率值，可能进一步增强SPAN的有效性。作者的消融研究还强调了在下采样金字塔结构和移位窗口机制在稀疏注意力模型对WSI分析的效能中的关键作用。这些元素对SPAN的性能至关重要，并可能指导该领域未来的创新。
参考
[1].Memory-Efficient Sparse Pyramid Attention Networks for Whole Slide Image Analysis.
点击上方卡片，关注 「AI视界引擎」 公众号