扩大组织病理学自监督学习：介绍 Phikon
机器学习和计算机视觉的最新创新为人工智能在医学领域的应用开辟了新机遇。数字病理学是图像表示和分类可以带来新研究突破和更高效疾病诊断的领域之一，这两者都有助于改善患者预后。
我们很高兴地宣布 Phikon，这是一个由 Owkin 使用基于自监督 Transformer 的框架 iBOT 开发的模型。在本文中，我们将详细介绍用于构建模型的数据集、训练技术和框架。我们还构建了一个 Colab notebook，以便社区可以轻松地利用该模型提取特征或使用 Hugging Face Transformers 和 PEFT 库对自己的数据集进行微调。
有用链接：
背景
组织病理学是癌症诊断和治疗计划的重要技术。这项技术要求病理学家分析载玻片上的病变组织，使用显微镜识别疾病的模式和标记。能够使用深度学习自动化或增强这些工作流程的部分，对患者来说具有巨大影响——可能使癌症诊断更快、更准确。
然而，正如其他领域一样，实施人工智能的瓶颈往往是高质量数据的可用性。训练这些模型所需的图像必须由训练有素的病理学家以耗费成本和劳动密集型的方式进行标注。
过去，这个问题通过使用在 ImageNet 上训练的模型进行迁移学习来解决。在 ImageNet 上训练的卷积神经网络 (CNN) 表现良好，可作为组织病理学图像的强大特征提取器，但受限于域外预训练的局限性。组织病理学图像包含复杂的细胞结构，而不同的数据源在颜色、纹理和染色方面也各不相同。因此，在 ImageNet 上预训练的模型难以充分捕捉对疾病诊断至关重要的细节。
然而，最近自监督学习作为一种有前景的解决方案出现，可以利用大量的未标注数据。我们的研究表明了使用掩码图像建模 (MIM) 进行域内预训练的优势。MIM 是一种受 BERT 启发的技术，其中图像（补丁或像素）的部分被随机掩盖。然后，模型可以通过重建被掩盖的部分来学习有意义的表示。
训练
模型的预训练数据由来自癌症基因组图谱 (TCGA) 的公开数据集组成。我们考虑了 3 个不同的数据集：COAD4M（结直肠癌）、PanCancer4M 和 PanCancer40M（多癌种）。我们应用了我们的全切片图像 (WSI) 处理流程，包括物质检测、伪影去除和并行切片，以从 WSI 中提取切片以进行高效训练。
我们使用 iBOT 框架来训练从 WSI 中提取的切片模型。为了评估其可扩展性，我们构建了五种不同尺寸、架构和预训练数据的模型。有关我们训练的不同模型的信息可在下面找到。用于比较的 MoCoV2 模型是 Owkin 先前发布的一个模型。
然后，我们对 17 个不同的下游任务测试了该模型，其中之一是使用 Camelyon16 检测转移，Camelyon16 是一个来自淋巴结切片的 H&E 染色载玻片数据集。
下面的代码展示了使用该模型从组织病理学切片中提取特征的简单程度。
import requests
from PIL import Image
import torch
from transformers import AutoImageProcessor, ViTModel
# download an image from GTeX website
url = "https://biospecimens.cancer.gov/gtexbiobank/images/histology_inset.jpg"
image = Image.open(requests.get(url, stream=True).raw).resize((224, 224))
# load phikon
image_processor = AutoImageProcessor.from_pretrained("owkin/phikon")
model = ViTModel.from_pretrained("owkin/phikon", add_pooling_layer=False)
# process the image
inputs = image_processor(image, return_tensors="pt")
# get the features
with torch.no_grad():
outputs = model(**inputs)
features = outputs.last_hidden_state[:, 0, :] # (1, 768) shape
由于此特征提取可能计算成本高昂，我们分享了一个已使用该模型提取的 Camelyon16 特征数据集。这些特征在笔记本中用于较小模型的癌症分类。
微调
该模型的真正强大之处在于它如何通过微调来提高特定癌症亚型的性能。在我们的笔记本中，我们展示了如何通过迁移学习和使用 LoRA 的微调，将 Phikon 专门用于结直肠癌。我们在 NCT-CRC 数据集的子集上训练了该模型，我们也已将其上传到 Hugging Face。
仍然需要研究优化模型性能的最佳方法，以及是否应通过微调整个模型、最后的 Transformer 块或通过 LoRA 来实现。
我们鼓励任何对组织病理学感兴趣的人，利用准备好的笔记本对模型进行微调，以适应自己的用例，因为该模型在域外任务中表现良好。
结论
为了进一步评估该模型，我们探讨了域内预训练与域外预训练的优势，以及 iBOT 框架相对于纯粹对比方法的优势。
此外，我们还探讨了扩展模型架构对整体性能的影响。将 ViT-S (21.7M) 模型扩展到 ViT-B (85.8M) 显着提高了性能，所有任务的平均增益为 2.5%。然而，令人惊讶的是，结果表明，进一步将架构尺寸从 ViT-B (85.8M) 增加到 ViT-L (307M) 并没有提高性能，反而导致整体性能损失 0.2%。
这表明，如果数据集大小没有相应扩展，那么仅仅增加模型大小并不总是能带来性能提升，这是一个重要的例子。
然后，我们通过重新训练 ViT-B (85.8M) 模型并增加数据集大小和多样性来构建我们表现最好的 Phikon 模型。
全面训练和微调模型之间的权衡始终是一项具有挑战性的任务，尤其是在医疗领域。由于用例的重要性，性能的微小改进相对于其他应用而言被过高评估。如果您对训练和使用的超参数有更详细的兴趣，我们邀请您阅读完整论文。
我们的研究结果将受益于进一步的探索和验证，特别是在数据集整理方面。我们相信这项工作为组织病理学开源基础模型铺平了道路。