AFFMAE: Scalable and Efficient Vision Pretraining for Desktop Graphics Cards
This work addresses the problem of limited in-domain foundation model development for research labs by enabling efficient training on desktop graphics cards, though it is incremental in improving existing methods.
The paper tackled the challenge of high-resolution self-supervised pretraining for computer vision by introducing AFFMAE, a framework that reduces computational requirements, achieving up to 7x fewer FLOPs and halving memory usage while matching performance on electron microscopy segmentation.
Self-supervised pretraining has transformed computer vision by enabling data-efficient fine-tuning, yet high-resolution training typically requires server-scale infrastructure, limiting in-domain foundation model development for many research laboratories. Masked Autoencoders (MAE) reduce computation by encoding only visible tokens, but combining MAE with hierarchical downsampling architectures remains 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. By discarding masked tokens and performing dynamic merging exclusively over visible tokens, AFFMAE removes dense-grid assumptions while preserving hierarchical scalability. We developed numerically stable mixed-precision Flash-style cluster attention kernels, and mitigate sparse-stage representation collapse via deep supervision. On high-resolution electron microscopy segmentation, AFFMAE matches ViT-MAE performance at equal parameter count while reducing FLOPs by up to 7x, halving memory usage, and achieving faster training on a single RTX 5090. Code available at https://github.com/najafian-lab/affmae.