StainDiffuser: MultiTask Dual Diffusion Model for Virtual Staining

TL;DR

StainDiffuser introduces a multitask diffusion framework that jointly learns H&E-conditioned IHC staining and H&E-based cell segmentation to enable accurate, multiplex virtual staining in pathology. By employing two diffusion branches that interact through a shared H&E encoder with task-specific attention, the model achieves state-of-the-art results on CD3 and CK8/18 stains while requiring only coarse segmentation during training. Extensive experiments and ablations across paired and unpaired baselines demonstrate the benefits of multitask learning and data-efficient training in data-scarce regimes. The work provides a scalable path toward rapid, cost-effective multiplex virtual staining, with potential extensions to additional stains and faster inference through improved sampling strategies.

Abstract

Hematoxylin and Eosin (H&E) staining is widely regarded as the standard in pathology for diagnosing diseases and tracking tumor recurrence. While H&E staining shows tissue structures, it lacks the ability to reveal specific proteins that are associated with disease severity and treatment response. Immunohistochemical (IHC) stains use antibodies to highlight the expression of these proteins on their respective cell types, improving diagnostic accuracy, and assisting with drug selection for treatment. Despite their value, IHC stains require additional time and resources, limiting their utilization in some clinical settings. Recent advances in deep learning have positioned Image-to-Image (I2I) translation as a computational, cost-effective alternative for IHC. I2I generates high fidelity stain transformations digitally, potentially replacing manual staining in IHC. Diffusion models, the current state of the art in image generation and conditional tasks, are particularly well suited for virtual IHC due to their ability to produce high quality images and resilience to mode collapse. However, these models require extensive and diverse datasets (often millions of samples) to achieve a robust performance, a challenge in virtual staining applications where only thousands of samples are typically available. Inspired by the success of multitask deep learning models in scenarios with limited data, we introduce STAINDIFFUSER, a novel multitask diffusion architecture tailored to virtual staining that achieves convergence with smaller datasets. STAINDIFFUSER simultaneously trains two diffusion processes: (a) generating cell specific IHC stains from H&E images and (b) performing H&E based cell segmentation, utilizing coarse segmentation labels exclusively during training. STAINDIFFUSER generates high-quality virtual stains for two markers, outperforming over twenty I2I baselines.

Figures (4)

• Figure 1: StainDiffuser: (A) Block diagram of the StainDiffuser model. The two diffusion models are tasked to (a) denoise the noisy IHC and, (b) corresponding segmentation of the cells on H&E images. The two diffusion models (with separate parameters) interact through the common H&E encoder and the attention blocks for different tasks as well as the back-propagation of losses. (B) Block Diagram of StainDiffuser extension for Multi-staining task, with $N$ number of stains for generation.
• Figure 2: Qualitative Results. Qualitative Results shown for CK818 and CD3 virtual staining for different models. We can observe that all the proposed models highlight the correct cells for the CK818 marker, however, StainDiffuser diffuser is the best in terms of correct cell coloring for CD3 markers.
• Figure 3: Additional CK818 Qualitative Results Comparing Proposed Methods. The results show that all proposed methods perform reasonably well in achieving accurate staining. However, StainDiffuser stands out with the highest precision in exact color matching.
• Figure 4: Additional CD3 Qualitative Results Comparing Proposed Methods. StainDiffuser stands out by achieving the highest precision in accurately coloring the greatest number of correct cells.