Extending Video Masked Autoencoders to 128 frames

TL;DR

This work proposes an effective strategy for prioritizing tokens which allows training on longer video sequences (128 frames) and gets better performance than, more typical, random and uniform masking strategies.

Abstract

Video understanding has witnessed significant progress with recent video foundation models demonstrating strong performance owing to self-supervised pre-training objectives; Masked Autoencoders (MAE) being the design of choice. Nevertheless, the majority of prior works that leverage MAE pre-training have focused on relatively short video representations (16 / 32 frames in length) largely due to hardware memory and compute limitations that scale poorly with video length due to the dense memory-intensive self-attention decoding. One natural strategy to address these challenges is to subsample tokens to reconstruct during decoding (or decoder masking). In this work, we propose an effective strategy for prioritizing tokens which allows training on longer video sequences (128 frames) and gets better performance than, more typical, random and uniform masking strategies. The core of our approach is an adaptive decoder masking strategy that prioritizes the most important tokens and uses quantized tokens as reconstruction objectives. Our adaptive strategy leverages a powerful MAGVIT-based tokenizer that jointly learns the tokens and their priority. We validate our design choices through exhaustive ablations and observe improved performance of the resulting long-video (128 frames) encoders over short-video (32 frames) counterparts. With our long-video masked autoencoder (LVMAE) strategy, we surpass state-of-the-art on Diving48 by 3.9 points and EPIC-Kitchens-100 verb classification by 2.5 points while relying on a simple core architecture and video-only pre-training (unlike some of the prior works that require millions of labeled video-text pairs or specialized encoders).

Figures (4)

• Figure 1: Left: Proposed Long Video MAE Decoder Masking. We leverage a jointly trained adaptive tokenizer and importance module to define a decoder mask and token targets for a video MAE pre-training strategy. The resulting sparsification in tokens (only 15%) allows pre-training with long videos (128-frames) and results in substantial performance gains. Right: Decoder masking and memory in long-video (128 frames) pre-training. We report memory and FLOPs per-device for a batch size of 1 using different decoder mask ratios and ViT-B architecture.
• Figure 2: Illustration of Adaptive FSQ-MagViT Training. FSQ-MagViT adaptive tokenizer includes MAGVIT encoder and CNN-based token scorer with a differentiable top-$k$ selection layer designating importance of tokens. During tokenizer training unselected tokens zeroed out and video is reconstructed using MAGVIT decoder. We then freeze this adaptive tokenizer and use it to generate target tokens for scalable pre-training of video MAE.
• Figure 3: Our Adaptive tokenizer visualized. We visualize the tokens masks by masking the corresponding input video (repeating frames to match the latent temporal dimension).
• Figure 4: Token selection strategies visualized. We can see that flow based token selection can be dominated by large background motion. Randomly selected masks are unable to focus the tokens on the interesting parts of the video. In contrast, we see that the adaptively selected tokens reflect well what is changing in the video.