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

会议: NeurIPS2026
arXiv: 2606.11846
代码: https://github.com/deepnoid-ai/SheafStain
领域: 医学图像
关键词: 虚拟染色、层论、薛定谔桥、空间一致性、病理基础模型

一句话总结¶

SheafStain 在非配对薛定谔桥中加入邻域 VFM 空间条件、重叠区域一致性损失和病理监督，再通过自适应补块与加权拼接生成 H&E 对应的 IHC，在两个乳腺病理数据集的 1024×1024 拼接区域上显著减少接缝，但不提供诊断正确性的保证。

研究背景与动机¶

虚拟染色希望从苏木精-伊红（H&E）切片预测免疫组织化学（IHC）染色，降低多标志物检测的时间与组织消耗。难点不只是把颜色变成棕色：HER2、ER、PR、Ki-67 的信号分布必须保留组织结构与生物学差异。BCI 和 MIST 的 H&E–IHC 对来自连续的不同组织切片，即使经过配准，也不是同一物理平面的逐像素真值。因此，严格配对的像素回归既可能惩罚合理结构，也可能鼓励模糊输出。

另一个问题发生在推理尺度上。大切片通常分成小块独立生成，同一组织在相邻视野中可能获得不同染色强度、核纹理与边界。病理视觉基础模型（VFM）并不自动解决这个问题：全局自注意力会把当前块内的非重叠内容混入重叠位置的 token，导致同一个物理位置因所在窗口不同而得到不同表征。作者在 4,873 张 BCI 图像上测得，224×224 VFM 窗口以步长 192 移动时，水平重叠 token 的平均余弦相似度仅为 0.626；即使重叠率达到 85.7%，也只有 0.918，而不是完全一致。

这里的切入点是把“相邻块是否能拼成同一幅图”变成训练时的明确问题，而不只在生成后抹平接缝。层论提供局部截面、限制映射和粘合相容性的语言；真正落地的操作则是共享邻域条件、惩罚输出重叠差异，以及更合理的覆盖与融合。核心 idea：让局部生成既看到一致的邻域组织条件，又为同一物理区域给出接近的输出，从而把单块真实性与拼接后的空间一致性共同纳入虚拟染色。

方法详解¶

整体框架¶

输入是 H&E 图像，输出是指定标志物的 IHC 图像。生成骨干沿用 Unpaired Neural Schrödinger Bridge（UNSB），从 H&E 分布通向 IHC 分布，而不是把源图像仅作为从高斯噪声出发的附加条件。冻结的 Prov-GigaPath 同时提供输入组织上下文与输出染色特征监督；训练仍包含对抗分布匹配、桥目标和 PatchNCE。

训练时抽取一个 256×256 参考生成块和两个邻接生成块，使三者存在共同重叠区。参考块周围另取 5–16 个 224×224 窗口供 VFM 编码，这组条件窗口与三个生成块不是同一组对象。空间条件对齐后，三个共享权重的生成器输出接受两两一致性约束；参考输出还接受弱配对 IHC 提供的病理监督。推理时不需要真实 IHC，也不通过相邻输出之间的消息传递修复图像，而是独立生成重叠块再融合。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    I["H&E 与重叠窗口"] --> A["邻域空间条件"]
    A --> G["共享权重 SB 生成器"]
    G -.->|仅训练：三个输出| B["重叠相容性约束"]
    B -.-> C["跨染色病理监督"]
    Y["弱配对 IHC"] -.->|仅训练| C
    G -->|仅推理：独立生成| D["自适应覆盖与融合"]
    D --> O["1024×1024 IHC"]

关键设计¶

1. 邻域空间条件：把组织上下文保留到每个位置

把 VFM 的整个视野压成一个向量，容易丢掉“哪一处是腺体、哪一处是间质”的空间信息。SheafStain 将参考块周围的 224×224 条件窗口批量送入冻结的 VFM，每个窗口产生 14×14 个空间 token 和一个 CLS token。条件窗口来自八个方向，基础步长为 80，水平与垂直方向再加入 ±32 像素横向扰动，以增加边界覆盖；实际数量随组织覆盖在 5–16 之间变化。

空间 token 按原图位置散射到参考生成块的 16×16 网格；落到同一格的多个 token 取平均，少见的无覆盖格置零。这样形成维度为 \(d\times16\times16\) 的空间条件图，而不是把 224×224 图像缩放成 256×256。本文以 16 像素单元描述这套坐标对应，Prov-GigaPath 的特征维度为 1536；224 是 VFM 条件窗口大小，256 是生成窗口大小，两者不能混用。CLS token 则在邻域窗口间平均，提供共享的邻域摘要。

空间条件图上采样到生成器残差块的分辨率，经学习投影后与时间嵌入、CLS 投影一起加到中间特征。空间与 CLS 投影零初始化，使新增条件在训练开始时不改变原有骨干，随后逐渐学习其作用。训练中的邻接生成块通过 token 网格平移复用参考条件，三块共享 CLS，因此共同重叠位置具有相同条件，输出差异更能反映生成器自身的不一致。

这不是把冻结 VFM 的原始 token 强制修成严格的层：多视野表征仍可能不同，方法先通过位置汇聚构造公共条件，再让生成器学会对这些条件作近似一致的解释。CLS 概括整体组织，空间图提供逐位置形态上下文；输出侧的 DAB 对齐才进一步提供染色生物学约束，不能把 CLS 单独当作标志物真值。

2. 重叠相容性约束：比较生成块之间，而不是强行对齐邻切片

层论中的限制映射在这里可以直观理解为“只取同一物理重叠区域”。严格相容性要求两个局部截面在重叠区域完全相同，才能谈精确粘合。实际训练使用软惩罚：先比较重叠区域的逐通道均值以约束色调，再比较对应像素以约束结构。

\[ \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. \]

其中 \(\mathcal{O}\) 是共享区域，\(\mu\) 是逐通道空间均值，\(\alpha=1\)。比较对象是同一 H&E 区域在不同生成视野下的预测，而不是两个相邻物理切片的真实像素，所以这部分监督不需要逐像素配对 IHC。它能减少接缝色差与形态断裂，但两个输出一致地生成错误标志物，也可以获得很低的损失。

训练再引入非空三重重叠区。参考块与第一邻接块已有一个 sheaf 项，所谓 cocycle 损失补上“参考块—第二邻接块”和“两邻接块”两项，使三个输出都互相约束。附录 A 明确指出，差异由输出相减构成，三者的差异相加关系本来就因望远镜消去而成立。因此，这里并非学习一个额外的非平凡一阶上同调障碍，而是把三重区域的剩余两两差异压小，使输出近似满足 0-cocycle 的相容性条件。

形式上的严格粘合要求与神经网络的有限样本软正则需要区分。训练后仍有非零重叠能量；作者测量的是生成器在融合前的均方重叠差异，证明不一致性下降，而不是证明所有窗口、所有病例都能精确限制到唯一的全局预测。

3. 跨染色病理监督：同时约束染色强度、组织频谱和标志物布局

纯重叠一致性不能决定染什么、染在哪里，因此参考输出还接受三个互补的病理信号。DAB 强度损失先通过 Beer–Lambert 色彩解卷积提取棕色显色剂通道，再比较输出和弱配对 IHC 中强度最高 10% 像素的平均值。原文称其为 p90，但它不是单个第 90 百分位数，而是超过该分位数后的均值。它不要求棕色信号出现在完全相同的位置，能减轻位移的影响，却不能独自保证阳性细胞定位。

Fourier 边缘损失比较灰度输出与目标的高频对数幅度谱，径向高通掩膜保留超过最大半径 25% 的频率，强调核轮廓、腺体边缘与间质纹理。幅度谱对平移具有不变性，因此比逐像素边缘回归更能容忍弱配准；但不同切片的结构变化并不等价于纯平移，不能据此声称任意形变与生物差异都被消除。这里的训练高通损失也不同于推理时用于补块的中频 FFT 能量。

最后，跨染色 VFM 对齐分别提取生成图与目标的 DAB 通道，重新渲染成白底棕色 RGB 图，再比较冻结 Prov-GigaPath 的逐 token 空间特征。去除其余染色成分，可以让监督更集中于显色剂的空间布局，而不只是标量强度；目标分支停止梯度，生成分支的梯度经冻结编码器传回生成器。

输入条件与输出监督使用同一个在 H&E 和 IHC 上共同预训练的编码器，是这项设计的关键前提。但输出特征对齐仍是对弱配对邻切片的特征比较，而不是获得了真实逐细胞对应；它只是更有针对性的学习信号。跨染色对齐延迟到最后 50 个 epoch 启用，避免从一开始就加入全部监督。

4. 自适应覆盖与融合：在不增加生成块尺寸的情况下修复拼接区域

推理先用 256×256 参考窗口以步长 192 覆盖 1024×1024 图像，末行末列贴齐图像边界，形成 25 个参考块；相邻规则窗口有 64 像素重叠。然后根据中频 FFT 能量增加 9 个参考块，频带为每像素 \([1/16,1/4]\) 周期，并要求候选位置间至少相距 112 像素。优先选择能量排序的两端，不仅补高纹理组织，也补低纹理背景，因为二者都可能产生不稳定拼接。

每个参考块独立获得 VFM 条件并生成输出，合计 34 个块。附录 G 进一步让条件邻域大小随局部 FFT 复杂度变化，利用图像级能量的第 25 与第 75 百分位数在 5–16 个条件窗口间插值。作者将其解释为细化局部覆盖；实际操作是一个频率启发式采样规则，不是已经证明最优的覆盖算法。

融合时，规则参考块采用边缘 32 像素线性爬坡、中央 192×192 平台的权重；额外块使用中心高、边界低的升余弦窗。对同一位置的输出进行归一化加权，避免额外块本身引入硬边界。

\[ \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)}. \]

这种融合会吸收残余差异，但不意味着融合前各块已经严格一致，也不意味着平均后的染色强度必然真实。方法同时依赖训练一致性与推理覆盖，不能把最终提升全部归因于一个理论损失。

原文对 VFM 调用范围存在需要保留的差异：主文 3.3 称推理“对整图应用一次 VFM”，随后从全局 token 图裁切；3.4、附录 F/G 则描述逐参考块邻域编码，附录 K 明确统计每图 200 次 VFM 样本前向、按 32 个批处理，且未计入 9 个额外参考块。不能将“一次批量调用”解释为“一张整图只编码一次”，本笔记按局部邻域机制说明，实际复现仍需核对代码。

一个完整示例¶

考虑一个跨越两个生成窗口边界的乳腺腺体。若两个窗口独立预测，同一细胞附近可能在左块呈深棕色、右块呈浅棕色，虽然两块各自看起来像 IHC，拼接却出现强度跳变。

训练先围绕参考块收集 224×224 条件窗口，把空间 token 映射到 256×256 生成块的 16×16 条件网格。再选择两个与参考块存在共同区域的邻接块，平移空间条件并共享 CLS，使三者看到相同的重叠条件。生成后，sheaf 与 cocycle 项惩罚共同组织上的色调与像素差异，而 DAB、Fourier、跨染色对齐分别约束染色量、组织纹理和显色布局。

推理面对整个 1024×1024 区域时，25 个规则生成块加 9 个频率驱动补块覆盖它，再按不同窗函数融合。这个例子说明的是边界组织如何受到一致性约束，并不假设弱配对 IHC 中存在逐像素相同的腺体，也不是已经验证的临床诊断过程。

损失函数 / 训练策略¶

总生成器目标保留 UNSB 的 GAN、SB、NCE 项，并加入 sheaf、cocycle、Fourier、DAB 与跨染色对齐。判别器和辅助能量网络沿用基线，并随生成器联合训练；因此实现并非一个不含对抗学习的纯扩散模型。

权重设为 \(\lambda_{\text{SB}}=\lambda_{\text{NCE}}=\lambda_{\text{sheaf}}=1\)，\(\lambda_{\text{cocycle}}=\lambda_{\text{DAB}}=0.1\)，\(\lambda_{\text{fourier}}=0.5\)。跨染色对齐在 epoch 1–350 权重为 0，在 epoch 351–400 为 1。生成器为带时间条件的 9 块 ResNet，VFM 保持冻结；Adam 学习率 \(2\times10^{-4}\)，\(\beta=(0.5,0.999)\)，有效批量 192，使用 8 张 H200，训练 400 个 epoch，前 200 固定学习率、后 200 线性衰减。

附录的 27 组权重筛选只训练 100 个 epoch，而且不使用额外推理参考块；其 KID 为原始值，不是主表的放大 1000 倍形式。筛选结果不能与完整模型主结果直接混比。作者对 BCI 和四个 MIST 标志物使用同一最终权重，报告使用最终 epoch 权重而不按验证指标选择 checkpoint。

实验关键数据¶

主实验¶

BCI 的 HER2 有 3,896 对训练图与 977 对测试图。MIST 的 HER2、ER、PR、Ki-67 分别有 4,642、4,153、4,139、4,361 对训练图，各有 1,000 对验证图；MIST 没有独立测试集，论文将验证集用于报告结果。所有方法在 256×256 上生成，再在 1024×1024 聚合区域评估。

下表摘自主文表 3，仅保留均值。每个区域选择此前 FID 最低的方法作对照，并不表示该方法也是其他指标的最佳基线。KID 的表中单位为 \(\mathrm{KID}\times10^3\)；FID、KID、DISTS、TS 越低越好，DAB-r 越高越好。TS 衡量块接缝边界差异，DAB-r 为生成与参考 DAB 图的 Pearson 空间相关性。

数据集 / 标志物	方法	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 在五组任务的 FID、KID、DISTS、TS 上均领先，但不是所有指标都第一。BCI HER2 的 UNIStainNet 有更高 DAB-r（0.1267 对 0.1209）和更低 LPIPS（0.4577 对 0.4689）。TS 最接近的竞争者其实是 D-VST，例如 BCI HER2 为 0.0164，因此不能拿表中其他基线的 TS 差值代替最强拼接基线的差值。

推理协议并不完全对称：D-VST 与 SheafStain 使用各自原生推理，其余五个方法将 16 个不重叠输出直接按 4×4 拼接。比较反映完整系统效果，也混合了训练改进和覆盖融合策略的收益，尚不能隔离“给所有模型同样重叠融合后”的纯训练优势。

消融实验¶

下表为主文表 2 的 BCI 渐进消融，每行在上一行基础上新增一项，并非独立删除实验。KID 仍乘 1000。

配置	FID	KID×1000	LPIPS	DISTS	TS	DAB-r
UNSB 基线	227.7889	236.2310	0.6635	0.3347	0.1464	0.0589
+ 邻域空间条件	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 边缘	36.5335	5.1630	0.4774	0.2285	0.0145	0.1190
+ DAB 强度	37.0935	4.1570	0.4718	0.2187	0.0144	0.1183
+ 跨染色对齐（完整）	36.3626	4.2070	0.4689	0.2132	0.0146	0.1209

最大单步改善来自条件信息：FID 从 227.7889 降到 55.2489，TS 从 0.1464 降到 0.0143。后续正则改善分布与感知质量，但并非所有指标单调变好；例如 cocycle 后 TS 略升，DAB 强度后 FID 与 DAB-r 均略退，完整模型 TS 也高于只加 pixel sheaf 时。原文“小而一致的改善”应按具体指标限定，不能概括成每步全面提升。

关键发现¶

融合前的一致性也有证据。MIST HER2 最后 checkpoint 的重叠对均方差中位数为 0.0019，UNSB 为 0.0052；第 95 百分位数分别为 0.0126 与 0.0584。这支持生成器学到近似一致性，而非只靠后处理遮缝，但残余差异仍非零。
额外参考块并不保证每张图改善：在 1,000 张 MIST HER2 验证图中，678 张的 TS 降低，配对 Wilcoxon 检验 \(p<10^{-33}\)。总体有效与逐病例可靠是两种不同结论。
弱配对会压低空间相关性。附录用真实 DAB 图的平移副本作分析，7.4 微米位移时相关性约为 0.40；这个位移实验解释指标敏感性，不能证明任一生成结果接近真实生物学上限。

下游实验用各方法生成的 3,896 张 BCI 训练 IHC 训练 ImageNet 初始化的 ResNet-50，再在 977 张真实 IHC 测试图上评估；不是在合成图上直接做临床评分。标签为 Low（0、1+）与 High（2+、3+）。

分类器训练图像来源	Accuracy	F1	AUROC
UNSB 生成 IHC	0.524	0.640	0.460
UNIStainNet 生成 IHC	0.724	0.835	0.737
SheafStain 生成 IHC	0.766	0.844	0.794
真实 IHC（参考上限）	0.974	0.982	0.997

相对 UNIStainNet，Accuracy 增加 4.2 个百分点、AUROC 增加 0.057，但仍明显落后于真实 IHC。测试集有 273 个 Low 与 704 个 High，准确率需与 F1、AUROC 一起看。将 2+ 并入 High 只是研究任务定义，不能消除其临床上需要进一步确认的属性。

亮点与洞察¶

将上下文污染与接缝问题关联起来，比“再加一个全局病理向量”更具体。条件是否在重叠位置一致，与输出是否一致被拆成两个可检验的问题。
层论在这里最有价值的是监督组织方式，而不是保证生成正确。局部截面和限制映射把共享组织上的比较范围说清楚，也提醒读者不能把弱配对邻切片当作同位置真值。
在融合前测量输出重叠差异，是可以迁移到其他分块生成任务的审计思路。只看最终无接缝图像，容易把核函数平滑效果误认为网络学到了空间一致性。

局限与展望¶

验证仅覆盖 BCI、MIST 的乳腺病理区域，没有外部临床队列、跨器官泛化或前瞻性诊断评估。1024×1024 能暴露内部接缝，但不是完整 gigapixel WSI 的临床验证。
高一致性不能防止共同幻觉：邻接块可能一致地遗漏阳性细胞或生成错误棕色信号。应增加细胞级定位、分级一致性、病例级误差与不确定性分析，而不能只追求低 TS。
计算代价明显。附录在 Quadro RTX 5000 上报告约 18 秒/图，VFM 占约 90% 成本；200 次 VFM 样本前向的 44.7 TFLOPs 统计还排除了 9 个额外块。更快加速器的 2–3 秒估计是投影，不是本文实测结果。
原文“整图一次 VFM”与逐块邻域成本统计不一致；token 对齐也同时出现 floor-rounded offset 与 nearest-integer 的描述。复现应确认缓存复用、非整数偏移处理和实际调用数，不把这些实现差异自行修成统一版本。
Low/High 分类不是标准四级 HER2 诊断，而且主协议称仅用训练集，附录 ROC 图注却写训练与验证输出。这里保留主协议结果，样本使用范围仍应核对。

评分¶

新颖性：4/5。把层论相容性语言转成虚拟染色的空间条件与重叠监督，有明确应用切入点；严格理论结论应与软实现分开。
实验充分度：3/5。两个数据集、四种标志物、渐进消融与融合前分析较完整，但仍缺全 WSI、外部病例与统一融合预算对照。
写作质量：3/5。v2 附录详细解释了 cocycle 的实际含义；VFM 调用和 token 对齐表述仍有复现歧义。
价值：4/5。为分块医学图像生成提供了可复用的一致性设计与评估方法，当前定位仍是研究系统而非可替代真实 IHC 的诊断工具。