AFFMAE: Scalable Vision Pre-Training for High-Resolution Microscopy Segmentation on Desktop Hardware

Self-supervised pretraining has transformed computer visionby enabling data-efficient fine-tuning, yet high-resolution pretraining typ-ically requires server-scale infrastructure, limiting custom in-domain train-ing for many research laboratories. Masked Autoencoders (MAE) reducecomputation by encoding only visible tokens, but combining MAE withhierarchical downsampling architectures has remained structurally chal-lenging due to dense grid priors and mask-aware design compromises. Weintroduce AFFMAE, a masking-friendly hierarchical pretraining frame-work 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 nu-merically 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 MAEfinetuning performance on foot process width estimation with ViT back-bone at equal parameter counts while being 2x faster during pre-trainingand halving peak memory usage. Furthermore, AFFMAE achieves up to5x throughput speedups fine-tuning at the 1024px resolution, providinghigh-resolution model training on desktop hardware. Code available athttps://github.com/najafian-lab/affmae.