Search papers, labs, and topics across Lattice.
The paper introduces AFFMAE, a masking-friendly hierarchical pretraining framework for vision transformers that uses adaptive, off-grid token merging to remove dense-grid assumptions and enable efficient high-resolution training. By dynamically merging only visible tokens and incorporating deep supervision to prevent representation collapse during sparse stages, AFFMAE achieves significant computational savings. Experiments on high-resolution electron microscopy segmentation demonstrate that AFFMAE matches ViT-MAE performance with 7x FLOPs reduction, halved memory usage, and faster training on a single RTX 5090.
Train high-resolution vision models on a desktop GPU: AFFMAE slashes FLOPs by 7x and halves memory usage compared to ViT-MAE, without sacrificing performance.
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.
