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

Conference: NeurIPS2026
arXiv: 2609.30682
Area: Self-Supervised Learning / Scientific Image Segmentation
Keywords: masked autoencoders, quadtrees, adaptive tokenization, structure-conditioned masking, ultra-high-resolution images

TL;DR

SGMA combines content-adaptive quadtree tokenization with cross-scale structure-conditioned masking for MAE pretraining, preserving scientific microstructure within a fixed sequence budget for standard ViTs; SGMA-SAM reaches 95.68% Dice on SpringXCT and SGMA-SAM 2 reaches 83.21% on PAIP, while the reported maximum 24.8× inference speedup comes from a separate, shorter-sequence configuration.

Background & Motivation

Important objects in scientific imaging are often not large objects centered in an image, but small structures such as cell boundaries, material pores, and polymer connections. Transmission electron microscopy, whole-slide pathology, and X-ray CT can all produce ultra-high-resolution images, yet their labels require domain expertise and their scale makes pixel-level annotation difficult to expand. Pretraining on unlabeled images before fine-tuning with limited labels is therefore a natural approach. MAE learns representations by masking image patches and reconstructing their pixels, but its default uniform grid and random masking do not explicitly distinguish homogeneous backgrounds from detail-rich regions.

The issue is not only whether the pretraining task is appropriate, but whether the resulting model can process the original resolution. A 32,768 × 32,768 image divided into 32 × 32 patches produces approximately one million tokens. Global self-attention in a standard ViT scales quadratically with sequence length; uniformly enlarging patches reduces memory pressure but removes the microstructure needed for segmentation. The MAE encoder processes only visible patches, reducing pretraining costs, whereas downstream segmentation usually still requires the full input sequence, so this bottleneck does not disappear automatically. Swin, linear or sparse attention, and multilevel models offer alternatives, but change the attention mechanism or backbone and may not directly reuse standard MAE/ViT designs.

This paper instead intervenes in both the input representation and the pretraining task: allocate finer patches to structurally rich regions, use larger patches for simple regions, and make reconstruction aware of these cross-scale structures. The fixed token budget constrains the sequence passed to the Transformer rather than uniformly shrinking the original image. Core idea: jointly generate a bounded sequence of structural tokens and multiscale masking conditions on the same content-adaptive quadtree, concentrating computation on microstructure while learning dense-prediction representations through structure-aware reconstruction.

Method

Overall Architecture

SGMA takes a high-resolution scientific image and produces a pretrained ViT encoder for downstream segmentation. The Structure-Guided Feature Tokenizer (SGFT) first constructs a quadtree according to edge salience until a prescribed token budget is reached. Each tree refinement simultaneously triggers Damped Accumulation (DA), adding signal-dependent local Gaussian responses to a full-resolution structure canvas that provides multiscale conditions for mask sampling. Tree-Aligned Reconstruction and Segmentation then uses visible tokens to reconstruct masked patches during pretraining, and retains the tree representation during fine-tuning to align labels and predictions on the same patch partition.

Solid arrows show the pretraining processing chain; dashed arrows show clean reconstruction targets, segmentation supervision, or downstream use, not continued random masking at inference time.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Ultra-high-resolution image"] --> B["Structure-Guided Feature Tokenizer<br/>SGFT"]
    B --> C["Damped Accumulation<br/>DA"]
    C -->|Pretraining structure conditions| D["Tree-Aligned Reconstruction<br/>and Segmentation"]
    B -.->|Fine-tuning and inference tokens| D
    A -.->|Clean patch reconstruction targets| D
    Y["Segmentation labels"] -.->|Same-tree supervision| D
    D --> O["Pretrained encoder / segmentation"]

Compression occurs at the input-tokenization stage. SGMA does not turn a standard Transformer into linear attention: with \(N\) input tokens, global self-attention still costs \(O(N^2)\). Its benefit is that \(N\) can be budgeted instead of growing with the original pixel count under uniform patching; preprocessing, canvas construction, and spatial reconstruction still have their own costs.

Key Designs

1. Structure-Guided Feature Tokenizer: allocate limited tokens to detail-rich regions

A uniform grid assigns the same spatial granularity to backgrounds and fine pores, or must enlarge every patch together. SGFT begins with a root covering the entire image, treats current leaf nodes as candidate tokens, and computes salience from a Canny edge map, for example by summing edge responses within a patch. Highly salient patches are more likely to split into four children, while less salient regions remain coarse. The result is a mixed-scale patch sequence covering different spatial extents, rather than selecting a few regions of interest and discarding the rest of the image.

Selecting the highest-scoring patch at every step requires a sequential search that becomes a tokenizer bottleneck as the active set grows. The authors replace it with temperature-controlled softmax sampling, favoring salient regions while supporting a batched implementation. The selection probability is:

\[ \pi(n_j\mid\mathcal A_k,\tau)= \frac{\exp(V(n_j)/\tau)}{\sum_{n_i\in\mathcal A_k}\exp(V(n_i)/\tau)}. \]

Here, \(\mathcal A_k\) is the current set of active leaves, \(V\) is the salience score, and \(\tau\) is the temperature. Low temperatures approach greedy selection, while high temperatures approach uniform selection; experiments use \(\tau=1.0\). Refinement stops at the target budget, so increasing image resolution need not proportionally increase the Transformer sequence. However, the main text does not fully explain how quadtree leaf-count increments relate to the implementation of an exact fixed-length sequence; the reported value 8194 is retained rather than changed to 8192.

This design preserves the standard ViT interface, but edge richness is only a proxy for structural importance. Weak-boundary targets may receive too few tokens, while noisy edges may consume the budget. It is therefore not lossless compression and does not guarantee preservation of every small structure; gains depend on salience quality and an adequate budget.

2. Damped Accumulation: condition masking on both ancestral and local scales

The edge density of a final leaf alone does not describe the broader organization to which a microstructure belongs. DA updates a full-resolution structure canvas during tree construction: the root supplies an initial global response, and each split adds local Gaussian responses within the new child patches. These responses depend on image signals and accumulate along root-to-leaf paths, so the final conditions retain multiple spatial scales rather than only the deepest local texture.

To prevent excessive accumulation in deep nodes, the variance of each new response is damped with depth \(j\) by \(1/\sqrt{j+1}\). For independent Gaussian components, the final variance is the sum of component variances; damping prevents simply stacking equally strong responses at every level. The key mechanism is simultaneous tree refinement and structural-condition generation, not constructing a tree first and independently sampling an unrelated random mask. The cached expression for the complete signal-conditioned variance contains duplicated rendering and unclear notation; this note does not reconstruct its expansion or reinterpret the image mean as conventional statistical variance.

The structure canvas subsequently contributes to masking conditions, and the appendix states that masks are sampled at each quadtree level to preserve the masking ratio across scales. This makes dense-prediction pretraining aware of multiscale spatial organization rather than only learning semantic completion of arbitrary missing patches. However, the source both says that the structural noise field determines masks and describes randomly selecting an index subset; its readable text does not fully specify the mapping from canvas values to masking probabilities. It is therefore not justified to claim that stronger responses are necessarily more likely to be masked, or to supply an unreported threshold or ranking algorithm.

3. Tree-Aligned Reconstruction and Segmentation: retain the structural representation downstream

If token compression is used only during pretraining, returning to a very long uniform sequence downstream restores the deployment bottleneck. SGMA passes its adaptive patch sequence to a standard ViT; pretraining uses an MAE-style asymmetric encoder–decoder, with visible patches providing context and a lightweight decoder predicting clean pixels for masked patches. The objective computes mean squared error only over masked indices, focusing on how the remaining structure explains missing regions without adding segmentation annotations to pretraining.

The main text contains an input-description ambiguity that must be retained: one passage explicitly says the encoder processes a clean visible sequence after removing masked indices, while another calls this a corrupted input and refers to noised patches. The supported common interpretation is that the encoder uses visible context and the decoder reconstructs clean target patches. A noise canvas alone does not establish that the encoder necessarily receives noisy visible pixels, nor does it make this a standard diffusion denoising model. The loss below follows the readable part of source Equation (5), without inventing noise injection or decoder-input details.

Fine-tuning discards the pretraining reconstruction decoder, retains the encoder, and partitions images and ground-truth masks with the same tree operator. The model predicts patch-level segmentation logits and compares them with corresponding tree-partitioned labels, rather than directly maintaining a full-resolution output during back-propagation. SAM variants use symmetric depatching to reconstruct spatial predictions, while UNETR variants use their segmentation head. The SpringXCT SAM 2 configuration processes 2D slices and stacks their outputs; this should not be interpreted as direct global attention over an entire enormous 3D volume.

A Worked Example

Consider a PAIP whole slide measuring 32,768 × 32,768 pixels: uniform 32 × 32 patches yield 1,048,576 tokens, far beyond the main experiment's sequence budget. SGFT retains large patches over background and refines tissue boundaries; its accuracy-oriented configuration uses 16,384 tokens instead of uniformly downsampling the slide. DA accumulates cross-scale conditions while the tree is constructed, and pretraining masks 75% of the adaptive sequence by default, reconstructing clean masked patches from the visible portion. This illustrates the mechanism, not reported per-region token counts or exact mask positions for a particular slide.

During labeled fine-tuning, the lesion mask is partitioned by the same tree, and the model learns patch-level segmentation before restoring predictions to image space. A speed-oriented run can instead use 2048 tokens; this is a separate configuration that sacrifices detail coverage, so its timing cannot be attached to the best Dice from the 16,384-token configuration. Inference retains structural tokenization and segmentation, but does not require the MAE reconstruction-masking task.

Loss & Training

Pretraining minimizes clean-pixel reconstruction error over masked patches:

\[ \min_{\theta,\phi}\; \mathbb E_{I\sim\mathcal D,\,\mathcal I_{\mathrm{masked}}\sim\mathrm{Mask}} \left[ \sum_{i\in\mathcal I_{\mathrm{masked}}} \mathcal L_{\mathrm{MSE}} \left(g_\phi(f_\theta(\mathcal P'),i),p_i\right) \right]. \]

\(\mathcal P'\) is the visible patch sequence, \(p_i\) is a clean target patch, and \(f_\theta\) and \(g_\phi\) are the encoder and reconstruction decoder. Fine-tuning compares patch-level predictions with ground truth partitioned by the same tree; the segmentation objective may use Dice and cross-entropy, but the text does not specify their exact combination weights. The reconstruction objective retains the source's summation form, without adding unreported area weighting or normalization.

Experiments use PyTorch, DeepSpeed, and NVIDIA H100 nodes, with a default masking ratio of 0.75. The pretraining decoder has 8 Transformer blocks, a hidden dimension of 512, and 16 attention heads, and is discarded afterward. Optimization uses AdamW, a cosine schedule, 40 warmup epochs, 400 pretraining epochs, 100 fine-tuning epochs, a 10× smaller learning rate for fine-tuning, and bf16 throughout. The cached base-learning-rate rendering is damaged, so this note does not infer an exact value from the corrupted display text. These settings show that standard ViT training costs remain; tokenization is not free computation and does not imply that all configurations fit on one GPU.

Key Experimental Results

Main Results

Experiments cover HydrogelTEM-1K electron-microscopy networks, SpringXCT-8K material CT, and PAIP-32K liver-cancer pathology slides, using Dice (%). SpringXCT pretraining uses 308,000 real scan slices, while virtual microstructures and simulated imaging degradations provide labeled training data. The paper reports 2,457 PAIP WSIs, split 70% / 10% / 20% into training, validation, and test sets. The following pairs come from source Table 1 and compare MAE with SGMA on matching architectures; gains are percentage points, not relative percentages.

Dataset	Fine-tuned model	MAE Dice	SGMA Dice	Gain	MAE / SGMA sequence length
HydrogelTEM-1K	SAM	72.31	73.48	+1.17	16384 / 16384
HydrogelTEM-1K	SAM 2	73.12	74.23	+1.11	16384 / 16384
HydrogelTEM-1K	UNETR	74.56	75.33	+0.77	4096 / 16384
SpringXCT-8K	SAM	82.68	95.68	+13.00	4096 / 16384
SpringXCT-8K	SAM 2	85.98	93.77	+7.79	4096 / 16384
SpringXCT-8K	UNETR	86.12	91.93	+5.81	4096 / 8194
PAIP-32K	SAM	65.78	82.11	+16.33	1024 / 16384
PAIP-32K	SAM 2	66.37	83.21	+16.84	1024 / 16384
PAIP-32K	UNETR	77.24	81.23	+3.99	1024 / 8194

Matching architectures do not mean identical input configurations. For example, on SpringXCT, MAE-SAM uses uniform patches of size 128, whereas SGMA's table lists an effective patch size of 2 and a different sequence length. The large gains therefore support the complete structural-tokenization and pretraining approach, not an isolated replacement of random masking with structural masking. The 95.68 score belongs to SAM on SpringXCT; on PAIP, 82.11 belongs to SAM and 83.21 to SAM 2, so these are not conflicting measurements of one model.

The efficiency results below come from source Table 2; paired comparisons use the same GPU count, but hardware scale differs across datasets.

Dataset	GPUs	MAE / SGMA time (s/image)	MAE / SGMA sequence length	MAE / SGMA Dice	Reported speedup
HydrogelTEM-1K	4	0.38261 / 0.09913	16384 / 8194	72.31 / 72.56	3.86×
SpringXCT-8K	128	2.5168 / 0.3512	4096 / 1024	82.68 / 89.37	7.17×
PAIP-32K	512	8.9812 / 0.3663	16384 / 2048	62.34 / 76.08	24.8×

The PAIP efficiency configuration uses 2048 tokens and achieves 76.08 Dice, whereas the best-accuracy configuration uses 16384 tokens. Its efficiency baseline of 62.34 is not the main table's FT-MAE-SAM result of 65.78; gains should not be calculated across these tables. The PAIP efficiency row also contains values requiring verification: the main table lists 1024 tokens for its 1024-patch-size configuration, while the efficiency table lists 16384; 8.9812 / 0.3663 is approximately 24.52, whereas the source reports 24.8. Both tables and the reported speedup are retained without correction; reproduction requires checking configurations and timing conventions. The counts of 128 and 512 refer to cluster GPUs, so these results are not single-GPU efficiency measurements and do not establish an overall training-time speedup.

Ablation Study

The cached main text reports DA gains for PAIP-16K classification and segmentation, but does not include the complete appendix table it cites. Only readable differences are recorded below, without invented absolute scores; the other two entries summarize appendix hyperparameter observations.

Config / analysis	Verifiable source result	Evidence boundary
SGMA-ViT versus NoDA, PAIP-16K classification	Top-1 increases by 0.88 points	Absolute Top-1 absent from cache
SGMA-ViT versus NoDA, PAIP-16K segmentation	Dice increases by 3.00 points	Absolute Dice absent from cache
HydrogelTEM masking-ratio validation	Among 0.50, 0.60, 0.75, 0.85, 0.90, the best ratio is 0.75	No per-setting numerical table
Fixed token-budget analysis	A 1024-token budget on PAIP-32K harms fine-structure coverage	Specific Dice not reported for this configuration

Key Findings

• SGMA's advantage over paired MAE models grows with resolution, but the main table also changes patch granularity and sequence budgets, so it is not a pure masking ablation.
• DA's reported gain of +3.00 Dice for dense prediction exceeds its +0.88 Top-1 classification gain, consistent with multiscale spatial conditions benefiting segmentation, but not a complete causal test.
• The authors show zero-shot pore analysis on real SpringXCT samples and skeleton-connectivity analysis on HydrogelTEM; these visualizations are not a comprehensive quantitative zero-shot evaluation with real pixel-level ground truth.
• Appendix Figure 7 reports 94.79 Dice for SAM 2 on one sample with 8194 tokens; this should not replace the aggregate test result of 93.77 in Table 1.

Highlights & Insights

• Adaptive representation is retained at deployment. The method not only reduces visible patches during pretraining, but also bounds the input sequence during fine-tuning and inference, addressing the gap between efficient MAE training and long downstream sequences.
• The tree is more than a compression container. Its construction also generates multiscale reconstruction conditions, coupling spatial partitioning to the learning task; the transferable idea is to make tokenizer structure part of self-supervised learning.
• Scientific utility depends on topology, not only average pixel overlap. Pore connectivity and skeleton degree are sensitive to thin boundaries, so local segmentation improvements can affect scientific statistics; independent topology metrics and error analysis are still needed.

Limitations & Future Work

• Edge dependence. The authors explicitly acknowledge that Canny scores can misidentify task importance under weak gradients, strong noise, or texture-dominated structures; multiscale texture, confidence, or learned salience could help, while preserving usability without labels.
• Fixed budgets lose detail. Highly heterogeneous images require dataset-specific token budgets; adaptive budgets based on structural density may be preferable to assuming lossless processing at every resolution.
• Incomplete reconstruction specification. The exact structure-canvas-to-mask mapping and the relationship between clean visible input and noised/corrupted descriptions require clearer implementation details, rather than reconstruction by this note.
• Incomplete evidence and configuration conventions. Several complete appendix tables cited in the text are missing from the cache, and A.5 ends at its table headers; efficiency configurations and ratios require verification, with no variance or confidence intervals reported here.
• Deployment boundaries. Large-GPU-count experiments do not directly demonstrate resource-constrained deployment; real zero-shot scientific statistics and clinical use require independent labeled validation, cross-device evaluation, and professional review.

Related Work & Insights

• vs MAE: MAE uses uniform patches and random masking; SGMA changes both spatial token allocation and structure-aware pretraining. Comparisons should disentangle compression, budgets, and DA.
• vs Adaptive Patching / SHF: These methods already exploit hierarchical spatial partitions or symmetric reconstruction; SGMA adds probabilistic refinement and structure-conditioned reconstruction coupled to tree construction, rather than introducing quadtrees to vision for the first time.
• vs Swin / HIPT / linear attention: These methods change backbone hierarchies or attention computation; SGMA mainly changes input preprocessing and retains standard ViTs, but remains subject to quadratic attention costs over the fixed token sequence.
• Research direction: Compare random masking, level-wise masking, and DA-conditioned masking with identical token counts, patch representations, and segmentation heads, while measuring boundaries, connectivity, and preprocessing time to identify the sources of accuracy gains.

Rating

• Novelty: 4/5 — Joint adaptive tokenization and structure-conditioned self-supervision is meaningful, although quadtrees and hierarchical compression are established ideas.
• Experimental Thoroughness: 3/5 — Three imaging modalities and multiple backbone comparisons provide useful coverage, limited by configuration differences, missing ablation tables, and insufficient quantitative real zero-shot evidence.
• Writing Quality: 3/5 — The problem and overall pipeline are clear, but masking implementation, noisy-input descriptions, and efficiency configurations need clarification.
• Value: 4/5 — A practical direction for standard ViTs on ultra-high-resolution microstructure, with additional validation needed for clinical or resource-constrained deployment.