Structure-Guided Masked Autoencoders for Ultra-High Resolution Scientific Image Understanding

Enzhi Zhang; Du Wu; Rui Zhong; Cong Ma; Isaac Lyngaas; Amir Koushyar Ziabari; Xiao Wang; Peng Chen

cs.CVSep 25, 2026

Structure-Guided Masked Autoencoders for Ultra-High Resolution Scientific Image Understanding

Authors: Enzhi Zhang, Du Wu, Rui Zhong, Cong Ma, Isaac Lyngaas, Amir Koushyar Ziabari, Xiao Wang, Peng Chen,

Organizations: Hokkaido University, Japan · RIKEN Center for Computational Science (R-CCS), Japan · Oak Ridge National Laboratory, USA · A*STAR, Singapore ·

Abstract

Self-supervised pre-training with Vision Transformers, including Masked Autoencoders (MAE), is difficult to apply to gigapixel scientific images. Random masking is poorly matched to the structured, multi-scale morphology of scientific data, while uniform tokenization produces prohibitively long sequences that make $O(N^2)$ attention impractical. We propose SGMA, a structure-guided masked autoencoding framework for ultra-high-resolution scientific images. SGMA couples two components: a content-adaptive quadtree tokenizer that compresses gigapixel images into a fixed-length sequence, and a structure-conditioned masking process that biases reconstruction toward spatially informative regions. To stabilize this process across scales, we introduce Damped Accumulation (DA), which aggregates signal-dependent responses across the tree into a structure canvas used to guide masking. The resulting pre-training task preserves fine microstructure while remaining compatible with standard ViT encoders and MAE-style reconstruction. Across electron microscopy, whole-slide optical microscopy, and X-ray CT datasets, SGMA consistently outperforms MAE baselines. It achieves 95.68% Dice on the 8K x 8K x 28K SpringXCT dataset, improving over the same-architecture MAE baseline by +13.00 points, and 83.21% Dice on the 32K^2 WSI PAIP dataset, improving by +16.84 points, while providing up to a 24.8x inference speedup.

Figures & tables

Figure 1: Overview of SGMA . A structure-guided quadtree decomposes the image and accumulates per-level spatial responses by superimposing scale-dependent Gaussian noise. The resulting structure-conditioned noise field determines where masks are placed on the image. When applied to fine-tuning SAM [ 31 ] on HydrogelTEM-1K electron microscopy data, SGMA compresses the sequence length by a large margin while maintaining key spatial structures, enabling stronger fine-tuning performance.

Panel of Figure 2 — Figure 2: Adaptive structural patches and spatial information from transmission electron microscopy: HydrogelTEM-1K [ 34 ] , X-ray CT: SpringXCT-8K [ 56 ] , and WSI microscopy: PAIP-32K [ 30 ] datasets. These datasets span centimeter to nanometer scales and include biomedical, macroscopic, and microscopic samples. SGMA adapts to the multi-scale micro-structure of each domain for MAE pre-training.

Panel of Figure 3 — Figure 3: We compared the segmentation difference between MAE-SAM [ 31 ] and SGMA -SAM on TEM HydrogelTEM-1K, X-ray CT SpringXCT-8K, and WSI PAIP-32K datasets. At the same GPU compute budgets, SGMA -SAM can go down to patch size of $8\times 8$ (vs. $1,024\times 1,024$ at best for SAM before going OOM). As a result, SGMA -SAM can extract and express mask details better than MAE-SAM, with the gap in accuracy favoring SGMA -SAM as the resolution gets higher.

Dataset	Model	Patch Size	Sequence Length	Dice Score (%)	Performance Improvement
HydrogelTEM-1K	SAM 2 [ 44 ]	8	16384	69.87	-
	UNet [ 46 ]	-	-	71.03	-
	UNETR [ 20 ]	16	4096	70.59	-
	FT-MAE-SAM	8	16384	72.31	(Baseline)
	FT-MAE-SAM 2	8	16384	73.12	(Baseline)
	FT-MAE-UNETR	16	4096	74.56	(Baseline)

Table 1: Quantitative segmentation results (Dice Score %

\uparrow

) on all datasets. Our SGMA pre-training method is compared against standard baselines and MAE pre-training on identical architectures. The "Improvement" column shows the Dice Score gain of our SGMA model over its equivalent MAE-trained baseline (e.g. FT-SGMA-UNETR vs. FT-MAE-UNETR).

Datasets	GPUs	Model	Time (s/img)	Seq. Len.	Dice (%)	Speedup ( $\times$ )
HydrogelTEM-1K	4	MAE-SAM-8 [ 22 , 31 ]	0.38261	16384	72.31	1 $\times$
HydrogelTEM-1K	4	SGMA-SAM (ours)	0.09913	8194	72.56	3.86 $\times$
SpringXCT-8K	128	MAE-SAM-128 [ 22 , 31 ]	2.5168	4096	82.68	1 $\times$
SpringXCT-8K	128	SGMA-SAM (ours)	0.3512	1024	89.37	7.17 $\times$
PAIP-32K	512	MAE-SAM-1024 [ 22 , 31 ]	8.9812	16384	62.34	1 $\times$
PAIP-32K	512	SGMA-SAM (ours)	0.3663	2048	76.08	24.8 $\times$

Table 2: Inference efficiency under the same compute budget (identical GPU count per row) for MAE-SAM and our SGMA-SAM. The MAE row uses the longest sequence that fits within its GPU budget before OOM, while SGMA uses a deliberately shorter sequence (e.g.

N{=}2{,}048

on PAIP-32K vs.

N{=}16{,}384

in Table 1 ) to match wall-clock compute. This enables a clear comparison of SGFT speedup. SGMA still achieves higher Dice and faster inference, so the reported speedups (up to 24.8

\times

) are conservative lower bounds.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Masking ratio ablation. Each row fixes a ratio $r\in\{0.25,0.50,0.75\}$ and shows the input, the quadtree partition, the visible patches after structure-guided masking, and the SGMA reconstruction. Higher masking ratios yield a harder pre-training task while the adaptive quadtree keeps structurally important regions visible.

Panel of Figure 5 — Figure 5: Sequence length ablation. The SGFT tokenizer concentrates more tokens in detail-rich regions as the budget $N$ grows from $1024$ to $8192$ , while keeping uniform background coarsely partitioned.

Panel of Figure 6 — Figure 6: Model size ablation. Reconstruction quality improves from ViT-L to ViT-XL on a HydrogelTEM-1K patch under the same structure-guided masking and quadtree tokenization. The larger backbone recovers finer network topology, which motivates our use of larger backbones on higher-resolution datasets.

Panel of Figure 7 — Figure 7: Segmentation comparison on a SpringXCT-8K simulated slice (per-sample Dice reported). UNet [ 46 ] loses small-pore topology; the Adaptive Patching baseline [ 65 ] at $N{=}8{,}194$ tokens improves but still drops thin channels; MAE-pretrained SAM 2 [ 44 ] at the smallest feasible $128^{2}$ patch on an 8K image captures coarse structure only; our SGMA -SAM 2 at a $2{\times}2$ effective patch size with $N{=}8{,}194$ SGFT tokens recovers the fine pore topology. The aggregate test-set Dice for each model is reported in Table 1 ; the numbers here are for this specific sample.

Panel of Figure 8 — Figure 8: Segmentation on a real SpringXCT-8K sample with zero-shot inference. The pixel-level microstructure (e.g. void area) can be precisely extracted, which is hard for human experts.

Figure 9: Example of SpringXCT-8K downstream task: connected component analysis [ 67 ] , separation, and mapping the stones distribution in the 3D image.

Figure 10: Example of HydrogelTEM-1K downstream task: a hydrogel sample processed through zero-shot segmentation and skeletonization, followed by junction detection and degree-wise node annotation, illustrating the structural complexity of the hydrogel network.

Masked Autoencoders

Masked autoencoders learn representations by hiding parts of an input and training a model to reconstruct what is missing.

See all 101 papers

AFFMAE: Scalable Vision Pre-Training for High-Resolution Microscopy Segmentation on Desktop Hardware
Feb 18, 2026David Smerkous, Zian Wang, Behzad NajafianMasked Autoencoders Self-Supervised Pre-Training
Department of Laboratory Medicine & Pathology, University of Washington, Seattle WA, USA · Electrical Engineering & Computer Science, Oregon State University, Corvallis, OR, USA · Paul G. Allen School of Computer Science & Engineering, University of Washington, Seattle WA, USA
WildMAE: Structure-Aware Masked Autoencoding for Wildlife Image Reconstruction
Oct 20, 2026Yicheng Zhang, Yujie Jin, Ke Zhang +1Masked Image Modeling Masked Autoencoders
The Chinese University of Hong Kong, Shenzhen
From Semantics to Pixels: Coarse-to-Fine Masked Autoencoders for Hierarchical Visual Understanding
Mar 10, 2026Wenzhao Xiang, Yue Wu, Hongyang Yu +3Self-Supervised Visual Representation Learning Masked Autoencoders
Key Laboratory of Intelligent Information Processing, Institute of Computing Technology, Chinese Academy of Sciences, Beijing 100190, China · University of Chinese Academy of Sciences, Beijing 100049, China · Pengcheng Laboratory, Shenzhen 518108, China +1

Feb 18, 2026J/K move · Enter open · S save

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

David Smerkous, Zian Wang, Behzad Najafian

Department of Laboratory Medicine & Pathology, University of Washington, Seattle WA, USA · Electrical Engineering & Computer Science, Oregon State University, Corvallis, OR, USA · Paul G. Allen School of Computer Science & Engineering, University of Washington, Seattle WA, USA

Masked Autoencoders Self-Supervised Pre-Training Efficient ViTs Microscopy Image Analysis Visual Token Compression

Open paper AskPDF arXiv

Self-supervised pretraining has transformed computer vision by enabling data-efficient fine-tuning, yet high-resolution pretraining typically requires server-scale infrastructure, limiting custom in-domain training for many research laboratories. Masked Autoencoders (MAE) reduce computation by encoding only visible tokens, but combining MAE with hierarchical downsampling architectures has remained 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. AFFMAE removes dense-grid assumptions while preserving hierarchical scalability during pre-training and fine-tuning. To support this architecture, we developed numerically 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 MAE finetuning performance on foot process width estimation with ViT backbone at equal parameter counts while being 2x faster during pre-training and halving peak memory usage. Furthermore, AFFMAE achieves up to 5x throughput speedups fine-tuning at the 1024px resolution, providing high-resolution model training on desktop hardware. Code available at https://github.com/najafian-lab/affmae.

Abstract

Masked Autoencoders

Explore similar work

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

WildMAE: Structure-Aware Masked Autoencoding for Wildlife Image Reconstruction

From Semantics to Pixels: Coarse-to-Fine Masked Autoencoders for Hierarchical Visual Understanding