MaskViT: Masked Visual Pre-Training for Video Prediction

• Video prediction must model complex, stochastic futures while remaining fast enough to support robot actions.Source: §1 · p. 1
• MaskViT uses masked visual modeling with spatial and spatiotemporal window attention to improve computational efficiency.Source: §1 · p. 1, Figure 1 · p. 2
• Variable-ratio masking during training and iterative decoding during inference reduce the mismatch between training and complete future-frame prediction.Source: §1 · p. 2
• 256 × 256 videos are supported, and iterative decoding makes MaskViT up to 512× faster than autoregressive methods.Source: §1 · p. 2
• The method enables planning for real-robot manipulation tasks through faster video prediction.Source: §1 · p. 2

• Video prediction generates future videos conditioned on past frames, optionally with language descriptions or motor commands.Source: §2 · p. 3
• Existing approaches include GANs, VAEs, invertible networks, autoregressive models, and diffusion models.Source: §2 · p. 3
• MaskViT reduces the inference-time drawback of autoregressive video generation through iterative decoding.Source: §2 · p. 3
• Masked autoencoders learn representations by reconstructing original inputs from corrupted or masked inputs.Source: §2 · p. 3

• MaskViT first discretizes video frames into visual tokens with a dVAE or VQ-GAN, then trains a bidirectional transformer with masked visual modeling.Source: §3 · p. 3, §3.1 · p. 3
• Per-frame compression produces a 16 × 16 token grid and preserves conditioning on arbitrary context frames.Source: §3.1 · p. 3
• Masked visual modeling keeps context-frame tokens intact and masks a variable number of future-frame tokens for prediction.Source: §3.2 · pp. 3–4
• The variable masking ratio is sampled from r ∈[0.5, 1) and reduces the gap between pre-training and inference.Source: §3.2 · p. 4
• Window-restricted attention uses spatial windows and spatiotemporal windows to reduce the costs of global self-attention.Source: §3.3 · p. 4
• Iterative decoding starts with all future tokens masked, predicts tokens repeatedly, and retains previously predicted and most confident tokens under a decreasing mask schedule.Source: §3.4 · p. 4

Comparison with Prior Work

• MaskViT is competitive or outperforms prior work across BAIR, KITTI, and RoboNet while using fewer parameters.Source: Table 1 · p. 5, §4.2 · pp. 5–6
• On BAIR, goal conditioning improves FVD by approximately 18%, while action conditioning improves FVD by almost 25%.Source: §4.2 · p. 5
• KITTI results are significantly better across all metrics, and MaskViT produces 256 × 256 predictions despite the dataset’s long-horizon, limited-data setting.Source: §4.2 · p. 6

Ablation Studies

• A 16 × 4 × 4 spatiotemporal window achieves better accuracy with 60% less memory and 3.3× faster training than full self-attention.Source: §4.3 · pp. 6–7
• 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.Source: §4.3 · p. 7

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.Source: §4.4 · pp. 7–8

• MaskViT uses masked visual modeling and window attention to obtain competitive video prediction results with efficient computation.Source: §5 · p. 8
• Iterative decoding is significantly faster than autoregressive decoding and enables planning for real robot manipulation.Source: §5 · p. 8
• Per-frame quantization can produce flicker artifacts, especially in videos with static backgrounds such as RoboNet.Source: §5 · p. 8
• Scaling video prediction remains challenging for scenarios with substantial camera motion, including self-driving and egocentric videos.Source: §5 · p. 8

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.Source: §A.1 · p. 14
• GAN losses are introduced after reconstruction loss converges because discriminator training can be unstable.Source: §A.1 · p. 14

Transformer

• Transformer blocks use spatial and spatiotemporal window attention with learnable space-time positional embeddings.Source: §A.1 · p. 14

Evaluation

• Evaluation parameters are selected by grid search over the cosine or square mask schedule, temperature, and decoding iterations.Source: §A.1 · p. 14

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.Source: §A.2 · p. 15

Visual-MPC

• Visual MPC minimizes weighted ℓ2 differences between predicted frames and a goal image over an action horizon.Source: §A.2 · p. 15
• CEM optimizes H = 10 future actions by sampling M = 256 action sequences and iteratively refitting their Gaussian distribution.Source: §A.2 · p. 15
• Correlated action noise biases sampled trajectories toward smoother sequences, while the gripper dimension is rounded to −1 or 1.Source: §A.2 · p. 15

Evaluation

• The evaluation covers table setting and sweeping, with five task variations and three trials per variation.Source: §A.2 · p. 15
• Five decoding iterations require approximately 6.2 seconds for a batch of 256 samples using eight RTX 3090 GPUs.Source: §A.2 · p. 16

Video prediction

• Additional qualitative video predictions are presented for BAIR, KITTI, and RoboNet.Source: §B.1 · p. 16

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.Source: §B.1 · p. 16
• RoboNet-only training fails because it cannot reconstruct the background and robot arm with high fidelity.Source: §B.1 · p. 16
• Domain-specific finetuning remains necessary for reasonable predictions in this real-robot setting.Source: §B.1 · p. 16

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.Source: §B.2 · p. 16
• Each task variant is evaluated with three trials for every model or policy.Source: Table 7 · p. 17
• Success requires the object center to lie within 8cm of the goal position after 15 steps.Source: Table 7 · p. 17

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.Source: Figure 5 · p. 18, §B.2 · pp. 18–19, Figure 6 · p. 19
• Additional qualitative visualizations report results on BAIR at 64 × 64 and KITTI and RoboNet at 256 × 256.Source: Figure 7 · p. 20, Figure 8 · p. 21, Figure 9 · p. 22