AFFMAE: Scalable Vision Pre-Training for High-Resolution Microscopy Segmentation on Desktop Hardware
David Smerkous, Zian Wang, Behzad Najafian
ECCV 2026
cs.CV
2026-02-18
UW's hierarchical MAE with off-grid token merging hits 0.633 mIoU at 1024px on an RTX 5090, matching ViT-MAE at 5x fine-tune throughput and half the memory.
Pathology labs sit on large unlabeled high-resolution electron-microscopy archives and still cannot pretrain in-domain. Two blockers dominate: high-resolution pretraining assumes multi-GPU servers, and moving the images off-site often trips HIPAA and IRB review. ImageNet weights are a weak substitute here. On kidney EM, official MAE ImageNet-1K (ViT-B, 86M parameters, 1600 epochs) fine-tunes to 0.494 mIoU; training from scratch is 0.476; the same MAE pretrained on in-domain EM reaches 0.606. A larger model, a larger dataset, and 4× the epochs do not close that domain gap.
MAE already drops masked tokens from the encoder. At microscopy resolution a ViT-MAE is still slow and memory-heavy on a desktop GPU. Hierarchical backbones such as Swin cut tokens with grid-aligned windows and patch merging, so discarding tokens breaks the operators. SimMIM, GreenMIM, and HiViT either keep mask tokens in the encoder or add a pretrain-only packing step. Pretrain and finetune then diverge, and thin structures suffer.
The downstream task is unforgiving. Foot process width (FPW) needs both filtration slits a few pixels wide and global glomerular geometry, usually at inputs larger than 512×512. Grid-aligned downsampling dilutes those slits into neighboring patches.
AFFMAE grafts AutoFocusFormer's off-grid adaptive merging onto MAE's mask-agnostic encoder.
Each token carries a feature and a 2D coordinate. Local attention runs inside equal-size coordinate clusters. Downsampling is not a 2× grid cut. A one-layer MLP scores importance; a retention ratio ds (for example 0.4) keeps the top tokens as anchors and merges the rest into nearby anchors, coordinates included. Empty texture is merged hard; slit neighborhoods stay dense.
The encoder never sees mask tokens, so pretraining and finetuning differ only in the input mask. The decoder is point-based deformable cross-attention: learnable mask queries predict an offset, gather the K nearest visible tokens at each encoder stage, and blend them with an exponential distance kernel written as a stable softmax. Four decoder stages mirror the four encoder stages. The same head is reused as a segmentation decoder, so UperNet is optional.
Two systems pieces make the FLOP cut show up on a wall clock. Cluster attention is rewritten in Triton in the FlashAttention style: no materialized score matrix, fp32 softmax accumulators, fp16 outputs. Decoder KNN is replaced by an H×W×K lookup table over the original patch grid, making neighbor retrieval effectively O(1) after the table is built (1.5× forward and 1.4× backward in the point decoder).
Two training choices keep the sparse stages alive. Without extra losses, the deepest tokens collapse toward positional grids and effective rank falls. Auxiliary reconstruction heads on intermediate decoder stages (deep supervision) keep normalized effective rank above 0.7 across stages. Random masking at high resolution is too local; the model can inpaint from neighbors. Perlin noise produces contiguous holes whose power spectrum tracks the EM 1/f decay. The masking-ratio peak is 50%, not MAE's usual 75%.
Pretraining uses 187,270 unlabeled glomerular EM images (7500–15000×) on one RTX 5090 at 512px, effective batch 256, 65M parameters, 400 epochs. Finetuning is the FPW set (570 train / 235 test) at 512, 768, and 1024, four random seeds.
| Method | Pretrain GFLOPs / mem / img/s | 1024 finetune mIoU / img/s / mem |
| ViT-MAE | 274.5 / 29.7 GB / 76 | 0.630 / 4 / 25.4 GB |
| SimMIM | 119.3 / 27.7 GB / 111 | 0.628 / 21 / 18.0 GB |
| AFFMAE | 58.7 / 14.5 GB / 151 | 0.633 / 20 / 13.4 GB |
Slits IoU at 512px: hierarchical baselines 0.399–0.415, MAE 0.447, AFFMAE 0.459; at 1024px AFFMAE reaches 0.514. FPW pixel error is 15.97 versus MAE's 19.26. HiViT is faster at pretraining (318 img/s) but plateaus at 0.623 mIoU and weaker slits.
Ablations on a 23M model, 300 epochs: ds = 0.5 gives 0.6009 mIoU; ds = 0.4 is 0.5908 at 16% less compute. Removing deep supervision drops to 0.5734. Perlin beats random 0.6009 vs 0.5947. At batch 16, MAE and SimMIM OOM past 640px on a 32 GB card; AFFMAE still fits 22.3 GB at 896px.
Public EM sets in the supplement: Lucchi++ MAE 0.8840 vs AFFMAE 0.8829; Kasthuri++ AFFMAE 0.8870 vs MAE 0.8818.
In-domain pretraining is the actual lever, larger than swapping in a bigger ImageNet checkpoint. AFFMAE makes that pretraining fit a desktop GPU: about 2× faster pretrain, half the memory, about 5× 1024 finetune throughput, with segmentation on par with ViT-MAE. Labs that cannot ship kidney EM to a cluster can train at home.
For hierarchical MIM, adaptive off-grid merging is a cleaner way to keep MAE's visible-token efficiency without packing tokens back onto a lattice. Code is at najafian-lab/affmae.
Irregular tokens make low-level efficiency messy. FLOPs fall; wall-clock does not always follow. Dense lookup tables can outgrow cache at very high resolution, and past 1200px grid indexing may win again. The study is 2D TEM; 3D is a sketch.
The comparison is also friendly. Every MIM method in the main table is pretrained on the same EM archive, so the test is architecture, not "ImageNet on a desktop card." The headline numbers are tied to one FPW annotation protocol; public-set gaps sit in the third decimal. An RTX 5090 is a flagship consumer GPU. A 24 GB card will not reprint these throughputs.