Semi-Supervised Domain Adaptation with Latent Diffusion for Pathology Image Classification

TL;DR

This work tackles domain shift in computational pathology by introducing a semi-supervised domain adaptation framework that leverages unlabeled data from both source and target domains through a latent diffusion model (LDM). The LDM is conditioned on morphology- and domain-relevant embeddings to generate target-aware, morphology-preserving synthetic images, which are combined with labeled source data to train a downstream classifier. Evaluated on LUAD prognosis classification, the approach improves target-domain performance (TCGA held-out set) while maintaining source-domain performance, with the best results achieved when combining LDM-generated data with stain augmentations. The findings demonstrate that target-aware diffusion-based augmentation effectively bridge cross-domain gaps, offering a promising path for robust generalization in computational pathology and potentially extending to other tissue-imaging tasks.

Abstract

Deep learning models in computational pathology often fail to generalize across cohorts and institutions due to domain shift. Existing approaches either fail to leverage unlabeled data from the target domain or rely on image-to-image translation, which can distort tissue structures and compromise model accuracy. In this work, we propose a semi-supervised domain adaptation (SSDA) framework that utilizes a latent diffusion model trained on unlabeled data from both the source and target domains to generate morphology-preserving and target-aware synthetic images. By conditioning the diffusion model on foundation model features, cohort identity, and tissue preparation method, we preserve tissue structure in the source domain while introducing target-domain appearance characteristics. The target-aware synthetic images, combined with real, labeled images from the source cohort, are subsequently used to train a downstream classifier, which is then tested on the target cohort. The effectiveness of the proposed SSDA framework is demonstrated on the task of lung adenocarcinoma prognostication. The proposed augmentation yielded substantially better performance on the held-out test set from the target cohort, without degrading source-cohort performance. The approach improved the weighted F1 score on the target-cohort held-out test set from 0.611 to 0.706 and the macro F1 score from 0.641 to 0.716. Our results demonstrate that target-aware diffusion-based synthetic data augmentation provides a promising and effective approach for improving domain generalization in computational pathology.

Figures (5)

• Figure 1: Overview of the proposed synthetic-augmentation semi-supervised domain adaptation pipeline. The latent diffusion model (LDM) is trained on unlabeled source-domain images (e.g., NLST) and unlabeled target-domain images (e.g., TCGA). Synthetic images are generated from the UNI features of the real, labeled images from NLST. These synthetic NLST tiles are combined with real labeled NLST tiles to train a downstream classifier. The downstream classifier is evaluated on a held-out set of labeled TCGA images, which does not contain patients used for LDM training.
• Figure 2: Examples of real and synthetic NLST images. For each real NLST image in the top row, a corresponding synthetic tile (bottom row) was generated from the LDM. The synthetic images preserve morphology while introducing variations in image appearance and tissue structures. Additional examples are provided in Supplementary Fig. S1.
• Figure 3: t-SNE visualization of real NLST, real TCGA, and synthetic NLST tiles. Plots are generated using CONCH feature embeddings. (a) shows the t-SNE feature space of real NLST and real TCGA tiles. (b) shows the t-SNE with real NLST, real TCGA, and synthetic NLST tiles. While real NLST and real TCGA images do not exhibit completely separated clusters, there still exist gaps between the two domains, as highlighted by the red circle. Synthetic NLST tiles generated by the LDM occupies an intermediate space between NLST and TCGA distributions, indicating that using unlabeled TCGA during diffusion training influences the generative prior and introduces target-domain characteristics.
• Figure 4: Box plots showing the distribution of F1 scores from 5-fold cross-validation when model using different data augmentation approaches are tested on the TCGA held-out test set. Statistical significance is represented by * ($p<0.05$), ** ($p<0.01$), *** ($p<0.001$), or ns ($p>0.05$).
• Figure 5: t-SNE visualization of CONCH feature embeddings for NLST and TCGA tiles in the 'good prognosis' class. A substantial overlap exists between NLST and TCGA data points, indicating a smaller cross-cohort domain shift in this class.