[论文解读] HoloHisto: End-to-end Gigapixel WSI Segmentation with 4K Resolution Sequential Tokenization

Yucheng Tang, Yufan He+12 more|arXiv (Cornell University)|Jul 3, 2024

CCD and CMOS Imaging Sensors被引用 4

一句话总结

HoloHisto 提出了一种用于全切片图像（WSI）分割的端到端框架，采用 4K 分辨率的序列化标记化方法，实现全 WSI 的直接输入/输出，无需分块处理。该方法结合 4K 分块采样器与两阶段 ViT 架构（采用 ReLU 线性注意力机制），在新型 WSI 级肾病病理基准上实现最先进性能，Dice 分数最高提升 8.8%。

ABSTRACT

In digital pathology, the traditional method for deep learning-based image segmentation typically involves a two-stage process: initially segmenting high-resolution whole slide images (WSI) into smaller patches (e.g., 256x256, 512x512, 1024x1024) and subsequently reconstructing them to their original scale. This method often struggles to capture the complex details and vast scope of WSIs. In this paper, we propose the holistic histopathology (HoloHisto) segmentation method to achieve end-to-end segmentation on gigapixel WSIs, whose maximum resolution is above 80,000$ imes$70,000 pixels. HoloHisto fundamentally shifts the paradigm of WSI segmentation to an end-to-end learning fashion with 1) a large (4K) resolution base patch for elevated visual information inclusion and efficient processing, and 2) a novel sequential tokenization mechanism to properly model the contextual relationships and efficiently model the rich information from the 4K input. To our best knowledge, HoloHisto presents the first holistic approach for gigapixel resolution WSI segmentation, supporting direct I/O of complete WSI and their corresponding gigapixel masks. Under the HoloHisto platform, we unveil a random 4K sampler that transcends ultra-high resolution, delivering 31 and 10 times more pixels than standard 2D and 3D patches, respectively, for advancing computational capabilities. To facilitate efficient 4K resolution dense prediction, we leverage sequential tokenization, utilizing a pre-trained image tokenizer to group image features into a discrete token grid. To assess the performance, our team curated a new kidney pathology image segmentation (KPIs) dataset with WSI-level glomeruli segmentation from whole mouse kidneys. From the results, HoloHisto-4K delivers remarkable performance gains over previous state-of-the-art models.

研究动机与目标

• 为解决基于分块的 WSI 分割方法因受限上下文而无法捕捉细微组织关系的问题。
• 实现无需预定义分块或裁剪策略的全太比特级 WSI 的端到端训练与推理。
• 开发一种可扩展、高效的架构，能够对富含上下文与局部细节的超高分辨率（4K）图像进行建模。
• 创建一个新的金标准 WSI 级分割数据集（KPIS），用于病理学模型的基准测试。
• 证明 4K 分辨率序列化标记化与基于 ReLU 的注意力机制在病理影像密集预测中的优越性。

提出的方法

• HoloHisto 采用 4K 分辨率分块采样器，相比标准 512×512 分块，可捕获多 31 倍的像素，从而提供更丰富的视觉上下文。
• 利用预训练图像标记化器将 4K 图像特征转换为离散的标记网格，降低维度的同时保留感知结构。
• 采用具有多尺度注意力与基于 ReLU 的线性注意力机制的两阶段视觉 Transformer（ViT）架构，替代标准的 softmax 注意力机制，以降低计算成本。
• 框架集成 cuCIM 与 MONAI，实现实时、动态的 WSI 加载与多倍率下的在线数据增强。
• 模型在全 WSI 上进行端到端训练与推理，直接输出太比特级分割掩码，无需重建过程。
• 提出一种新型随机 4K 采样器，通过从 WSI 中任意感兴趣区域采样，实现高效的数据加载与训练。

实验结果

研究问题

• RQ1 与基于分块的方法相比，整体性端到端方法是否能在太比特级 WSI 上实现更优的分割性能？
• RQ2 4K 分辨率序列化标记化是否显著优于标准 256×256 或 512×512 分块，在上下文建模与特征表示方面？
• RQ3 基于 ReLU 的线性注意力机制是否能在计算成本更低的前提下，有效替代超高分辨率图像建模中的标准 softmax 注意力机制？
• RQ4 采用 4K 分块策略的直接 WSI 输入/输出在密集小目标任务（如肾小球检测）中，是否显著提升分割精度？
• RQ5 新的 WSI 级基准数据集（KPIS）是否能有效支持下一代病理学分割模型的开发与评估？

主要发现

• HoloHisto-4K 在 4K 分块上的 Dice 分数达到 88.85，比次优方法（SAM ViT-H 多尺度）高出 0.87 分。
• 在全 WSI 级别，HoloHisto-4K 的 Dice 分数达到 86.11，显著优于依赖分块与重建的基线模型。
• 消融实验表明，将标记化器与 ReLU 注意力结合，相比基线（Linear Proj + MHSA），在 4K 分块上可使 Dice 分数提升 1.85 分。
• 使用随机 4K 采样器替代分块策略，可使 WSI 级 Dice 分数提升 3.67 分，证明了端到端学习的优势。
• 在 NEPTUNE 数据集上，HoloHisto 在所有组织尺度下均获得高于基线的 Dice 分数，尤其在 3000×3000 分辨率下保持一致的性能增益。
• 模型可在单次前向传播中成功生成完整的太比特级分割掩码，充分证明其真正的端到端能力。