Source-linked AI summary
MaskViT: Masked Visual Pre-Training for Video Prediction
Agrim Gupta, Stephen Tian, Yunzhi Zhang, Jiajun Wu, Roberto Martín-Martín, Li Fei-Fei
TL;DR
Video prediction for embodied agents must capture complex, stochastic futures while remaining fast enough for action and planning. MaskViT addresses this with masked visual modeling, window-restricted attention, variable training masks, and iterative decoding, achieving strong video-prediction results and up to 512× faster inference. The method also supports 256 × 256 videos and real-robot planning, while scaling to substantial camera motion remains challenging.
Problem
Video prediction for embodied agents must model complex and stochastic future observations while maintaining inference speed that supports robot actions.
Method
MaskViT pre-trains bidirectional transformers with masked visual modeling, spatial and spatiotemporal window attention, variable masking ratios, and iterative decoding.
Results
MaskViT achieves competitive or state-of-the-art results across several real-world datasets, supports 256 × 256 videos, and is up to 512× faster than autoregressive methods.
Takeaways & Limitations
Iterative decoding enables MaskViT to support planning for real-robot manipulation tasks.
Takeaways & Limitations
Per-frame quantization can cause flicker artifacts, and scaling video prediction to scenarios with significant camera motion remains challenging.
Abstract
from arXiv · showhide
The ability to predict future visual observations conditioned on past observations and motor commands can enable embodied agents to plan solutions to a variety of tasks in complex environments. This work shows that we can create good video prediction models by pre-training transformers via masked visual modeling. Our approach, named MaskViT, is based on two simple design decisions. First, for memory and training efficiency, we use two types of window attention: spatial and spatiotemporal. Second, during training, we mask a variable percentage of tokens instead of a fixed mask ratio. For inference, MaskViT generates all tokens via iterative refinement where we incrementally decrease the masking ratio following a mask scheduling function. On several datasets we demonstrate that MaskViT outperforms prior works in video prediction, is parameter efficient, and can generate high-resolution videos (256x256). Further, we demonstrate the benefits of inference speedup (up to 512x) due to iterative decoding by using MaskViT for planning on a real robot. Our work suggests that we can endow embodied agents with powerful predictive models by leveraging the general framework of masked visual modeling with minimal domain knowledge.
1 Introduction
MaskViT addresses the computational and train–test challenges of video prediction with masked visual modeling, efficient window attention, and iterative decoding. It achieves competitive or state-of-the-art results while supporting high-resolution prediction and substantially faster inference.
- Video prediction must model complex, stochastic futures while remaining fast enough to support robot actions.
- MaskViT uses masked visual modeling with spatial and spatiotemporal window attention to improve computational efficiency.
- Variable-ratio masking during training and iterative decoding during inference reduce the mismatch between training and complete future-frame prediction.
- 256 × 256 videos are supported, and iterative decoding makes MaskViT up to 512× faster than autoregressive methods.
- The method enables planning for real-robot manipulation tasks through faster video prediction.
2 Related Work
Prior video-prediction methods condition generated videos on past frames, motor commands, or other context and use several generative-model families. MaskViT targets the large inference time associated with autoregressive generation by using iterative decoding.
- Video prediction generates future videos conditioned on past frames, optionally with language descriptions or motor commands.
- Existing approaches include GANs, VAEs, invertible networks, autoregressive models, and diffusion models.
- MaskViT reduces the inference-time drawback of autoregressive video generation through iterative decoding.
- Masked autoencoders learn representations by reconstructing original inputs from corrupted or masked inputs.
3 MaskViT: Masked Video Transformer
MaskViT uses a two-stage pipeline: per-frame visual tokenization followed by bidirectional masked visual pre-training with efficient window attention. During inference, iterative non-autoregressive decoding progressively unmasks future-frame tokens.
- MaskViT first discretizes video frames into visual tokens with a dVAE or VQ-GAN, then trains a bidirectional transformer with masked visual modeling.
- Per-frame compression produces a 16 × 16 token grid and preserves conditioning on arbitrary context frames.
- Masked visual modeling keeps context-frame tokens intact and masks a variable number of future-frame tokens for prediction.
- The variable masking ratio is sampled from r ∈[0.5, 1) and reduces the gap between pre-training and inference.
- Window-restricted attention uses spatial windows and spatiotemporal windows to reduce the costs of global self-attention.
- Iterative decoding starts with all future tokens masked, predicts tokens repeatedly, and retains previously predicted and most confident tokens under a decreasing mask schedule.
4 Experimental Evaluation
MaskViT is evaluated across video prediction datasets, ablations, inference efficiency, and real-robot control. It achieves strong dataset performance, benefits from efficient window attention and variable masking, and supports fast planning on a Sawyer robot.
- Comparison with Prior Work: MaskViT is competitive or outperforms prior work across BAIR, KITTI, and RoboNet while using fewer parameters.
- Comparison with Prior Work: On BAIR, goal conditioning improves FVD by approximately 18%, while action conditioning improves FVD by almost 25%.
- Comparison with Prior Work: KITTI results are significantly better across all metrics, and MaskViT produces 256 × 256 predictions despite the dataset’s long-horizon, limited-data setting.
- Ablation Studies: A 16 × 4 × 4 spatiotemporal window achieves better accuracy with 60% less memory and 3.3× faster training than full self-attention.
- Ablation Studies: Variable masking performs best, while excessive masking degrades FVD; concave scheduling functions outperform linear and convex alternatives.The scheduling study also finds that additional decoding steps improve FVD only up to a point.
- Visual Model Predictive Control with MaskViT on a Real Robot: Iterative decoding enables approximately 6.5 seconds per CEM iteration for real-robot planning, which is orders of magnitude more efficient than autoregressive models.The robot experiments use MaskViT for visual model-predictive control after training with RoboNet and additional setup data.
5 Conclusion
MaskViT combines masked visual modeling with windowed transformers for efficient video prediction and faster iterative decoding. The approach achieves competitive results while exposing limitations from per-frame quantization and challenging high-motion scenarios.
- MaskViT uses masked visual modeling and window attention to obtain competitive video prediction results with efficient computation.
- Iterative decoding is significantly faster than autoregressive decoding and enables planning for real robot manipulation.
- Per-frame quantization can produce flicker artifacts, especially in videos with static backgrounds such as RoboNet.
- Scaling video prediction remains challenging for scenarios with substantial camera motion, including self-driving and egocentric videos.
A.1 Training MaskViT
MaskViT training combines dataset-specific VQ-GAN tokenization, windowed transformer modeling, and tuned decoding parameters. Training and evaluation use staged procedures intended to stabilize reconstruction and support efficient inference.
- VQ-GAN: VQ-GAN models downsample each frame into 16 × 16 latent codes, with compression factors of 4 for 64 × 64 and 16 for 256 × 256 frames.
- VQ-GAN: GAN losses are introduced after reconstruction loss converges because discriminator training can be unstable.
- Transformer: Transformer blocks use spatial and spatiotemporal window attention with learnable space-time positional embeddings.
- Evaluation: Evaluation parameters are selected by grid search over the cosine or square mask schedule, temperature, and decoding iterations.
A.2 Real Robot Experiments
The real-robot system collects low-resolution trajectories and uses visual model-predictive control with MaskViT predictions inside cross-entropy action optimization. Evaluation measures task completion across table-setting and sweeping variations.
- Data collection: The robot dataset contains 4000 trajectories of 30 timesteps recorded with a Sawyer arm and resized from 480 × 480 crops to 64 × 64.
- Visual-MPC: Visual MPC minimizes weighted ℓ2 differences between predicted frames and a goal image over an action horizon.
- Visual-MPC: CEM optimizes H = 10 future actions by sampling M = 256 action sequences and iteratively refitting their Gaussian distribution.
- Visual-MPC: Correlated action noise biases sampled trajectories toward smoother sequences, while the gripper dimension is rounded to −1 or 1.
- Evaluation: The evaluation covers table setting and sweeping, with five task variations and three trials per variation.
- Evaluation: Five decoding iterations require approximately 6.2 seconds for a batch of 256 samples using eight RTX 3090 GPUs.
B.1 Qualitative Results
Qualitative results cover BAIR, KITTI, and RoboNet, while real-robot examples show that broadly trained MaskViT can guide actions toward visual goals. RoboNet-only training fails to reconstruct important scene elements adequately in the tested setting.
- Video prediction: Additional qualitative video predictions are presented for BAIR, KITTI, and RoboNet.
- Real robot experiments: The planner brings a blue bowl or soft red hat close to the goal position using the model trained on all data.
- Real robot experiments: RoboNet-only training fails because it cannot reconstruct the background and robot arm with high fidelity.
- Real robot experiments: Domain-specific finetuning remains necessary for reasonable predictions in this real-robot setting.
B.2 Quantitative Results
The real-robot evaluation measures MaskViT-based planning across table-setting and sweeping task variants, using per-task success rates and visualized predicted plans. Success is assessed by whether the pushed object reaches a defined target region after 15 steps.
- Real robot experiments: The evaluation includes two task types—table setting and sweeping—with five variations each.Table-setting variations use unseen bowls, while sweeping variations use toys previously seen during finetuning.
- Real robot experiments: Each task variant is evaluated with three trials for every model or policy.
- Real robot experiments: Success requires the object center to lie within 8cm of the goal position after 15 steps.
- Qualitative results: The qualitative results visualize predicted plans for table-setting and sweeping tasks, including pushing a blue bowl left and moving a hat to the bottom-right corner.
- Qualitative results: Additional qualitative visualizations report results on BAIR at 64 × 64 and KITTI and RoboNet at 256 × 256.