Weakly supervised training of pixel resolution segmentation models on whole slide images

TL;DR

This paper introduces a weakly supervised method for training pixel-level segmentation models on whole slide images, using confidence-based patch sampling and a noise-robust divergence to handle noisy labels, demonstrated on cancer data.

Contribution

It proposes a novel training approach combining confidence-guided patch sampling and a noise-robust KL divergence extension for weakly supervised segmentation.

Findings

Successful tumor segmentation with morphological consistency

Effective handling of noisy labels in weak supervision

Promising results on CAMELYON 16 dataset

Abstract

We present a novel approach to train pixel resolution segmentation models on whole slide images in a weakly supervised setup. The model is trained to classify patches extracted from slides. This leads the training to be made under noisy labeled data. We solve the problem with two complementary strategies. First, the patches are sampled online using the model's knowledge by focusing on regions where the model's confidence is higher. Second, we propose an extension of the KL divergence that is robust to noisy labels. Our preliminary experiment on CAMELYON 16 data set show promising results. The model can successfully segment tumor areas with strong morphological consistency.

Equations

$$Q_{i,j,k}=\frac{e^{L_{i,j,k}}}{e^{L_{i,j,0}}+e^{L_{i,j,1}}}$$

$${\text{l}}_0=\mathbb{E}[L_{.,.,0}]$$

$${\text{l}}_1=max(L_{.,.,1})$$

$${\text{l}}_1^{\eta }=\mathbb{E}[L_{.,.,1}|L_{.,.,1}>P_n]$$

$${\mathcal{L}}_{\mathbf{\beta }}(\mathbf{p}\Vert \mathbf{q})=-\sum _i\frac{p_i}{\text{boldmath}{\beta }_i}\left[{\left(\frac{q_i}{p_i}\right)}^{{\beta }_{\mathbf{i}}}-1\right]$$

$$\underset{\beta \to 0}{\lim }\frac{p}{\text{boldmath}\beta }\left[{\left(\frac{q}{p}\right)}^{\beta }-1\right]=p\log \left(\frac{q}{p}\right)$$

$$\frac{\partial {\mathbb{E}}_X[{\mathcal{L}}_{\mathbf{\beta }}(\mathbf{p}\Vert \mathbf{q})]}{\partial {\theta }_0}=(1-r)\gamma {\theta }_0^{{\beta }_1-1}-r(1-{\theta }_0{)}^{{\beta }_0-1}$$

$${\beta }_1=\frac{1}{\log {\theta }_0}\left[({\beta }_0-1)\log (1-{\theta }_0)-\log \left(\frac{(1-r)\gamma }{r}\right)\right]+1$$

$$P(i,j|S,\theta ,\alpha )=\frac{Q_{i,j,1}^{\alpha }}{{\sum }_{u,v}{Q}_{u,v,1}^{\alpha }}$$

$$P(S=s^{\ast }|\text{benign})=\frac{{\sum }_{i,j}{Q}_{i,j,1}|s^{\ast },\text{benign}}{{\sum }_{s,i,j}{Q}_{i,j,1}|s,\text{benign}}$$

Full text

1 Introduction

Pathologists generally base their diagnosis on the presence of localized morphological bio-markers in histopathological images. Training deep neural networks on Whole Slide Images (WSIs) for automatic disease detection and localization is the main approach for the development of decision support tools. However it requires large annotated datasets that are difficult to construct because annotations at the pixel resolution are often received through a costly, noisy, and low bit-rate channel. Weakly supervised approaches aim to tackle the problem of data acquisition by relying on a supervision available on the WSI level. For cancer detection task, the WSIs are labeled as $benign$ or $malign$ depending on the cancer presence. This information is cheaper to acquire, more robust, and allows construction of larger datasets. A model is trained using this high level labeling to retrieve localized, pixel level, disease information.

We propose to use a segmentation model that produces pixel level unnormalized log evidences (or logits) that represent localized cancer presence probabilities in WSIs. We form WSI level cancer presence probability from aggregation of pixel level logits. We can train the model to maximise the mutual information between this probability and the WSIs label. Our assumption is that this serves as a proxy to maximize the pixel level mutual information between pixel wise logits and the (unknown) cancer localization.

However this training scheme requires the segmentation model to work on all the pixels of a WSI at a time. Since a WSI often contain an order of $10^{9}$ pixels, that makes this approach intractable. To tackle this problem, we propose to train the model on patches randomly sampled from the WSIs. A patch is given a label that is inherited from the slide it is sampled from to form training samples. The drawback is that this can lead patches to be incorrectly labeled. For instance, the cancerous region of a malignant WSI can represent less than $0.1\%$ of the tissue. Therefore, there is high chance to sample patches from benign regions, they will be incorrectly labeled as malignant. This produces noisy supervision that would prevent the model to converge. We propose to address this with two things. First, we propose an extension of the Kullback–Leibler (KL) divergence that is robust to noisy labels. Second, we sample patches dynamically during training. The sampling is made using a distribution defined by the pixel wise cancer probability map given by the model. Patches are sampled more frequently from regions were the model gives higher cancer probability. In the beginning of the training, the sampling distribution is uniform over the whole slide. As the model converges, the patches are more often sampled from regions containing cancer, thus reducing labeling noise and further improving convergence.

We demonstrate our approach on the CAMELYON 16 dataset [1] and show that the model is able to segment tumors at pixel resolution. Even more, the segmentations are strongly correlated with the morphological structures of the tissues. We evaluated the ability of the model to classify slides as malign or benign. The model reached a ROC AUC score of $0.82$ which is a promising score given that only one set of hyper-parameters have been tested.

2 Related work

Image segmentation at pixel-level under weak supervision has been studied in [2, 3, 4]. These approaches propose to work under the Multiple Instance Learning framework (MIL). The images are seen as bag of pixels or patches. The elements of the bags have an underlying label that is not available. However a label at the bag level is available and relate to the label of its elements. A negative (i.e. $benign$) bag contains only negative elements, while a positive (i.e. $malign$) bag contains at least one positive element. A convolutional neural network (CNN) is used to produce pixel (patch) level features that are aggregated through a pooling function to form the bag level class prediction. However these approaches need the image to fit into memory and do not scale on WSIs that contain giga pixels.

Weakly supervised approaches on WSIs divide the slide into a grid of tiles [5, 6, 7]. Each tile is processed through a CNN to produce a score that allows them to be ranked. The top/bottom tiles are then used to train a WSI classifier. The gradients are propagated into the scoring CNN up to some depth allowing to learn the scores. The scores of tiles are then interpreted as disease localization confidences. These approaches have several constraints. First, they produce disease location at tile resolution and not at pixel resolution. Second, they are limited in their ability to train complex models. In [6] the tiles are, prior to the training, mapped into a fixed feature space using a deep neural network trained on ImageNet [8]. A simpler parameterized function is trained over this feature space to produce the scores, consequently limiting the representation power of the learned model. In [5], the scoring CNN is applied on all the tiles at the beginning of each epochs, the training is made only on the top tile of each slide thereafter. This leads to slow convergence that requires larger dataset.

In [9] the authors propose to scale the training of a CNN on WSIs using the streaming stochastic gradient descent method [10]. This approach does not allow to train segmentation model architectures because it requires upper feature maps layers to fit in memory which prevents the usage of up-convolution layers. The authors propose to extract a saliency map using the gradients at the input pixels.

In [11] the authors propose the use of a deep recurrent attention model that classifies a WSI using information provided from a limited number of patches (or glimpses) using visual attention method. Disease localisation is inferred from the glimpses location.

Our approach combines the ability the train pixel resolution disease segmentation models on WSIs without constraining complexity of the learned CNN. Moreover, we do not train on pre-tiled slides, instead we learn on patches dynamically extracted during training according to the current knowledge state of the model.

3 Model

Our model is defined using a segmentation model represented by a parameterized function $M_{\theta}:[0,1]^{h\times w\times c}\rightarrow\mathbb{R}^{h\times w\times 2}$. $M_{\theta}$ can be an instance of any kind of state-of-the-art segmentation model such as U-net[12] or DeepLab[13]. For any given image $X$ of dimension $h\times w$ with $c$ channels, the model gives $L:=M_{\theta}(X)\in\mathbb{R}^{h\times w\times 2}$ representing the un-normalized log probabilities (logits) for benign and malign classes for each pixels. We get the classes probabilities of a pixel $X_{i,j,.}$ using the softmax function:

[TABLE]

In a weakly supervised setup the pixel level annotation is missing. Instead, the supervision is available on the WSI level. A WSI can either be benign or malign. A benign WSI have all its pixels belonging to the benign class while a malign WSI has at least one of its pixel being malignant. Hence we can translate the segmentation problem to a constrained WSI binary classification problem. The model outputs the WSI’s class by computing the slide level benign and malign logits (we note respectively $l_{0}$ and $l_{1}$) with:

[TABLE]

A WSI is malignant if at least one of its pixels is malignant. We can translate that relationship with the following expression of $l_{1}$:

[TABLE]

However, such strategy results into slow and noisy learning of the model because the gradients would flow through only one pixel at a time. Therefore, we generalize to a top K expression:

[TABLE]

where $P_{\eta}$ is the $100-\eta$ percentile. For low $n$ we retrieve the max definition of $\textit{l}_{1}$ and for $\eta=100$ we get the same average expression of $l_{0}$. As for the pixels probabilities, we compute the WSI’s classes probabilities using the softmax function. We note $q_{0}$ and $q_{1}$ the probability of the WSI to be respectively benign and malignant.

If the dimensions $h\times w$ are not too large, the model can be trained using the cross-entropy loss function on the WSI level probabilities. However, in general $h\times w$ is large and the training is intractable.

3.1 Training on patches using dynamic sampling

Instead of training on the full WSI which is intractable, we propose to train on dynamically extracted patches. The model is trained by classifying the patches instead of the WSIs. The patches inherit the class of the slide they are extracted from. In that setup a patch extracted from a benign slide have all its pixels being benign. However, a patch extracted from a malign slide have high chance to contain no malign pixels. This breaks the assumption that at least one pixel have to be malign within $X$ and makes the patch level supervision noisy. We propose to tackle this problem by using an extended version of the KL divergence loss that is robust to noise:

[TABLE]

where $\mathbf{p}$ is the true distribution and $\mathbf{q}$ the model’s distribution. Since:

[TABLE]

we can properly extend $\mathcal{L}_{\bm{\beta}}$ for the terms where $\beta_{i}=0$, and particularly we have $\mathcal{L}_{\bm{0}}(\mathbf{p}\|\mathbf{q})=D_{KL}(\mathbf{p}\|\mathbf{q})$.

If $\forall i,\beta_{i}<1$, then we have $\mathcal{L}_{\bm{\beta}}(\mathbf{p}\|\mathbf{q})=0$ iff $\mathbf{p}=\mathbf{q}$. The vector $\bm{\beta}$ allows to control the shape of the loss for the different classes.

3.1.1 Tuning $\beta$ under noisy supervision

Let $\gamma$ be the probability that a patch $X$ from a malign slide does not contain any malign pixel. These patches are indistinguishable from patches extracted from benign slides and will confuse the model. Since they contain no malign pixel they should be classified as benign. We can leverage the noise induced by the label inconsistency by adjusting $\bm{\beta}$.

Let’s suppose we have a data-set made of $malign$ and $benign$ slides. Let’s suppose that we can extract one pixel patches from a slide such that $X\in\{0,1\}$ where $X=0$ is a benign pixel such that we have $P(X=0|malign)=\gamma$ and $P(X=0|benign)=1$. Let’s define a trivial model $\Psi_{\bm{\theta}}$ with $\bm{\theta}=\begin{pmatrix}\theta_{0}\\ \theta_{1}\end{pmatrix}\in[0,1]^{2}$, such that $P(malign|X=0,\bm{\theta})=\theta_{0}$ and $P(malign|X=1,\bm{\theta})=\theta_{1}$. Ideally, we would like the model to classify the benign pixel as benign with probability one, i.e. we would like $P(malign|X=0,\bm{\theta})=\theta_{0}=0$.

Optimization of this model using $\mathcal{L}_{\bm{\beta}}$ would give the following gradient over $\theta_{0}$:

[TABLE]

where $r=P(benign)$ is the ratio of benign slides.

In the case of the KL divergence, obtained with $\bm{\beta}=\bm{0}$, under a balanced data-set $r=\frac{1}{2}$, the minimum would be reached on $\theta_{0}=\frac{1}{1+\frac{1}{\gamma}}$. For instance, a noise of $\gamma=33\%$ would give $\theta_{0}\approx 0.25$.

We can leverage the noise and obtain lower value of $\theta_{0}$ by tuning $\bm{\beta}$. The minimum on $\theta_{0}$ is obtained when $\frac{\partial\mathbb{E}_{X}[\mathcal{L}_{\bm{\beta}}(\mathbf{p}\|\mathbf{q})]}{\partial\theta_{0}}=0$ which gives the following value of $\beta_{1}$:

[TABLE]

If we have $r=\frac{1}{2}$ and set $\beta_{0}=0$, if $\gamma=33\%$ and if we wish to get $\theta_{0}=5\%$, we can set $\beta_{1}\approx 0.61$. The figure 1 shows how the parameter $\bm{\beta}$ shapes the loss for two Bernoulli distributions. The false positives are less penalized than the false negatives. This affects the gradient by reducing the slope of the loss for the miss-classified examples. Note that our approach is different than scaling the gradients differently for each class. In fact, with this approach, also called balanced cross entropy in [14], the gradient ratio between a well classified and a miss-classified example would remain constant. In our approach this ratio depends on the parameter $\bm{\beta}$. Also, we differ from the focal loss [14] as this approach is designed to down weight well classified examples to focalize more on the miss-classified ones. This relies on the assumption that the label information is not noisy.

3.1.2 Training with dynamic sampling of patches

We train $M_{\theta}$ by sampling patches from slides according to a distribution defined by the probability map of the model on the slides. Given a WSI $S\in[0,1]^{h\times w\times c}$, the probability that the pixel $S_{i,j,.}$ is $malign$ according to the model is $Q_{i,j,1}$. The probability to sample a patch centered on the pixel $(i,j)$ is given by:

[TABLE]

where $\alpha\geq 0$ controls the entropy of the distribution. With $\alpha=0$, the distribution is uniform and all pixels are sampled with equal probability, while higher values narrow the distribution on pixels having higher classification confidence.

The training is performed using two processes. The mapping process computes the probability maps on slides using up-to-date model weights. The use of fully convolutional segmentation models for $M_{\theta}$ allows to speedup the computation of the probability maps by allowing to infer the model on larger patches. The training process samples patches from the latest probability map to feed a shuffle buffer from which training batches are built and used to update the model’s parameters. The two processes are running synchronously in order to optimize the up-to-dateness of the probability maps from which the patches are sampled; and in parallel allowing the training to run without bottleneck. See the figure 2 for an illustration of the training pipeline.

Slides sampling strategy

The mapping process samples malign and benign slides with the same rate. Benign slides are sampled with an emphasis on those containing high segmentation error according to the following probability:

[TABLE]

The probability maps $Q$ are taken from the latest map computed by the mapping process. A default value is used for slides that haven’t been computed yet. We ensure that every slides are sampled at least once every $N$ epochs.

4 Experiments

Dataset

We trained a model on the CAMELYON 16 dataset [1], which consists of 400 WSIs. The dataset has been developed for the task consisting of detection of metastases in H&E stained WSIs of lymph node sections. The tumors have been annotated at the pixel level by pathologists. For training we only used the slide level cancer presence information to train the model in a weakly supervised setup. We used the pixel level annotations for evaluation.

Segmentation model

We instantiated $M_{\theta}$ with a modified U-net[12] network. We replaced the convolutions with separable convolutions. We added residual skip connections and batch-normalization. We set the starting number of filters at 32 in the first layer. The model is trained on a resolution of 1 micrometer per pixel.

Pre-training

We pre-trained $M_{\theta}$ by initializing it with an auto-encoder trained with the same network architecture but without the lateral skip connections. It has been trained to reconstruct the input using the $L_{2}$ loss.

Data augmentation

We augmented the patches with random rotation, mirroring, elastic deformation and color jittering.

Hyperparameters

We have set $\beta_{0}=0$, $\beta_{1}=1$, $\eta=50\%$ and $\alpha=2$.

Results

We measured the ability of the model to discriminate malignant and benign slides. We scored the slides with their maximum cancer pixel probability value. Given such scoring, the model could reach a classification performance of $0.89$ ROC AUC.

The quality of the segmentation has been measured using the same metric as the one proposed for the CAMELYON 16 challenge. It is the average sensitivity at 6 predefined false positive rates: $1/4$, $1/2$, $1$, $2$, $4$, and $8$ FPs per whole slide image. The model reached an average sensitivity of $0.43$. The figure 3 shows examples of successfully detected metastases. We can see that the segmentations are consistent with the morphology and have strong agreement with the ground truth.

5 Conclusion

We have proposed a novel method to train a segmentation model at pixel resolution in a weakly supervised setup. We have shown that the model trained with scarce WSI level supervision was able to retrieve cancer localization at pixel resolution. Preliminary experiment have shown promising results in term of classification and segmentation performances inviting for further investigation.