跳转至

Structure-Guided Masked Autoencoders for Ultra-High Resolution Scientific Image Understanding

会议: NeurIPS2026
arXiv: 2609.30682
领域: 自监督表示学习 / 科学图像分割
关键词: 掩码自编码器、四叉树、自适应 token 化、结构条件掩码、超高分辨率图像

一句话总结

SGMA 将内容自适应四叉树 token 化与跨尺度结构条件掩码共同用于 MAE 预训练,让标准 ViT 在固定序列预算下保留科学图像微结构;SpringXCT 的 SGMA-SAM 达到 95.68% Dice,PAIP 的 SGMA-SAM 2 达到 83.21%,而最高 24.8 倍推理加速来自另一个较短序列配置。

研究背景与动机

科学成像中的重要对象往往不是占据画面中心的大物体,而是细胞边缘、材料孔隙、聚合物连接等细小结构。 透射电镜、病理全切片和 X 射线 CT 都能产生超高分辨率图像,但标签依赖领域专家,数据规模也使逐像素标注难以扩展。 因此,先利用未标注图像预训练,再用少量标签微调,是自然的路线。 MAE 通过遮挡图像块并重建其像素学习表示,却默认采用规则网格和随机掩码,没有明确区分均匀背景与细节密集区域。

问题不仅在于预训练任务是否合适,还在于训练好的模型是否能处理原始分辨率。 对 32,768 × 32,768 图像使用 32 × 32 图像块,就会产生约一百万个 token。 标准 ViT 的全局自注意力随序列长度平方增长;统一放大图像块虽然能缓解显存压力,却同时抹去分割最需要的微结构。 MAE 编码器只处理可见块,能够降低预训练开销,但下游分割通常仍要处理整条输入序列,这一问题不会自动消失。 Swin、线性或稀疏注意力以及多层级模型各有解决办法,却需要改变注意力或骨干结构,不一定能直接复用标准 MAE/ViT。

本文选择从输入表示和预训练任务两端入手:让结构丰富的区域获得更细的图像块,让简单区域用更大的块表示,并让重建任务感知这些跨尺度结构。 固定 token 预算约束的是送入 Transformer 的序列长度,而不是把原图统一缩小。 核心 idea:在同一棵内容自适应四叉树上联合生成有限长度的结构 token 和多尺度掩码条件,使计算预算主要用于微结构,同时通过结构相关重建学习适合密集预测的表示。

方法详解

整体框架

SGMA 输入一张高分辨率科学图像,输出可用于下游分割的预训练 ViT 编码器。 结构引导特征 tokenizer(Structure-Guided Feature Tokenizer,SGFT)先依据边缘显著性构造四叉树,直到达到预定 token 预算。 树的每次细分同时触发阻尼累积(Damped Accumulation,DA),把信号相关的局部高斯响应叠加到原分辨率结构画布中,为掩码采样提供多尺度条件。 随后进行树对齐重建与分割:预训练用可见 token 重建被遮挡块,微调时沿用树表示,并让标签与预测在相同的块划分上对齐。

图中实线表示预训练处理链;虚线表示干净重建目标、分割监督或下游使用,不表示推理时仍执行随机遮挡。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["超高分辨率图像"] --> B["结构引导特征 tokenizer<br/>SGFT"]
    B --> C["阻尼累积<br/>DA"]
    C -->|预训练结构条件| D["树对齐重建与分割"]
    B -.->|微调与推理 token| D
    A -.->|干净块重建目标| D
    Y["分割标签"] -.->|同树划分监督| D
    D --> O["预训练编码器/分割结果"]

这里的压缩发生在输入 token 化层面。 SGMA 没有把标准 Transformer 改成线性注意力:若送入编码器的 token 数为 \(N\),全局自注意力仍为 \(O(N^2)\)。 好处是 \(N\) 可以按预算控制,而不必与原图像素数按规则网格同步膨胀;预处理、画布构建和最终还原图像仍有各自开销。

关键设计

1. 结构引导特征 tokenizer:把有限 token 分配给细节密集区域

规则网格对背景和细微孔隙分配相同空间粒度,或者只能整体增大块尺寸。 SGFT 从覆盖整图的根节点开始,把当前叶节点视为候选 token,使用 Canny 边缘图构造显著性分数,例如块内边缘响应之和。 显著性高的块更可能继续分成四个子块,显著性低的区域则保持粗粒度。 最终得到的是覆盖不同空间范围的混合尺度块序列,而不是只选择若干感兴趣区域并抛弃剩余图像。

直接每次选择最高分块需要串行搜索,随着活跃节点增加会成为 tokenizer 的瓶颈。 作者改用带温度的 softmax 概率选择,使高显著性区域更常被细分,同时支持批量化实现。 其选择概率为:

\[ \pi(n_j\mid\mathcal A_k,\tau)= \frac{\exp(V(n_j)/\tau)}{\sum_{n_i\in\mathcal A_k}\exp(V(n_i)/\tau)}. \]

其中 \(\mathcal A_k\) 是当前活跃叶节点集合,\(V\) 是显著性分数,\(\tau\) 是温度。 温度很低时接近贪心选择,很高时接近均匀选择;实验固定为 \(\tau=1.0\)。 分裂达到目标预算后停止,因此原图分辨率提高并不要求 Transformer 序列按同等比例增长。 不过,四叉树叶节点增量与实际定长序列实现的细节在正文中没有完全展开,表中的 8194 按原值保留,不改写为 8192。

这种设计保留了标准 ViT 的接口,但“边缘丰富”只是结构重要性的代理。 弱边界目标可能没有获得足够 token,噪声边缘反而可能消耗预算。 所以它不是无损压缩,也不能保证所有细小结构都被保存;收益取决于显著性质量和预算是否足够。

2. 阻尼累积:让掩码条件同时包含祖先尺度和局部尺度

只知道最终叶节点的边缘密度,还不足以表达一个微结构属于怎样的全局组织。 DA 在建树过程中同步更新全分辨率结构画布:根节点提供初始全局响应,每次分裂又在新子块覆盖的位置添加局部高斯响应。 这些响应的尺度由图像信号决定,并沿着根到叶的祖先路径累积,因此最终条件保留了多个空间尺度,而不只是最深层的局部纹理。

为避免深层节点累积过强,新增响应方差随深度 \(j\) 按 \(1/\sqrt{j+1}\) 阻尼。 独立高斯分量相加后,其方差为各层分量方差之和;阻尼使深度增长不会简单地造成同强度响应逐层堆积。 这段机制的重点是“树细分”和“结构条件生成”同时发生,而不是先独立构造一棵树,再额外随机采样一个无关掩码。 缓存中信号条件方差的完整表达存在重复排版与符号不清,本文不猜补其展开式,也不把图像均值擅自解释成常规统计方差。

结构画布随后参与掩码条件生成,附录说明掩码在四叉树各层采样,以维持各尺度上的遮挡比例。 这让密集预测预训练能感知多尺度空间组织,而不是只学习任意缺块的语义补全。 但原文一方面称掩码由结构噪声场决定,另一方面又写随机选择索引子集;它没有在可读文字中完整给出从画布值到掩码概率的规则。 因此,不能进一步断言“响应越强就一定越容易被遮挡”,也不能指定一个未报告的阈值或排序算法。

3. 树对齐重建与分割:把结构表示贯穿预训练和下游任务

token 压缩如果只在预训练使用,下游仍恢复为超长规则序列,部署瓶颈就会回来。 SGMA 将自适应块序列交给标准 ViT;预训练采用 MAE 式非对称编码器—解码器,可见块提供上下文,轻量解码器预测被遮挡块的干净像素。 训练目标只在被遮挡索引上计算均方误差,因此它关注的是如何利用其余结构解释缺失区域,而不是直接为分割引入额外人工标签。

正文对输入有需要保留的表述歧义:一处明确说编码器处理移除掩码索引后的干净可见序列,另一处又称其为 corrupted input,并描述 noised patches。 可确认的共同部分是“编码器使用可见上下文,解码器以干净块为重建目标”。 不能仅凭噪声画布就声称编码器必然接收加噪可见像素,也不能把它解释为标准扩散去噪模型。 下文损失采用原文式(5)的可读部分,不补造噪声注入或解码器输入细节。

微调时丢弃预训练重建解码器,保留编码器,并用同一树算子划分图像和真值掩码。 模型在块级输出分割 logits,与对应树块的标签计算分割损失,而不是直接在反向传播中维护完整原分辨率输出。 SAM 版本借用对称 depatching 将块级预测还原为空间结果,UNETR 版本使用其分割头。 SpringXCT 的 SAM 2 方案处理二维切片并堆叠输出;不能因此把它当成对整幅超大三维体数据直接执行全局注意力。

一个完整示例

以 PAIP 的 32,768 × 32,768 全切片为例,规则 32 × 32 图像块产生 1,048,576 个 token,远超过主实验的序列预算。 SGFT 让背景保留较大块,组织边界附近继续细分;其精度优先配置使用 16,384 个 token,而非把整张切片均匀缩小成低分辨率图像。 树生成时 DA 同时积累跨尺度条件,预训练默认遮挡 75% 的自适应序列,并利用可见部分重建干净被遮挡块。 这是机制示例,不代表原文公布了该切片各区域的实际 token 数或精确掩码位置。

进入有标签微调后,病灶掩码按照同一树划分,模型学习块级分割并还原到图像空间。 若更重视速度,可以改用 2048 token;这是另一个实验配置,会牺牲细节覆盖,不能把其速度结果绑定到 16,384-token 配置的最佳 Dice 上。 推理沿用结构 token 化和分割路径,但不再需要 MAE 的重建遮挡任务。

损失函数 / 训练策略

预训练损失为被遮挡块的干净像素重建误差:

\[ \min_{\theta,\phi}\; \mathbb E_{I\sim\mathcal D,\,\mathcal I_{\mathrm{masked}}\sim\mathrm{Mask}} \left[ \sum_{i\in\mathcal I_{\mathrm{masked}}} \mathcal L_{\mathrm{MSE}} \left(g_\phi(f_\theta(\mathcal P'),i),p_i\right) \right]. \]

\(\mathcal P'\) 是可见块序列,\(p_i\) 是干净目标块,\(f_\theta\) 和 \(g_\phi\) 分别是编码器与重建解码器。 微调将块级预测与同树划分的真值比较,分割损失可采用 Dice 与交叉熵;正文没有给出二者的精确组合权重。 这里保留重建损失的求和形式,不自行添加原文未写明的面积加权或归一化方式。

实验采用 PyTorch、DeepSpeed 与 NVIDIA H100 节点,默认遮挡比例为 0.75。 预训练解码器为 8 个 Transformer block、512 隐藏维度、16 个注意力头,训练后丢弃。 优化使用 AdamW、余弦学习率、40 个 warmup epoch,预训练 400 epoch、微调 100 epoch,微调学习率缩小 10 倍,全程 bf16。 缓存中的基础学习率排版有损,本文不从破损显示文本推断精确值。 这些设置说明作者仍承担标准 ViT 训练成本;token 化并非免费计算,也不意味着单卡即可完成所有配置。

实验关键数据

主实验

实验覆盖 HydrogelTEM-1K 电镜网络、SpringXCT-8K 材料 CT 和 PAIP-32K 肝癌病理切片,指标为 Dice(%)。 SpringXCT 使用 308,000 张真实扫描切片预训练,再通过虚拟微结构及成像退化模拟构造有标签训练数据。 PAIP 按论文报告包含 2,457 张 WSI,按 70% / 10% / 20% 划分训练、验证、测试。 以下为原文 Table 1 的同架构 MAE—SGMA 配对;提升单位是百分点,不是相对百分比。

数据集 微调模型 MAE Dice SGMA Dice 提升 MAE / SGMA 序列长度
HydrogelTEM-1K SAM 72.31 73.48 +1.17 16384 / 16384
HydrogelTEM-1K SAM 2 73.12 74.23 +1.11 16384 / 16384
HydrogelTEM-1K UNETR 74.56 75.33 +0.77 4096 / 16384
SpringXCT-8K SAM 82.68 95.68 +13.00 4096 / 16384
SpringXCT-8K SAM 2 85.98 93.77 +7.79 4096 / 16384
SpringXCT-8K UNETR 86.12 91.93 +5.81 4096 / 8194
PAIP-32K SAM 65.78 82.11 +16.33 1024 / 16384
PAIP-32K SAM 2 66.37 83.21 +16.84 1024 / 16384
PAIP-32K UNETR 77.24 81.23 +3.99 1024 / 8194

这里“同架构”不等于输入配置完全相同。 例如 SpringXCT 的 SAM 比较中,MAE 使用 128 的规则块尺寸,SGMA 表列有效块尺寸为 2,序列长度也不同。 因此,大幅提升说明完整结构 token 化与预训练方案的效果,不能全部归因于把随机掩码换成结构掩码。 95.68 属于 SpringXCT 的 SAM;PAIP 的 82.11 属于 SAM,83.21 属于 SAM 2,不是同模型相互冲突的数字。

效率表来自原文 Table 2,同行比较使用相同 GPU 数,但不同数据集之间的硬件规模不同。

数据集 GPU 数 MAE / SGMA 时间(秒/图) MAE / SGMA 序列长度 MAE / SGMA Dice 原文加速比
HydrogelTEM-1K 4 0.38261 / 0.09913 16384 / 8194 72.31 / 72.56 3.86×
SpringXCT-8K 128 2.5168 / 0.3512 4096 / 1024 82.68 / 89.37 7.17×
PAIP-32K 512 8.9812 / 0.3663 16384 / 2048 62.34 / 76.08 24.8×

PAIP 的效率配置是 2048 token、76.08 Dice,最佳精度配置则使用 16384 token。 效率表中的 62.34 不是主表 FT-MAE-SAM 的 65.78,不应跨表混算提升。 另外,PAIP 效率行的 MAE 序列长度及加速比与其他原值存在核验问题:主表对应 1024 块尺寸列为 1024 token,而效率表列为 16384;8.9812 / 0.3663 约为 24.52,原表写 24.8。 这里保留两表和加速比原值,不自行修正;复现实验需要进一步核对配置及计时口径。 128 与 512 是集群 GPU 数,以上不能解读为单卡效率,也不能据此给出完整训练时间加速结论。

消融实验

缓存正文给出了 DA 在 PAIP-16K 分类与分割上的增益描述,但所引用的完整附录表未包含在当前文本中。 因此只记录可读的差值,不编造绝对分数;另外两项是附录的超参数观察。

配置 / 分析 原文可核实结果 证据边界
SGMA-ViT 相对 NoDA,PAIP-16K 分类 Top-1 增加 0.88 点 绝对 Top-1 未见于缓存
SGMA-ViT 相对 NoDA,PAIP-16K 分割 Dice 增加 3.00 点 绝对 Dice 未见于缓存
HydrogelTEM 遮挡比例验证 0.50、0.60、0.75、0.85、0.90 中,0.75 最佳 没有逐设置数值表
固定 token 预算分析 PAIP-32K 的 1024-token 小预算损害细结构覆盖 未报告该配置具体 Dice

关键发现

  • 随分辨率提高,SGMA 相对配对 MAE 的优势明显增大,但主表也同时改变了块粒度与序列预算,不能视作纯掩码消融。
  • DA 对密集预测的已报告增益为 +3.00 Dice,高于分类的 +0.88 Top-1;这与跨尺度空间条件更适合分割的解释一致,但不是完整的因果验证。
  • 作者展示了真实 SpringXCT 的零样本孔隙分析及 HydrogelTEM 骨架连接分析;这些可视化不是带真实像素真值的整套零样本定量评测。
  • 附录 Figure 7 的 SAM 2 为单样本 94.79 Dice、8194 token,不能拿它替换 Table 1 的整体测试结果 93.77。

亮点与洞察

  • 部署阶段也保留自适应表示。 不仅减少预训练编码器看到的可见块,还直接限制微调和推理的输入序列,解决了 MAE 训练高效但下游序列仍过长的落差。
  • 树不是单纯的压缩容器。 建树过程同时生成多尺度重建条件,让空间划分与学习任务关联;可迁移的启发是让 tokenizer 的结构成为自监督任务的一部分。
  • 科学用途关注拓扑而非只看平均像素重合。 孔隙连接和骨架度数对细边界非常敏感,说明提高局部分割质量可能影响最终科学统计;但仍需要独立拓扑指标与误差分析支持。

局限与展望

  • 边缘依赖。 作者明确承认 Canny 分数在弱梯度、强噪声或纹理主导场景中可能错配任务重要性;可考虑多尺度纹理、置信度或学习式显著性,但需保留无标签可用性。
  • 固定预算会丢细节。 高异质性图像需要按数据集调节 token 数;未来可根据结构密度动态选择预算,而不是认为固定长度就能无损处理所有分辨率。
  • 重建实现说明不足。 结构画布到掩码的具体映射,以及干净可见输入与 noised/corrupted 描述的关系,需要更明确的实现说明,不能由笔记替作者补齐。
  • 证据不完全与配置口径。 缓存缺少正文提到的若干完整附录表,A.5 也仅剩表头;效率表存在配置与比值核验疑点,缺少方差或置信区间,需结合代码和完整材料复核。
  • 真实部署边界。 大规模 GPU 实验不能直接证明资源受限部署可行;真实零样本科学统计与临床使用还需要独立标签验证、跨设备评测和专业审查。

相关工作与启发

  • vs MAE:MAE 使用规则块与随机遮挡;SGMA 同时改变空间 token 分配和结构相关预训练。比较时应分离压缩、预算和 DA 各自贡献。
  • vs Adaptive Patching / SHF:前者已利用层级空间划分或对称还原;SGMA 增加概率化细分和与树构建耦合的结构条件重建,而不是首次提出四叉树用于视觉。
  • vs Swin / HIPT / 线性注意力:这些方法改变骨干层级或注意力计算;SGMA 主要改输入预处理,保留标准 ViT,但仍受固定 token 序列的二次注意力成本约束。
  • 研究启发:可在相同 token 数、相同块表示和相同分割头下比较随机掩码、分层掩码及 DA 条件掩码,并同时测量边界、连通性和预处理耗时,以厘清准确率提升来源。

评分

  • 新颖性: 4/5 — 自适应 token 化与结构条件自监督任务的联合设计有意义,但四叉树和层级压缩并非全新。
  • 实验充分度: 3/5 — 三种成像模态和多骨干对照较丰富,仍受配置差异、缺失消融表及真实零样本定量证据不足限制。
  • 写作质量: 3/5 — 问题和主流程清楚,但掩码实现、输入噪声表述及效率配置口径需要澄清。
  • 价值: 4/5 — 为标准 ViT 处理超高分辨率微结构提供了实用方向,临床或资源受限部署仍需额外验证。