Self-supervised Vision Transformer are Scalable Generative Models for Domain Generalization

TL;DR

This work tackles the problem of domain generalization in digital histopathology, where cross-domain staining and imaging variations hinder robust generalization. It introduces a self-supervised generative approach based on a Vision Transformer encoder that disentangles patch-wise anatomy ($z^a$) from image characteristics ($z^c$), and uses an image synthesizer to reconstruct images and generate diverse synthetic samples by mixing anatomy with cross-sample characteristics. The method employs three loss terms—anatomical consistency $L^a_C$, characteristic consistency $L^c_C$, and self-reconstruction $L_R$—within a unified objective, enabling training without target-domain labels. Empirically, it outperforms state-of-the-art domain generalization methods on two histopathology benchmarks (e.g., Camelyon17-wilds and epithelium-stroma) and demonstrates scalability by leveraging unlabeled data and deeper ViT backbones, with the authors providing code for reproducibility. The approach offers a flexible, scalable path toward robust generalization in medical imaging and potentially other modalities, reducing reliance on domain-specific annotations and target-domain data.

Abstract

Despite notable advancements, the integration of deep learning (DL) techniques into impactful clinical applications, particularly in the realm of digital histopathology, has been hindered by challenges associated with achieving robust generalization across diverse imaging domains and characteristics. Traditional mitigation strategies in this field such as data augmentation and stain color normalization have proven insufficient in addressing this limitation, necessitating the exploration of alternative methodologies. To this end, we propose a novel generative method for domain generalization in histopathology images. Our method employs a generative, self-supervised Vision Transformer to dynamically extract characteristics of image patches and seamlessly infuse them into the original images, thereby creating novel, synthetic images with diverse attributes. By enriching the dataset with such synthesized images, we aim to enhance its holistic nature, facilitating improved generalization of DL models to unseen domains. Extensive experiments conducted on two distinct histopathology datasets demonstrate the effectiveness of our proposed approach, outperforming the state of the art substantially, on the Camelyon17-wilds challenge dataset (+2%) and on a second epithelium-stroma dataset (+26%). Furthermore, we emphasize our method's ability to readily scale with increasingly available unlabeled data samples and more complex, higher parametric architectures. Source code is available at https://github.com/sdoerrich97/vits-are-generative-models .

Figures (4)

• Figure 1: Schematic Visualization of our self-supervised generative approach. A single ViT encoder ($E$) is used to separate anatomy from image-characteristic features of distinct images which are subsequently intermixed among each other and processed by an image synthesizer ($IS$) to generate synthetic images.
• Figure 2: Examples from the histopathology datasets used for evaluating domain generalization. Left: Camelyon17-wilds for which the domains are hospitals. Right: Combined epithelium-stroma dataset for which the domains are datasets.
• Figure 3: Qualitative evaluation of our method's reconstruction capability on the Camelyon17-wilds dataset.
• Figure 4: Qualitative evaluation of the method's generative capabilities on the Camelyon17-wilds dataset by means of synthetic images created through its anatomy-characteristics intermixing.