让显微镜图像高分辨率训练在普通电脑上跑起来,速度翻倍还省内存。
AFFMAE: Scalable Vision Pre-Training for High-Resolution Microscopy Segmentation on Desktop Hardware
- 用自适应非网格令牌合并,解决高层级预训练的结构难题
- 在1024px分辨率下,微调速度提升5倍,峰值内存减半
- 适合想在桌面设备上做高分辨率显微图像分割的研究者
自监督预训练推动了计算机视觉的发展,但高分辨率预训练通常需要服务器级算力,限制了众多实验室的领域定制训练。掩码自编码器(MAE)通过仅编码可见标记降低计算量,但与分层下采样架构结合时,因密集网格假设和遮蔽感知设计妥协而难以实现。我们提出AFFMAE,一种基于自适应、非网格令牌合并的遮蔽友好分层预训练框架,消除密集网格假设,同时保持预训练与微调中的层级可扩展性。为此,我们开发了数值稳定的混合精度Triton内核和轻量点式解码器,可直接作为分割头复用。在高分辨率显微图像分割任务中,AFFMAE在参数量相当的情况下,达到与ViT骨干的MAE微调相当的脚掌宽度估计性能,预训练速度提升2倍,峰值内存减半。此外,在1024px分辨率下微调速度最高提升5倍,实现了桌面硬件上的高分辨率模型训练。代码已开源:https://github.com/najafian-lab/affmae。
原文摘要 · Abstract (English)
Self-supervised pretraining has transformed computer vision by enabling data-efficient fine-tuning, yet high-resolution pretraining typically requires server-scale infrastructure, limiting custom in-domain training for many research laboratories. Masked Autoencoders (MAE) reduce computation by encoding only visible tokens, but combining MAE with hierarchical downsampling architectures has remained structurally challenging due to dense grid priors and mask-aware design compromises. We introduce AFFMAE, a masking-friendly hierarchical pretraining framework built on adaptive, off-grid token merging. AFFMAE removes dense-grid assumptions while preserving hierarchical scalability during pre-training and fine-tuning. To support this architecture, we developed numerically stable mixed-precision Triton kernels and a lightweight, point-based decoder that can be directly repurposed as a segmentation head. On high-resolution microscopy segmentation, AFFMAE matches MAE finetuning performance on foot process width estimation with ViT backbone at equal parameter counts while being 2x faster during pre-training and halving peak memory usage. Furthermore, AFFMAE achieves up to 5x throughput speedups fine-tuning at the 1024px resolution, providing high-resolution model training on desktop hardware. Code available at https://github.com/najafian-lab/affmae.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。