Title: D-VST: Diffusion Transformer for Pathology-Correct Tone-Controllable Cross-Dye Virtual Staining of Whole Slide Images
Venue: NeurIPS 2025
Accepted: yes
Keywords: Computational Pathology

Whole Slide Images

Virtual Staining

Diffusion Transformer
Abstract: Diffusion-based virtual staining methods of histopathology images have demonstrated outstanding potential for stain normalization and cross-dye staining (e.g., hematoxylin-eosin to immunohistochemistry). However, achieving pathology-correct cross-dye virtual staining with versatile tone controls poses significant challenges due to the difficulty of decoupling the given pathology and tone conditions. This issue would cause non-pathologic regions to be mistakenly stained like pathologic ones, and vice versa, which we term “pathology leakage.” To address this issue, we propose diffusion virtual staining Transformer (D-VST), a new framework with versatile tone control for cross-dye virtual staining. Specifically, we introduce a pathology encoder in conjunction with a tone encoder, combined with a two-stage curriculum learning scheme that decouples pathology and tone conditions, to enable tone control while eliminating pathology leakage. Further, to extend our method for billion-pixel whole slide image (WSI) staining, we introduce a novel frequency-aware adaptive patch sampling strategy for high-quality yet efficient inference of ultra-high resolution images in a zero-shot manner. Integrating these two innovative components facilitates a pathology-correct, tone-controllable, cross-dye WSI virtual staining process. Extensive experiments on three virtual staining tasks that involve translating between four different dyes demonstrate the superiority of our approach in generating high-quality and pathologically accurate images compared to existing methods based on generative adversarial networks and diffusion models. Our code and trained models will be released.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents D-VST (Diffusion Virtual Staining Transformer), a new diffusion-based method for virtual staining of histopathology images, with a focus on accurate pathology preservation and flexible tone control. Traditional methods often suffer from pathology leakage, where incorrect staining occurs due to entanglement of pathology and tone information. To address this, the authors propose:\n\n1. A dual-encoder architecture (pathology and tone encoders) that separates content and style in the staining process.\n2. A two-stage curriculum learning strategy to effectively train the model on these disentangled conditions.\n3. A frequency-aware adaptive patch sampling technique for efficient and scalable inference on ultra-high-resolution whole slide images, enabling zero-shot generalization.\n\nThe method demonstrates strong performance across three stain translation tasks, outperforming existing GAN-based and diffusion-based approaches in both visual quality and pathological fidelity.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. First Diffusion-Based Model with Tone Control for Cross-Dye Staining\n\nThe proposed Diffusion Virtual Staining Transformer (D-VST) is the first diffusion-based model to support versatile tone control in cross-dye virtual staining, marking a meaningful advance in the virtual staining literature.\n\n2. Pathology and Tone Representations Decoupling via Curriculum Learning\n\nD-VST introduces a two-stage curriculum learning strategy to disentangle pathology and tone representations, effectively addressing the critical issue of pathology leakage, where pathological and non-pathological regions are misrepresented during virtual staining.\n\n3. Frequency-Aware Adaptive Sampling for High-Quality Zero-Shot Virtual Staining for WSIs\n\nTo mitigate mosaic artifacts and enhance zero-shot generation quality on whole slide images (WSIs), the authors introduce a frequency-aware adaptive patch sampling strategy. By prioritizing sampling in low-frequency regions, where artifacts are more likely, this method enables efficient and high-fidelity inference on billion-pixel WSIs.\n\n4. Comprehensive Experimental Evaluation\n\nThe method is validated on three cross-dye staining tasks involving four stain types, with additional downstream classification tasks and ablation studies, demonstrating strong performance against recent approaches.\n\nWeaknesses:\n\n1. Insufficient Explanation of HER2 Labels in Figure 1\n\nFigure 1 includes references to HER2 0, 1+, 2+, 3+, but lacks explanation of what these labels mean clinically. This omission reduces clarity for readers unfamiliar with HER2 scoring systems.\n\n2. Evidence for Key Claims Should Be Included in Figure 1\n\nThe caption of Figure 1 claims the method is \"Fast\", \"Pathology-Correct\", and “No Mosaic Artifacts”, but no visual or quantitative evidence is provided to substantiate these assertions in this figure.\n\n3. Missing Runtime Comparison\n\nWhile the paper asserts that the method is computationally efficient, there is no runtime or inference speed comparison with other models.\n\n4. Requirement for Cross-Stain Paired Training Data\n\nD-VST relies on paired cross-stain training data, which can be challenging and costly to obtain in real-world medical imaging workflows. This reliance limits the method’s scalability and broader applicability, yet the paper does not discuss this practical constraint or propose solutions for unpaired or weakly supervised settings.\n\nQUESTIONS:\n1. Clarification and Supporting Evidence for Figure 1 Claims\n\nFigure 1 highlights several core advantages of D-VST, such as being “Fast,” “Pathology-Correct,” and free of “Mosaic Artifacts.” However, these claims are not directly supported within the figure.\n\nSuggestion: Please include visual examples (e.g., zoom-in patches with and without mosaic artifacts) and/or quantitative metrics that demonstrate the absence of artifacts and the correctness of pathology preservation. While including some visual or metric-based evidence in Figure 1 would be ideal, alternatively, you could explicitly reference other figures, tables, or sections in the paper that substantiate these claims. For example, show where “no mosaic artifact” is evidenced, or where “fast” is quantified.\n\nEvaluation impact: Providing this evidence would significantly improve the clarity and credibility of the figure, and would strengthen the justification for your claims.\n\n2. Runtime and Efficiency Comparison Missing\n\nThe method is described as fast, but no runtime or efficiency comparisons are provided.\n\nSuggestion: Please provide a runtime or inference time comparison with baseline models (e.g., GAN-based and other diffusion-based methods), ideally on WSIs and/or standard image sizes.\n\nEvaluation impact: Demonstrating competitive or superior efficiency would support real-world deployment claims and positively influence the significance of the work.\n\nLIMITATIONS:\nNo. The authors have addressed several important limitations. To further strengthen the discussion, it would be helpful to acknowledge the reliance on cross-dye paired training data, which may be challenging to obtain in practice. Highlighting this would enhance the clarity and practical relevance of the work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes a method for integrating tone control for virtual staining using diffusion transformer models. This integration is done to reduce false virtual staining of regions in tissue images (e.g., staining of non-cancer regions as cancer regions). The method is evaluated on several datasets and virtual staining tasks.\n\nSTRENGTHS AND WEAKNESSES:\nThe use of tone control to improve virtual staining performance is different and interesting. The proposed approach implements a two-stage training strategy to incorporate pathology and tone information. It also employs an adaptive patch-sampling technique to improve image generation quality in whole slide images. \n\nExperimental evaluation using several datasets and virtual staining tasks (HE-to-IHC, FFPE-to-HE and HE-to-mIHC) shows performance improvements by the proposed method in several metrics (FID, KID, SSIM, etc) compared with other methods. \n\nThe experimental evaluation also includes a downstream image classification task. This is a useful experiment as it shows the performance improvements are not just in image generation metrics (KID, FID, etc), but also have discernable impact in pathology focused tasks.\n\nQUESTIONS:\nPage 7 (lines 224-225): “use a random patch from the same WSI but not overlapping with the target image for tone conditioning”. Per this sentence, it sounds like a separate patch is selected for a source image. Is this realistic? In a real setting, there may not be more than one target image (e.g., an IHC image in case of HE-to-IHC image generation). Clarification is needed for how tone conditioning patches are selected during the training and test phases in the experiments.\n\nPage 8 (lines 274-276): “Next, for each HE image in the sub-test set, we randomly select an IHC image in the sub-test set that is not paired with the specific HE image as the tone condition for virtual staining.” Selecting a different tone condition image for each HE image may affect the experimental results. In a real application setting, there may not be multiple IHC images. Please clarify if multiple IHC patches were selected for tone conditioning in the experiments. If so, a set of experiments should be executed in which only one IHC conditioning patch is selected.\n\nLIMITATIONS:\n.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a diffusion transformer for virtual staining, e.g IHC staining from HE images, or HE staining from FFPE (unstained tissue). It has two main contributions: 1) it uses a two-stage curriculum learning scheme to separate tone and pathology features, solving the pathology leakage problem, and ii) it uses a frequency-aware adaptive patch sampling strategy to mitigate mosaic artefacts during high-resolution inference. They evaluate their model on three datasets and show improved or competitive performance compared to state-of-the-art virtual staining methods.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1: The paper is well written with clearly explained motivation, methodology and experimental design. \n\n2. The methodology is logical and well-designed. Although the essential concept like disentanglement of pathology features and tones, is not necessarily new (e.g. used in histology stain normalisation/augmentation), but the design of the presented model architecture is novel, as far as I can see. \n\n3. The experimental evaluation is also comprehensive, covering three different datasets. Particularly, the evaluation of downstream HER classification and scoring provides quantitative evidence of the proposed method over other methods in pathology leakage prevention (Table 3 back up qualitative results in Fig. 2) . \n\n4. Ablation studies show that two-stage curriculum learning is indeed the key of pathology structure preservation in terms of prediction accuracy of HER classification (from around 0.77 to 0.94).\n\nWeakness:\n1. The theoretical contribution of the paper is limited as it is pretty much a very applied paper (not necessarily a weakness of the paper itself, but theoretical contribution is something I like for NeurIPS conference)\n2. Sampling strategy seems only marginal better (according to Table 1). Is it worth the computational cost?\n3. Not entirely clear about model generalisation (see Questions for details)\n\nQUESTIONS:\n1. According to Session: Datasets and Evalaution Metrics (line 218-265), authors trained three models, one for each dataset and evaluated the model on the testing split of the corresponding dataset. This already raised questions on the model generalisation. As for clinical usage, a model need to be evaluated on an independent patient cohort. \n2. Another question is for tone image selection. It is written that tone image is “ a random patch from the same WSI but not overlapping  with the target image for tone conditioning.” For model inference, do users need to select one tone image for the same WSI of a different stain? That leads to a practical problem, if users have the stained the same image with a different stain, why do they need the generated one? Is it possible to select the tone image from the training set. \n3. There is rapid development for pathology foundation model recently,  why not use pathology foundation models trained on multi stains as conditioning like UNI, https://huggingface.co/MahmoodLab/UNI2-h\n4. Ablation study on the autoencoder size would also be interesting, since authors only use a small and lightweight one.\n5. Ref [49] is a very relevant work as it does something similar but on TMAs and using GNN. Could authors discuss their methodological differences?\n\nLIMITATIONS:\nMaybe indeed authors can add a few sentence about the status of their model and potential limitations of using their model in clinical settings.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors present a diffusion-based virtual staining (VS) model specifically designed to address potential tone variation and pathology leakage in the generated images. The three main contributions of the framework are: (1) diffusion-based staining and tone control, (2) a curriculum learning method to prevent pathology leakage, and (3) a new frequency domain content-based sampling technique to produce high-quality VS images.  The analysis includes comparisons against GAN-based methods, showing superior performance compared to methods in the literature while also decreasing the computational complexity\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n- The authors presented a systematic approach to the problem. They tailored their solutions according to the needs of the application and drawbacks in the current methods. \n- The experiments are well designed and illustrate the performance of the model well. \n- Ablation studies are presented with illustrate the performance boost introduced by different components\n- Patch selection strategy is well justified, provide good efficacy\n\n\n\n**Weaknesses** \n\n- It is not clear why the authors particularly picked SynDiff and StainFuser for comparisons. There are several other methods available in the literature (such as \"Pixel super-resolved virtual staining of label-free tissue using diffusion models\" or DeepLiif - rather than cycleGAN)\n- Conclusions section is too short and does not contain anything more than the summary of the paper\n- It would have been nice to give a WSI example in the paper, and also some timing comparison between DVST and the competing methods for virtual staining of WSI images.\n\nQUESTIONS:\nCan the authors give some timing information for large WSI generation and compare it against at least StainFuse and SynDiff\n\nDid the authors observe any significant difference between the VS images when they used different tone domain image examples for the same sample? How robust is the model against variation of the tone domain images?\n\nLIMITATIONS:\nPerformance comparisons are limited to FID, PSNR, SSIM metrics, which can be diagnostically deceptive.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a new method called D-VST. It is for cross-dye virtual staining of pathology images. The method uses a Diffusion Transformer. It tackles two main problems. One is \"pathology leakage\". This is when pathology features from a reference image wrongly appear in the result. The other problem is the \"mosaic artifact\". This happens when making very large Whole Slide Images (WSIs). D-VST uses two separate encoders for pathology and tone. It also uses a two-stage training plan. This helps to separate pathology from tone information. For large WSIs, the paper suggests a new patch sampling strategy. It focuses on low-frequency areas to avoid artifacts. This is also computationally efficient. The authors test their method on three tasks. The results show D-VST is better than other methods. It creates high-quality images and keeps the pathology correct. A classification task confirms this.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1.\tD-VST is the first diffusion model to achieve tonal control for cross-staining virtual staining, filling a gap in existing methods.\n2.\tThe paper uses separate encoders to extract pathology structure and staining tone. A two-stage curriculum learning first focuses on pathology, then introduces tone control to prevent “pathology leakage.” A frequency-aware adaptive patch sampling enables efficient inference on gigapixel whole-slide images and removes mosaic artifacts.\n3.\tThe adaptive sampling method is a good innovation for large images. It solves the quality versus cost problem. This makes the tool practical.\n\n**Weaknesses**\n\n1.\tThe VAE used was pre-trained on a natural image (Stable Diffusion 1.5), which is not suitable for pathology images due to the obvious domain shift. You note in the appendix that the fine-tuning failed, but you do not analyze why, which makes the choice of VAE seem arbitrary. I suggest adding a direct performance comparison with pathology-specific VAEs (e.g., VAEs using the PathCLIP encoder) to quantify the impact of domain shifting. I would also suggest analyzing the distribution of reconstruction errors to show which specific pathology-specific structures are poorly reconstructed by current VAEs, rather than just reporting average scores.\n2.\tThe frequency-aware adaptive sampling is a key contribution, but the choice of hyperparameters α=1 and β=8 appears inadequately justified. The paper mentions an appendix for the grid search of n, but a similar rigorous ablation study for α (which controls the sampling probability distribution ) and β (which governs the cost-quality trade-off ) is absent. The robustness and generalizability of the sampling strategy depend heavily on these parameters, and a dedicated analysis is needed to strengthen this contribution.\n3.\tThe authors rightly note the difficulty in quantitatively evaluating billion-pixel WSIs due to the lack of suitable datasets. While understandable, this remains a limitation, especially since WSI processing is a core claim. You should suggest a way to test it. For example, you could measure consistency between patches. Or you could ask pathologists to rate the image quality blindly.\n4.\tThe paper lacks theoretical analysis or proof and relies solely on empirical evidence. This may limit a deeper understanding of model behavior and principles, especially when explaining why two-stage learning and frequency sampling are effective.\n5.\tIt is recommended to test the model on more diverse pathology datasets and staining types to assess its robustness and generalization capabilities. For example, images of different organs or pathology types could be included.\n6.\tOn Ext-FFPE2HE, lower PSNR/SSIM vs. task-specific [27] attributed to \"poor pixel alignment\" without evidence.\n7.\tThere is insufficient discussion of the direction of generation and the direction of medical treatment. It would have been better for the authors to present more methods [A-E].\n\n**Minor issues**\n\n1. Simplify long sentences for readability.\n2. Check figure/table numbering and ensure all subfigures have clear captions.\n3. Standardize reference formats.\n\n**Reference**\n\n[A] Improving generative adversarial network inversion via fine-tuning GAN encoders[J]. TPAMI, 2024.\n\n[B] Multi-marginal wasserstein gan[J]. Advances in Neural Information Processing Systems, 2019, 32. \n\n[C] Closed-loop matters: Dual regression networks for single image super-resolution[C]CVPR. 2020.\n\n[D] Cross-ray neural radiance fields for novel-view synthesis from unconstrained image collections[C] ICCV 2023\n\n[E] Segment anything model for medical image segmentation: Current applications and future directions[J]. Computers in Biology and Medicine, 2024\n\nQUESTIONS:\nNA\n\nLIMITATIONS:\nSee the Weaknesses", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-04"}