MaskViT: Masked Visual Pre-Training for Video Prediction
Authors: Agrim Gupta
Stephen Tian
Yunzhi Zhang
Jiajun Wu
Roberto Martín-Martín
Li Fei-Fei
Venue: ICLR 2023 poster (2023)
Keywords: Video Prediction
Masked Visual Modeling
Visual MPC
Transformers

Abstract: 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 ($256 \times $256). Further, we demonstrate the benefits of inference speedup (up to $512 \times$) 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. 

Published as a conference paper at ICLR 2023
MASKVIT: MASKED VISUAL PRE-TRAINING FOR VIDEO PREDICTION
Agrim Gupta', Stephen Tian', Yunzhi Zhang', Jiajun Wu', Roberto Martin-Martin 1.2.3 Li Fei-Fei
1 Stanford University, UT Austin, SSaeesooree Al
ABSTRACT
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- pre-training transformers via masked visual modeling. Our approach, named Mask 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 variable percentage of tokens instead of fixed mask ratio. For inference, MaskViT generates all to- kens via iterative refinement where we incrementally decrease the masking ratio following a mask scheduling function. On several datasets we demonstrate that Mask outperforms prior works in video prediction, parameter efficient, gen- erates high-resolution videos (256 x 256) and can be easily adapted to perform goal-conditioned video prediction. Further, we demonstrate the benefits of in- ference 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.
INTRODUCTION
Evidence from neuroscience suggests that human cognitive and perceptual capabilitiess are supported by predictive mechanism to anticipate future events and sensory signals (Tanji & Evarts, 1976; Wolpert al., 1995). Such mental model of the world can be used to simulate, evaluate, and select among different possible actions. This process is fast and accurate, even under the computational limitations of biological brains (Wu et al., 2016). Endowing robots with similar predictive capabil- ities would allow them to plan solutions to multiple tasks in complex and dynamic environments, e.g., via visual model-predictiv control (Finn & Levine, 2017; Ebert et al., 2018).
Predicting visual observations for embodied agents is however challenging and computationally demanding: the model needs to capture the complexity and inherent stochasticity of future events while maintaining an inference speed that supports the robot's actions. Therefore, recent advances in autoregressive generative models, which leverage Transformers (Vaswani et al., 2017) for building neural architectures and learn good representations via self-supervised generative pretraining (De- vlin et al., 2019), have not benefited video prediction or robotic applications. We in particular identify three technical challenges. First, memory requirements for the full attention mechanism in Transformers scale quadratically with the length the input sequence, leading to prohibitively large costs for videos. Second, there is inconsistency between the video prediction task and au- toregressive masked visual pretraining- while the training process assumes partial knowledge of the ground truth future frames, at test time the model has to predict a complete sequence future frames from scratch, leading to poor video prediction quality (Yan et al., 2021; Feichtenhofer et al., 2022). Third, the common autoregressive paradigm effective other domains would be too slow for robotic applications.
To address these challenges, we present Masked Video Transformers (Mask (iT): a simple, effective and scalable method for video prediction based on masked visual modeling. Since using pixels
"Correspondence agri wttrrrimtttttffd..mmm

---

Published as a conference paper at ICLR 2023
L-T-1
1=:
t=0
t=T-1
(context)
(oontext)
(context)
Bidirectional Window Transformer
(a) Training
T-T-1
1-0
1-1
T-
Maeked Istent token
3D window attention
(contexd)
(context]
(coneext)
Iteration
Iteration 4
Neration 16
(b) Inference
Figure 1: Mask (a) Training: We encode the video frames into latent codes via VQ-GAN. A variable number of tokens in future frames masked, and the network is trained to predict the masked tokens. A block in MaskViT consists of two layers with window-restricted attention: spatial and spatiotemporal. (b) Inference: Videos are generated via iterative refinement where we incrementally decrease the masking ratio following a mask scheduling function. Videos available at this project page.
directly frame tokens would require inordinate amount of memory, autoencoder (dVAE) (Van Den Oord & Vinyals, 2017; Esser et al., 2021) that compresses frames into smaller grid of visual tokens. We opt for compression in the spatial (image) domain instead of the spatiotemporal domain (videos), as preserving the correspondence between each original and tokenized video frame allows for flexible conditioning on any subset of frames initial (past), final (goal), and possibly equally spaced intermediate frames. However, despite operating on tokens, representing 16 frames at 256 tokens per frame still requires 4, 096 tokens, incurring prohibitive memory requirements for full attention. Hence, to further reduce memory, MaskViT is composed of alternating transformer layers with non-overlapping window-restricted (Vaswani al., 2017) spatial and spatiotemporal attention.
discrete variational
To reduce the inconsistency between the masked pretraining and the video prediction task and to speed up inference, we take inspiration from non-autoregressive. iterative decoding methods in gen- erative algorithms from other domains (Sohl-Dickstein al., 2015; Ho et al., 2020; Nichol & Dhari- wal, 2021; Ghazvininejad et al., 2019; Chang et al., 2022). We propose novel iterative decoding scheme for videos based on mask scheduling function that specifies, during inference, the number of tokens to be decoded and kept each iteration. In contrast to autoregressive decoding, which involves predicting tokens one by one, our iterative decoding scheme is faster as the number of decoding iterations significantly less than the number of tokens. A few initial tokens are pre- dicted over multiple initial iterations, and then the majority of the remaining tokens can be predicted rapidly over the final few iterations. This brings us closer the ultimate video prediction task, where only the first frame is known and all tokens for other frames must be inferred. To further close the training-tes test gap, during training we mask variable percentage of tokens, instead of using fixed masking ratio. This simulates the different masking ratios Mask VIT will encounter during iterative decoding the actual video prediction task.
Through experiments on several publicly available real-world video prediction datasets (Ebert et al., 2017: Geiger et al., 2013; Dasari et al., 2019), we demonstrate that Mask ViT achieves competitive or state-of-the-art results in variety of metrics. Moreover, Mask can predict considerably higher resolution videos (256 x 256) than previous methods. We also show the flexibility of Mask ViT by adapting it to predict conliconditi video frames. In addition, thanks to iterative decoding, Mask VIT is up to 512x faster than autoregressive methods, enabling its application for planning on real robot ($ 4.5). These results indicate that endow embodied agents with powerful
2

---

Published as a conference paper at ICLR 2023
predictive models by leveraging the advances in self-supervised learning in language and vision, without engineering domain- -epeuiii--seiiiieeee solutions.
2 RELATED WORK
Video prediction. The video prediction task refers to the problem of generating videos conditioned on past frames (Ranzato et al., 2014; Lotter et al., 2016), possibly with an additional natural lan- guage description (Li et al., 2018; Gupta et al., 2018; Pan et al., 2017; Wu et al., 2021b) and/or motor commands (Finn et al., 2016; Villegas et al., 2019; Wu et al., 2021a; Babaeizadeh et al., 2021). Multiple classes of generative models have been utilized to tackle this problem, such as Gen- erative adversarial networks (GANs) (Clark et al., 2019; Tulyakov et al., 2018; Luc et al., 2020), Variational Autoencoders (VAEs) (Villegas et al., 2019; Wu et al., 202 la; Babacizadeh et al., 2021; 2018; Denton & Fergus, 2018; Akan al., 2021: 2022), invertible networks (Dorkenwald et al., 2021), autoregressive (Yan et al., 2021; Rakhimov et al., 2020; Nash et al., 2022) and diffusion (Ho al., 2022: Voleti al., 2022) models. Our work focuses on predicting future frames conditioned on past frames or motor commands and belongs to the family of two-stage methods that first encode the videos into a downsampled latent space and then use transformers to model an autoregressive prior (Yan et al., 2021; Rakhimov et al., 2020). A common drawback of these methods is the large inference time due to autoregressive generation. Mask ViT overcomes this issue by using an iterative decoding scheme, which significantly reduces inference time.
Masked autoencoders. Masked autoencoders are a type of denoising autoencoder (Vincent et al., 2008) that learn representations by (re)generating the original input from corrupted (i.e., masked) inputs. Masked language modeling (MLM) was first proposed in BERT (Devlin et al., 2019) and has revolutionized the field of natural language processing, especially when scaled to large datasets and model sizes (Brown et al., 2020; Radford et al., 2019). MLM with fixed and low masking ratio has also been extended for multi-modal representation learning (Lu et al., 2019; Sun et al., 2019; Zellers et al., 2021). The success in NLP has also been replicated in vision by masking a fixed but high ratio of patches of pixels (He et al., 2021; Dosovitskiy et al., 2020) or masking tokens generated by a pretrained dVAE (Bao et al., 2022; Chen et al., 2020). Recently, these works have also been extended to video domains to learn good representations for action recognition (Tong et al., 2022; Feichtenhofer et al., 2022). Unlike them, we apply masked visual modeling for video prediction, and we use a variable masking ratio during training to reduce the difference between masked pretraining and video prediction. Another related line of work is leveraging good visual representations Iearnt via self supervised learning methods (Laskin et al., 2020; Nair et al., 2022; Parisi et al., 2022) including masked autoencoders (Xiao al., 2022) for motor control.
Non-autoregeessive decoding. A key limitation of autoregressive decoding is that each token is predicted sequentially. Hence, the decoding time is proportional the dimensionality of the data. An appealing alternative is iterative decoding, where all tokens are predicted simultaneously and then refined for fixed number of steps independent of the data dimensionality (Gu et al., 2017: Ghazvininejad et al., 2019; Saharia et al., 2020). Ghazvininejad et al. (2019) proposed a mask- predict decoding paradigm for machine translation by training a bidirectional decoder instead of a causal decoder. During inference, entire translated sequences are predicted in parallel within a constant number of steps. Chang et al. (2022) extend this framework to the image domain and use mask scheduling functions during training and decoding. We extend the framework of mask- predict (Ghazvininejad et al., 2019) and MaskGiT (Chang et al., 2022) for video prediction. We pre-train bidirectional window transformer model via masked visual modeling where we mask a variable percentage of tokens independent of the mask scheduling function. Importantly, unlike MaskGiT, instead of keeping the most confident tokens during each decoding step, we add temper- ature annealed Gumbel noise to the token confidence to produce diverse outputs.
MASKVIT: MASKED VIDEO TRANSFORMER
Mask VIT is the result of two-stage training procedure (Van Den Oord & Vinyals, 2017; Razavi et al., 2019): First. learn an encoding of the visual data that discretizes images into tokens based on a discrete variational autoencoder (dVAE). Next, we deviate from the common autoregres- sive training objective and pre-train bidirectional transformer with window-restricted attention via
3

---

Published as a conference paper at ICLR 2023
KITI
1-0
1
1: 1::10
1-
1+20
t- 26
29
Ground ruuh
(context)
Predicted
1+3
1-6
1+9
Ground truth
(conted)
Predicted
Figure 2: High resolution video prediction. Video prediction results on test of KITTI and RoboNet at 256 x 256 resolution. See $ B.2 for additional qualitative results.
masked visual modeling (MVM). In the following section, describe image tokenizer, bidirec- tional transformer, masked visual pre-training, and iterative decoding procedure.
3.1 LEARNING VISUAL TOKENS
Videos contain too many pixels to be used directly as tokens a transformer architecture. Hence, to reduce dimensionality, we first train a VQ-VAE (Van Den Oord & Vinyals, 2017) for individual video frames so that we can represent videos as sequences of grids of discrete tokens. VQ-VAE con- sists of an encoder E(x:) that encodes input image IR xxx3 into series of latent vectors. The vectors are discretized through a nearest neighbour look up in a codebook of quantized embed- dings, 2 {=x h== C Rm= decoder D trained to predict reconstruction of the image, & from the quantized encodings. In our work, we leverage VQ-GAN (Esser et al., 2021) l), which improves upon VQ-VAE by adding adversarial (Goodfellow al., 2014) and perceptual losses (Johnson al., 2016; Zhang et al., 2018). Each video frame is individually tokenized into a 16 x 16 grid of tokens, regardless of their original resolution (Fig. 1, left). Instead of using 3D extensions of VQ-VAE which perform spatiotemporal compression of videos (Yan et al., 2021). our per-frame compression enables us to condition on arbitrary context frames: initial, final, and possibly intermediate ones.
3.2 MASKED VISUAL MODELING (MVM)
Inspired by the success of masked language (Devlin al., 2019) and image (Bao al., 2022; He et al., 2021) modeling, and in the spirit of unifying methodologies across domains, we pre-train MaskViT via MVM for video prediction. Our pre-training task and masking strategy are straightfor- ward: we keep the latent codes corresponding to context frames intact and mask random number of tokens corresponding future frames. The network is trained predict masked latent codes conditioned on the unmasked latent codes.
Concretely, we assume access to input context frames for T. time steps, and our goal is to pre- dict T2 frames during test time. We first quantize the entire video sequence into latent codes Z I T Let Zp [2 denote the latent tokens corresponding to future video frames, where N = Tp x h x w. Unlike prior work on MVM (Bao et al., 2022; He et al., 2021) that uses a fixed masking ratio, we propose to use variable masking ratio that reduces the gap be- tween pre-training task and inference leading to better evaluation results (see $ 3.4). Specifically, during training, for each video in batch. first select masking ratio E [0.5, 1) and then domly select and replace Ir N tokens in Z, with a [MASK] token. The pre-training objective to minimize the negative log-likelihood of the visual tokens given the masked video input:
4

---

Published as a conference paper at ICLR 2023
MVM -- E Emenas logp(z;
Zo)] where D is the training dataset, N represents ran-
ZED
domly masked positions, and Z denotes the output of applying the mask to 2p, and Z. are latent tokens corresponding to context frames. The MVM training objective is different from the causal autoregressive training objective: LAR =
EE Eviez, logp(z; Rjecs Zo) The key difference ZED
is that for MVM the conditional dependence is bidirectional: all masked tokens are predicted con-
ditioned on all tokens.
3.3 BIDIRECTIONAL WINDOW TRANSFORMER
Transformer models composed entirely of global self-attention modules incur significant compute and memory costs, especially for video tasks. To achieve more efficient modeling, we propose to compute self-attention in windows, based on two types of non-overlapping configurations: 1) Spatial Window (SW): attention is restricted to all the tokens within subframe of size 1 x h wo (the first dimension is time); 2) Spatiotemporal Window (STW): attention is restricted within 3D window of size T h' w'. We sequentially stack the two types of window configurations to gain both local and global interactions in a single block (Fig. 1, center) that we repeat T. times. Surprisingly, we find that a small window size of h' w' 4 is sufficient to learn a good video prediction model while significantly reducing memory requirements (Table 2b). Note that our proposed block enjoys global interaction capabilities without requiring padding or cyclic-shifting like prior works (Liu et al., 2021a;b), nor developing custom CUDA kernels for sparse attention (Child et al., 2019) as both window configurations can be instantiated via simple tensor reshaping.
3.4 ITERATIVE DECODING
Decoding tokens autoregressively during inference is time-consuming, as the process scales linearly with the number of tokens, and this can be prohibitively large (e.g.. 4, 096 for video with 16 frames and 256 tokens per frame). Our video prediction training task allows us to predict future video frames via novel iterative non-autoregressive decoding scheme: inspired by the forward diffusion process in diffusion models (Ho et al., 2020; Nichol & Dhariwal, 2021) and the iterative decoding generative models e.g., mask mnask-predict Ghazvininejad et al. (2019) and MaskGiT Chang et al. (2022), we predict videos in T steps where T << N, the total number of tokens to predict.
Coneretely, let >(),, where E ( ( J..., I,I), be a mask scheduling function (Fig. 3) that computes the mask ratio for tokens as a function of the decoding steps. We choose (t) such that it monotonically decreasing with respect to t, and holds that (0) + and 7(1) -> to ensure that our method converges. At t = 0, we start with Z = [Z.: ZpI where all the tokens in Zp are [MASK] tokens. At each decoding iteration, we predict all the tokens conditioned on all the previously predicted tokens. For the next iteration, we mask out = [++)) tokens by keeping all the previously predicted tokens and the most confident token predictions in the current decoding step. We use the softmax probability our confidence measure.
Empirically we found that selecting the most confident tokens while performing iterative decoding generates videos with little or no motion. Intuitively, given only few (1 - 2) context frames, selecting the most confident tokens results in videos where the context frames are copied across all time steps. Hence, we add temperature annealed Gumbel noise to the token confidence to encourage the model to produce more diverse outputs. Concretely, let C E R. be the vector of softmax probabilities of the sampled tokens Z/ decoding iteration L. We select the most confident tokens to keep at iteration t from C₁ where, C2 Ct + Gumbel 1) (1- +)
4 EXPERIMENTAL EVALUATION
In this section, we evaluate our method on three different datasets and compare its performance with prior state-of-the-ee methods, using four different metrics. We also perform extensive ablation studies of different design choices, and showcase that the speed improvements due to iterative de- coding enable real-tima planning for robotic manipulation tasks. For qualitative results, see $ B.2 and videos our project website.
5

---

Published as a conference paper at ICLR 2023
RoboNet
param.
FVDL
P
SSIM
T.PISS]
SVG
298M
123.2
23.9
87.8
0.060
BAIR
parum.
FVDJ
GHVAR
599M
95.2
24.7
89.1
0.036
SV2P (Bahoeizudeh: at. 2018)
262.5
FirVid
302M
62.5
28.2
89.3
0.024
LVT (Kakhimov al, 2020)
125.8
Mi
257M
133.5
23.2
80.5
0.042
SAVP (1 et 2018)
116.4
VIS (256)
228M
211.7
20.4
67.1
0.170
DVD-GAN-FFP (Claris at.l,,
VideoGPT (Yin
109.8
103.3
KITTI
param
FVDJ
PSNR1
SSIM
LPIPS
VT (Weissenbom 2020)
THVD--AN-FF (f al, 2020)
373M
103.3
94.0
SVG GHVAE FirVid
298M
1217.3
15.0
41.9
0.327
FilVid (Babacizzadh et al,, 2021)
302M
93.6
599M
552.9
15.8
51.2
0.286
MaskViT (ours)
189M
93.7
302M
884.5
17.1
49.1
0.217
MaskVIT (ours, goal cond.)
255M
76.9
MiskVIT
18IM
27.2
58.1
MiskViT (ours, act cond.)
255M
70.5
MIskVIT (256)
228M
446.1
26.2
40.7
0.270
Table 1:
Comparison with prior work. We evaluate Vis on BAIR, RoboNet and KITTI
datasets. Our method is competitive or outperforms prior work while being more parameter efficient.
4.1 EXPERIMENTAL SETUP
Implementation. Our transformer model stack of L blocks, where each block consists of two transformer layers with attention restricted to the window size of 1 x 16 x 16 (spatial window) and T' x (spatiotemporal window), unless otherwise specified. We use learnable positional embeddings, which are the sum of space and time positional embeddings. See $ A.1 for architecture details and hyperparameters.
Metrics. We use four evaluation metrics to compare our method with prior work: Fréchet Video Distance (FVD) (Unterthiner et al., 2018), Peak Signal-to-noise Ratio (PSNR), Structural Simi- larity Index Measure (SSIM) (Wang et al., 2004) and Learned Perceptual Image Patch Similarity (LPIPS) (Zhang et al., 2018). To account for the stochastic nature of video prediction, we follow prior work (Babaeizadeh et al., 2021; Villegas et al., 2019) and report the best SSIM, PSNR, and LPIPS scores over 100 trials for each video. For FVD, we use all 100 with a batch size of 256. We only conducted trial per video for evaluating performance the BAIR dataset (Ebert al., 2017).
4.2 COMPARISON WITH PRIOR WORK
BAIR. We first evaluate our model on the BAIR robot pushing dataset (Ehert et al., 2017), one of the most studied video modeling datasets (Fig. 7). We follow the evaluation protocol prior works and predict 15 video frames given 1 context frame and no actions. The lack of action conditioning makes this task extremely challenging and tests the model's ability to predict plausible future robot trajectories and object interactions. MaskViT achieves similar performance to FitVid (Babaeizadeh et al., 2021) while being more parameter efficient, and it outperforms all other prior works (Table l). Finally, predict cittonncondiiioned future frames, linearly project the action vectors and add them to Z. As expected, action conditioning performs the best, with 25% improvement in FVD.
KITTL. The KITTI dataset (Geiger et al., 2013) is relatively small dataset of 57 training videos. We follow the evaluation protocol of prior work (Villegas et al., 2019) and predict 25 video frames given 5 context frames. Compared other datasets in our evaluation, KITTI especially challenging, as it involves dynamic backgrounds, limited training data, and long-horizon predictions. We use color jitter and random cropping data augmentation for training VQ-GAN and do not use any data augmentation for training the second stage. Across all metrics, we find that Mask is significantly better than prior works while using fewer parameters (Table 1). Training a transformer model with full self-attention would require prohibitively large GPU memory due to the long prediction horizon (30 x 16 16 7680). However, Vis can attend to all tokens because windows significantly reduce the size of the attention context. We also report video prediction results for the KITTI dataset at 256 x 256 resolution (Fig. 2, 8), a higher resolution that prior work was not able obtain.
spatiotemporal
RoboNet. RoboNet (Dasari et al., 2019) is a large dataset of 15 million video frames of 7 different robotic arms interacting with objects and provides 5 dimensional robot action annotations. We follow the evaluation protocol of prior work (Babaeizadeh et al., 2021) and predict 10 video frames given 2 context frames and future actions. At 64 x 64 resolution, Mask is competitive but does

---

Published as a conference paper at ICLR 2023
si window
FVDJ
train mem.
train time
mask ralio
FVD
blocks cmbd.
FVD4
16 x
100.6
5.4 GB
12.3 hr
0.75
189.3
768
96.6
96.6
7.0 GB
12.5 hr
0.90
124.
1024
94.2
8
93.7
7.9 GB
14.2
110.9
768
99.3
16
x x 16
16
96.6
11.6 GB
27.9 hr
0.98
214.4
1024
99.5
full self atin.
98.2
16.4 GB
40.3 hr
0.5-1
96.6
(a) Model size. Increasing em- hedding dim improves FVD.
(b) Spatiotemporal window size.
(c) Mask ratio. able masking works best.
Vari-
Smaller window size is faster, memory efficient, and achieves lower FVD scores.
ratio
Table 2: MaskVTT ablation experiments BAIR. We compare FVD scores ablate important design deci- sions with the default setting: blocks, 768 embedding dimension (embd. dim), x x 16 spatial window, 16 4 4 saptiotemporal (st) window, and variable masking ratio. Default setrings marked in blue
not outperform prior works (Table 1). FVD of the VQ-GAN reconstructions is lower bound for MaskViT. We found flicker artifacts in the VQ-GAN reconstructions, probably due to per-frame latents, resulting in a high FVD score of 121 for the VQ-GAN reconstructions. Mask ViT achieves FVD scores very close to this lower bound but performs worse than prior works due to temporally inconsistent VQ-GAN reconstructions. Finally, we also report video prediction results for the RoboNet dataset at 256 256 resolution (Fig. 2, 9).
of
4.3 FLEXIBLE CONDITIONING
Planning for long horizon tasks via Visual MPC can quickly become computationally intractable as the search complexity grows exponentially with the number planning steps. Goal conditioned planning can be especially helpful in reducing the search space by breaking down tasks into sub- goals (Nasiriany et al., 2019; Pertsch et al., 2020). Methods like FitVid (Babaeizadeh et al., 2021) and prior two-stage methods (Yan al., 2021), which rely on compression in both temporal and spa- tial domain, cannot predict goal conditioned video frames. In contrast, we can easily adapt Mask ViT predict goal-conditioned video frames by including the last frame in Ze- Goal conditioning signif- icantly (18% FVD improvement) improves performance without requiring any architectural change (Table 1). We believe our method could also be extended perform video interpolation but leave that investigation to future work.
4.4 ABLATION STUDIES
We ablate Mask VIT to understand the contribution of each design decision with the default settings: 6 blocks, 768 embedding dimension, 1 16 x 16 spatial window, 16 4 x 4 spatiotemporal win- dow, and variable masking ratio (Table 2). See $ B.1 and $ B.2 for additional ablation studies and discussion.
Model hyperparameters. We compare the effect of the number of blocks and the embedding dimension Table 2a. We find that having larger embedding dimension improves the performance slightly, whereas increasing the number of blocks does not improve the performance.
Spatiotemporal window (STW). An important design decision of Mask ViT is the size of the STW (Table 2b). We compare three different window configurations and Mask VIT with full self-attention all layers. Note that training model with full self-attention requires using gradient checkpointing. which significantly increases the training time. In fact, the STW size of 16 x 4 x achieves better accuracy, while requiring 60% less memory and speeds up training time by 3.3x.
Masking ratio during training. We find that fixed masking ratio results in poor video prediction performance (Table 2c). A large masking ratio until a maximum of 95% decreases FVD. Further increase in masking ratio significantly deteriorates the performance. A variable masking ratio per- forms best, as it best approximates the different masking ratios encountered during inference.
Mask scheduling, The choice of mask scheduling function and the number of decoding iterations during inference has a significant impact on image (Chang et al., 2022) and video generation quality (Fig. 3). We compare three types of scheduling functions: concave (cosine, square, cubic, exponen-

---

Published as a conference paper at ICLR 2023
Exponential
Cubic
Square
Cosine
Linear
Square root
1.0
500
400
0.8
400
300
0.6
300
RVO
4.5
3.0
200
200
0.2
100
100
0.2
0.4
0.6
0.8
1.0
10
15
20 25
10
15
20
25
30
t/T
Decoding iterations
Decoding iterations
Figure 3: Mask scheduling functions. Left: categories of mask scheduling functions: concave (cosine, square, cubic, exponential), linear, and convex (square root). Middle: FVD scores for differ- ent mask scheduling functions and decoding iterations. Concave functions perform the best. Right: FVD score vs. decoding iterations for different temperature values. Lower and higher temperature values lead to poor FVD scores, with sweet spot temperature value of 3 and 4.5.
tial), linear, and convex (square root). Concave functions performed significantly better than linear and convex functions. Videos have a lot of redundant information due temporal coherence, and consequently with only 5% of unmasked tokens (Table 2c) the entire video can be correctly pre- dicted. The critical step is to predict these few tokens very accurately. We hypothesize that concave functions perform better they capture this intuition by slowly predicting the initial tokens over multiple iterations and then rapidly predicting the majority of remaining tokens conditioned on the (more accurate) initial tokens in the final iterations. Convex functions operate in opposite manner and thus perform significantly worse. Across all functions, FVD improved with increased numbers of decoding steps until certain point. Further increasing the decoding steps did not improve FVD. Additionally, we found that selecting the most confident tokens while performing iterative decod- ing led video predictions with little noise to the token confidence to encourage the model to produce more diverse outputs (Fig. 3). Empirically, we found that temperature value of 4.5 works best across different datasets.
motion. Hence, add temperature annealed Gumbel
4.5 VISUAL MODEL PREDICTIVE CONTROL WITH MASKVIT ON A REAL ROBOT
In this section, we test how our gener- ally designed framework performs when applied real robotic tasks. We eval- uate the capability of our method on control of embodied agents through Sawyer robot arm. We first train our model on the RoboNet dataset along with small collection of random interaction trajectories in our setup. We then leverage MaskViT to perform visual model-predictive control, and evaluate the robot's performance on several ma- nipulation tasks.
auto reg. fwel. pass
ours fwdd
dataset
pred frames
speed up
BAIR
3,840
24
160x
BAIR w/ act.
3,840
12
320x
KITTI
6,400
48
133x
experimental evaluation
RobuNet
10
2,560
5
512x
Table 3: Inference speedup of MaskViT over auto- regressive generation as measured by the number of for- ward passes. Iterative decoding in Mask can predict video frames in significantly fewer forward passes, espe- cially when conditioned actions.
Setup and data collection. We found that the VQGAN model trained only on RoboNet data did not generalize variations in lighting, and background, or to novel objects. Therefore, we au- tonomously collect 120K frames of additional data with random policy to finetune VQGAN to our setup by augmenting RoboNet. We hypothesize that this dependency on domain-specific data could be substantially reduced with large-scale pre-training of VQ-VAEs on internet-scale data as shown by prior work (Ramesh et al., 2021), 1), but beyond the scope of our experiment. At each timestep, the robot takes 5-dimensional action representing desired change in end-effector pose: [x, y, gripper position, 0 yaw angle, and a binary gripper open/close command. During data collection, the robot interacts with diverse collection of random household objects (Fig. 4).
8

---

Published as a conference paper at ICLR 2023
method
success
Mis (all data)
67%
FitVid (all data)
63%
Misk (finetuned)
53%
MiskViT (RN only)
6%
Random policy
3%
Figure 4: Left: Third person view of real-world data collec- tion. Right: An example evaluation task. The overlaid white arrow depicts the goal location of the green bowl.
Table 4: Control evaluation results. We perform 30 trials for each method, and report aggregated success rates.
Model predictive control using MaskViT. We evaluate the planning capabilities of Mask VIT using visual foresight (Finn & Levine, 2017; Ebert al., 2018) on series of robotic pushing tasks. Our control evaluation contains two task types: table e sett i ing which involves pushing a bowl a specified location, and sweeping, where objects are moved into an unseen dustpan. For each task, the robot is given a 64 x 64 goal image and we perform planning based on MaskViT by optimizing a sequence of actions using the cross-entropy method (CEM) (Boer al., 2004), using the &g pixel eтror between the last predicted image and the goal as the cost. We compare with two baseline methods: a Vid model trained on combined RoboNet and the domain-speciffe dataset, and random Gaussian policy. We evaluate three variants of our model: one trained on the combined RoboNet and -emain-speccfic datasets (all data), one pretrained using RoboNet and then finetuned with domain-specific data ffinetunedo and one trained only RoboNet (RN only). See $ A.2 for additional details and hyperparameters.
Results. As shown in Table 4, our model improves performance compared FitVid on real robot experiments. Our model achieves slightly better performance when trained jointly on all data opposed to pre- pre-training found that domain-specific data significantly improves the scene fidelity and arm motion prediction accuracy for model, resulting in better planning performance.
RoboNet and then finetuning real robot data. Qualitatively,
We highlight two advantages our method when applied to real robot tasks compared to prior methods. The first advantage the computation efficiency during inference time. Prior two-stage methods (Rakhimov et al., 2020; Yan et al., 2021) cannot be used for real-time planning due to extremely slow inference speed. Furthermore, compared to autoregressive models, ViT is orders of magnitude more efficient as shown in Table 3. Indeed, our method achieves ~ 6.5 seconds per CEM iteration. Second, compared to CNN-based architectures such as FitVid. which is the state- of-the-art method on this task, our method does not require any Nomain-speciff inductive biases and allows for a more general inference procedure, without sacrificing any performance.
5 CONCLUSION
In this work, we explore Mask a simple but powerful and versatile method for video prediction that leverages masked visual modeling as pre-training task and transformers with window atten- tion as a backbone for computation efficiency. We showed that by masking variable number of tokens during training, we can achieve competitive video prediction results. Our iterative decoding scheme is significantly faster than autoregressive decoding and achieves state-of-the-art performance in planning for real robot manipulation tasks.
Limitations and future work. While our results are encouraging, we found that using per frame quantization can lead to flickering artifacts, especially in videos that have a static background like in RoboNet. Although MaskViT is efficient in terms of memory and parameters, scaling up video prediction, especially for scenarios that have significant camera motion (e.g., self-driving (Geiger et al., 2013) and egocentric videos (Grauman et al., 2021)) remains challenging. Finally, im- portant future avenue of exploration is scaling up the complexity of robotic tasks (Srivastava et al., 2022) integrating our video prediction method in complex planning algorithms.

---

Published as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
We use an open-source implementation of VQ-GAN (ottps://githnn) com/ caom/CoopVisss for all our experiments. We use the default parameters specified in original implementation unless specified in Table A.I. We use PyTorch (Paszke et al., 2019) 1.7 library for implementing MaskViT. Architecture details and hyperparameters are described in $ 3 and Table A. respectively. The inference parameters used for iterative decoding for different datasets provided in Table A.1. Finally, please refer to the detailed description of our real robot data collection setup and experiments in $ A.2.
caming-trans
ACKNOWLEDGMENTS
This work is in part supported by ONR MURI N00014-22-1-2744 and the Stanford Human-Centered Institute for AI (HAI).
REFERENCES
Adil Kaan Akan, Erkut Erdem, Aykut Erdem, and Fatma Güney. Slamp: Stochastic latent ap- pearance and motion prediction. In Proceedings the IEEE/CVF International Conference Computer Vision, pp. 14728-14737, 2021.
Adil Kaan Akan, Sadra Safadoust, Erkut Erdem, Aykut Erdem, and Fatma Güney, Stochastic video prediction with structure and motion. arXiv preprint arXiv:2203.1055 2022.
Mohammad Babaeizadeh, Chelsea Finn, Dumitru Erhan, Roy H. Campbell, and Sergey Levine. Stochastic variational video prediction. In International Conference on Learning Representations, 2018
Mohammad Babaeizadeh, Mohammad Taghi Saffar, Suraj Nair, Sergey Levine, Chelsea Finn, and Dumitru Erhan. Fitvid: Overfitting in pixel-level video prediction. arXiv preprint arXiv:2106.15 13195, 2021.
Hangbo Bao, Li Dong. Songhao Piao, and Furu Wei. BEit: BERT pre-training of image transform- In International Conference Learning Representations, 2022.
De Boer, Kroese, S. Mannor, and R. Y. Rubinstein. A tutorial on the cross-entropy method. Annals of Operations Research, 134(1):19-65 67, 2004.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. In Advances in Neural Information Processing Systems, 2020.
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. arXiv preprint Stti::20020022000 2022.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and llya Sutskever. Generative pretraining from pixels. In International Conference on Machine Learning, pp. 1691- 1703. PMLR. 2020.
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: /904. 10509, 2019.
Aidan Clark, Jeff Donahue, and Karen Simonyan. Adversarial video generation complex datasets. arXiv preprint arXiv:/907.066
2019.
Sudeep Dasari, Frederik Ebert, Stephen Tian, Suraj Nair, Bernadette Bucher, Karl Schmeckpeper, Siddharth Singh, Sergey Levine, and Chelsea Finn. Robonet: Large-scale multi-robot learning. arXiv preprint arXiv: 1910. 11215, 2019.
Pieter- Tjerk De Boer, Dirk P Kroese, Shie Mannor, and Reuven Rubinstein. tutorial on the cross-entropy method. Annals of operations research, 134(1):19-67, 2005.
10

---

Published as a conference paper at ICLR 2023
Emily Denton and Rob Fergus. Stochastic video generation with a learned prior. In International Conference on Machine Learning, pp. 1174-1183. PMLR, 2018.
Jacob Devlin, Ming-War Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the Conference of the North American Chapter of the Association for Compulational Linguistics: Human Language Technologies, 2019.
Michael Dorkenwald, Timo Milbich, Andreas Blattmann, Robin Rombach, Konstantinos G Derpa- nis, and Bjorn Ommer. Stochastic image-to-video synthesis using cinns. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 3742-3753 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 11999, 2020.
Frederik Ebert, Chelsea Finn, Alex X and Sergey Levine. Self-supervised visual planning with temporal skip connections. In Conference on Robot Learning, pp. 344-356, 2017.
Frederik Ebert, Chelsea Finn, Sudeep Dasari, Annie Xie, Alex Lee, and Sergey Levine. Visual fore- sight: Model-based deep reinforcement learning for vision-based robotic control. arXiv preprint /8Xiv: 2018.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recog- nilion, pp. 12873-12883, 2021.
Christoph Feichtenhofer, Haoqi Fan, Yanghao Li, and Kaiming He. Masked autoencoders as spa- tiotemporal learners. arXiv preprint arXiv:2205.00 09.13, 2022.
Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. In IEEE Inter- national Conference on Robotics and Automation, 2017.
Chelsea Finn, Ian Goodfellow, and Sergey Levine. Unsupervised learning for physical interaction through video prediction. In Advances in Neural Information Processing Systems, 2016.
Andreas Geiger, Philip Lenz, Christoph Stiller, and Raquel Urtasun. Vision meets roboties: The kitti dataset. The International Journal of Robotics Research, 32(11):1231-2237,22 2013.
Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoye. Mask-predict: Parallel de- coding of conditional masked language models. In Proceedings of the 2019 Conference on Em- pirical Methods in Natural Language Processing and the 9th International Joint Conference Natural Language Processing (EMNLP-JJCNP)), pp. 61 1211122, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.1 1t..1653/v1/D1....
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Infor- mation Processing Systems, 2014.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint htt.::/70.0227.7 2017.
Kristen Grauman, Andrew Westbury, Eugene Byrne, Zachary Chavis, Antonino Furnari, Rohit Gird- har, Jackson Hamburger, Hao Jiang, Miao Liu, Xingyu Liu, et al. Ego4d: Around the world in 3,000 hours of egocentric video. arXiv preprint arXiv:21 3, 2021.
Jiatao Gu, James Bradbury, Caiming Xiong, Victor OK Li, and Richard Socher. Non-autoregresssvv neural machine translation. arXiv preprint 1e1Iiv11711000
2017.
Tanmay Gupta, Dustin Schwenk, Ali Farhadi, Derek Hoiem, and Aniruddha Kembhavi. Imagine this! scripts to compositions to videos. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 598 613, 2018
11

---

Published as a conference paper at ICLR 2023
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. arXiv preprint arXiv:2111.06377,2000 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in
Neurul Information Processing Systems, 33:6840-6851, 2020.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. arXiv preprint arXiv:2204.03458, 2022.
Justin Johnson, Alexandre Alahi, and Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In European Conference on Computer Vision, pp. 694-711. Springer, 2016.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR (Poster), 2015. URL http: wttpvv.viii..///ev//...e/...aava www.iii.ooo org/ abs 1112. 69800.80.......
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representa- tions for reinforcement learning. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume I 119 of Proceedings of Machine Learning PMLR, 13-18 Jul 2020.
Research, pp.
Alex X Lee, Richard Zhang, Frederik Ebert, Pieter Abbeel, Chelsea Finn, and Sergey Levine. Stochastic adversarial video prediction. arXiv preprint arXiv: 2018.
Yitong Li, Martin Min, Dinghan Shen, David Carlson, and Lawrence Carin. Video generation from text. In Proceedings of the AAAI Conference on Artificial Intelligence, 2018
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 10012-10022, 2021a.
Ze Liu, Jia Ning, Yue Cao, Yixuan Wei, Zheng Zhang, Stephen Lin, and Han Hu. Video swin transformer. arXiv preprint arXiv:2106. 13230, 2021b.
William Lotter, Gabriel Kreiman, and David Cox. Deep predictive coding networks for video pre- diction and unsupervised learning. arXiv preprint arXiv: 1at..1000000.4.. 2016.
Jiasen Lu, Dhruv Batra, Devi Parikh, and Stefan Lee. Vilbert: Pretraining task-agnostic visiolin- guistic representations for vision-and-langee lan-uang tasks. Advances in neural information processing systems, 32, 2019.
Pauline Luc, Aidan Clark, Sander Dieleman, Diego de Las Casas, Yotam Doron, Albin Cassirer, and Karen Simonyan. Transformation- based adversarial video prediction on large-scale data, arXiv preprint arXiv:2003.0403: 2020.
Anusha Nagabandi, Kurt Konoglie, Sergey Levine, and Vikash Kumar. Deep Dynamics Models for Learning Dexterous Manipulation. In Conference on Robot Learning (CoRL), 2019.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: A univer- sal visual representation for robot manipulation. arXiv preprint arXiv:2203./2601, 2022.
Charlie Nash, João Carreira, Jacob Walker, lain Barr, Andrew Jaegle, Mateusz Malinowski, and Peter Battaglia. Transframer: Arbitrary frame prediction with generative models. arXiv preprint 2:X1::2003.09494,, 2022.
Soroush Nasiriany, Vitchyr Pong, Steven Lin, and Sergey Levine. Planning with goal-conditioned policies. Advances in Neural Information Processing Systems, 32, 2019.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 8162-8171. PMLR, 2021.
Yingwei Pan, Zhaofan Qiu, Ting Yao, Houqiang Li, and Tao Mei. To create what you tell: Generat- ing videos from captions. In Proceedings of the 25th ACM International Conference on Multime- dia, pp. 1789-1795 2017.
12

---

Published as a conference paper at ICLR 2023
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsurprising effectiveness of pre-trained vision models for control. arXiv preprint arXiv:2203.03550, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury. Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Karl Pertsch, Oleh Rybkin, Frederik Ebert, Shenghao Zhou, Dinesh Jayaraman, Chelsea Finn, and Sergey Levine. Long-horizon visual planning with goal-conditioned hierarchical predictors. Ad- vances in Neural Information Processing Systems, 33:17321 1112211-113,,,3 2020.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Ruslan Rakhimov, Denis Volkhonskiy, Alexey Artemov, Denis Zorin, and Exgeny Burnaev. Latent video transformer. arXiv preprint arXiv:2006. 10704, 2020.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning. pp.
PMLR, 2021.
MarcAurelio Ranzato, Arthur Szlam, Joan Bruna, Michael Mathieu, Ronan Collobert, and Sumit Chopra. Video (language) modeling: baseline for generative models of natural videos. arXiv preprint arXiv: 141.::1411.6604 2014.
Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq- In Advances in Neural Information Processing Systems, 2019.
Chitwan Saharia, William Chan, Saurabh Saxena, and Mohammad Norouzi. Non-autoreeresssiv machine translation with latent alignments. arXiv preprint arXiv:2004.07437 2020.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learn- ing. pp. 2256-2265. PMLR, 2015.
Sanjana Srivastava, Chengshu Li, Michael Lingelbach, Roberto Martín-Martín, Fei Xia, Kent Elliott Vainio, Zheng Lian, Cem Gokmen, Shyamal Buch, Karen Liu, et al. Behavior: Benchmark for everyday houschold activities in virtual, interactive, and ecological environments. In Conference on Robot Learning, pp. 477-490, PMLR, 2022.
Chen Sun, Austin Myers, Carl Vondrick, Kevin Murphy, and Cordelia Schmid. Videobert: A joint model for video and language representation learning. In Proceedings of the IEEE/CVF Interna- tional Conference on Computer Vision, pp. 7464-7473, 2019.
JUN Tanji and Edward V Evarts. Anticipatory activity of motor cortex neurons in relation to direc- tion of an intended movement. Journal of neurophysiology. 1065 1976.
Zhan Tong, Yibing Song, Jue Wang, and Limin Wang. Videomae: Masked autoencoders are data- efficient learners for self-supervised video pre- pre-training. arXiv preprint arXiv:2203.12602. 2022.
Sergey Tulyakov, Ming-Yi Yu Liu, Xiaodong Yang, and Jan Kautz. Mocogan: Decomposing motion and content for video generation. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 1526-1535, 2018.
Thomas Unterthiner, Sjoerd van Steenkiste, Karol Kurach, Raphael Marinier, Marein Michalski. and Sylvain Gelly. Towards accurate generative models of video: A new metric & challenges. arXiv preprint arXiv:18/2.017: 2018
Aaron Van Den Oord and Oriol Vinyals. Neural discrete representation learning. In Advances in Neural Information Processing Systems, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances Neural Infor- mation Processing Systems, 2017.
13

---

Published as a conference paper at ICLR 2023
Ruben Villegas, Arkanath Pathak, Harini Kannan, Dumitru Erhan, Quoe V Le, and Honglak Lee. High fidelity video prediction with large stochastic recurrent neural networks. Advances in Neural Information Processing Systems, 32, 2019.
Pascal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre-Antoine Manzagol. Extracting and composing robust features with denoising autoencoders. In Proceedings of the 25th International Conference on Machine Learning, pp. 1096-1103, 2008.
Vikram Voleti, Alexia for prediction, generation, and interpolation. arXiv preprint arXiv:2205.098 099533, 2022.
and Christopher Pal. Masked conditional video diffusion
Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE Transactions on Image Processing, 13(4:)00 612, 2004.
Dirk Weissenbomn, Oscar Täckström, and Jakob Uszkoreit. Scaling autoregressive video models. In International Conference on Learning Representations, 2020.
Daniel M Wolpert, Zoubin Ghahramani, and Michael Jordan. An internal model for sensorimotor integration. Science, 269(5232) 1995.
Bohan Wu, Suraj Nair, Roberto Martin-Martin, Fei-Fei, and Chelsea Finn. Greedy hierarchi- cal variational autoencoders for scage video prediction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 2318-2328, 2021a.
Chenfei Wu, Jian Liang, Lei Ji, Fan Yang, Yuejian Fang, Daxin Jiang, and Nan Duan. NUWA: Visual synthesis pre-training for neural visual world creation. arXiv preprint arXi::2111.12417, 2021b.
Tingting Wu, Alexander J Dufford, Melissa-Ann Mackie, Laura J Egan, and Jin Fan. The capacity of cognitive control estimated from a perceptual decision making task. Scientific reports, 6(1); 1-11, 2016.
Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik. Masked visual pre-training for motor control. arXiv preprint arXiv:2203.06173. 2022.
Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. Videogpt: Video generation using vq-vae and transformers. arXiv preprint arXiv:2 10157, 2021.
Rowan Zellers, Ximing Lu, Jack Hessel, Youngjae Yu, Jae Sung Park, Jize Cao, Ali Farhadi, and Yejin Choi. Merlot: Multimodal neural script knowledge models. Advances in Neural Information Processing Systems, 14.23634-23663 2021.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 586-595, 2018.
14

---

Published as a conference paper at ICLR 2023
A IMPLEMENTATION DETAILS
A.1 TRAINING MASKVIT
VQ-GAN. We train a VQ-GAN (Esser et al., 2021) model for each dataset which downsamples each frame into 16 X 16 latent codes, i.e., by factor of 4 for frames of size 64 x 64 frames and 16 for frames of size 256 x 256. Table 5 summarizes our settings for all three datasets. Training VQ-GAN with discriminator loss can lead to instabilities. Hence, suggested by (Esser et al., 2021) we start GAN losses after the reconstruction loss has converged. We also found that GAN losses were not always helpful, especially at lower input resolutions for BAIR and RoboNet.
Transformer: Our transformer model stack L blocks, where each block consists of two transformer layers with attention restricted to the window size of 1 x 16 16 (spatial window) and T x (spatiotemporal window), unless otherwise specified. We use learnable positional embeddings, which are the of space and time positional embeddings. Following (Liu et al., 2021a), we adopt relative position biases in our layers. We use the Adam (Kingma & Ba, 2015) optimizer with linear warmup (Goyal et al., 2017) and cosine decay learning rate schedule, Table summarizes our settings for all three datasets.
Evaluation. We find the optimal evaluation parameters by doing grid search of the following parameters: (cosine, square), temperature (3, 4.5) and decoding iterations depending on the pre- diction horizon length. We use top-p value of 0.95 for the BAIR dataset only. Table summarizes our evaluation settings.
Dataset
BAIR
KITTI
KITTI
RoboNet
RoboNet
Image resolution
64
64
256
64
256
Context frames
VQ-GAN
Channel
160
128
128
192
128
K
1024
1024
1024
1024
1024
256
256
256
256
256
Batch size
320
1120
112
720
112
Training steps
3e5
5e4
3e5
3e5
3e5
Learning rate
le-4
1e-3
le-4
5e-4
le-4
Disc. start
2e4
1.5e5
1.5e5
Transformer
Spatial window Spatiotemporal window Blocks
x x
1 16 x 16
1 16 16
1 16 16
1 16
16
16
8
16 x 4 4
16
x4x4
16
4 4
16 x
4
4
Attention heads
4
4
4
Embedding dim.
768
768
1024
768
1024
Feedforward dim.
3072
3072
4096
3072
4096
Dropout
0.0
0.0
0.0
0.0
0.0
Batch size
64
32
32
224
224
Learning rate
3e-4
3e-4
3e-4
3e-4
3e-4
Training steps
le5
le5
le5
3e5
3e5
Evaluation
Mask scheduling func. Decoding iters.
square 18
cosine 48
cosine
cosine
cosine 16
64
Temperature
4.5
3.0
4.5
Table 5: Training and evaluation hyperparameters.
15

---

Published as a conference paper at ICLR 2023
A.2 REAL ROBOT EXPERIMENTS
Data collection. Our robot setup consists of a Sawyer robot arm with a Logitech C922 PRO con- sumer webcam for recording video frames at 640 480 resolution. All raw image observations are center-cropped to 480 x 480 resolution before being resized to 64 x 64 for model training and control in our experiments. For the finetuning dataset, we autonomously collect 5000 trajectories of 30 timesteps. At cach step the robot takes 5-dimensional action representing a change in state of the end-effector: a delta translation in Cartesian space, [xr, y, 2]. for the gripper position in meters, change in 0 yaw angle of the end-effector, and a binary gripper open/close command. Following the action space used to collect the RoboNet dataset, the pitch and roll of the end-effector are kept fixed such that the gripper points with the fingers towards the table surface. Each action within a trajectory is selected independently and each dimension of the action vector is independent of the others, sampled from diagonal Gaussian distribution, except the gripper open/close command that closes automatically when the z-position of the end-effector reaches below a certain thresh- old to increase the rate of object interaction. The random action distribution is parameterized by (0, diag([0.035, 0.035, 0.08, T/ /88 2)). During data collection, we provide the robot with a di- verse set of training objects to interact with. During evaluation, we test on tasks which require the robot to manipulate bowls in one setting and to push training items into an unseen dustpan in another.
Visual-MPC. Our control strategy is visual MPC (Finn & Levine, 2017; Ebert et al., 2018) procedure. Given a start and goal image Lo, Lg E R64 the objective is to find an optimal sequence of actions to reach the goal observation from the start. The planning objective can be written as: minay, "2..... CalffTT 01 .. T where f(Tn, a1 (1:--01)) represents the ith predicted frame by the learned video prediction model (Mask ViT in our case), and e; are a sequence of constant hyperparameters that determine the importance of the difference between the predicted frame and the goal for each time step.
We use the cross-entropy method (CEM) (De Boer al., 2005) to optimize a sequence of H = 10 future actions for this objective. In each planning iteration, we first sample M = 256 sequences of random actions. We then provide these sequences, together with two consecutive context frames (the previous and current step observations), and one context action (the action taken at the previous step) to Mask Action sequences are sampled according to multivariate Gaussian distribution. To bias action sampling towards smoother trajectories, the noise samples for actions in given random trajectory are correlated across time as in Nagabandi et al. (2019). Specifically, given a
correlation coefficient hyperparameter 3, we first compute u/, u², ... M i.i.d (0, Σ.), where D; is the variance of the action at timestep i in the current optimization iteration. The noise at timestep i for the jth random trajectory. re/, is then computed as a weighted combination of the new noise sample and the noise sample the previous timestep, that is, ni (1 B) uf 3 * ni_ After all noise samples are computed, they are summed with means μ4; for each timestep, which are also iteratively updated. The final random trajectories are formed by rounding the elements in the last action dimension (gripper action) to - or 1, whichever is closer.
Next, compare the predictions to the goal image by computing the la error and summing over time as described by the objective above. We weight the cost on the final timestep by 10x, but still include the costs on intermediate timesteps in the summation task quickly. The best action sequences based on this score are used to refit the sampling distribu- tion mean and variance for the next optimization iteration. After K = optimization iterations, we execute the best scoring action sequence on the robot for the first 3 steps before performing replanning.
encourage the robot to solve the
The robot uses a total of 15 steps to solve the task, including one initial action [0, 0, -0.08, 0.1, 0] which is executed at the beginning of every trajectory. This ensures that at least two context images provided for planning. The planning hyperparameters are summarized in Table 6.
FitVid model training details. For the FitVid comparison, we train FitVid on a combined dataset consisting of RoboNet and our finetuning data, just when training our model in the all data case. Like in the original paper, we train the model to predict 10 future frames given 2 context frames, and use the Adam optimizer with learning rate of 3e - 4. To be consistent with the training setup for our model, we do not use data augmentation. We train for 245K gradient steps using a batch size of 32 with 4 NVIDIA TITAN RTX GPUs
16

---

Published as a conference paper at ICLR 2023
Hyperparameter
Value
Total trajectory length (I')
15
Planning horizon
10
Number of steps between replanning
3
Action dimension
5
# of samples per CEM iteration (M)
256
of CEM iterations (K)
3
Weights on each timestep in cost (c;)
1 if i = O....; 10 i
Initial sampling distribution mean Initial sample distribution std.
[0. 0. 0.5,
[0.05, 0.05, 0.08, x/18,2
Sampled noise correlation coefficient (1))
0.3
CEM fraction of clites
0.05
ViT mask scheduling function
Cosine
MaskVIT decoding iterations
5
Table 6: Hyperparameters for visual-MPC.
Evaluation. We perform control evaluation on two categories of tasks: table setting and sweeping. For each task, we test 5 different variations with 3 trials each. A trial is considered successful if the center of the object of interest is within 8 cm of the goal position after the trajectory complete. Model inference for real robot control is performed using NVIDIA RTX 3090 GPUs with a batch size of 16 per GPU. We use decoding iterations, which yields a forward pass time of approximately 6.2 seconds for a batch of 256 samples.
B ADDITIONAL RESULTS
B. QUANTITATIVE RESULTS
Factorized attention. An important decision decision of MaskViT factorizing full-self attention into a block comprised of spatial and aptio-temporal window layers. This design ensures that each token can attend to all other tokens just as in a full-self-attention layer but without huge memory requirements. In Table 8, we compare Misk with Transformer model composed entirely of STW layers. For all window sizes, the FVD scores are worse than our default block design. We note that a higher spatial resolution helps, which showcases the importance of our block design.
VQ-GAN training dataset. A key benefit of operating on individual frames is that we can poten- tially use VQ-VAF trained internet scale datasets (Ramesh al., 2021). As a step towards that goal, we train a VQ-GAN model jointly on all three datasets considered in this work, i.c. BAIR, RoboNet and KITTI. During training, we create batch consisting of equal number of frames from each dataset. We use the parameeersss for the RoboNet VQ-GAN model (Table 5). As shown Table 9, across the three data sets, the performance Mask VIT trained with a shared VQ-GAN is better or similar to Mask ViT trained with VQ-GAN specific to the dataset.
Real robot experiments. Table shows task success rates for our real-world robotic control evaluation. Each of our two task types (table setting, sweeping) has 5 variations, each of which in- volves different objects to push (unseen bowls for table setting, toys previously seen in the finetuning data for sweeping) and different target locations.
B.2 QUALITATIVE RESULTS
Video prediction. We present additional qualitative video prediction results for BAIR (Fig. 7), KITTI (Fig. 8) and RoboNet (Fig. 9).
Real robot experiments. Fig. 5 and Fig. 6 depict sample predictions for MaskVIT (all data)
and MiskVIT trained only on RoboNet (RN only) for two example control tasks. We observe that with our model, the planner is able to find sequences of actions which bring the blue bowl or the soft red hat close to the position specified in the goal image. However, with a model which is trained only the RoboNet dataset, planning fails. We qualitatively that the model trained in
17

---

Published as a conference paper at ICLR 2023
Task
MaskVIT (all data)
MaskVIT (finctuned)
MaskVIT (RN only)
Random
FirVid
table setring (bowl color; destination)
blue; front-left
1/3
1/3
0/3
0/3
2/3
green: back-right
1/3
0/3
0/3
0/3
0/3
bluc: front-right
3/3
2/3
red; front
3/3
3/3
0/3
0/3
3/3
green: left
3/3
2/3
0/3
0/3
2/3
sweeping (object; destination)
toys; back-right
3/3
hat; front-right
2/3
2/3
0/3
1,3
3/3
hat; front-left
1/3
1/3
1/3
0/3
2/3
toys; back-left
1/3
1/3
0/3
0/3
1/3
toys; front-left
2/3
1/3
0/3
0/3
1/3
Aggregated
20/30
16/30
2/30
1/30
19/30
Table 7: Per-task quantitative results for robotic control evaluation. We evaluate each of variants of the two tasks using 3 trials with each model or policy. Success is determined by the center of the object being within 11em of the goal position at the end of 15 steps.
dataset
VQ-GAN
FVD
BAIR
joint 99.8
st window
window
FVD4
BAIR
indy.
96.6
16 x 1
442.8
KITTI
joint 388.4
16 x 4 x
242.9
KITTI
indy. 401.9
16 x 8 8
126.6
RoboNet
joint 169.6
16 x 4 1x 16
16
96.6
RoboNet
indv.
133.5
Table 9: VQ-GAN training dataset. We com- pare using VQ-GAN models trained on indi- vidual datasets and a single VQ-GAN models trained BAIR, KITTI and RoboNet dataset. Across all three datasets the performance of Mask trained with shared VQ-GAN is bet- ter or similar to Mask trained with dataset specific VQ-GAN.
Table 8: Block design ablation on BAIR. We compare FVD scores for blocks consisting only spatiotemporal window (STW) with our block design. All STW sizes perform worse due to lack of global interactions across all tokens.
RoboNet-only, even when solely performing reconstruction of the first two context images using the VQ-GAN component, is unable reconstruct the background and robot arm with high fidelity. Despite the diversity of the RoboNet dataset, finetuning on domain-specific data still required to produce reasonable predictions our setting.
Discussion. We performed extensive experiments comparing different design choices: model size, window design, masking ratio, and mask scheduling functions. Qualitatively, found that all model sizes and different window designs work equally well, as also indicated by similar FVD scores. However, when our backbone consisted only of spatio-temporal windows, we found the performance to be much worse (Table 8). A key failure mode was lack of global consistency in the predicted video frames, e.g. the robot arm would generally disappear after the first few frames. We believe that this failure could be attributed to the lack of global interactions i.e. token could no longer attend to all other tokens. Models trained with a low masking ratio or without adding annealed noise had high FVD score because they showed very little arm. Finally, in comparison with other works, we found that the biggest qualitative difference was in the KITTI dataset. KITTI presents unique challenge, as we have predict 25 video frames. As shown in Fig. 2 and Fig. 8. Mask ViT can predict future video frames without perceptually visible degradation, whereas prior work predict frames that degrade rapidly.
motion of the robot
18

---

Published as a conference paper at ICLR 2023
B.3 LIMITATIONS AND FUTURE WORK
Although Mask addresses important limitations of current work in video prediction, we believe the field has yet witnessed progress similar text-to-image generation (Ramesh et al., 2021). Key areas of improvement are short-term temporal consistency, long-term temporal consistency, increasing the prediction horizon, and moving beyond narrow cemainn-peccfic datasets to "in-the- wild" settings like everyday egocentric videos (Grauman et al., 2021).
One possible way to improve short-term consistency is incorporating temporal consistency losses for VQ-GAN training. Despite using bidirectional window attention mechanism to reduce memory requirements, we are still limited to predicting 25 frames in the future. Possible ways to predict long- horizon videos include training multiple networks at different temporal resolutions and incorporating advances increasing the context length of Transformers. Finally, scaling both data and model parameters might be necessary to move more realistic and complicated everyday scenes.
19

---

Published as a conference paper at ICLR 2023
1-14
Figure 5: Visualizations for task set table: push blue bowl left.
t=6
t-
Goal image
Visunitration of predicted plan att founn GEM Horations
Visuliization of best predicted found after 3 CEM Ierallone
found
1-2
1=5
MaskV/T (RoboNet only)
MaskVIT (all data)
recondiructed context
reconsirucboo cloo context
context
Starting imae
Visuelization
Full
20

---

Published as a conference paper at ICLR 2023
1-14
push hat to bottom- right corner.
t=
t-
Figure 6: Visualizations for task sweep:
Goal image
Visunizration of predicted plan att Horations
found iterations
Visuelization of best predicted found after 3 CEM nerallone
1-2
1=5
MaskV/T (RoboNet only)
recondiructed context
reconsirucdod cdod context
context
MaskVIT (all data)
Starting imge
Full
21

---

Published as a conference paper at ICLR 2023
t= 15
t=
1
13
t= 13
Figure Qualitative results: BAIR (action free, 64 64)
1=11
t=11
t=11
1=11
t=
t=9
1=
t=
t=
t=
1=
t= 5
t=
t=3
t=3
t=1
t=1
(context)
(oontext)
(context)
(context)
t=
t= 0
t=
1= 0
Predicted
Predicted
Predicted
Predicted
Ground truth
Ground truth
Ground truth
Ground truth
22

---

Published as a conference paper at ICLR 2023
t= 29
t=
t 29
=
t
27
t=
t:
1
Figure 8: Qualitative results: KITTI (256 256)
t=
21
11
t 18
18
t=
t= 15
1 15
1 15
1
t= 12
t= 12
t= 12
1=9
t
1=9
tw6
t=6
1=4
1=4
(context)
(contexx)
(context)
(context)
0
t=
t=
t=0
Predicted
Predicted
Predicted
Predicted
Ground truth
Ground truth
Ground truth
Ground truth
23

---

Published as a conference paper at ICLR 2023
t= 11
1=11
1=11
1=11
t= 10
t=
t 10
t 10
t= 9
t=
t
Figure 9: Qualitative results: RoboNet (256 256)
[ 0
t
t=8
t=8
t=7
1=7
t=
t
t=5
t=4
1=4
t= 4
==3
t
==3
=
1#2
1=2
1#1
1=1
(context)
(context)
(context)
(context)
0
t=0
==0
Predicted
Predicted
Predicted
Predicted
Ground truth
Ground truth
Ground truth
Ground truth
24


Reviews:
[{"review_id": "oBkTLHIfA4l", "reviewer": "Reviewer_f78f", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work proposes a new framework for video prediction based on masked discrete token modeling. This work extends masked image modeling (MIM) into masked video modeling with two modifications. \n- First, reducing the complexity of the full attention between discrete tokens within the bidirectional transformer.\nThis model uses decoupled attention for spatial and temporal domains, to sparsify full attention for visual tokens.\n\n- Second, this work introduces a masking policy during training and a decoding policy for sampling to mitigate the inconsistency between training and sampling distribution. Especially for decoding policy, this model sample causally in time direction to support generating long video sequence.\n\nRaised concerns in the paper are verified in the experiment section. The proposed method shows comparable performance to the state-of-the-art methods. Interestingly, this work also demonstrates video prediction in a real robot control domain and achieves better performance compared to the off-the-shelf video prediction method.", "summary_of_the_review": "The proposed methods extend masked image modeling into the video domain by tackling raised problems during extension.\nThe proposed method mainly design two sub-components; mixing spatial & spatio-temporal attention for bidirectional transformer and iterative decoding policy for video prediction. The choices of architecture and method design are natural and intuitive. However, I conjecture that the rationales behind the choices are not explored well in the main experiments.", "strength_and_weaknesses": "## Strength\n### 1. a straightforward extension of masked image modeling into the video domain\n\nMasked image modeling (MIM) shows good performance on both representation learning via self-supervised learning and image generation. This work extends MIM into masked video modeling and effectively tackles raised problems caused by the nature of video domain.\n\n### 2. benchmark on real robot control\n\nCompared to the other video synthesis tasks such as video generation (no conditioned on frames), video interpolation or video temporal super-resolution, video prediction can be served as a good representation for setting like control problem or reinforcement learning. In that sense, I believe that demonstrating this kind of task is a contribution to the community of video synthesis research. \n\n## Weaknesses\n### 1. ablation study on spatial & spatial-temporal window size\n\nThe ablation study on the configuration of attention is weak. \n\n**comparison to only spatio-temporal attention (the most important ablation)**\n\nHow about stacking single spatio-temporal attention, instead of stacking spatial and spatio-temporal attention alternatively as in the proposed architecture? I am curious about those settings because adopting homogeneous attention block is already explored in the action recognition literatures. \n\n**more diverse setting in spatio-temporal window sizes**\nAbation study on Table 2 (b) only tested the window size from (16 x 4 x 4) to (16 x 16 x 16). \nSo, I'm curious about the case in below,\n- only temporal window (T x 1 x 1), this setting is crucial because it fully decompose temporal and spatial effect.\n- different temporal window sizes while fixing the entire token numbers\n   - For example, comparison with three settings (16 x 4 x 4), (4 x 8 x 8), (32 x 2 x 2) --> (last one is not available for BAIR benchmark)\n\n### 2. a paragraph for iterative decoding policies in detail\nI think that related work section (2) or iterative decoding subsection (3.4) needs a paragraph for iterative decoding policies in detail, such as MaskGIT (CVPR 22) or other relevant works. I think that comparison to relevant methods is a good way to explain the proposed one.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "cXJZtKSbQe", "reviewer": "Reviewer_G4sC", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper tackles the problem of video prediction: given a video, the aim is to predict the future frames. This is a challenging problem as it requires learning the dynamics of the scene content in addition to the content itself. The authors argue that the classical autoregressive paradigm is both inconsistent with masked pre-training and too slow for this task. \nThey propose a novel approach based on: 1) a  VQGAN encoding to reduce dimensionality, 2) alternating Transformer layers based on a combination of spatial and spatiotemporal windows, 3) iterative decoding to decode latent tokens to frames, 4) variable masking during training. \n\nThe approach is evaluated on two main datasets: KITTI (driving scenes), and RoboNet (robotic manipulation), and is demonstrated to produce better predictions than the state-of-the-art with a considerably faster inference time. The approach is also used for robot predictive control. ", "summary_of_the_review": "I enjoyed reading this paper. It tackles an interesting task and describes a clear progress on this task. The paper is well written, although I would have liked more discussion of the chosen methods and of the results and their implications (possibly some quantitative results could be instead pushed in appendices), and the results are convincing. \n", "strength_and_weaknesses": "### Strengths\n- This paper addresses an extremely challenging and important problem for computer vision and robotics\n- The evaluation is convincing and thorough, the improvements over the state of the art follow from a clear rationale and are well discussed. \n- The ablation study is well done. \n- The proposed method is a clear improvement over previously published approaches. \n- Applying the method to a robotic scenario, in addition to classical benchmark provides additional confidence that the method is robust. \n### Weaknesses\n- Although the paper is generally well written and clear, some critical aspects of the system could have been described and discussed in more details. For example 3.2 and 3.4 are essential contributions and would benefit from more explanation. Additional figures might have helped.\n- The paper lacks a discussion of success and failure cases, in comparison to other approaches (and with respect to the ablation study). Because of this the paper provides little intuition on how well the method performs, beyond the (impressive) quantitative results.  ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "9zA-BNHbSHj", "reviewer": "Reviewer_dKZs", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents a video generation model that iteratively decodes patches in future frames. The paper also presents novel architectural components. The results are generally strong and the application of the method to a robotics task is convincing. Ablations are also strong.", "summary_of_the_review": "I generally feel positively about the paper. The method is motivated by a need to reduce memory and improve inference time, the method seems methodologically novel, and the key contributions of the paper are empirically justified.\n\nI currently recommend for weak acceptance of the paper.\n\nI am most concerned about the training of different models for different datasets (see *W3*). I am also interested in the performance of using a pre-trained VQ-GAN instead of training one specifically for each video dataset (see *W6*). I feel these experiments will help provide context for the generality and scalability of the approach.\n\nI am willing to revisit my evaluation during the rebuttal period.\n\nPOST REBUTTAL:\nThe authors have answered all of my concerns and added a worthwhile investigation of a shared VQ-GAN encoder. Additionally, they made it more clear that their work extends original ideas that were applied to static images. The robot experiments greatly elevate this paper. Hence, I elect to raise my score from a 6 to an 8.\n", "strength_and_weaknesses": "**Strengths**\n\n*S1.* Key tricks that are required to get the technique to work are presented, which is extremely useful to a reader trying to learn from this paper.\n\n*S2.* Ablations are extensive and evaluate key design decisions.\n\n*S3.* Appx. give appropriate details when necessary hence improving the the quality of the main paper without sacrificing clarity on key details.\n\n*S4.* Real world robot experiments are a major strength, highlighting the applicability and usefulness of the method.\n\n*S5.* Experiments on goal conditioning on the last frame are eliminating.\n\n*S6.* Generality of the method to be applied to different datasets.\n\n**Weaknesses**\n\n*W1.* Missing some related work. I consider the masking techniques used in this paper as related to BERT-style masking that has been previously explored in the context of videos. Some relevant references to consider are Lu et al., 2019 (https://arxiv.org/abs/1908.02265), Sun et al., 2019 (https://arxiv.org/abs/1904.01766), and Zellers et al., 2021 (https://arxiv.org/abs/2106.02636). I recommend a more comprehensive review of recent related work to better contextualize the paper.\n\n*W2.* It is not clear early in the paper why \"iterative\" is necessarily better than \"auto-regressive.\" Auto-regression may also be considered an iterative process. I suggest being more precise here to give better intuition on why the proposed method is faster than existing auto-regressive methods.\n\n*W3.* Different models (e.g., VQ-GANs) are trained for different datasets, which each represent fairly narrow domains. I would be interested to see if the technique could be extended to training a single model on a large corpus of video data or at least the union of the datasets considered. How well does this model perform compared to the single models? I feel this experiment is important to elucidate the scalability of the method.\n\n*W4.* The intuition for why keeping the most confident frames leads to static videos is not clear to me. Given motion in the video ground truth, won't such predictions lead to very high loss? \n\n*W5.* Why is the evaluation protocol different for BAIR than the other video datasets (L194-196)?\n\n*W6.* Why train VQ-GAN from scratch, why not use the encoder trained on internet scale data mentioned in L291-292. Would using this tokenizer improve performance? I do not consider this to be outside the scope of the paper.\n\n**Minor**\n\n*M1.* I suggest changing the title as ViT usually refers to the Dosovitskiy et al., 2021 work (https://arxiv.org/abs/2010.11929), where \"Vi\" does not mean video. Hence the title could be a little misleading.\n\n*M2.* The paper claims that \"there is an inconsistency between the video prediction task and autoregressive masked visual pretraining – while the training process assumes partial knowledge of the ground truth future frames, at test time the model has to predict a complete sequence of future frames from scratch.\" However, this explanation might not be clear for someone without a lot of background knowledge on autoregressive techniques in video. I suggest giving an example or providing a reference here to make things more clear.\n\n*M3.* I feel a major strength of the method is to be able to condition prediction on arbitrary frames in a video. This naturally allows for goal conditioning for robot manipulation, which seems like a major plus. I recommend emphasizing this feature more in the writing earlier on, otherwise it may get lost.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4DuGtfu89I", "reviewer": "Reviewer_V2jW", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper addresses the video prediction problem where the model needs to predict the future frames (or tokens) conditioned on the first frame (context). This paper specifically focuses on addressing the high memory cost and inference efficiency. The method builds upon VQGAN to encode frames into quantized tokens and learn a transformer for predicting the missing tokens. The core components of the methods are \n1) spatial attention (attend only to the context frame) and spatially localized spatiotemporal attention window, \n2) iterative token decoding with variable mask ratios.\nExtensive results showcase improvement over the state-of-the-art on multiple datasets.", "summary_of_the_review": "I think the method is technically sound, the results are solid, and the writing is clear. However, all these improvements are simply validating the applicability of prior work MaskGit to video data. Without careful analysis and comparison against a simple extension of MaskGit for video, it's hard to evaluate the novelty/significance of the proposed method. ", "strength_and_weaknesses": "Strength:\n+ The paper is easy to read.\n+ The inference speedup and training memory reduction from iterative decoding and window attention make sense. The method is technically sound.\n+ The experimental results are extensive. The results demonstrate state-of-the-art performance on BAIR, RoboNet, and KITTI datasets.\n+ The results show inference speedup over models with autoregressive generation (Table 3).\n\nWeakness:\n- The primary concern I have about this paper lies in its technical novelty. The core component that significantly improves the inference speed is the iterative decoding with a bidirectional transformer. This is primarily based on the work of MaskGit [Chang et al. 2022]. The MaskGit also considers multiple masking designs (Sec 3.3 in the MaskGit paper). The paper's exploration of mask scheduling (Figure 4) confirms similar findings. The variable masking ratio is also explored in MaskGit as part of the mask schedule design.\n\n- MaskGit also considers the \"conditional generation\" tasks (e.g., image inpainting or extrapolation) where tokens are in the known regions. From this perspective, the proposed MaskViT is a simple adaptation of MaskGit on video data. I thus consider the novelty to be limited.\n\n- Aside from iterative decoding, another claimed contribution is the window attention. This is interesting, but it appears to be a simple hyperparameter change with a trade-off on the training time vs. quality. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster