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

Conference: NeurIPS2026
arXiv: 2606.11846
Code: https://github.com/deepnoid-ai/SheafStain
Area: Medical Imaging
Keywords: virtual staining, sheaf theory, Schrödinger bridge, spatial consistency, pathology foundation model

TL;DR¶

SheafStain adds neighborhood VFM spatial conditioning, overlap-consistency losses, and pathology supervision to an unpaired Schrödinger bridge, then uses adaptive extra patches and weighted stitching to generate IHC from H&E, substantially reducing seams on 1024×1024 assembled regions from two breast pathology datasets without guaranteeing diagnostic correctness.

Background & Motivation¶

Virtual staining predicts immunohistochemistry (IHC) from hematoxylin and eosin (H&E) sections to reduce the time and tissue consumed by multi-marker testing. The challenge is not merely to turn colors brown: HER2, ER, PR, and Ki-67 signals must retain tissue structure and biological differences. H&E–IHC pairs in BCI and MIST come from different consecutive tissue sections. Even after registration, they are not pixel-exact ground truth from the same physical plane. Strict paired pixel regression can therefore penalize reasonable structures and encourage blurred predictions.

A second problem arises at inference scale. Large slides are typically generated as independent small patches, allowing the same tissue to receive different staining intensities, nuclear textures, and boundaries in neighboring fields of view. A pathology vision foundation model (VFM) does not automatically solve this: global self-attention mixes non-overlapping content into tokens at overlapping positions, so the same physical location receives different representations in different windows. Across 4,873 BCI images, the authors find a mean horizontal overlap-token cosine similarity of only 0.626 for 224×224 VFM windows at stride 192; even at 85.7% overlap, it reaches only 0.918 rather than exact agreement.

The approach makes “can neighboring predictions assemble into one coherent image?” an explicit training question instead of merely smoothing seams afterward. Sheaf theory supplies the language of local sections, restriction maps, and gluing compatibility; the implemented operations are shared neighborhood conditioning, output-overlap penalties, and improved coverage and blending. Core Idea: expose local generation to consistent neighborhood tissue conditions and encourage nearby views to predict similar outputs for the same physical region, bringing both patch realism and assembled spatial consistency into virtual staining.

Method¶

Overall Architecture¶

The input is an H&E image and the output is an IHC image for a specified marker. The generative backbone follows Unpaired Neural Schrödinger Bridge (UNSB), connecting the H&E distribution to the IHC distribution rather than using the source image solely as additional conditioning on a path starting from Gaussian noise. Frozen Prov-GigaPath provides both input tissue context and output stain-feature supervision; training retains adversarial distribution matching, the bridge objective, and PatchNCE.

Training samples one 256×256 reference generation patch and two adjacent generation patches with a nonempty common overlap. Separately, 5–16 surrounding 224×224 windows are encoded by the VFM; these conditioning windows are not the same objects as the three generation patches. After spatial alignment of their conditions, the three weight-shared generator outputs receive pairwise consistency constraints. The reference output also receives pathology supervision from weakly paired IHC. Inference needs no real IHC and does not repair images through messages between neighboring outputs: overlapping patches are generated independently and then blended.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    I["H&E and overlapping windows"] --> A["Neighborhood spatial<br/>conditioning"]
    A --> G["Weight-shared SB generator"]
    G -.->|Training only: three outputs| B["Overlap compatibility<br/>constraints"]
    B -.-> C["Cross-stain pathology<br/>supervision"]
    Y["Weakly paired IHC"] -.->|Training only| C
    G -->|Inference only: independent generation| D["Adaptive coverage<br/>and blending"]
    D --> O["1024×1024 IHC"]

Key Designs¶

1. Neighborhood spatial conditioning: retain tissue context at each position

Compressing an entire VFM field of view into one vector can discard spatial distinctions between glands and stroma. SheafStain batches the 224×224 conditioning windows around a reference patch through the frozen VFM. Each window produces 14×14 spatial tokens and one CLS token. Windows are extracted in eight directions at base stride 80, with additional ±32-pixel lateral jitter for horizontal and vertical directions to improve boundary coverage; their count varies from 5–16 with tissue coverage.

Spatial tokens are scattered by image position into the reference generation patch's 16×16 grid. Tokens landing in the same cell are averaged, and the rare uncovered cells default to zero. The result is a spatial conditioning map of size \(d\times16\times16\), not a 224×224 image resized to 256×256. The paper describes this coordinate correspondence using 16-pixel cells, with a Prov-GigaPath feature dimension of 1536; 224 denotes the VFM conditioning-window size, whereas 256 denotes the generation-window size. They must not be conflated. CLS tokens are averaged across neighborhood windows to provide a shared neighborhood summary.

The spatial map is upsampled to the resolution of each generator residual block, projected, and added to intermediate features together with time and CLS projections. Spatial and CLS projections are zero-initialized, so the new conditioning initially leaves the backbone unchanged and gradually acquires influence. Adjacent generation patches reuse reference conditioning through token-grid shifts and share the CLS vector during training. Their common overlap therefore receives identical conditions, making output discrepancies more directly attributable to generator inconsistency.

This does not force the frozen VFM's original tokens into a strict sheaf: representations from different views may still differ. The method first constructs common conditions through position-wise aggregation, then trains the generator to interpret them approximately consistently. CLS summarizes tissue context, while the spatial map supplies position-specific morphology; output-side DAB alignment provides an additional staining-related biological constraint. CLS alone is not a biomarker ground truth.

2. Overlap compatibility constraints: compare generated patches rather than force alignment between consecutive sections

A restriction map here can be understood as selecting the same physical overlap region. Strict compatibility requires two local sections to agree exactly on their overlap before exact gluing is possible. The implemented training uses a soft penalty: it compares per-channel overlap means for tone agreement and corresponding pixels for structural agreement.

\[ \mathcal{L}_{\text{sheaf}}(G_i,G_j)=\big\|\mu(G_i|_{\mathcal{O}})-\mu(G_j|_{\mathcal{O}})\big\|_1+\alpha\big\|G_i|_{\mathcal{O}}-G_j|_{\mathcal{O}}\big\|_1. \]

Here, \(\mathcal{O}\) is the shared region, \(\mu\) is the per-channel spatial mean, and \(\alpha=1\). The comparison is between predictions for the same H&E region under different generation views, not between real pixels from adjacent physical sections, so this supervision does not require pixel-paired IHC. It can reduce seam tone differences and morphological breaks, but two outputs that consistently predict an incorrect marker pattern can also achieve a low loss.

Training further introduces a nonempty triple overlap. The reference and first adjacent patch already receive one sheaf term; the so-called cocycle loss adds the reference–second-adjacent and adjacent–adjacent terms, constraining all three outputs against one another. Appendix A explicitly notes that discrepancies are constructed by subtracting outputs, so their additive relation already holds through telescoping cancellation. This is not the learning of an additional nontrivial first-cohomology obstruction. It suppresses the remaining pairwise discrepancies on the triple region so that outputs approximately satisfy the compatibility condition of a 0-cocycle.

Strict formal gluing must be distinguished from finite-sample neural soft regularization. Nonzero overlap energy remains after training. The authors measure mean squared overlap discrepancies before blending, providing evidence that generator inconsistency decreases rather than proving that every window and every case restricts exactly to a unique global prediction.

3. Cross-stain pathology supervision: constrain stain intensity, tissue spectra, and marker layout together

Overlap agreement alone cannot determine what should be stained or where. The reference output therefore receives three complementary pathology signals. The DAB intensity loss uses Beer–Lambert color deconvolution to extract the brown chromogen channel and compares the mean intensity of the strongest 10% of pixels in the output and weakly paired IHC. The paper calls this p90, but it is the mean above the 90th percentile, not the percentile value itself. It does not require brown signals to occupy exactly the same positions, reducing sensitivity to displacement while being insufficient by itself to guarantee positive-cell localization.

The Fourier edge loss compares high-frequency log-magnitude spectra of grayscale output and target images. Its radial high-pass mask retains frequencies beyond 25% of the maximum radius, emphasizing nuclear contours, gland boundaries, and stromal texture. Magnitude spectra are translation-invariant, making this loss more tolerant of weak registration than pixelwise edge regression. However, structural differences between sections are not equivalent to pure translation; arbitrary deformation and biological variation are not thereby removed. This training-time high-pass loss also differs from the mid-band FFT energy used to place extra patches at inference.

Finally, cross-stain VFM alignment extracts DAB from generated and target images, re-renders each as brown RGB on a white background, and compares frozen Prov-GigaPath per-token spatial features. Removing other stain components focuses supervision on chromogen layout rather than scalar intensity alone. The target branch is detached, while gradients from the generated branch pass through the frozen encoder to the generator.

Using the same encoder, co-pretrained on H&E and IHC, for input conditioning and output supervision is an important premise. Output-feature alignment is still a comparison against a weakly paired neighboring section, not true cell-by-cell correspondence; it is a more targeted learning signal. Cross-stain alignment is enabled only in the final 50 epochs rather than introducing every supervision term from the outset.

4. Adaptive coverage and blending: refine the assembled region without enlarging generation patches

Inference first covers a 1024×1024 image with 256×256 reference windows at stride 192, snapping the trailing row and column to the image boundary. This yields 25 reference patches, with 64-pixel overlap between stride-adjacent windows. It then adds 9 reference patches using mid-band FFT energy at \([1/16,1/4]\) cycles per pixel and a minimum candidate spacing of 112 pixels. Both extremes of the energy ranking receive priority: extra coverage targets low-texture background as well as high-texture tissue because both may produce unstable stitching.

Each reference patch obtains VFM conditioning and is translated independently, giving 34 generation patches in total. Appendix G further adapts conditioning-neighborhood size to local FFT complexity, interpolating between 5–16 conditioning windows using image-wide 25th and 75th energy percentiles. The authors interpret this as refinement of the local cover. Operationally, it is a frequency-based sampling heuristic, not a proven optimal covering algorithm.

For blending, regular reference patches use a linear ramp over the 32-pixel border and a central 192×192 plateau. Extra patches use a raised-cosine window that peaks centrally and tapers toward the boundary. Outputs at each location are combined by a normalized weighted average to avoid introducing hard edges around the extra patches.

\[ \hat{Y}(x,y)=\frac{\sum_{i:(x,y)\in U_i}w_i(x,y)s_i(x,y)}{\sum_{i:(x,y)\in U_i}w_i(x,y)}. \]

Blending absorbs residual discrepancies, but it does not imply exact agreement before fusion or biologically correct stain intensity after averaging. The method relies on both training consistency and inference coverage; its final improvements cannot all be attributed to one theoretical loss.

The source contains differing descriptions of the VFM invocation scope. Section 3.3 states that inference applies the VFM once over the image and then crops a global token map. Section 3.4 and Appendices F/G describe neighborhood encoding for each reference patch, while Appendix K explicitly counts 200 VFM sample forwards per image in batches of 32, excluding the 9 extra reference patches. “One batched invocation” must not be interpreted as “encoding the entire image only once.” This note explains the local-neighborhood mechanism; reproduction still requires checking the implementation.

A Worked Example¶

Consider a breast gland crossing the boundary between two generation windows. Independent predictions may render the same cell neighborhood dark brown in the left patch and light brown in the right. Each patch may look plausible as IHC, yet stitching creates an intensity discontinuity.

Training gathers 224×224 conditioning windows around a reference patch and maps their spatial tokens onto a 16×16 conditioning grid for the 256×256 generation patch. It then selects two adjacent patches sharing a common region with the reference, shifts the spatial conditions, and shares CLS so that all three receive identical overlap conditions. After generation, sheaf and cocycle terms penalize tone and pixel differences on shared tissue, while DAB, Fourier, and cross-stain alignment constrain stain amount, tissue texture, and chromogen layout, respectively.

At inference over the full 1024×1024 region, 25 regular generation patches plus 9 frequency-driven extra patches provide coverage, followed by blending with different window functions. This example illustrates how boundary tissue receives consistency constraints. It does not assume a pixel-identical gland exists in the weakly paired IHC reference and is not a validated clinical diagnostic procedure.

Loss & Training¶

The full generator objective retains UNSB's GAN, SB, and NCE terms and adds sheaf, cocycle, Fourier, DAB, and cross-stain alignment losses. The discriminator and auxiliary energy network are inherited from the baseline and trained jointly with the generator; the implementation is therefore not a pure diffusion model without adversarial learning.

Weights are \(\lambda_{\text{SB}}=\lambda_{\text{NCE}}=\lambda_{\text{sheaf}}=1\), \(\lambda_{\text{cocycle}}=\lambda_{\text{DAB}}=0.1\), and \(\lambda_{\text{fourier}}=0.5\). Cross-stain alignment has weight 0 during epochs 1–350 and 1 during epochs 351–400. The generator is a time-conditioned 9-block ResNet, and the VFM remains frozen. Adam uses learning rate \(2\times10^{-4}\) and \(\beta=(0.5,0.999)\), with effective batch size 192 on 8 H200 GPUs over 400 epochs: 200 at a fixed learning rate and 200 with linear decay.

The appendix's 27-setting weight screening uses only 100 epochs and no extra inference reference patches. Its KID values are raw rather than multiplied by 1000 as in the main tables, so its results are not directly comparable to full-model results. The authors use one final weight configuration for BCI and all four MIST markers and report final-epoch weights rather than selecting checkpoints by validation metrics.

Key Experimental Results¶

Main Results¶

BCI HER2 contains 3,896 training pairs and 977 test pairs. MIST HER2, ER, PR, and Ki-67 contain 4,642, 4,153, 4,139, and 4,361 training pairs, respectively, with 1,000 validation pairs each. MIST has no independent test set, so the paper reports results on its validation sets. All methods generate at 256×256 and are evaluated on 1024×1024 assembled regions.

The following selection from main-text Table 3 retains means only. Each block uses the prior method with the lowest FID as its comparator; this does not mean that method is the best baseline on every other metric. KID is reported as \(\mathrm{KID}\times10^3\). Lower FID, KID, DISTS, and TS are better; higher DAB-r is better. TS measures boundary discrepancies at patch seams, and DAB-r is the Pearson spatial correlation between generated and reference DAB maps.

Dataset / Marker	Method	FID	KID×1000	DISTS	TS	DAB-r
BCI / HER2	UNIStainNet	67.6322	29.3000	0.2279	0.0448	0.1267
BCI / HER2	SheafStain	36.3626	4.2070	0.2132	0.0146	0.1209
MIST / HER2	PSPStain	50.1078	10.6750	0.2466	0.0763	0.0841
MIST / HER2	SheafStain	34.5080	2.2350	0.2056	0.0292	0.1187
MIST / ER	PSPStain	45.9200	16.0910	0.2534	0.0752	0.0906
MIST / ER	SheafStain	29.0824	2.5000	0.1934	0.0292	0.1248
MIST / PR	PSPStain	48.6565	19.2100	0.2744	0.0797	0.0905
MIST / PR	SheafStain	29.7240	2.1910	0.1962	0.0272	0.1184
MIST / Ki-67	UNSB	38.1956	10.1780	0.2514	0.0857	0.0326
MIST / Ki-67	SheafStain	25.3978	3.9220	0.2124	0.0261	0.0493

SheafStain leads FID, KID, DISTS, and TS across all five tasks, but not every metric. On BCI HER2, UNIStainNet has higher DAB-r (0.1267 versus 0.1209) and lower LPIPS (0.4577 versus 0.4689). The closest TS competitor is actually D-VST, with 0.0164 on BCI HER2, for example; differences against other baselines in this selected table must not be presented as margins over the strongest stitching baseline.

Inference protocols are not fully symmetric. D-VST and SheafStain use their native procedures, whereas the other five methods place 16 disjoint outputs into a direct 4×4 assembly. The comparison captures complete-system performance and mixes training improvements with coverage and blending benefits. It does not isolate the training advantage under identical overlapping fusion for all methods.

Ablation Study¶

The following is the progressive BCI ablation from main-text Table 2. Each row adds one component to the preceding row rather than independently removing it. KID remains multiplied by 1000.

Config	FID	KID×1000	LPIPS	DISTS	TS	DAB-r
UNSB baseline	227.7889	236.2310	0.6635	0.3347	0.1464	0.0589
+ Neighborhood spatial conditioning	55.2489	16.5750	0.5538	0.2960	0.0143	0.0823
+ pixel sheaf	41.5356	6.8270	0.5225	0.2721	0.0131	0.1059
+ cocycle	38.0140	4.6070	0.5176	0.2679	0.0132	0.1091
+ Fourier edge	36.5335	5.1630	0.4774	0.2285	0.0145	0.1190
+ DAB intensity	37.0935	4.1570	0.4718	0.2187	0.0144	0.1183
+ Cross-stain alignment (full model)	36.3626	4.2070	0.4689	0.2132	0.0146	0.1209

The largest single-step improvement comes from conditioning: FID drops from 227.7889 to 55.2489 and TS from 0.1464 to 0.0143. Subsequent regularizers improve distributional and perceptual quality, but not every metric improves monotonically. TS rises slightly after cocycle, both FID and DAB-r regress slightly after DAB intensity, and the full model's TS exceeds that of the pixel-sheaf-only addition. The source's “small but consistent gains” should be qualified by metric, not generalized to universal improvement at each step.

Key Findings¶

Consistency before blending is also supported. On MIST HER2 at the final checkpoint, the median pairwise mean squared overlap discrepancy is 0.0019 versus UNSB's 0.0052; the 95th percentiles are 0.0126 and 0.0584. This supports learned approximate consistency rather than seam suppression exclusively through post-processing, while residual disagreement remains nonzero.
Extra reference patches do not improve every image. TS decreases on 678 of 1,000 MIST HER2 validation images, with paired Wilcoxon \(p<10^{-33}\). Aggregate effectiveness and per-case reliability are different conclusions.
Weak pairing compresses spatial correlation. The appendix analyzes shifted copies of real DAB maps and reports correlation of about 0.40 at 7.4 micrometers displacement. This explains metric sensitivity but does not establish that a generated result approaches a true biological ceiling.

The downstream experiment trains an ImageNet-initialized ResNet-50 using 3,896 generated BCI training IHC images from each method and evaluates on 977 real IHC test images. It is not direct clinical scoring of synthetic images. Labels are Low (0, 1+) and High (2+, 3+).

Source of classifier training images	Accuracy	F1	AUROC
UNSB-generated IHC	0.524	0.640	0.460
UNIStainNet-generated IHC	0.724	0.835	0.737
SheafStain-generated IHC	0.766	0.844	0.794
Real IHC (reference ceiling)	0.974	0.982	0.997

Relative to UNIStainNet, Accuracy increases by 4.2 percentage points and AUROC by 0.057, but performance remains substantially below real IHC. The test set contains 273 Low and 704 High cases, so accuracy should be interpreted alongside F1 and AUROC. Assigning 2+ to High is only a research-task definition; it does not remove the clinical need for further confirmation of that grade.

Highlights & Insights¶

Connecting context contamination to seams is more specific than adding another global pathology vector. Whether conditioning agrees at overlap positions and whether outputs agree become separate, testable questions.
The most useful contribution of sheaf theory here is the organization of supervision, not a guarantee of correct generation. Local sections and restriction maps clarify where shared tissue should be compared and why weakly paired neighboring sections are not same-position ground truth.
Measuring output-overlap discrepancies before blending is a transferable audit strategy for other patch-based generation tasks. Inspecting only seamless final images can mistake kernel smoothing for learned spatial consistency.

Limitations & Future Work¶

Validation covers only breast pathology regions in BCI and MIST, without an external clinical cohort, cross-organ generalization, or prospective diagnostic evaluation. A 1024×1024 region exposes internal seams but is not clinical validation on complete gigapixel WSIs.
Strong agreement cannot prevent shared hallucinations: neighboring patches may consistently omit positive cells or generate incorrect brown signals. Cell-level localization, grading concordance, case-level errors, and uncertainty should complement low TS.
Compute cost is substantial. The appendix reports roughly 18 seconds per image on a Quadro RTX 5000, with the VFM accounting for about 90% of cost; the 44.7-TFLOP accounting for 200 VFM sample forwards excludes the 9 extra patches. The 2–3-second estimate for faster accelerators is a projection, not a measured result in this paper.
The source's “one VFM pass over the image” conflicts with its per-patch neighborhood cost accounting. Token alignment is also described using both floor-rounded offsets and nearest-integer mapping. Reproduction should check cache reuse, noninteger-offset handling, and actual invocation counts rather than silently reconciling these differences.
Low/High classification is not standard four-grade HER2 diagnosis. The main protocol uses only the training split, but the appendix ROC caption refers to training and validation outputs. This note retains the main-protocol results; the actual sample scope still needs verification.

vs UNSB: SheafStain retains an unpaired bridge, adversarial matching, and the same backbone family while adding neighborhood conditions, overlap losses, pathology supervision, and inference coverage. The UNSB comparison reflects the complete increment, not a classification gain attributable to cocycle alone.
vs D-VST: D-VST also uses frequency-aware sampling to reduce assembly artifacts. SheafStain links coverage refinement to pairwise and triple-overlap consistency during training. Its TS advantage must be interpreted with each method's inference budget.
vs UNIStainNet / HistDiT: Pathology VFM guidance is not unique to this paper. The distinction is position-specific spatial information, a shared neighborhood summary, and output supervision through DAB-only feature alignment. UNIStainNet still leads some BCI metrics, showing that more spatial constraints do not automatically win every biological metric.
vs Neural Sheaf Diffusion: Learnable restriction maps in graph neural networks are not the same implementation as image-position indexing and pixel-discrepancy regularization here. A useful extension would equalize coverage budgets and compare globally cached conditions with local neighborhood conditions to separate contextual structure from extra compute.

Rating¶

Novelty: 4/5. Translates sheaf compatibility into spatial conditioning and overlap supervision for a specific virtual-staining problem; strict theoretical claims must remain separate from the soft implementation.
Experimental Thoroughness: 3/5. Two datasets, four markers, progressive ablations, and pre-blending analysis provide substantial evidence, but full WSI evaluation, external cases, and equal-fusion-budget controls remain absent.
Writing Quality: 3/5. The v2 appendices explain the operational meaning of cocycle in detail; VFM invocation and token alignment retain reproduction ambiguities.
Value: 4/5. Offers reusable consistency designs and evaluation methods for patch-based medical image generation, while remaining a research system rather than a diagnostic replacement for real IHC.