AFFMAE: Scalable Vision Pre-Training for High-Resolution Microscopy Segmentation on Desktop Hardware
David Smerkous, Zian Wang, Behzad Najafian
ECCV 2026
cs.CV
2026-02-18
华盛顿大学把离网格自适应合并接到层次化MAE,单卡RTX 5090上预训练约快2倍、显存减半,1024微调吞吐20对4 img/s,mIoU 0.633对齐ViT-MAE的0.630。
实验室堆着大量未标注高分辨率电镜图,想做域内自监督预训练,却卡在两件事上:算力默认按服务器多卡来,数据又常常出不了院。ImageNet 权重救急不够。肾电镜上,官方 MAE ImageNet-1K(ViT-B,86M 参数,1600 epoch)微调 mIoU 只有 0.494,从零训练是 0.476,同域预训练直接到 0.606。更大的模型、更大的数据集、四倍 epoch,都没挡住域差。
MAE 只编码可见 token,省显存。显微镜分辨率上,ViT-MAE 在消费级卡仍然又慢又肥。Swin 这类分层骨干靠密网格做窗口注意力和 patch merging,丢掉 mask token 会把核心算子拆掉。SimMIM、GreenMIM、HiViT 要么把 mask token 留在编码器里,要么为预训练单独改接口,预训练和微调对不齐,细长结构尤其吃亏。
任务本身也苛刻。足突宽度(FPW,foot process width)要同时看滤过裂隙这种像素级细缝和肾小球全局几何,输入通常要大于 512×512。网格对齐的下采样会把细缝稀释进邻域。
AFFMAE 把 AutoFocusFormer 的离网格自适应合并接到 MAE「编码器不看 mask token」的设定上。
每个 token 带着特征向量和二维坐标。局部注意力在等大小的邻域簇上做,位置用坐标差。下采样不按 2 倍网格砍,而用一层 MLP 打重要性分:保留率 ds 可以是 0.4 这种任意比例,高分当锚点,其余并进邻近锚,坐标跟着走。纹理空的地方狠并,裂隙附近留密。
编码器对 mask 无感,预训练和微调除了输入遮挡几乎同一套。解码器是点式可变形交叉注意力:可学习的 mask query 预测偏移,在各层用 K 近邻可见 token 做指数距离加权(写成带 max-logit 的 softmax,混合精度才稳),四层对齐编码器四阶段。这颗头微调时直接当分割头,不必再上 UperNet。
工程上两件实事。簇注意力用 Triton 按 FlashAttention 思路重写,分数矩阵不物化,softmax 用 fp32 累加。解码器 KNN 预建成 H×W×K 查找表,查询近似 O(1);点解码器前向 1.5×、反向 1.4×。
另外两处设计选择。深层 token 会塌成位置编码网格,有效秩掉下去,所以在解码器中间阶段加辅助重建。加上之后,各层有效秩能维持在 0.7 以上,和下采样率几乎无关。高分辨率随机遮挡太碎,模型靠邻域插值就能交差,于是改用 Perlin 噪声造连片遮挡,频谱更接近电镜图的 1/f 衰减。遮挡率峰值在 50%,不是 MAE 常说的 75%。
预训练:187,270 张未标注肾小球电镜(7500–15000×),单卡 RTX 5090,512px,有效 batch 256,约 65M 参数,400 epoch。微调:FPW 570 张训练 / 235 张测试,分辨率 512/768/1024,四随机种子。
| 方法 | 预训练 GFLOPs / 显存 / 吞吐 | 1024 微调 mIoU / 吞吐 / 显存 |
| ViT-MAE | 274.5 / 29.7 GB / 76 img/s | 0.630 / 4 img/s / 25.4 GB |
| SimMIM | 119.3 / 27.7 GB / 111 | 0.628 / 21 / 18.0 GB |
| AFFMAE | 58.7 / 14.5 GB / 151 | 0.633 / 20 / 13.4 GB |
512px 裂隙 IoU:分层基线 0.399–0.415,MAE 0.447,AFFMAE 0.459;1024px AFFMAE 到 0.514。FPW 像素误差 15.97,低于 MAE 的 19.26。HiViT 预训练吞吐更高(318 img/s),1024 微调 mIoU 停在 0.623,裂隙更弱。
消融(23M,300 epoch):ds=0.5 时 mIoU 0.6009,0.4 时 0.5908 但算力少约 16%;去掉 Deep Supervision 掉到 0.5734;Perlin 对随机是 0.6009 对 0.5947。batch 16 时 MAE 和 SimMIM 超过 640px 就 OOM,AFFMAE 在 896px 仍占 22.3 GB。
公开电镜集在补遗:Lucchi++ 上 MAE 0.8840、AFFMAE 0.8829;Kasthuri++ 上 AFFMAE 0.8870,略超 MAE 的 0.8818。
同域预训练的收益远大于换更大 ImageNet 模型。这篇把这件事做到一张消费卡上:预训练大约快 2 倍、显存砍半,1024 微调吞吐约 5 倍,分割不掉点。肾脏电镜这种出不了院的数据,可以在实验室自己预训练。
分层骨干想吃 MAE 的可见 token 效率、又不想把细结构并没,离网格合并比「把可见 token 塞回窗口」更干净。代码在 najafian-lab/affmae。
作者承认:不规则 token 让底层效率变脏,FLOPs 降了,墙钟不一定同比例降。查找表在很高分辨率会撑爆缓存,大约 1200px 以上网格索引可能重新更快。工作停在二维透射电镜,三维只是路线图。
对照设置偏友好。主表里所有 MIM 方法都在同域电镜上预训练,比的是架构,不是「能不能在桌面卡上从 ImageNet 出发」。主结果绑在足突宽度这一套标注上,公开集差距在百分位。RTX 5090 是旗舰消费卡,实验室如果还是 24 GB 卡,数字要自己打折。