Unsupervised Latent Stain Adaption
for Digital Pathology

Abstract

In digital pathology, deep learning (DL) models for tasks such as segmentation or tissue classification are known to suffer from domain shifts due to different staining techniques. Stain adaptation aims to reduce the generalization error between different stains by training a model on source stains that generalizes to target stains. Despite the abundance of target stain data, a key challenge is the lack of annotations. To address this, we propose a joint training between artificially labeled and unlabeled data including all available stained images called Unsupervised Latent Stain Adaption (ULSA). Our method uses stain translation to enrich labeled source images with synthetic target images in order to increase supervised signals. Moreover, we leverage unlabeled target stain images using stain-invariant feature consistency learning. With ULSA we present a semi-supervised strategy for efficient stain adaption without access to annotated target stain data. Remarkably, ULSA is task agnostic in patch-level analysis for whole slide images (WSIs). Through extensive evaluation on external datasets, we demonstrate that ULSA achieves state-of-the-art (SOTA) performance in kidney tissue segmentation and breast cancer classification across a spectrum of staining variations. Our findings suggest that ULSA is an important framework towards stain adaption in digital pathology.

1 Introduction

Recent advances in DL for digital pathology have shown promising results for a wide range of applications, from cancer and biomarker detection to tissue structure segmentation [3]. However, large-scale studies have indicated that the effectiveness of DL techniques in histology is heavily dependent on the availability of labeled data [15]. Despite its theoretical promise, acquiring a sufficient number of expert annotations remains challenging. In the realm of digital pathology, image datasets often consist of sequential slides stained with various techniques, each providing distinct insights into the same region of interest. Despite variations in staining protocols, these slides frequently share a significant amount of consistent information. However, expert annotations may be available for one type of staining but may be lacking for others, which are often accessible in large quantities without labels. Generating expert annotations for multiple staining techniques for the same analysis tasks would be exceedingly time-consuming. In the era of foundational models [11], we also prefer generalized DL models robust to data shifts instead of domain experts. In this paper, we question how to tailor a DL model trained for a specific task to handle variations in staining within the distribution of target stains, for which no annotations are available. This can be accomplished by incorporating unlabeled data during the training phase. The aspect of stain adaptation across different inter-staining techniques has not been sufficiently explored so far. Despite efforts to develop stain-to-stain translation techniques, their effectiveness is typically evaluated either visually by experts or through translation metrics [19]. Prior research has not focused on directly incorporating unlabeled target stain images into the training process yet, only using them for translation [2, 7]. Here, we present ULSA, a semi-supervised strategy designed for joint training of all staining data for the first time. We introduce a framework that integrates unlabeled target stain images into supervised training by maintaining the supervised learning signal for synthetic target stainings generated through cycle GAN (cGAN) inference [18]. Feature-wise stain-adaption enables using unlabeled target data and enforces feature consistency across stains. Combining these key ingredients, we propose a new method for efficient stain adaption, that outperforms current SOTA approaches. Our novelties can be summarized as

• (1) Unsupervised stain adaption. ULSA leverages target stain data in a supervised and unsupervised fashion with only annotated source stains. We propose a framework for training stain-invariant models for digital pathology.

• (2) Feature consistency learning. We maximize cosine-similarity between hierarchical features across stains to achieve stain-invariance on feature level.

• (3) Task agnostic framework. ULSA is applicable for classification and segmentation training of stain-invariant models.

• (4) Outperforming SOTA. Our approach outperforms methods from stain-translation, DL based augmentation and semi-supervised learning slightly for source stains and by a large margin for target stains. Our approach only needs 10% of labels to reach the same performance as SOTA trained with all data.

2 Method

Stain adaption aims to minimize the generalization error in task performance between a source staining $s\in S$ and a target staining $t\in T$ (Fig. 1a). In particular, a parameterized model $m_{\theta}$ trained on labeled source staining data $x_{S}^{L}$ should ideally maintain task performance on other unlabeled target stainings $x_{T}^{U}$ where no labels are available. We propose to address this challenge by incorporating unlabeled target stainings through (i) a cGAN model to augment labeled images into target stains inheriting the same annotation, (ii) unsupervised stain adaption (USA) to jointly train on all stains with supervised and unsupervised objectives including all stains, followed by (iii) stain-invariant feature consistency learning (FCL) by unsupervised matching of latent representations between stainings. The overall method is outlined in Fig. 2.

Figure 1: (a). Problem statement of unsupervised stain adaption. (b). Stain-invariant feature consistency learning. (c). Artificial images generated by cGAN.

2.0.1 (i) cGAN stain augmentation.

After pretraining we used cGANs, which we define as stain translation function $\mathcal{G}(x_{S}^{L})=x_{S\cup T}^{L}$ to synthetically augment source training data into target stainings. This process is structure preservering [18], thus each target stain image inherits the label corresponding to the associated source image $x_{S}^{L}$ used for translation. Fig. 1c shows exemplary translation results. This strategy aims to increase the labeled training dataset and thus the supervisory signal to achieve stain-invariance on prediction level.

Figure 2: ULSA model. Labeled source stains are translated into synthetic target stain data to obtain supervision for image-wise stain-invariance. We extract features for a real and stain translated noised image of target stained data, where we maximize cosine similarity to achieve unsupervised feature-wise stain-invariance in latent space.

2.0.2 (ii) Unsupervised stain adaption.

Since labeled data are given only in source staining $x^{L}\sim p_{S}(x^{L})$, we would ideally desire a map** from source to target samples. We approximate this map** by using $|S|\cdot|T|$ distinct cGAN augmenter to obtain additional labeled target samples $x^{L}\sim\hat{p}_{t\in T}(x^{L}\mid x^{U})$. With $|\cdot|$ we denote the cardinality of a set of stains. Note that all inferred samples $x^{L}$ inherit the label $y$ associated with $x^{L}$. We approximate $p_{S\cup T}$ by a mixture

$$\hat{p}_{S\cup T}\left(x^{L}\right)=\frac{1}{|S|+|T|}\left(\sum_{s\in S}p_{s}\left(x^{L}\right)+\sum_{t\in T}\hat{p}_{t}\left(x^{L}\right)\right).$$

In addition, we leverage unlabeled data by unsupervised learning. Given an unlabeled image $x_{t\in T}^{U,1}$, we translate it to $\tilde{x}_{t\in T}^{U,1}=\mathcal{I}\left(\mathcal{R}\left(x_{t\in T}^{U,1}\mid x_{t\in T}^{U,2}\right)\right)$ by Reinhard translation $\mathcal{R}$ and noise injection $\mathcal{I}$. Note that $\mathcal{R}$ receives another unlabeled random sampled image $x_{t\in T}^{U,2}$ from target stains $t\in T$ as reference for subsequent translations. Fig. 1b include example inputs. We used Reinhard normalization as Macenkos method has much higher runtime leading to computational overhead, see supplementary material (SM). Finally, we enforce the model to embed images invariant to stain translations by maximizing cosine similarity in unsupervised loss $\mathcal{L}_{U}$, see (iii). For this reason, we used light gaussian blurring for noise injection, see SM for other choices. In summary, we define our objective

$$\min_{\theta}\mathcal{L}\left(\theta\mid x^{L},y,x^{U}\right)=\underset{x^{L}\sim\hat{p}_{S\cup T}\left(x^{L}\right)}{\mathbb{E}}\left[\mathcal{L}_{S}\left(\theta\mid x^{L},y\right)\right]+\lambda\underset{x^{U}\sim p_{S\cup T}\left(x^{U}\right)}{\mathbb{E}}\left[\mathcal{L}_{U}\left(\theta\mid x^{U}\right)\right]$$

where $\mathcal{L}_{S}$ is a supervised loss. With minimizing $\mathcal{L}_{S}$ using labeled images of synthetic target stains, we aim for stain-adaption on prediction level. We use multi-class and binary cross entropy loss for segmentation and classification, respectively. We set equal weight $\lambda=1$, obtained by performing a grid search as described in SM. In each iteration we compute $\mathcal{L}_{S}$ on a batch of labeled data and compute feature consistency with $\mathcal{L}_{U}$ on a batch of unlabeled data. Next, we detail the unsupervised loss term associated with feature consistency learning.

2.0.3 (iii) Stain-invariant feature consistency learning.

For unsupervised learning, we aim to calculate similar feature representations for an input $x=x_{t\in T}^{U,1}$ and a stain translated noised version $\tilde{x}=\tilde{x}_{t\in T}^{U,1}$. To enforce this for every downsampling block in a model, we apply non-parametric 2D adaptive average pooling. In Einstein notation, this tensor operation leads to $B,C_{i},H_{i},W_{i}\rightarrow B,C_{i}$ where $B$, $C$, $H$, $W$ and $i$ refers to batch-size, channels, height, width and the model block, respectively (Fig. 1b). We aim to maximize the cosine similarity between all hierarchical features from a model $m_{\theta}$. In this way, we enforce feature similarity by updating the model such that extracted image features in latent space become similar. We define features as model outputs until block $i$ after 2D adaptive pooling, $f_{\bar{\theta},i}=m_{\bar{\theta},i}(x)$, $\tilde{f}_{\theta,i}=m_{\theta,i}(\tilde{x})$. For brevity, index $\theta,i$ also refers to parameter optimization from input layer until block $i$. Following the literature [10, 17], we do not backpropagate gradients for non-augmented input, denoted with $\bar{\theta}$. Hence, we define our unsupervised objective $\mathcal{L}_{u}$ as

$$\mathcal{L}_{U}(\theta)=-\frac{1}{b}\sum_{i\in b}\frac{\mathbf{f}_{\bar{\theta},i}\cdot\tilde{\mathbf{f}}_{\theta,i}}{\lVert\mathbf{f}_{\bar{\theta},i}\rVert_{2}\cdot\lVert\tilde{\mathbf{f}}_{\theta,i}\rVert_{2}},$$

where $\lVert\cdot\rVert_{2}$ denotes the Euclidean norm and $b$ the number of downsampling blocks. With minimizing $\mathcal{L}_{U}$, we force stain-invariance on feature level.

3 Experiments

We applied our proposed method to kidney tissue segmentation and cancer classification. We measured the performance using the dice score and area under the receiver operating characteristic curve (AUROC), respectively. In the following we describe our datasets, more information about statistics can be found in SM.

3.1 Datasets and Comparable Methods

For slide tiling, we used a modified version of the CLAM preprocessing pipeline [8] and manually selected color thresholds as different stains require adjustments in tissue detection. All image data were processed with a size of $512\times 512$ px.

3.1.1 Kidney segmentation datasets.

Our internal annotated train $n_{train}^{PAS}=2100$ and validation $n_{val}^{PAS}=160$ datasets consist of annotated patches extracted from periodic acid-Schiff (PAS) stained WSIs with $20\times$ magnification. Annotation masks contain the classes tubule, glomerulus, glomerular tuft, artery, arterial lumen, and vein. For external testing we included annotated glomerulus images from PAS stained WSIs provided by HuBMAP consortium [4]. After processing the raw data we obtained $n_{ext,hubmap}^{PAS}=2670$ samples. Moreover, we included test data from the NEPTUNE [1] study. This dataset contains $n_{ext,neptune}^{HE}=402$, $n_{ext,neptune}^{PAS}=1176$, $n_{ext,neptune}^{SIL}=688$ Silver (SIL) and $n_{ext,neptune}^{TRI}=817$ Trichome (TRI) stain samples. The classes glomerulus, glomerular tuft, artery and tubule were annotated and images were taken from WSIs with $40\times$ magnification, which we rescaled to $20\times$ magnification. We further processed unlabeled WSIs at $20\times$ magnification from the KPMP database [6]. We obtained $n_{unlabeled}^{HE}=385670$, $n_{unlabeled}^{PAS}=409554$, $n_{unlabeled}^{SIL}=538412$, $n_{unlabeled}^{TRI}=415822$ tiles.

3.1.2 Breast cancer classification datasets.

Our cancer classification datasets contain $n_{train}^{HE}=1950$, $n_{val}^{HE}=287$ images obtained from HE stained WSIs and tiled at $40\times$ magnification level. Our test set contains Cytokeratin (CK5) and Cluster of differentiation (CD8) stains with $n_{test}^{CK5}=1000$ and $n_{test}^{CD8}=1000$ samples. Sample sizes and binary label ratios were equalized for test sets after preprocessing and excess data were held out of experiments to avoid data leakage at patient level. Additionally we have unlabeled datasets containing $n_{unlabeled}^{HE}=117369$, $n_{unlabeled}^{CK5}=41735$, $n_{unlabeled}^{CD8}=52492$ samples. We splitted data on slide-level such that each patient appears exclusively in one dataset.

3.1.3 Comparable Methods.

We selected comparable methods from the domains of stain translation, unsupervised augmentation, and semi-supervised consistency training. (1) Baseline. For comparison, we include a naive approach where we train a model on source stains without access to target stains, thus obtain a lower bound for stain adaptation. (2) Reinhard. Reinhard’s method is a stain normalization technique that adjusts the color appearance of histopathology images by aligning them with a reference color space [12]. (3) Macenko. Macenko’s method [9] is a stain normalization technique that standardizes the color appearance of images by map** them to a reference color space. It has been shown that this method provides the best performance across various colorization techniques for downstream tasks [7]. (4) cGAN Augmentation. This method generates synthetic images with diverse staining variations in order to translate between stainings and train stain-invariant models[2]. (5) FixMatch. A semi-supervised learning approach that combines labeled and unlabeled data by enforcing consistency between predictions made on unlabeled samples based on confidence scores [13]. (6) Unsupervised Data Augmentation (UDA). UDA is a semi-supervised learning technique designed for classification tasks that leverages augmentations on unlabeled data to improve model performance by minimizing probability distributions between two versions of an image. Originally proposed for natural images [17] it has been adapted to histology images [5].

3.2 Implementation

3.2.1 cGAN pretraining.

We initially started with hyperparameters following the literature [2] and performed a grid search with details noted in SM. By visually evaluating stain-translation results, we selected a model trained with $300$ epochs and a learning rate of $1.5e-4$. For each stain translation we used $10,000$ unlabeled images from KPMP. For source stains, we used unlabeled HE data for cancer classification, and our labeled PAS stained training set for segmentation. Each stain translation training and inference task took around 8 hours.

3.2.2 Segmentation and Classification.

For all experiments we used a ResNet-50 as classification model and encoder model for U-Net in segmentation. We tuned hyperparameters on validation sets and set a learning rate decay of $1e-10$ and early stop** for 5 and 10 consecutive epochs with no decrease in validation loss. We employed AdamW as optimizer and used a starting learning rate of $1e-04$ and weight decay $1e-05$. Images were resized into $224\times 224$ px scale. We used an overall batch-size of 128 for all experiments. For semi-supervised learning we batched data to 32 and 96 for labeled and unlabeled data, respectively. All models were initialized with ImageNet-pretrained weights and experiments were performed on a single Nvidia A100 GPU. We measured a maximum training time across all methods of 10 and 7 hours for segmentation and classification, respectively (except Macenko, see SM).

4 Results and Discussion

In this section, we report results and question the necessary number of labels for stain adaption in kidney tissue segmentation and breast cancer classification.

Table 1: Dice and AUROC scores for segmentation and classification tasks, respectively. We report mean and standard deviation across three consecutive runs. All methods except for Baseline access additional unlabeled target stains in trainings phase.

Method	Segmentation					Classification
	Intra-stain	Inter-stain				Inter-stain
	PAS	TRI	HE	SIL	Overall	CK5	CD8	Overall
Baseline	87.4 (0.66)	59.4 (4.19)	43.2 (4.01)	76.4 (1.67)	62.1 (3.24)	86.6 (9.61)	90.5 (7.10)	88.6 (8.35)
Reinhard [12]	87.3 (0.50)	64.8 (6.50)	40.2 (3.23)	77.9 (2.57)	64.3 (4.39)	89.8 (3.65)	94.1 (1.48)	91.9 (2.32)
Macenko [7]	85.0 (0.80)	71.6 (2.60)	48.3 (4.70)	81.6 (0.50)	70.3 (2.29)	89.5 (2.90)	93.4 (2.06)	91.4 (2.48)
cGAN [2]	84.2 (0.92)	69.9 (1.44)	46.0 (4.71)	79.0 (1.63)	68.1 (2.18)	87.1 (6.24)	88.0 (3.32)	87.5 (3.67)
FixMatch [13]	87.2 (0.51)	64.8 (2.67)	40.3 (7.28)	78.2 (1.02)	64.5 (3.05)	88.3 (2.25)	94.3 (1.37)	91.3 (1.78)
UDA [5]	87.2 (0.54)	64.9 (1.84)	45.8 (2.49)	77.5 (0.60)	65.4 (1.53)	89.7 (2.19)	92.8 (1.97)	91.3 (1.47)
Ours	87.9 (0.33)	74.1 (0.94)	53.6 (1.46)	81.8 (0.61)	72.6 (0.93)	91.6 (4.63)	94.6 (1.40)	93.1 (3.01)

4.0.1 Measuring stain adaption.

In the area of stain adaption, we aim to at least maintain performance on source stains (Intra-stain, Tab. 1) and maximize performance on target stains (Inter-stain, Tab. 1). For segmentation on source data our method is on par with other methods, showing we also slightly increased performance. Note that all other methods decrease performance on source stains compared to the baseline. More importantly, we increase the dice score for target stains (Inter-stain, Overall) by more than $10$ and $2.3$ compared to naive training and Macenko colorization, the best comparable method. However, note that despite Macenko provides the second best stain adaption result, source performance is decreased by $-2.4$ and $-2.9$ compared to baseline and ULSA. In the case of classification, we had no further annotated intra-stain data for testing but report results for target data. Our method increase AUROC by 4.5 and 1.2 compared to naive and best comparable method. Overall, these findings demonstrate efficient stain adaption by increasing target while not only maintaining but increasing source performance (Fig. 1a). Additionally, we trained all methods using different fractions of labeled data and tested for performance on targets in segmentation (Fig. 3a). Interestingly, ULSA shows great stain adaption even for very less annotated data compared to all other methods. ULSA is on par with the second best comparable method using only 10% of labeled data.

4.0.2 Ablation study.

We measured the influence of different components of our method in segmentation (Fig. 3b). By drop** either cGAN or FCL components, overall target performance decrease. Using hierarchical features (ULSA) instead of features from last block (LB FCL) increase scores. Apart from that, we also initialized our ULSA method with pretrained weights from a foundation model for histology (FM ULSA), obtained by large-scale learning of HE images [16]. With this setup we demonstrate that ImageNet pretrained weights yield better performance. This hints that building foundation models should incorporate stain variations to avoid catastrophic forgetting about learned morphologies.

Figure 3: (a). Performance for different fractions of labeled data. (b). Ablation study.

5 Conclusion

We proposed ULSA, a novel SOTA strategy to decrease stain generalization errors in digital pathology tasks. Our semi-supervised learning strategy leverages annotated data for both source and artificial target stains. Moreover, we incorporate unlabeled data for stain-invariant feature consistency learning. Finally, joint optimization of supervised and unsupervised objectives enables efficient stain adaption. We empirically demonstrated that ULSA training increases performance on unlabeled target stains in patch level segmentation and classification. This suggests that ULSA is a task agnostic framework. We further showed ULSA achieves efficient stain adaption even in settings with scarce labels. A potential limitation of our approach is that even if the performance for unseen interstains is increased, augmentation strategies may not translate the correct marker information from other stains. A possible example is immunohistochemical (IHC) staining, where specific immune cells are highlighted. This could potentially affect downstream applications in certain scenarios not seen in this study. Future work could compare and replace cGAN with other GAN-based augmentation strategies such as HistAuGAN [14].

5.0.1 Acknowledgements.

This work was supported by Deutsche Forschungsgemeinschaft (DFG Project number 445703531). The authors gratefully acknowledge the computational and data resources provided by the Leibniz Supercomputing Centre (www.lrz.de).

5.0.2 Disclosure of Interests.

The authors declare that they have no conflicts of interest related to this work.