Stain-AgLr: Stain Agnostic Learning for Computational Histopathology Using Domain Consistency and Stain Regeneration Loss

Stain color variations between Whole Slide Images (WSIs) is a key challenge in the application of Computational Histopathology. Deep learning-based algorithms are susceptible to domain shift and degrade in performance on the WSIs captured from a different source than the training data due to stain color variations. We propose a training methodology Stain-AgLr, that achieves high invariance to stain color changes on unseen test data. In addition to task loss, Stain-AgLr training is supervised with a consistency regularization loss that enforces consistent predictions for training samples and their stain altered versions. An additional decoder is used to regenerate stain color from feature representation of the stain altered images. We compare the proposed approach to state-of-the-art strategies using two histopathology datasets and show significant improvement in model performance on unseen stain variations. We also visualize the feature space distribution of test samples from multiple diagnostic labs and show that Stain-AgLr achieves a significant overlap between the distributions.