Exploring Long-Sequence Masked Autoencoders

TL;DR

This work examines how input specifications, particularly sequence length, influence Masked Autoencoder pre-training for vision. By decoupling mask size from patch size and introducing block-wise masking, the authors create a minimally altered long-sequence MAE that scales with $L=(I/p)^2$. Across detection, segmentation, and some classification tasks, longer pre-training sequences yield consistent gains, especially on scene-rich datasets like COCO and Open Images, with larger ViT models deriving the most benefit. While longer sequences increase pre-training cost, they do not add transfer-time computation, suggesting sequence-length as a practical axis for scaling vision models with real-world impact on localization tasks.

Abstract

Masked Autoencoding (MAE) has emerged as an effective approach for pre-training representations across multiple domains. In contrast to discrete tokens in natural languages, the input for image MAE is continuous and subject to additional specifications. We systematically study each input specification during the pre-training stage, and find sequence length is a key axis that further scales MAE. Our study leads to a long-sequence version of MAE with minimal changes to the original recipe, by just decoupling the mask size from the patch size. For object detection and semantic segmentation, our long-sequence MAE shows consistent gains across all the experimental setups without extra computation cost during the transfer. While long-sequence pre-training is discerned most beneficial for detection and segmentation, we also achieve strong results on ImageNet-1K classification by keeping a standard image size and only increasing the sequence length. We hope our findings can provide new insights and avenues for scaling in computer vision.

Figures (3)

• Figure 1: Input specifications for MAE. An $I{\times} I$ input image is converted to a patch sequence of length $L=(I/p)^2$, according to the patch size $p{\times} p$. The encoder processes a random subset of $L_e=L\cdot(1-\gamma)$ patches based on the mask ratio $\gamma$, whereas the decoder produces all $L_d=L$ patches.
• Figure 2: Decoupled mask size from patch size. Besides mask ratio ($\gamma=0.75$), mask size $m$ is another important factor that controls task difficulty, which was incidentally coupled with patch size in MAE he2021masked. For long-sequence MAE, we reinstate $m$: if $m=1{\times} 1$, the model can mostly use near-by low-level texture details for reconstruction; $m=2{\times} 2$ renders a more semantically meaningful task.
• Figure 3: Scaling trends of sequence length $L$ for pre-training on COCO or ImageNet-1K; and evaluating on COCO object detection with fixed transferring length. We make three observations: 1) Increasing sequence length $L$ generally leads to a healthy trend of better AP$^{b}$ on COCO object detection. 2) The trend for ViT-L is better than ViT-B -- the former continues to scale even at $L{=}1024$, whereas the latter saturates and even starts to decline after a certain point. 3) COCO pre-training benefits more from longer sequences than ImageNet-1K.