SheafStain: Sheaf-Theoretic Schrödinger Bridge for Spatially and Biologically Coherent Virtual Staining

1Department of Biomedical Informatics, Korea University College of Medicine2DEEPNOID Inc.

NeurIPS 2026

    SheafStain reinterprets pathology vision foundation model (VFM) features as sheaf-like sections, so virtual IHC translated patch by patch stays spatially and biologically coherent.

    TL;DR

    Virtual staining of gigapixel whole-slide images runs patch by patch, and independent patches fail to keep spatial continuity. Pathology VFMs offer rich representations, but their self-attention gives the same physical region inconsistent embeddings under different global contexts. We formalize this context contamination as a sheaf-theoretic problem and propose SheafStain, which integrates the VFM class and patch tokens into a Schrödinger bridge as sheaf-like sections. Evaluated on stitched 1024×1024 outputs across HER2, ER, PR and Ki-67, SheafStain shows promising results against six prior methods1,2,3,4,5,6 while mitigating patch-boundary stitching artifacts.

    Motivation

    An IHC stain can be applied to only one tissue section, so virtual staining predicts it from the H&E of the same tissue. Whole-slide images force this translation onto small patches, and each patch is produced without its neighbors. Reassembling the patches then shows three failures: patch-boundary discontinuities, stain-tone drift, and chromatin-distribution mismatch.

    VFMs are a natural source of conditioning, but self-attention ties each region's embedding to its surrounding context. Viewed over overlapping patches, the embeddings form a presheaf whose sections disagree on overlaps, so no global section restricts to them. SheafStain treats the VFM features as sheaf-like sections and trains the generator to agree on the overlaps.

    Method at a glance

    SheafStain training pipeline: (a) overview with a reference patch and two adjacent targets sharing a triple overlap, (b) the spatial map built from VFM tokens of overlapping patches, (c) pixel sheaf and cocycle losses on shared regions.
    Training. (a) A reference patch and two adjacent targets sharing a triple overlap are sampled. The VFM-derived spatial map M and neighborhood class token cn condition three weight-shared generator passes. (b) VFM tokens of overlapping patches are aggregated into M, a sheaf-consistent section over the reference patch. (c) Pixel sheaf and cocycle losses penalize disagreement between adjacent outputs on shared regions.
    SheafStain inference pipeline: a stride-driven reference grid over the H&E, extra reference patches from the FFT energy map, per-patch translation under VFM conditioning, and soft-kernel weighted blending.
    Inference. The input H&E is covered by a stride-driven reference grid, supplemented with additional reference patches drawn from the high- and low-energy extremes of the mid-band FFT energy map. Each reference patch is translated under the VFM-derived conditioning, and the patches are reassembled by soft-kernel weighted blending.

    Results

    On BCI7 and MIST8 datasets, we translate at 256×256 and evaluate the stitched 1024×1024 outputs.

    Stitched 1024×1024 evaluation on BCI · HER2
    MethodFID ↓KID×103 ↓LPIPS ↓DISTS ↓TS ↓DAB-r ↑
    SheafStain (ours)36.36264.20700.4689±0.05840.2132±0.04790.0146±0.00720.1209
    Pix2Pix1172.7932138.90500.4800±0.07220.2753±0.06980.0533±0.02140.0559
    CycleGAN296.349650.80200.5448±0.06590.2920±0.05990.0475±0.02620.0668
    PSPStain3185.6991173.73600.5695±0.06400.2924±0.04940.0885±0.03540.0317
    D-VST487.310062.76700.5107±0.07070.2570±0.05960.0164±0.00700.0736
    UNIStainNet567.632229.30000.4577±0.07110.2279±0.05310.0448±0.03200.1267
    UNSB6227.7889236.23100.6635±0.06770.3347±0.04520.1464±0.01610.0589

    Best in bold, second best underlined. Mean ± standard deviation over test images for the per-image metrics.
    TS is the Tiling Score9 at patch seams. DAB-r is the DAB Pearson correlation with the ground-truth IHC.

    Qualitative comparison on BCI and MIST HER2: input H&E, outputs of prior methods and SheafStain, and the ground-truth IHC.
    Qualitative comparison on BCI and MIST HER2. More samples are in the gallery.
    References (9)
    • 1P. Isola, J.-Y. Zhu, T. Zhou, and A. A. Efros, Image-to-image translation with conditional adversarial networks, CVPR 2017.
    • 2J.-Y. Zhu, T. Park, P. Isola, and A. A. Efros, Unpaired image-to-image translation using cycle-consistent adversarial networks, ICCV 2017.
    • 3F. Chen, R. Zhang, B. Zheng, Y. Sun, J. He, and W. Qin, Pathological semantics-preserving learning for H&E-to-IHC virtual staining, MICCAI 2024.
    • 4S. Yang, D. Wei, Y. Hu, Q. Peng, H. Liu, Y. Huang, X. Wu, Y. Zheng, and L. Wang, D-VST: Diffusion transformer for pathology-correct tone-controllable cross-dye virtual staining of whole slide images, NeurIPS 2025.
    • 5J. R. Saurav, T. L. H. Pham, P. Mukherjee, P. Yi, B. A. Orr, and J. M. Luber, UNIStainNet: Foundation-model-guided virtual staining of H&E to IHC, arXiv preprint 2026. arXiv:2603.12716
    • 6B. Kim, G. Kwon, K. Kim, and J. C. Ye, Unpaired image-to-image translation via neural Schrödinger bridge, ICLR 2024.
    • 7S. Liu, C. Zhu, F. Xu, X. Jia, Z. Shi, and M. Jin, BCI: Breast cancer immunohistochemical image generation through pyramid pix2pix, CVPR Workshops 2022.
    • 8F. Li, Z. Hu, W. Chen, and A. Kak, Adaptive supervised PatchNCE loss for learning H&E-to-IHC stain translation with inconsistent groundtruth image pairs, MICCAI 2023.
    • 9O. Madar and O. Fried, Tiled Diffusion, CVPR 2025.

    BibTeX

    If you find our work useful, please cite it using the BibTeX below.

    @inproceedings{lim2026sheafstain,
  title     = {SheafStain: Sheaf-Theoretic Schr\"odinger Bridge for Spatially and Biologically Coherent Virtual Staining},
  author    = {Lim, Hyeongyeol and Yoon, Hongjun and Jang, Eunjin and Jeong, Daeky and Cho, Won June and Lee, Hwamin},
  booktitle = {The Fortieth Annual Conference on Neural Information Processing Systems},
  year      = {2026},
  url       = {https://arxiv.org/abs/2606.11846}
}