[論文レビュー] Cluster-to-Conquer: A Framework for End-to-End Multi-Instance Learning for Whole Slide Image Classification

Yash Sharma, Aman Shrivastava+4 more|arXiv (Cornell University)|Mar 19, 2021

AI in cancer detection参考文献 26被引用数 51

ひとこと要約

C2C は、スライドごとにパッチをクラスタリングし、クラスタからサンプルを採取し、 KL-divergence 正則化を用いた適応型アテンションを用いてスライドレベルの予測を改善する、WSI分類のエンドツーエンド MIL フレームワークを提示します。

ABSTRACT

In recent years, the availability of digitized Whole Slide Images (WSIs) has enabled the use of deep learning-based computer vision techniques for automated disease diagnosis. However, WSIs present unique computational and algorithmic challenges. WSIs are gigapixel-sized ($\sim$100K pixels), making them infeasible to be used directly for training deep neural networks. Also, often only slide-level labels are available for training as detailed annotations are tedious and can be time-consuming for experts. Approaches using multiple-instance learning (MIL) frameworks have been shown to overcome these challenges. Current state-of-the-art approaches divide the learning framework into two decoupled parts: a convolutional neural network (CNN) for encoding the patches followed by an independent aggregation approach for slide-level prediction. In this approach, the aggregation step has no bearing on the representations learned by the CNN encoder. We have proposed an end-to-end framework that clusters the patches from a WSI into ${k}$-groups, samples ${k}'$ patches from each group for training, and uses an adaptive attention mechanism for slide level prediction; Cluster-to-Conquer (C2C). We have demonstrated that dividing a WSI into clusters can improve the model training by exposing it to diverse discriminative features extracted from the patches. We regularized the clustering mechanism by introducing a KL-divergence loss between the attention weights of patches in a cluster and the uniform distribution. The framework is optimized end-to-end on slide-level cross-entropy, patch-level cross-entropy, and KL-divergence loss (Implementation: https://github.com/YashSharma/C2C).

研究の動機と目的

• gigapixel WSI の課題に対処するため、MIL ベースの WSI 分類におけるエンドツーエンド学習を実現する。
• クラスタベースのサンプリングを活用して、モデルを多様で識別力のあるパッチ特徴にさらす。
• パッチエンコーディング、アテンションベースの集約、および KL-divergence 正則化を統合して、パッチとスライド表現の共同学習を改善する。
• 胃腸疾患および乳癌データセットで、二段階 MIL アプローチと比べて競争力のあるまたは優れた性能を示す。

提案手法

• 各エポックごとに patch embeddings 上の k-means によって各 WSI からパッチを k 個のクラスタにクラスタリングする。
• 各クラスタから k' 個のパッチをサンプリングして、WSI ごとに管理可能なトレーニングサブセットを形成する。
• CNN エンコーダーでパッチをエンコードし、パッチ表現 h を計算して 2 層のアテンションモジュールを適用してインスタンス重み a_n を得る。
• アテンションプーリングによってパッチ表現を WSI 表現 z に集約し、スライドレベルのラベルを予測する。
• 結合損失 L = alpha*L_WSI + beta*L_Patch + gamma*L_KLD を用いてエンドツーエンドで学習する；L_KLD は intra-cluster アテンション分散を正則化する。
• 各クラスタ内のパッチアテンション分布と一様分布との間に KL-divergence を適用して正則化する。

実験結果

リサーチクエスチョン

• RQ1 クラスタベースのサンプリングは、WSI の MIL における learned patch 表現の多様性と質にどのような影響を与えるのか？
• RQ2 アテンションベースの集約を用いたエンドツーエンド学習は、二段階 MIL 手法よりスライドレベルの精度を向上させるか？
• RQ3 KL-divergence 正則化はアテンション分布とモデル性能にどのような影響を持つのか？
• RQ4 C2C は GI 生検 WSI および CAMELYON16 乳癌データにおいて、完全監督および二段階のアプローチと比べてどの程度の性能を示すか？

主な発見

方法	精度	適合率	再現率	F1-スコア
Campanella-MIL	82.8	94.9	74.5	83.5
Campanella-MIL RNN	74.7	75.4	84.3	79.6
Two-Stage Mean	81.6	87.3	80.3	83.7
C2C (w WSI Loss)	81.6	80.7	90.1	85.2
C2C (w WSI+KLD Loss)	83.9	84.9	86.3	85.4
C2C (w WSI+Patch Loss)	85.1	86.5	88.2	87.4
C2C (w WSI+Patch+KLD Loss)	86.2	85.5	92.2	88.7

• C2C は GI データセットのセリアック病 vs 正常分類で二段階 MIL ベースラインを上回る。
• WSI 損失、パッチ損失、および KL-divergence 損失を組み込むと F1 スコアが向上し、適合率/再現率も競合的になる。
• C2C は slide レベルの監視なしで、ResNet-18 バックボーンを用いて CAMELYON16 で高い性能を達成する。
• KL-divergence 正則化は陽性インスタンスクラス間のアテンションを安定化させる（例：MNIST バッグ実験）。
• クラスタベースのサンプリングは多様な識別パッチへの露出を高め、エンドツーエンド学習を支援する。