HistDiT: A Structure-Aware Latent Conditional Diffusion Model for High-Fidelity Virtual Staining in Histopathology
arXiv:2604.08305v1 [eess.IV] 9 Apr 2026
Aasim Bin Saleem1 , Amr Ahmed1 , Ardhendu Behera1 , Hafeezullah Amin1 , Iman Yi Liao2 , Mahmoud Khattab3 , Pan Jia Wern4 , and Haslina Makmur4 Edge Hill University, Ormskirk, L39 4QP, United Kingdom {aasimr∗ ,ahmeda,beheraa,aminh}@edgehill.ac.uk 2 University of Nottingham Malaysia, 43500 Semenyih, Malaysia [email protected], [email protected] 3 University of Southampton Malaysia, 79100 Iskandar Puteri, Malaysia 4 Cancer Research Malaysia, , 47500 Subang Jaya, Malaysia {jiawern.pan,haslina.makmur}@cancerresearch.my 1
Abstract. 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. Code and trained model weights will soon be available at HistDiT . Keywords: Virtual staining · Breast cancer · Diffusion transformer · HistDiT · IHC staining · HER2 biomarker · stain-to-stain translation.
2
R. Aasim et al.
1
Introduction
Breast cancer remains a significant global health concern, with over 2.3 million cases every year and nearly 0.7 million deaths, thus demanding continuous advancements in diagnostic methodologies [1]. Precise treatments involve the evaluation of specific immune biomarkers, including estrogen/progesterone receptors (ER/PR), CD4/CD8, and Human Epidermal Growth Factor Receptor 2 (HER2) being the only FDA-approved biomarker. Accurate assessment of HER2 expression, scored at scales of 0, 1+, 2+, or 3+, dictates patient’s eligibility for targeted therapies like Trastuzumab, significantly improving survival rates among women [2]. Over the years, Hematoxylin and Eosin (H&E) staining serves as the cornerstone of histological analysis, but it lacks the specificity to accurately identify these biomarkers. Therefore, pathologists rely on advanced molecular techniques to obtain specialized immunohistochemical (IHC) stains. They require extensive manual processing of samples with specialized laboratory infrastructure and skilled pathologists, which is time-consuming, generates chemical waste and often prone to human errors that can alter tissue homeostasis and affect diagnostic accuracy [3]. These limitations necessitate the development of digital staining alternatives for early diagnosis and treatment of breast cancer. Virtual staining utilizes deep learning to generate synthetic IHC images from respective H&E inputs, predicting specific biomarker expressions without altering the histological workflow. Several state-of-the-art virtual staining algorithms have been developed, but face significant challenges in generating accurate stains. Traditional approaches still rely on Generative Adversarial Networks (GANs), such as cCycleGAN [4], PyramidPix2Pix [6], and AdaptiveSupervisedPatchNCE(ASP) [7], suffer from mode collapse and often fail to capture the full diversity of pathological textures (staining). Denoising Diffusion Probabilistic Models (DDPMs) [8] though are promising in high-quality image generation but struggle in stain-to-stain translation due to architectural limitations. Recent works, including SynDiff [9], StainDiffuser [10] and PST-Diff [11], primarily use Convolutional U-Nets. Although U-Net is excellent in local feature extraction but struggles to capture global context, often resulting in spatial misalignments and “hallucinations” in high-expression regions [12]. Moreover, training high-capacity diffusion models on small-scale histopathology datasets without adequate conditioning often leads to overfitting and poor generalization. To address these limitations, we present our Histopathology Diffusion Transformer (HistDiT), a transformer-based architecture that integrates long-range dependencies in tissue structure and the efficient global attention mechanism for precise virtual staining. We incorporate independent conditional guidance to strictly preserve complex cellular structures and pathology-specific embeddings to provide high-level semantic context, ensuring diagnostic consistency. The core innovation of HistDiT lies in its dual-stream conditioning strategy, the multi-objective loss function and the use of Structural Correlation Metric (SCM), as briefed below: a. We introduce a novel Diffusion Transformer (DiT) architecture that utilizes a dual-conditioning mechanism to modulate the denoising process. Unlike
HistDiT
3
standard concatenation, this explicitly aligns the generated noise with both the global semantic phenotypes (extracted via UNI Model) and the spatial tissue morphology (encoded as a structural blueprint via pre-trained VAE). b. We propose a multi-objective robust loss function that combines an auxiliary L1 term with the standard MSE to mitigate the inherent blurring effects caused by imperfect registration between serial tissue sections. This formulation effectively sharpens high-frequency details, like nucleus and membrane boundaries, that are typically smoothed out by pure Gaussian objectives. c. We identify that standard SSIM has a mathematical bias towards brightfield microscopy, where luminance masks structural errors. Therefore, we suggest the Structural Correlation Metric (SCM), that isolates the structure component to provide a more rigorous assessment of histological fidelity. We demonstrate that HistDiT qualitatively and quantitatively outperforms state-of-the-art GAN and diffusion baselines on the BCI and MIST bench-marks. Crucially, we validate the diagnostic viability of generated IHC stains by expert pathologists via blind study, confirming superior structural fidelity compared to existing approaches.
2
Related Work
The task of stain-to-stain translation is defined as learning the mapping function f : XH&E → YIHC . Unlike natural image translation (e.g., transforming horseto-zebra), virtual staining requires rigorous preservation of biological structures. A generated cell nucleus must occupy the exact spatial location as in the source image; any deviation or spatial hallucination might lead to critical misdiagnosis. 2.1
Generative Adversarial Networks (GANs) in Virtual Staining
Early research predominantly utilized GANs [13] for virtually stained IHC image generation. Prominent work on H&E to IHC translation was made by Xu et al. [4], where they proposed a conditional CycleGAN (cCGAN) for unpaired translation, introducing structural similarity losses patch-wise labels to maintain tissue details, but the method shows hows class-dependent inconsistencies and sensitivity to staining variations. In 2018, Isola et al. [5] introduced Pix2Pix, a widely adopted paired translation model; however, its restrictive pixel-level losses often limit the generation of complex pathological variations. To address this, Liu et al. [6] developed PyramidPix2Pix for HER2 characterization, incorporating multi-scale processing with Gaussian filtering and introducing the BCI benchmark dataset. Overall this improved PSNR and SSIM but it struggles with high-expression regions (level 3+) being washed out. Recently, Duan et al. [14] separated structural content and staining style using attention mechanisms, but their reliance on PatchGAN that assumes pixel-level independence among patches, lead to global staining inconsistencies. Ultimately, GANs remain difficult to train and prone to mode collapse, thus limiting the diversity in image generation and their clinical usage [15].
4
2.2
R. Aasim et al.
Diffusion Models in Stain-to-Stain Translation
Denoising Diffusion Probabilistic Models (DDPMs) [16] have recently surpassed GANs in high-resolution image synthesis and stain translation by iteratively refining noisy images. Latent Diffusion Models (LDMs) further improved efficiency by operating in a compressed latent space while maintaining high visual quality [17]. In histopathology for cancer diagnosis, DDPMs are emerging as a valuable tool for generating diverse synthetic data. Moghadam et al. [18] first explored that diffusion models could generate diverse tissue textures, avoiding GAN-based instability. Later, Kataria et al. [10] proposed StainDiffuser, a dual-diffusion architecture that simultaneously performs cell-specific segmentation and IHC staining on H&E images. While effective for specific markers (CK8/18, CD3) but study highlights a critical gap between quantitative metrics and diagnostic use. Xuanhe et al. [19] proposed conditional DDPM for HER2 expression levels assessment, utilizing an attention U-Net conditioned on H&E with classifierfree guidance (CFG). Reported better image quality metrics while operating at 64 × 64. He et al. [11] also introduced PST-Diff, a latent diffusion method for HER2 assessment using asymmetric attention, dynamic variance scheduling and latent-transfer modules. Improving metrics over GANs but failing to match PyramidPix2Pix in structural similarity (SSIM) on BCI benchmark. To address spatial misalignment, Liu et al. [20] proposed Star-Diff, combining deterministic restoration path to preserve structure and stochastic path for molecular variability. Similarly, Großkopf et al. introduced HistDiST [21], which utilizes a DDIM sampler with a cosine schedule to balance structure and diversity. However, HistDiST reported lower SSIM scores for HER2, suggesting that the accelerated, deterministic sampling of DDIM may compromise the reconstruction of high-frequency membrane textures (essential for grading) as the model has fewer opportunities to correct alignment errors. These findings indicate that while diffusion models offer superior diversity, but existing CNN-based architectures struggle with the precision required for clinical diagnostics. 2.3
Diffusion Transformers (DiTs)
While GANs and DDPMs traditionally rely on CNN-based backbones like UNets, recent advancements suggest that Transformer backbones can outperform CNNs in capturing high-resolution representations and long-range dependencies. State-of-the-art models like GPT-4 [22], DALL-E 3 [23], and Stable Diffusion 3.5 [24] have successfully integrated Transformers. Peebles et al. [25] formally introduced Diffusion Transformers (DiTs), treating diffusion as a sequence-tosequence problem. Unlike U-Nets, which are biased towards local processing, DiTs use global attention and Adaptive Layer Normalization (adaLN) to integrate conditioning information effectively [26]. Our proposed HistDiT adapts this scalable architecture for virtual staining by extending the standard DiT conditioning mechanism to support both spatially dense inputs (the H&E structural map) and semantically rich embeddings (the UNI phenotype vector) simultaneously, a novel configuration in the context of medical image synthesis.
HistDiT
2.4
5
Foundation Models as Clinical Priors
Foundation models like UNI [27] and CTransPath [28], trained on massive wholeslide image (WSI) datasets using self-supervised learning, have revolutionized feature extraction. UNI [27], a ViT-H model, captures high-level semantic concepts—such as tumor grade and lymphocytic infiltration—that are invariant to local deformations. In HistDiT, we utilize UNI as a “semantic prior provider”, injecting diagnostic context into the generative process to help model easily distinguishing cellular structures (e.g. lymphocytes from tumor nuclei) that lowlevel pixel data alone cannot resolve.
3
Proposed Methodology
This paper introduces our HistDiT (Histopathology Diffsuion Transformer), a novel virtual staining method for generating realistic and accurate IHC stained images used for HER2 assessment in breast cancer diagnostics. HistDiT utilises a dual-conditioned Diffusion Transformer to translate H&E images into specialized IHC stains with high structural and semantic fidelity. The overall framework is illustrated in Fig. 1. We construct our approach upon conditional DDPMs
Fig. 1. The HistDiT Architecture replaces the standard U-Net with a Transformer backbone that integrates two distinct conditioning streams: (i) Global Semantic Guidance uses frozen UNI embeddings injected via Adaptive Layer Norm (adaLN) to enforce diagnostic consistency; (ii) Spatial Structural Guidance uses VAE-encoded H&E latents injected via Cross-Attention to strictly preserve tissue morphology.
[29], which learn to approximate a data distribution q(x0 ) by reversing a gradual noising process. The forward process q(xt |x0 ) iteratively adds Gaussian noise ϵ ∼ N (0, I) according to a variance schedule βt , such a sample √ that √ at timestep t can be expressed in closed form as xt = ᾱt x0 + 1 − ᾱt ϵ, where ᾱt is the cumulative signal after noise addition. The reverse process
6
R. Aasim et al.
pθ (xt−1 |xt ) is then parametrized as learned Gaussian shifts where a neural network ϵθ (xt , t, c) predicts the noise component to denoise xt conditioned on context c. The model is trained to minimize the variational lower bound, simplified to the mean squared error (MSE) between the true noise ϵ and the predicted noise ϵθ : LM SE = Ex0 ,t,ϵ ∥ϵ − ϵθ (xt , t, c)∥22 . 3.1
Latent Diffusion with HistDiT
To reduce computational complexity while maintaining high-resolution fidelity, we operate in the compressed latent space of a pre-trained Variational AutoEncoder (VAE). We utilize the Stable Diffusion VAE (E, D), which compresses W H the input I ∈ RH×W ×3 into a latent representation z = E(I) ∈ R 8 × 8 ×4 [30]. Diffusion is performed on the latent vector z, and the final image is reconstructed via I˜ = D(z). To align the latent distribution with the standard normal prior expected by the diffusion model, we normalize the latent space to unit variance following standard latent diffusion protocols [17]. Dual-Stream Transformer Architecture. Unlike traditional U-Net backbones that rely on simple channel concatenation (which often dilutes structural guidance), HistDiT strategically handles two distinct streams of conditioning: Semantic and Spatial. This separation is critical for virtual staining, where the model must alter the biochemical appearance (stain intensity) without hallucinating or distorting the physical tissue structure. The noisy latent zt is first “patchified” into a sequence of N tokens, where N = (h/p) × (w/p) and p is the patch size (p = 2). The DiT block consists of standard multi-head self-attention (MHSA), layer normalization (LN), and pointwise feedforward (FF) layers, but innovates in how conditions are injected (Fig. 1: flagged branches). i. Global Semantic Guidance via adaLN: To resolve diagnostic ambiguity (e.g., differentiating HER2 level 1+ from 2+ based on subtle stain intensity), we utilize the “UNI Foundation Model” as a semantic prior [27]. The input H&E image xH is processed by the frozen UNI encoder to extract a global phenotype embedding csem ∈ R1536 . This vector added with the timestep embedding temb is injected into transformer block using Adaptive Layer Normalization (adaLN)[25]. Specifically, a simple MLP regresses MLP(ccombined ) into dimension-wise scale (γ) and shift (β) parameters, modulating the normalized latent features Zt as: adaLN(Zt , ccombined ) = γ(ccombined ) ⊙ LN(Zt ) + β(ccombined )
(1)
Here, ⊙ denotes element-wise multiplication. This allows high-level pathological concepts (like tumor grade, tissue subtype, biomarker information) to globally influence the generative statistics (the style and intensity) without disrupting local structure [25]. This ensures that the biochemical expression levels are consistent with the tissue phenotype. ii. Spatial Structural Guidance via Cross-Attention: To ensure the generated IHC image perfectly respects the cellular morphology of the source H&E, we encode the H&E image xH into a spatial latent map cspatial = E(xH ) ∈ Rh×w×4
HistDiT
7
which is then flattened into a sequence of spatial tokens. This structural blueprint is injected using Multi-Head Cross-Attention layers within the DiT blocks [32]. The intermediate representation of noisy latent Zt serves as the Query (Q), while cspatial acts as the Key (K) and Value (V ). This mechanism forces the model to attend to specific spatial locations in the H&E blueprint when synthesizing high-frequency details (e.g., nuclei). Unlike simple concatenation, cross-attention allows the model to dynamically attend to relevant structural features at different stages of the denoising process, thereby preventing the hallucinations. iii. Classifier-Free Guidance: We utilize Classifier-Free Guidance (CFG) to include adherence to this dual-conditioning constraints. During training, both semantic (csem ) and spatial (cspatial ) conditions are randomly dropped with probability pdrop = 0.11 to learn an unconditional prior. During inference, the noise prediction is extrapolated to amplify the conditional signal. Defined as, ϵ̃θ = ϵθ (zt , ∅) + scale · (ϵθ (zt , csem , cspatial ) − ϵθ (zt , ∅))
(2)
We empirically set the guidance scale to 3.0. This pushes the generation towards high diagnostic fidelity without sacrificing the staining diversity [29]. 3.2
Multi-Objective Loss Function
Standard diffusion models minimize the MSE between the predicted and actual noise. However, histopathology datasets use serial sections that introduce inevitable spatial misalignment. Pure MSE heavily penalizes slight pixel shifts, causing the model to produce blurry “average” images to minimize variance. To address this, we introduce an auxiliary L1 objective on the noise prediction. The L1 norm is less sensitive to outliers (misaligned edges) and encourages sparsity in the error, resulting in sharper boundaries. The total loss Ltotal is: Ltotal = λM SE .Ez0 ,ϵ,t ∥ϵ − ϵθ (zt , t, c)∥22 + λL1 Ez0 ,ϵ,t [∥ϵ − ϵθ (zt , t, c)∥1 ]
(3)
Empirically, we set λM SE = 0.7 and λL1 = 0.3. This hybrid loss is a critical contributor to the model’s ability to generate sharp cellular membranes despite imperfect ground truth alignments, as shown in Fig. 4.
4
Experiments and Evaluation
We evaluate HistDiT on the benchmark BCI dataset [6] and the Multi-IHC Stain Translation (MIST) dataset [7], publicly available collection of paired H&E and IHC images specifically designed for histopathological image translation on HER2 assessment. Each dataset comprises approximately 5,000 paired patches (1024×1024) sourced from several WSIs. They cover all clinically relevant HER2 expression levels (0, 1+, 2+, 3+). The datasets are split between 4k pairs for training and 1k for testing, consistent with standard split used in prior works.
8
4.1
R. Aasim et al.
Implementation Details
Due to computational constraints and the patch size of DiT, images were resized to 512 × 512 and normalized to [−1, 1]. We used the DiT-B/2 backbone configuration [25]. This choice was empirically determined to balance computational efficiency (Base model) with the high spatial resolution (patch size of 2) required to resolve fine-grained histological structures like nuclei boundaries. We utilized the pre-trained VAE (sd-vae-f t-mse) to compress inputs into the latent space, applying the standard scaling factor to ensure unit variance [17]. The UNI model (UNI2-h) is used in a frozen state to extract 1536-dimensional embeddings. The model was trained on an NVIDIA RTX 6000 (48GB) GPU using the AdamW optimizer with a learning rate of 3 × 10−5 and a batch size of 8 for 1000 epochs. Notably, we use the scaled-Linear Noise Schedule (βstart = 0.0001, βend = 0.02); that prevents rapid drop in signal-to-noise ratio (SNR) during early timesteps, critical for preserving key morphological features in histopathology. 4.2
Quantitative Metrics
We use standard metrics including Mean Square Error (MSE), Peak Signalto-Noise Ratio (PSNR) [33], Structural Similarity Index (SSIM) [34], Learned Perceptual Image Patch Similarity (LPIPS) [35], and vanilla Fréchet Inception Distance (FID) [36]. Additionally, to address the limitations of standard SSIM in histopathology, we report the Structural Correlation Metric (SCM). The standard SSIM is heavily biased by the white background (high intensity values), mathematically inconsistent for Virtual Staining because the high luminance of the white background dominates the score, masking structural errors. SSIM is a product of luminance, contrast and structure. For a generated IHC y ′ to be compared with real IHC y, it is, SSIM (y, y ′ ) = f (l(y, y ′ ), c(y, y ′ ), s(y, y ′ )) SSIM (y, y ′ ) =
2µy µ′y + C1 2σyy′ + C2 σyy′ + C3 . . µ2y + µ2y′ + C1 σy2 + σy2′ + C2 σy σy′ + C3
(4)
The luminance term compares mean brightness; however, a generated image dominated by bright pixels (even with poor cellular structures) will often make a mean brightness (µ) close to the ground truth. This results in a luminance score near 1.0, which acts as a multiplier that bloats the final SSIM value while masking a poor structural score. In histopathology, where preserving tiny, complex structures is critical, evaluation metrics must prioritize structural correlation. Therefore, we isolate structure (s) component to provide a honest assessment. The Structural Correlation Metric (SCM) is a split of the structure component of multi-scale SSIM [37] with a window-size of 11 × 11, given by, M
N
1 XX SCM (y, y ) = M N i=1 j=1 ′
σyy′ (i, j) + C σy (i, j) σy′ (i, j) + C
(5)
HistDiT
9
This metric focuses purely on the correlation of variance (texture and edges), ignoring mean luminance shifts inherent in staining differences.
5
Results and Discussion
We compared HistDiT against established GAN and diffusion baselines on the MIST and BCI test datasets. Our method presents quality staining and successfully achieves superior performance across perceptual and structural metrics. 5.1
Experimental Results on BCI Benchmark
GT IHC
Ours
pP2P[6] cDiff [19]
H&E
The BCI dataset serves as a rigorous benchmark for evaluating breast cancer due to its complex tissue morphology and subtle HER2 staining variations. It demands high fidelity in preserving pathological details, critical for accurate diagnosis. Test data contains 977 paired H&E, IHC samples across varying HER2 expression levels. A qualitative comparison of the generated samples, with state-of-the-art algorithms, is shown in Fig. 2, highlighting the model’s superior performance in handling complex HER-2 expression patterns. Visual results
Level 0 Level 1+ Level 1+ Level 2+ Level 2+ Level 3+ Level 3+ Level 3+
Fig. 2. Qualitative comparison on BCI Dataset across HER2 expression levels, ranging from negative (Level 0) to strongly positive (Level 3+). Rows compare the input H&E and baseline methods against our HistDiT and Ground Truths. HistDiT demonstrates higher fidelity and accurate stain intensity, particularly in high-grade regions (2+, 3+).
demonstrate that HistDiT effectively distinguishes between subtle staining variations, particularly in high-expression (Level 3+) samples where baselines often
10
R. Aasim et al.
produce “washed out” artifacts. We observed few cases where generated samples diverge from Ground Truth. This is attributed to serial sectioning artifacts inherent to the dataset, where the Input H&E and Ground Truth IHC are physically different tissue slices cut 4–5µm apart. Due to this depth disparity, cellular structures visible in the H&E may terminate or shift position in the subsequent IHC slice. In these instances, HistDiT strictly adheres to the morphological structures present in the source H&E for cellular staining, prioritizing diagnostic safety over hallucination. The quantitative results in Table 1, demonstrate better improveTable 1. Quantitative comparison with State-of-the-Art methods on BCI dataset. The best are highlighted in bold, and the second-best are underlined. Symbol ρ indicates pixel space operation and z means latent space. (↓ lower is better, ↑ higher is better)
Method/Model MSE↓ PSNR(dB)↑ SSIM↑ SCM↑ LPIPS↓ FID↓ Cycle GAN [4]ρ 1892.00 15.98 0.372 0.447 0.624 188.33 Pix2Pix [5]ρ 1560.26 17.08 0.303 0.391 0.769 160.39 Pyramid Pix2Pix [6]ρ 1348.54 19.61 0.397 0.473 0.466 167.40 — — 0.5032 — — — ASP [7] ρ Syn Diff [8]z — 14.28±2.52 0.32±0.1 — — — — 16.75±4.20 0.38±0.1 — — — PST-Diff [9]z Conditional Diff. [19]z 3011.36 15.76 0.4050 0.494 0.591 194.82 Star-Diff (No Visual)[20]z — 21.30±0.01 0.5301 — — — HistDiST [21]z — — 0.4693 — — — Proposed [HistDiT]z 891.53 21.43 0.4769 0.540 0.412 49.15 Improvement to SoTA -457.01 0.14 -0.026 0.046 -0.054 -111.2 %age Improvement 33.89% 0.66% 5.2% 9.31% 11.59% 69.4%
ment in terms of MSE, PSNR, FID, outperforming state-of-the-art GAN and diffusion-based methods on BCI data. The significantly lower FID (49.15) indicates that our dual-stream architecture generates textures indistinguishable from Table 2. Level-wise comparison of HER-2 Biomarker generation on the BCI Dataset. HER-2 Scale IQA Metrics cycleGAN [4] Pyr. Pix2Pix[6] cond. Diff [19] HistDiT PSNR ↑ 15.92 17.96 16.75 22.49 Level 0 SSIM ↑ 0.3587 0.3758 0.4595 0.5172 (38) LPIPS ↓ 0.5210 0.4683 0.5161 0.3945 PSNR ↑ 16.84 19.92 15.89 23.22 Level 1+ SSIM ↑ 0.3712 0.4099 0.4027 0.5046 (235) LPIPS ↓ 0.5185 0.4467 0.5790 0.3836 PSNR ↑ 16.21 20.39 15.70 20.4 Level 2+ SSIM ↑ 0.3759 0.3984 0.3746 0.4645 (446) LPIPS ↓ 0.5025 0.4524 0.5908 0.4376 PSNR ↑ 15.64 18.23 15.41 18.33 Level 3+ SSIM ↑ 0.3583 0.3859 0.3654 0.4052 (258) LPIPS ↓ 0.5394 0.5052 0.5887 0.4664
HistDiT
11
real pathology. The PSNR of (21.43) sets a new benchmark for signal fidelity. Notably, while standard SSIM score is slightly lower than some pixel-aligned baselines (e.g., ASP) but SCM is superior. This discrepancy, highlights a limitation of the SSIM in masking biological structures. In Table 2, we separately evaluate image quality across different HER-2 expression levels (0, 1+, 2+, 3+) by splitting the BCI test dataset. Our proposed HistDiT demonstrates superior performance across all levels of HER2 biomarker, particularly the challenging level 2+ and 3+ categories, where it successfully preserves the structural fidelity (SCM) and perceptual quality required for accurate medical analysis. 5.2
Experimental Results on MIST Dataset
GT IHC
Ours
pP2P[6] cDiff [19]
H&E
The MIST dataset presents a significant challenge for generative models due to its inherent complexity and lack of standardized structural alignment. It is characterized by high intra-class variability and complex textural details that are difficult for conventional GANs to synthesize without artifacts. This makes it an ideal benchmark for testing a model’s ability to maintain structural fidelity and perceptual realism under unconstrained conditions. In the test dataset we have 1000 input/output pairs for HER2 biomarker assessment. Visual quality analysis presented in Fig 3, reveals a superior staining quality of HistDiT.
Fig. 3. Visual comparison on the MIST dataset. Rows compare the input H&E and baselines against our HistDiT and Ground Truth. HistDiT demonstrates superior stain restoration, generating sharp morphological details.
We identified multiple instances, where the Ground Truth IHC images suffered from acquisition artifacts, like defocus blur or scanning noise. In these cases,
12
R. Aasim et al.
HistDiT generated crispier, high-contrast images that were visually superior to the Ground Truth itself Fig. 4. By conditioning strictly on the sharp H&E input, our model effectively restores the intended biological structures rather than over-fitting to the degraded quality of the target labels. Similarly, the results in Table 3 demonstrate that HistDiT achieves state-of-the-art performance across all reported metrics. Notably, HistDiT achieves SSIM of 0.211 and an FID of 59.3, outperforming competitors like HistDiST and PixCell. This indicates, when the data is noisy and unconstrained, our model’s reliance on VAE structural blueprint provides a good stability that baseline approaches lack. Table 3. Quantitative comparison on the MIST Dataset. The best results are bold and second-best are underlined. (↑ higher is better, ↓ lower is better).
Method/Model MSE↓ PSNR(dB)↑ SSIM↑ SCM↑ LPIPS↓ FID↓ ρ Cycle GAN [4] 4190.52 11.91 0.174 0.209 0.553 125.7 Pix2Pix [5]ρ 3485.20 12.74 0.150 0.194 0.614 128.1 2986.44 14.24 0.165 0.214 0.543 107.4 Pyramid Pix2Pix [6]ρ ASP [7]ρ — — 0.1945 — — 54.28 Conditional Diff. [19]z 6252.98 11.13 0.168 0.254 0.564 82.6 — — 0.2059 — — — HistDiST [21]z PixCell (no LoRA) [38]z — — 0.1880 — — 67.68 Proposed [HistDiT]z 3396.88 14.26 0.211 0.302 0.489 59.30 Improvement to SoTA -410.44 0.02 0.0051 0.048 -0.054 5.02 %age Improvement 13.74% 0.14% 2.5% 18.9% 9.94% 9.25%
These results suggest that proposed HistDiT learns the underlying concept of the stain distribution rather than merely memorizing pixel-to-pixel mappings, thus offering potential utility as a quality enhancement tool in clinical workflows. 5.3
Expert Evaluation of Visual Fidelity
To evaluate the perceptual realism and staining fidelity of our generated samples, we conducted a blind qualitative assessment with domain experts. The study involved immunologists from Cancer Research Malaysia (CRMY), who have extensive experience in IHC analysis. The experts were presented with anonymized random set of image patches containing both Real IHC (Ground Truth) and HistDiT-generated stains, paired with their corresponding H&E. The experts were asked to distinguish the real biological samples from the virtually stained ones, serving as a visual Turing test. The feedback indicated that the experts were unable to consistently differentiate between the two, confirming that HistDiT successfully captures the complex staining patterns and structural coherence while minimizing the artifacts. While this validates the model’s structural similarity and textural realism, we acknowledge that HER2 scoring is a specialized clinical task; therefore, strict diagnostic grading validation with board-certified pathologists is reserved for future clinical feasibility studies.
HistDiT
5.4
13
Ablation Studies
To validate the effectiveness of our dual-stream conditioning and the hybrid objective function, we conducted a component-wise analysis. The architectural choices are detailed in Table 4, and impact of loss functions is shown in Fig. 4. Table 4. Comparison for different architectural choices on BCI Dataset. Best are bold. Model Configuration PSNR(dB)↑ SSIM↑ SCM↑ LPIPS↓ FID↓ Pyramid Pix2Pix [6]ρ 19.61 0.397 0.473 0.466 167.40 HistDiT (Semantics only)z 16.39 0.423 0.492 0.553 81.51 18.58 0.416 0.560 0.452 62.38 HistDiT (Spatial only with Cross-Attention)z HistDiT (Spatial [Concatenation], Semantics)z 20.63 0.449 0.521 0.438 56.80 HistDiT (Spatial [Cross-Attention], Semantics)z 21.43 0.477 0.540 0.412 49.15
Impact of Dual-Stream Conditioning: Using UNI embeddings (Semantics Only) produced perceptually realistic textures but caused severe structural hallucinations, generating high-density tissue artifacts. Conversely, removing the UNI guidance (Spatial Only) preserved tissue morphology but failed to resolve stain intensity between HER2 levels, often washing out dense expression levels. Furthermore, injecting spatial VAE-latents via cross-attention significantly outperformed standard channel-wise concatenation, confirming that dynamic attention is crucial for aligning structural blueprints with high-level semantic context.
Method/Approach PSNR↑ SSIM↑ FID↓ M SE only 16.86 0.438 67.88 L1 only 19.44 0.450 93.31 0.9M SE + 0.1L1 20.44 0.443 56.18 21.43 0.477 49.15 0.7M SE + 0.3L1 MSE only L1 only
MSE+L1 Real IHC
Fig. 4. Ablation Study on Objective Functions: (Left) Quantitative comparison of loss components on the BCI dataset. M SE+L1 maintains high metric scores. (Right) Visual samples demonstrating that the combined objective produces sharper structures (column 3) compared to the smoothing artifacts seen in MSE only (column 1).
Optimization of Training Objectives: We investigated the impact of training loss towards the reconstruction of high-frequency details (Fig. 4). Training with “MSE only” resulted in over-smoothed textures, a known limitation where mean-square suppresses high-frequency variance. While the “L1 ” loss improved structural contrast, it degraded visual fidelity (FID 93.31), likely due to the statistical mismatch between the Laplacian distribution of L1 and the added Gaussian noise in forward diffusion. This forces the model to learn a suboptimal approximation of the noise distribution, degrading color fidelity. Our “Hybrid Objective” provides the optimal balance, utilizing L1 to sharpen structures while
14
R. Aasim et al.
retaining MSE for statistical consistency required for accurate stain-translation. The introduced hybrid objective helps in generating indistinguishable textures.
6
Conclusion
This study introduces HistDiT, a dual-stream diffusion transformer for virtual immunohistochemistry. By integrating the structural precision of VAE-based spatial conditioning with the semantic richness of the UNI model, the architecture overcomes the limitations of previous GAN-based and U-Net diffusion approaches. The rigorous mathematical formulation of incorporating a scaled linear noise schedule and a robust hybrid loss function ensures training stability and visual fidelity. Additionally, the Structural Correlation Metric (SCM) rectifies the luminance bias of standard SSIM in histopathology. We compared HistDiT against established GAN and diffusion baselines on the BCI and MIST datasets, demonstrating superior performance across structural and perceptual quality metrics. Our findings confirm that HistDiT not only outperforms stateof-the-art methods but is capable of preserving diagnostic integrity required for clinical decision-making. This work positions DiTs, supported by domain-specific Foundation Models, as new standard for generative computational pathology. Acknowledgements This research is fully funded by the Edge Hill University’s GTA studentship, and is in collaboration with Cancer Research Malaysia (CRMY) and University of Nottingham, Malaysia (UNM).
References 1. Sedeta, E.T., Jobre, B., Avezbakiyev, B.: Breast cancer: Global patterns of incidence, mortality, and trends. J. Clin. Oncol. 41(16-suppl) (2023) 2. Golestan, A., et al.: Unveiling promising breast cancer biomarkers. BMC Cancer 24(1), 155 (2024). https://doi.org/10.1186/s12885-024-11913-7 3. Bai, B., et al.: Deep learning-enabled virtual histological staining. Light Science & Applications 12(1), 57 (2023) 4. Xu, Z., et al.: GAN-based virtual re-staining. arXiv:1901.04059 (2019) 5. Isola, P., Zhu, J., Zhou, T., Efros, A.A.: Image-to-image translation with conditional adversarial networks. In: CVPR, pp. 1125–1134. IEEE (2017) 6. Liu, S., et al.: BCI: Breast cancer immunohistochemical image generation through pyramid pix2pix. In: CVPR Workshops, pp. 1815–1824. IEEE (2022) 7. Li, F., et al.: Adaptive supervised PatchNCE loss for learning H&E-to-IHC stain translation. In: MICCAI, pp. 41–51. Springer (2023) 8. Dhariwal, P., Nichol, A.: Diffusion models beat GANs on image synthesis. In: NeurIPS, vol. 34, pp. 8780–8794 (2021) 9. Özbey, M., et al.: Unsupervised medical image translation with adversarial diffusion models. IEEE Transaction on Medical Imaging 42(12), 3524–3539 (2023) 10. Kataria, T., Knudsen, B., Elhabian, S.Y.: StainDiffuser: Multi-task dual diffusion model. arXiv:2403.11340 (2024)
HistDiT
15
11. He, Y., et al.: PST-Diff: Achieving high-consistency stain transfer by diffusion models with pathological and structural constraints. IEEE TMI (2024) 12. Kazerouni, A., et al.: Diffusion models in medical imaging: A comprehensive survey. Medical Image Analysis 88, 102846 (2023) 13. Goodfellow, I., et al.: Generative adversarial nets. In: NeurIPS, vol. 27 (2014) 14. Duan, G., et al.: A virtual staining method for immunohistochemical images of breast cancer. In: CISP-BMEI, pp. 1–5. IEEE (2023) 15. G., Ainur: GAN mode collapse explanation. Towards AI Medium (2023). https: //pub.towardsai.net/gan-mode-collapse-explanation-fa5f9124ee73 16. Ho, J., Jain, A., Abbeel, P.: Denoising diffusion probabilistic models. In: NeurIPS, vol. 33, pp. 6840–6851 (2020) 17. Rombach, R., et al.: High-resolution image synthesis with latent diffusion models. In: CVPR, pp. 10684–10695. IEEE (2022) 18. Moghadam, P.A., et al.: A morphology focused diffusion probabilistic model for synthesis of histopathology images. In: WACV, pp. 2000–2009. IEEE (2023) 19. Er, X., et al.: Conditional diffusion-based virtual staining. In: ICPR, pp. 193–207. Springer, Cham (2024) 20. Liu, J., et al.: From pixels to pathology: Restoration diffusion for diagnosticconsistent virtual IHC. Comput. Biol. Med. 198, 111264 (2025) 21. Großkopf, E., et al.: HistDiST: Histopathological diffusion-based stain transfer. arXiv:2505.06793 (2025) 22. OpenAI, et al.: GPT-4 Technical Report. arXiv:2303.08774 (2023) 23. Lin, K., et al.: DEsignBench: Exploring and benchmarking DALL-E 3. arXiv:2310.15144 (2023) 24. Stability AI: Stable Diffusion 3.5 Medium. Hugging Face (2024). 25. Peebles, W., Xie, S.: Scalable diffusion models with transformers. In: ICCV, pp. 4195–4205. IEEE (2023) 26. Wu, J., et al.: PTQ4DiT: Post-training quantization for diffusion transformers. In: NeurIPS, vol. 37 (2024) 27. Chen, R.J., et al.: Towards a general-purpose foundation model for computational pathology. Nature Medicine 30(3), 850–862 (2024) 28. Wang, X., et al.: Transformer-based unsupervised contrastive learning for histopathological image classification. Medical Image Analysis 81, 102559 (2022) 29. Nichol, A.Q., Dhariwal, P.: Improved denoising diffusion probabilistic models. In: ICML, PMLR, vol. 139, pp. 8162–8171 (2021) 30. Stability AI: SD-VAE-FT-MSE. Hugging Face (2024) 31. Hang, T., et al.: Improved noise schedule for diffusion training. In: ICCV (2025) 32. Dosovitskiy, A., et al.: An image is worth 16x16 words: Transformers for image recognition at scale. In: ICLR (2021) 33. Hore, A., Ziou, D.: Image quality metrics: PSNR vs. SSIM. In: ICPR, pp. 2366– 2369. IEEE (2010) 34. Wang, Z., Bovik, A.C., Sheikh, H.R., Simoncelli, E.P.: Image quality assessment: From error visibility to structural similarity. IEEE TIP 13(4), 600–612 (2004) 35. Zhang, R., et al.: The unreasonable effectiveness of deep features as a perceptual metric. In: CVPR, pp. 586–595. IEEE (2018) 36. Parmar, G., Zhang, R., Zhu, J.-Y.: On aliased resizing and surprising subtleties in GAN evaluation. In: CVPR, pp. 11410–11420. IEEE (2022) 37. Venkataramanan, A.K., et al.: A hitchhiker’s guide to structural similarity. IEEE Access 9, 28872–28896 (2021) 38. Yellapragada, S., et al.: PixCell: A generative foundation model for digital histopathology images. arXiv:2506.05127 (2025)