AFFMAE: Scalable and Efficient Vision Pretraining for Desktop Graphics Cards

TL;DR

AFFMAE introduces a scalable, efficient vision pretraining framework that enables high-resolution training on desktop GPUs by using adaptive token merging and optimized attention kernels, matching state-of-the-art performance with reduced computational costs.

Contribution

It presents AFFMAE, a novel hierarchical pretraining method that overcomes dense-grid limitations, enabling efficient high-resolution vision model training on consumer hardware.

Findings

Matches ViT-MAE performance on electron microscopy segmentation

Reduces FLOPs by up to 7x compared to baseline

Halves memory usage and accelerates training on a single GPU

Abstract

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.