UNIStainNet: Foundation-Model-Guided Virtual Staining of H&E to IHC

Jillur Rahman Saurav; Thuong Le Hoai Pham; Pritam Mukherjee; Paul Yi; Brent A. Orr; Jacob M. Luber

cs.CVMar 13, 2026

UNIStainNet: Foundation-Model-Guided Virtual Staining of H&E to IHC

Authors: Jillur Rahman Saurav, Thuong Le Hoai Pham, Pritam Mukherjee, Paul Yi, Brent A. Orr, Jacob M. Luber

Organizations: Department of Computer Science, University of Texas at Arlington, Arlington, Texas, USA · Department of Radiology, St. Jude Children’s Research Hospital, Memphis, Tennessee, USA · Department of Pathology, St. Jude Children’s Research Hospital, Memphis, Tennessee, USA

Abstract

Virtual immunohistochemistry (IHC) staining from hematoxylin and eosin (H&E) images can accelerate diagnostics by providing preliminary molecular insight directly from routine sections, reducing the need for repeat sectioning when tissue is limited. Existing methods improve realism through contrastive objectives, prototype matching, or domain alignment, yet the generator itself receives no direct guidance from pathology foundation models. We present UNIStainNet, a SPADE-UNet conditioned on dense spatial tokens from a frozen pathology foundation model (UNI), providing tissue-level semantic guidance for stain translation. A misalignment-aware loss suite preserves stain quantification accuracy, and learned stain embeddings enable a single model to serve multiple IHC markers simultaneously. On MIST, UNIStainNet achieves state-of-the-art distributional metrics on all four stains (HER2, Ki67, ER, PR) from a single unified model, where prior methods typically train separate per-stain models. On BCI, it also achieves the best distributional metrics. A tissue-type stratified failure analysis reveals that remaining errors are systematic, concentrating in non-tumor tissue. Code is available at https://github.com/facevoid/UNIStainNet.

Figures & tables

Figure 1. UNIStainNet architecture. (a) Overview: The H&E image is split into $4\!\times\!4$ sub-crops and processed by a frozen UNI ViT-L/16 to produce multi-scale spatial maps. A CNN encoder compresses the H&E input through a self-attention bottleneck; the SPADE+FiLM decoder then receives UNI spatial maps via SPADE modulation, edge features via concatenation, and a stain embedding via FiLM, with encoder skip connections preserving fine-grained detail. An unconditional multi-scale discriminator provides adversarial training. (b) SPADE+FiLM block detail: at each decoder stage, edge and skip features are concatenated and convolved, then modulated by UNI spatial maps (spatially-varying $\gamma_{s},\beta_{s}$ ) and the stain embedding (channel-wise $\gamma_{c},\beta_{c}$ ), combined additively after instance normalization.

Dataset	Stain	Train	Test
BCI ( Liu et al., 2022 )	HER2	3,896	977
MIST ( Li et al., 2023 )	HER2	4,642	1,000
	Ki67	4,361	1,000
	ER	4,153	1,000
	PR	4,139	1,000

Table 1. Dataset summary. Both datasets pair H&E and IHC from consecutive tissue sections.

	HER2			Ki67			ER			PR
Method	FID $\downarrow$	KID $\times$ 1k $\downarrow$	SSIM $\uparrow$	FID $\downarrow$	KID $\times$ 1k $\downarrow$	SSIM $\uparrow$	FID $\downarrow$	KID $\times$ 1k $\downarrow$	SSIM $\uparrow$	FID $\downarrow$	KID $\times$ 1k $\downarrow$	SSIM $\uparrow$
ODA-GAN ( Wang et al., 2025 )	68.0	8.6	0.189	—	—	—	—	—	—	—	—	—
PASB ( Qiu et al., 2025 )	46.0	21.0	0.274	—	—	—	40.5	21.6	0.254	—	—	—
SIM-GAN ( Guan et al., 2025 )	39.6	—	0.154	31.6	—	0.182	34.6	—	0.176	36.1	—	0.183
ASP ( Li et al., 2023 )	51.4	12.4	0.200	51.0	19.1	0.241	41.4	5.8	0.206	44.8	10.2	0.216
USI-GAN ( Peng et al., 2025 )	37.8	—	0.187	27.4	—	0.232	33.1	—	0.202	34.6	—	0.216

Table 2. Quantitative comparison on MIST (1,000 test images per stain). UNIStainNet uses a single unified model serving all four IHC markers via learned stain embeddings; compared methods typically train a separate model per stain. All numbers are taken from the original papers. Cells marked — indicate the metric or stain was not reported. Best and second best highlighted per column.

Figure 2. Unified multi-stain generation on MIST. Four randomly sampled validation images per stain (HER2, Ki67, ER, PR). Columns: H&E input, ground truth IHC, and UNIStainNet output. A single model produces stain-specific expression patterns: membrane (HER2), punctate nuclear (Ki67), and diffuse nuclear (ER/PR).

Figure 3. Representative qualitative comparison on MIST HER2. Additional randomly sampled comparisons are provided in the supplementary material (Section D).

Method	FID $\downarrow$	KID $\times$ 1k $\downarrow$	SSIM $\uparrow$
ASP † ( Li et al., 2023 )	65.1	9.9	0.503
ODA-GAN † ( Wang et al., 2025 )	59.3	10.1	0.473
PSPStain ‡ ( Chen et al., 2024a )	45.7	30.4	0.433
PASB ( Qiu et al., 2025 )	43.6	9.6	0.426
UNIStainNet (ours)	34.6	6.5	0.541

Table 3. Quantitative comparison on BCI test set (977 images). All numbers are taken from the original papers. † Evaluated on a 500-image test split. ‡ As reported by ( Qiu et al., 2025 ) . Best and second best highlighted.

Figure 4. Qualitative comparison on BCI (one sample per HER2 class). Columns: H&E input, ground truth IHC, ASP ( Li et al., 2023 ) , PSPStain ( Chen et al., 2024a ) , and UNIStainNet. Our method preserves tissue morphology and produces correctly graded DAB staining across all HER2 levels.

Model	#Models	Train. Params	Avg FID $\downarrow$	Avg P- $r$ $\uparrow$	Avg DAB KL $\downarrow$
Per-stain specialists	4	170M	29.8	0.930	0.171
Unified (dim 64)	1	42M	30.0	0.937	0.159

Table 4. Unified model vs. per-stain specialists on MIST (macro-averaged across HER2, Ki67, ER, PR). The unified model matches specialist-level stain accuracy with

4\times

fewer total parameters. Per-stain breakdown in supplementary material (Section B).

Dataset	Res.	FID $\downarrow$	P- $r$ $\uparrow$	DAB KL $\downarrow$
MIST	$512$	30.0	0.937	0.159
MIST	$1024$	40.3	0.961	0.099
BCI	$512$	34.6	0.867	0.482
BCI	$1024$	34.1	0.892	0.392

Table 5.

512

vs.

1024

generation. MIST values are macro-averaged across 4 stains (unified model).

1024

improves stain accuracy on both datasets with minimal parameter overhead.

Figure 5. Tissue-type stratified failure analysis via CONCH ( Lu et al., 2024 ) zero-shot classification. (a,b) Failure rate and mean DAB KL by tissue type on MIST (4,000 images, 4 stains, macro-averaged) and BCI (977 images). (c,d) Representative examples: invasive carcinoma produces lower error, while adipose and necrotic regions produce higher error.

Figure 6. Worst failure cases on MIST per stain, ranked by whole-image DAB KL divergence (top-4 per stain). Most failures occur in non-tumor tissue and where consecutive-section misalignment is more pronounced.

Config	FID $\downarrow$	P- $r$ $\uparrow$	DAB KL $\downarrow$
Full model	30.0	0.937	0.159
$-$ Edge encoder	31.7	0.939	0.162
$-$ Feat. matching	29.9	0.932	0.171
$-$ DAB loss	30.4	0.927	0.184
$-$ LPIPS	31.4	0.931	0.189
CLS $4\!\times\!4$	33.7	0.938	0.150

Table 6. Ablation study on MIST (macro-averaged across 4 stains, unified model). Best and second best highlighted.

Virtual Staining

Virtual staining uses AI to predict histochemical or immunohistochemical appearances from unstained or differently stained tissue images.

See all 18 papers

HistDiT: A Structure-Aware Latent Conditional Diffusion Model for High-Fidelity Virtual Staining in Histopathology
Apr 9, 2026Aasim Bin Saleem, Amr Ahmed, Ardhendu Behera +5Virtual Staining Computational Pathology
Edge Hill University, Ormskirk, L39 4QP, United Kingdom · University of Nottingham Malaysia, 43500 Semenyih, Malaysia · University of Southampton Malaysia, 79100 Iskandar Puteri, Malaysia +1
PGVMS: A Prompt-Guided Unified Framework for Virtual Multiplex IHC Staining with Pathological Semantic Learning
Feb 26, 2026Fuqiang Chen, Ranran Zhang, Wanming Hu +6Virtual Staining Computational Pathology
Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences and University of Chinese Academy of Sciences, Beijing, China · Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences, Shenzhen, China · Department of Pathology, Sun Yat-sen University Cancer Center, State Key Laboratory of Oncology in South China, Guangzhou, China +2
Towards Reliable Stain Transfer: An Iterative Data-Model Co-Optimization Framework Based on Multimodal Expert-Guided Assessment
Jul 28, 2026Siyuan Xu, Yan Wang, Haofei Song +7Computational Pathology Image Style Transfer
East China Normal University Shanghai, China · Ruijin Hospital, Shanghai Jiao Tong University School of Medicine Shanghai, China · Hangzhou Hyperspectral Imaging Technology Co., Ltd. Hangzhou, China +1

Apr 9, 2026J/K move · Enter open · S save

HistDiT: A Structure-Aware Latent Conditional Diffusion Model for High-Fidelity Virtual Staining in Histopathology

Aasim Bin Saleem, Amr Ahmed, Ardhendu Behera, Hafeezullah Amin+4

Edge Hill University, Ormskirk, L39 4QP, United Kingdom · University of Nottingham Malaysia, 43500 Semenyih, Malaysia · University of Southampton Malaysia, 79100 Iskandar Puteri, Malaysia+1

Virtual Staining Computational Pathology Conditional Diffusion Models Diffusion Transformer

Open paper AskPDF arXiv

Immunohistochemistry (IHC) is essential for assessing specific immune biomarkers like Human Epidermal growth-factor Receptor 2 (HER2) in breast cancer. However, the traditional protocols of obtaining IHC stains are resource-intensive, time-consuming, and prone to structural damages. Virtual staining has emerged as a scalable alternative, but it faces significant challenges in preserving fine-grained cellular structures while accurately translating biochemical expressions. Current state-of-the-art methods still rely on Generative Adversarial Networks (GANs) or standard convolutional U-Net diffusion models that often struggle with "structure and staining trade-offs". The generated samples are either structurally relevant but blurry, or texturally realistic but have artifacts that compromise their diagnostic use. In this paper, we introduce HistDiT, a novel latent conditional Diffusion Transformer (DiT) architecture that establishes a new benchmark for visual fidelity in virtual histological staining. The novelty introduced in this work is, a) the Dual-Stream Conditioning strategy that explicitly maintains a balance between spatial constraints via VAE-encoded latents and semantic phenotype guidance via UNI embeddings; b) the multi-objective loss function that contributes to sharper images with clear morphological structure; and c) the use of the Structural Correlation Metric (SCM) to focus on the core morphological structure for precise assessment of sample quality. Consequently, our model outperforms existing baselines, as demonstrated through rigorous quantitative and qualitative evaluations.

Abstract

Virtual Staining

Explore similar work

HistDiT: A Structure-Aware Latent Conditional Diffusion Model for High-Fidelity Virtual Staining in Histopathology

PGVMS: A Prompt-Guided Unified Framework for Virtual Multiplex IHC Staining with Pathological Semantic Learning

Towards Reliable Stain Transfer: An Iterative Data-Model Co-Optimization Framework Based on Multimodal Expert-Guided Assessment