# SheafStain: Sheaf-Theoretic Schr"odinger Bridge for Spatially and Biologically Coherent Virtual Staining

arXiv: https://arxiv.org/abs/2606.11846 · HTML: https://mlsanity.com/paper/2606.11846/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 22 of its 39 references that are in the map.
- First posted on arXiv 10 Jun 2026.

## Builds on (most cited first, 22 of 22)

- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Image-to-Image Translation with Conditional Adversarial Networks](https://mlsanity.com/paper/1611.07004/) (2016)
- [Demystifying MMD GANs](https://mlsanity.com/paper/1801.01401/) (2018)
- [Image Quality Assessment: Unifying Structure and Texture Similarity](https://mlsanity.com/paper/2004.07728/) (2020)
- [A survey of the Schr"odinger problem and some of its connections with optimal transport](https://mlsanity.com/paper/1308.0215/) (2013)
- [Neural Sheaf Diffusion: A Topological Perspective on Heterophily and Oversmoothing in GNNs](https://mlsanity.com/paper/2202.04579/) (2022)
- [Unpaired Image-to-Image Translation via Neural Schr"odinger Bridge](https://mlsanity.com/paper/2305.15086/) (2023)
- [BCI: Breast Cancer Immunohistochemical Image Generation through Pyramid Pix2pix](https://mlsanity.com/paper/2204.11425/) (2022)
- [Adaptive Supervised PatchNCE Loss for Learning H&E-to-IHC Stain Translation with Inconsistent Groundtruth Image Pairs](https://mlsanity.com/paper/2303.06193/) (2023)
- [Sheaf Hypergraph Networks](https://mlsanity.com/paper/2309.17116/) (2023)
- [Pathological Semantics-Preserving Learning for H&E-to-IHC Virtual Staining](https://mlsanity.com/paper/2407.03655/) (2024)
- [Comparison of Consecutive and Re-stained Sections for Image Registration in Histopathology](https://mlsanity.com/paper/2106.13150/) (2021)
- [Cooperative Sheaf Neural Networks](https://mlsanity.com/paper/2507.00647/) (2025)
- [StainDiffuser: MultiTask Dual Diffusion Model for Virtual Staining](https://mlsanity.com/paper/2403.11340/) (2024)
- [Structural Cycle GAN for Virtual Immunohistochemistry Staining of Gland Markers in the Colon](https://mlsanity.com/paper/2308.13182/) (2023)
- [Sheaf theory: from deep geometry to deep learning](https://mlsanity.com/paper/2502.15476/) (2025)
- [Tiled Diffusion](https://mlsanity.com/paper/2412.15185/) (2024)
- [Learning Sheaf Laplacian Optimizing Restriction Maps](https://mlsanity.com/paper/2501.19207/) (2025)
- [Generating Seamless Virtual Immunohistochemical Whole Slide Images with Content and Color Consistency](https://mlsanity.com/paper/2410.01072/) (2024)
- [GANs vs. Diffusion Models for virtual staining with the HER2match dataset](https://mlsanity.com/paper/2506.18484/) (2025)
- [UNIStainNet: Foundation-Model-Guided Virtual Staining of H&E to IHC](https://mlsanity.com/paper/2603.12716/) (2026)
- [HistDiT: A Structure-Aware Latent Conditional Diffusion Model for High-Fidelity Virtual Staining in Histopathology](https://mlsanity.com/paper/2604.08305/) (2026)

## Abstract (arXiv)

Current virtual staining approaches offer the potential for time- and cost-efficient biomarker quantification in cancer diagnostics and prognostics. However, patch-wise inference for gigapixel whole-slide images (WSIs) fails to maintain spatial continuity, yielding artifacts that cause catastrophic mismatches with ground-truth images. Although pathology Vision Foundation Models (VFMs) offer rich representations, their self-attention causes varying global contexts to produce inconsistent embeddings for the same physical region. We formalize and validate this ``context contamination'' as a sheaf-theoretic problem where these embeddings form a presheaf whose sections disagree on overlaps, so no global section restricts to them. To address this, we propose SheafStain, a new approach that reinterprets VFM features as sheaf-like sections for spatially and biologically coherent virtual staining. Specifically, SheafStain integrates class and patch tokens into a Schr"odinger Bridge framework as sheaf-like sections. While the class token anchors biological consistency, patch tokens form a per-position spatial map. An encoder co-pretrained on Hematoxylin & Eosin (H&E) and Immunohistochemistry (IHC) yields cross-stain sections, so a single VFM feature space supervises both input conditioning and output stain alignment. Departing from prior work that evaluates on isolated 256 × 256 patches and either random-crops or resizes the 1024 × 1024 ground truth, we translate at 256 × 256 and evaluate on the stitched 1024 × 1024 outputs across HER2, ER, PR, and Ki-67. SheafStain demonstrates promising results against six prior methods while mitigating patch-boundary stitching artifacts. Code is available at https://github.com/deepnoid-ai/SheafStain.