MAGVIT: Masked Generative Video Transformer

Introduction

Recent years have witnessed significant advances in image and video content creation based on learning frameworks ranging from generative adversarial networks (GANs) , diffusion models , to vision transformers . Inspired by the recent success of generative image transformers such as DALL·E and other approaches , we propose an efficient and effective video generation model by leveraging masked token modeling and multi-task learning.

We introduce the MAsked Generative VIdeo Transformer (MAGVIT) for multi-task video generation. Specifically, we build and train a single MAGVIT model to perform a variety of diverse video generation tasks and demonstrate the model’s efficiency, effectiveness, and flexibility against state-of-the-art approaches. MAGVIT: Masked Generative Video Transformer(a) shows the quality metrics of MAGVIT on a few benchmarks with efficiency comparisons in (b), and generated examples under different task setups such as frame prediction/interpolation, out/in-painting, and class conditional generation in (c).

MAGVIT models a video as a sequence of visual tokens in the latent space and learns to predict masked tokens with BERT . There are two main modules in the proposed framework. First, we design a 3D quantization model to tokenize a video, with high fidelity, into a low-dimensional spatial-temporal manifold . Second, we propose an effective masked token modeling (MTM) scheme for multi-task video generation. Unlike conventional MTM in image understanding or image/video synthesis , we present an embedding method to model a video condition using a multivariate mask and show its efficacy in training.

We conduct extensive experiments to demonstrate the quality, efficiency, and flexibility of MAGVIT against state-of-the-art approaches. Specifically, we show that MAGVIT performs favorably on two video generation tasks across three benchmark datasets, including UCF-101 , BAIR Robot Pushing , and Kinetics-600 . For the class-conditional generation task on UCF-101, MAGVIT reduces state-of-the-art FVD from $332$ to $76$ ($\downarrow\!77\%$). For the frame prediction task, MAGVIT performs best in terms of FVD on BAIR ($84$ $\!\rightarrow\!62,\downarrow\!26\%)$ and Kinetics-600 ($16$ $\!\rightarrow\!9.9,\downarrow\!38\%)$.

Aside from the visual quality, MAGVIT’s video synthesis is highly efficient. For instance, MAGVIT generates a 16-frame 128$\times$128 video clip in $12$ steps, which takes $0.25$ seconds on a single TPUv4i device. On a V100 GPU, a base variant of MAGVIT runs at $37$ frame-per-second (fps) at 128$\times$128 resolution. When compared at the same resolution, MAGVIT is two orders of magnitude faster than the video diffusion model . In addition, MAGVIT is $60$ times faster than the autoregressive video transformer and $4$-$16$ times more efficient than the contemporary non-autoregressive video transformer .

We show that MAGVIT is flexible and robust for multiple video generation tasks with a single trained model, including frame interpolation, class-conditional frame prediction, inpainting, and outpainting, etc. In addition, MAGVIT learns to synthesize videos with complex scenes and motion contents from diverse and distinct visual domains, including actions with objects , autonomous driving , and object-centric videos from multiple views .

To the best of our knowledge, we present the first masked multi-task transformer for efficient video generation and manipulation. We show that a trained model can perform ten different tasks at inference time.

We introduce a spatial-temporal video quantization model design with high reconstruction fidelity.

We propose an effective embedding method with diverse masks for numerous video generation tasks.

We show that MAGVIT achieves the best-published fidelity on three widely-used benchmarks, including UCF-101, BAIR Robot Pushing, and Kinetics-600 datasets.

Preliminaries: Masked Image Synthesis

We take MaskGIT as an example. In the second stage, it applies a binary mask ${\textnormal{m}}_{i}\in\{x\!\rightarrow\!x,x\!\rightarrow\!\texttt{[MASK]}\}$ to each token to build a corrupted sequence $\overline{{\mathbf{z}}}={\mathbf{m}}({\mathbf{z}})$. Condition inputs, such as class labels, are incorporated as the prefix tokens ${\mathbf{c}}$. A BERT parameterized by $\theta$ is learned to predict the masked tokens in the input sequence $[{\mathbf{c}},\overline{{\mathbf{z}}}]$, where $[\cdot,\cdot]$ concatenates the sequences. The objective is to minimize the cross-entropy between the predicted and the ground-truth token at each masked position:

During training, MaskGIT randomly samples ${\mathbf{m}}$ from a prior distribution $p_{{\mathcal{U}}}$ where the mask ratio follows a cosine scheduling function $\gamma(\cdot)$ . Specifically, it first uniformly samples a per-token mask score ${\textnormal{s}}_{i}\sim{\mathcal{U}}(0,1)$ to form a sequence denoted as ${\mathbf{s}}$. Then it samples $r\sim{\mathcal{U}}(0,1)$ and computes a cut-off threshold $s^{*}$ as the $\lceil\gamma(r)N\rceil$-th smallest element in ${\mathbf{s}}$. Finally, a mask ${\mathbf{m}}$ is created such that ${\textnormal{m}}_{i}(x)=\texttt{[MASK]}$ if ${\textnormal{s}}_{i}\leq s^{*}$ and ${\textnormal{m}}_{i}(x)=x$ otherwise.

For inference, the non-autoregressive decoding method is used to synthesize an image . For example, MaskGIT generates an image in $K\!\!=\!\!12$ steps from a blank canvas with all visual tokens masked out. At each step, it predicts all tokens in parallel while retaining tokens with the highest prediction scores. The remaining tokens are masked and predicted in the next iteration until all tokens are generated. Similar to the training stage, the mask ratio is computed by the schedule function $\gamma$, but with a deterministic input as $\gamma(\frac{t}{K})$, where $t$ is the current step.

Masked Generative Video Transformer

Our goal is to design a multi-task video generation model with high quality and inference efficiency. We propose MAsked Generative VIdeo Transformer (MAGVIT), a vision transformer framework that leverages masked token modeling and multi-task learning. MAGVIT generates a video from task-specific condition inputs, such as a frame, a partially-observed video volume, or a class identifier.

The framework consists of two stages. First, we learn a 3D vector-quantized (VQ) autoencoder to quantize a video into discrete tokens. In the second stage, we learn a video transformer by multi-task masked token modeling.

Fig. 2 illustrates the training in the second stage. At each training step, we sample one of the tasks with its prompt token, obtain a task-specific conditional mask, and optimize the transformer to predict all target tokens given masked inputs. During inference, we adapt the non-autoregressive decoding method to generate tokens conditionally on the task-specific inputs, which will be detailed in Algorithm 1.

The VQ autoencoder is a crucial module as it not only sets a quality bound for the generation but also determines the token sequence length, hence affecting generation efficiency. Existing methods apply VQ encoders either on each frame independently (2D-VQ) or on a supervoxel (3D-VQ) . We propose different designs that facilitate MAGVIT to perform favorably against other VQ models for video (see Tab. 7).

We design a 3D-VQ network architecture to model the temporal dynamics as follows. The encoder and decoder of VQGAN consist of cascaded residual blocks interleaved by downsampling (average pooling) and upsampling (resizing plus convolution) layers. We expand all 2D convolutions to 3D convolutions with a temporal axis. As the overall downsampling rate is usually different between temporal and spatial dimensions, we use both 3D and 2D downsampling layers, where the 3D ones appear in the shallower layers of the encoder. The decoder mirrors the encoder with 2D upsampling layers in the first few blocks, followed by 3D ones. Appendix A.1 illustrates the detailed architecture. Note that a token is not only correlated to its corresponding supervoxel but depends on other patches due to the non-local receptive field.

Inflation and padding.

We initialize our 3D-VQ with weights from a 2D-VQ in a matching architecture to transfer learned spatial relationships , known as 3D inflation. We use inflation on small datasets such as UCF-101 . We use a central inflation method for the convolution layers, where the corresponding 2D kernel fills in the temporally central slice of a zero-filled 3D kernel. The parameters of the other layers are directly copied. To improve token consistency for the same content at different locations , we replace the same (zero) padding in the convolution layers with reflect padding, which pads with non-zero values.

Training.

We apply the image perceptual loss on each frame. The LeCam regularization is added to the GAN loss to improve the training stability. We adopt the discriminator architecture from StyleGAN and inflate it to 3D. With these components, unlike VQGAN, our model is trained stably with GAN loss from the beginning.

2 Multi-Task Masked Token Modeling

In MAGVIT, we adopt various masking schemes to facilitate learning for video generation tasks with different conditions. The conditions can be a spatial region for inpainting/outpainting or a few frames for frame prediction/interpolation. We refer to these partially-observed video conditions as interior conditions.

We argue that it is suboptimal to directly unmask the tokens corresponding to the region of the interior condition . As discussed in Section 3.1, the non-local receptive field of the tokenizer can leak the ground-truth information into the unmasked tokens, leading to problematic non-causal masking and poor generalization.

We propose a method, COnditional Masked Modeling by Interior Tokens (or COMMIT for short), to embed interior conditions inside the corrupted visual tokens.

We can decompose the loss in Eq. 3 into three parts according to Eq. 2: $\mathcal{L}_{\text{refine}}$ refines the task-specific condition tokens, $\mathcal{L}_{\text{mask}}$ predicts masked tokens , and $\mathcal{L}_{\text{recons}}$ reconstructs target tokens. Let $\overline{{\mathbf{c}}}=[\rho,{\mathbf{c}},\overline{{\mathbf{z}}}]$ for simplicity,

While $\mathcal{L}_{\text{mask}}$ is the same as the MTM loss in Eq. 1 and $\mathcal{L}_{\text{recons}}$ sometimes is used as a regularizer (e.g., in NLP tasks), $\mathcal{L}_{\text{refine}}$ is a new component introduced by COMMIT.

The COMMIT method facilitates multi-task video generation in three aspects. First, it provides a correct causal masking for all interior conditions. Second, it produces a fixed-length sequence for different conditions of arbitrary regional volume, improving training and memory efficiency since no padding tokens are needed. Third, it achieves state-of-the-art multi-task video generation results (see Tab. 5).

Video generation tasks.

We consider ten tasks for multi-task video generation where each task has a different interior condition and mask: Frame Prediction (FP), Frame Interpolation (FI), Central Outpainting (OPC), Vertical Outpainting (OPV), Horizontal Outpainting (OPH), Dynamic Outpainting (OPD), Central Inpainting (IPC), and Dynamic Inpainting (IPD), Class-conditional Generation (CG), Class-conditional Frame Prediction (CFP). We provide the detailed definitions in Appendix B.1.

Inference.

We use a non-autoregressive decoding method to generate video tokens from input conditions in $K$ steps (e.g., $12$). Each decoding step follows the COMMIT masking in Eq. 2 with a gradually reduced mask ratio. Algorithm 1 outlines the inference procedure.

Fig. 3 compares the non-autoregressive image decoding and our video decoding procedure. Different from the MTM decoding in which performs denoising from all [MASK], COMMIT decoding starts from a multivariate mask that embeds the interior conditions. Guided by this mask, Algorithm 1 performs a conditional transition process toward the output tokens by replacing a portion of newly generated tokens at each step. In the end, all tokens are predicted where the interior condition tokens get refined.

Experimental Results

We conduct extensive experiments to demonstrate the video generation quality (Section 4.2), efficiency (Section 4.3), and flexibility for multi-task generation (Section 4.4). We show a few generation results here, and refer to the web pagehttps://magvit.cs.cmu.edu for more examples.

We evaluate the single-task video generation performance of MAGVIT on three standard benchmarks, i.e., class-conditional generation on UCF-101 and frame prediction on BAIR Robot Pushing (1-frame condition) and Kinetics-600 (5-frame condition). For multi-task video generation, we quantitatively evaluate MAGVIT on BAIR and SSv2 on 8-10 tasks. Furthermore, to evaluate model generalizability, we train models with the same learning recipe on three additional video datasets: nuScenes , Objectron , and 12M Web videos. We show their generated videos in the main paper and quantitative performance in Appendix C.

Evaluation metrics.

We use FVD as our primary evaluation metric. Similar to , FVD features are extracted with an I3D model trained on Kinetics-400 . We also report the Inception Score (IS) calculated with a C3D model on UCF-101, and PSNR, SSIM , and LPIPS on BAIR. We report the mean and standard deviation for each metric calculated over four runs.

Implementation details.

We train MAGVIT to generate 16-frame videos at 128$\times$128 resolution, except for BAIR at 64$\times$64. The proposed 3D-VQ model quantizes a video into 4$\times$16$\times$16 visual tokens, where the visual codebook size is 1024. We use the BERT transformer to model the token sequence, which includes 1 task prompt, 1 class token, and 1024 visual tokens. Two variants of MAGVIT, i.e., base (B) with 128M parameters and large (L) with 464M, are evaluated. We train both stages with the Adam optimizer in JAX/Flax on TPUs. Appendix B.2 details training configurations.

2 Single-Task Video Generation

The model is given a class identifier in this task to generate the full video. Tab. 1 shows that MAGVIT surpasses the previous best-published FVD and IS scores. Notably, it outperforms Make-A-Video which is pretrained on additional 10M videos with a text-image prior. In contrast, MAGVIT is just trained on the 9.5K training videos of UCF-101.

Footnote 3 compares the generated videos to baseline models. We can see that CCVS+StyleGAN gets a decent single-frame quality, but yields little or no motion. TATS generates some motion but with artifacts. In contrast, our model produces higher-quality frames with substantial motion.

Frame prediction.

The model is given a single or a few frames to generate future frames. In Tab. 2, we compare MAGVIT against highly-competitive baselines. MAGVIT surpasses the previous state-of-the-art FVD on BAIR by a large margin ($84\xrightarrow[]{}62$). Inspired by , a “debiased” FVD is also reported in the parentheses to overcome the small validation set. See more discussion in Appendix B.3. In Tab. 3, it demonstrates better image quality.

On the large dataset of Kinetics-600, it establishes a new state-of-the-art result, improving the previous best FVD in from $16.2$ to $9.9$ by a relative $39\%$ improvement. The above results verify MAGVIT’s compelling generation quality, including on the large Kinetics dataset.

3 Inference-Time Generation Efficiency

Video generation efficiency is an important metric in many applications. We conduct experiments to validate that MAGVIT offers top speed in video generation. Fig. 5 shows the processing time for each frame on a single V100 GPU at different resolutions. We compare MAGVIT-B with an autoregressive transformer of the same size and a diffusion-based model . At 128$\times$128 resolution, MAGVIT-B runs at $37$ frames-per-second (fps). When running on a single TPUv4i , MAGVIT-B runs at $190$ fps and MAGVIT-L runs at $65$ fps.

Fig. 5 compares the sequence lengths and inference steps of these models. Diffusion models typically require 256-1000 diffusion steps with a 3D U-Net . Autoregressive models, such as TATS , decode visual tokens sequentially, which runs $60$ times slower than MAGVIT at 128$\times$128. Compared to the recent non-autoregressive model MaskViT , MAGVIT is $4$ to $16$ times faster due to more efficient decoding on shorter sequences.

4 Multi-task Video Generation

To demonstrate the flexibility in multi-task video synthesis, we train a single MAGVIT model to perform eight tasks on BAIR or ten tasks on SSv2. We do not intend to compare with dedicated models trained on these tasks but to demonstrate a generic model for video synthesis.

We perform a multi-task evaluation on BAIR with eight self-supervised tasks. Tab. 4 lists the “debiased” FVD for each task, where the third column computes the average. We compare the multi-task models (MT) with two single-task baselines trained on unconditional generation (UNC) and frame prediction (FP).

As shown in Tab. 4, the multi-task models achieve better fidelity across all tasks. Single-task models perform considerably worse on the tasks unseen in training (gray values in Tab. 4), especially on the tasks that differ more from the training task. Compared to the single-task models in their training task, MT performs better with a small gain on FP with the same model size.

Ten tasks on SSv2.

We evaluate on the large-scale SSv2 dataset, where MAGVIT needs to synthesize 174 basic actions with everyday objects. We evaluate a total of ten tasks, with two of them using class labels (CG and CFP), as shown on the right side of Tab. 4. We observe a pattern consistent with BAIR: multi-task models achieve better average FVD across all tasks. The above results substantiate model generalization trained with the proposed multi-task objective.

5 Ablation Study

We demonstrate the efficacy of COMMIT by comparing it with conventional MTM methods, including the latent masking in MaskGIT for image synthesis and the commonly-used prefix condition that prepends cropped condition tokens to the input sequence.

Tab. 5 compares these methods on the BAIR dataset where the same 3D-VQ tokenizer is used in all approaches. As discussed in Section 3.2, latent masking in , which directly unmasks tokens of the condition region at inference time, leads to poor generalization, especially for the multi-task setup. Prefix condition produces a long sequence of variable length, making it less tractable for multi-task learning. In contrast, COMMIT yields a fixed-length sequence and better generalizability for both single- and multi-task setups.

Training losses.

The bottom section of Tab. 5 shows the contribution of the training loss components in Eq. (4).

Decoding methods.

Tab. 6 compares Algorithm 1 with existing autoregressive (AR) and non-autoregressive (NAR) decoding methods. We consider two NAR baselines, i.e., MaskGIT for image and MaskViT for video synthesis. We use the same 3D-VQ tokenizer for MaskGIT, AR, and MAGVIT. As shown, the proposed decoding algorithm produces the best quality with the 3D-VQ and has a $4\times$ shorter sequence than the 2D-VQ. While the AR transformer obtains a reasonable FVD, it takes over $85\times$ more steps at inference time.

VQ architecture and training techniques.

We evaluate the design options of our 3D-VQ model in MAGVIT. Tab. 7 lists the reconstruction FVD and IS metrics on the UCF-101 training set, which are different from the generation metrics as they measure the intermediate quantization. Nevertheless, reconstruction quality bounds the generation quality.

Tab. 7 compares the proposed 3D architecture with existing 2D and 3D VQ architectures. We train the MaskGIT 2D-VQ and our 3D-VQ with the same protocol and evaluate the official TATS 3D-VQ model. We compare two inflation methods for our 3D-VQ model, i.e., average and central inflation.

The results show the following. First, 3D-VQ models, despite producing a higher compression rate, show better video reconstruction quality than 2D-VQ, even with fewer parameters. Second, the proposed VQ performs favorably against baseline architectures with a similar size and gets much better with a larger model. Third, ImageNet initialization boosts the performance for 2D and 3D models, where the central inflation outperforms the average inflation. The results demonstrate the excellent reconstruction fidelity of our tokenizer design.

Related Work

Early success in video synthesis has been made by GAN models . Training instability and lack of generation diversity are known issues of GAN models.

Autoregressive transformers.

Inspired by the success of GPT , autoregressive transformers have been adapted for image and video generation . A focus for video is autoregressive modeling of visual dynamics. Studies have switched from modeling the raw pixels to the discrete codes in a latent space . The state-of-the-art model TATS uses two hierarchical transformers to reduce the computation for long video generation, with tokens learned by a 3D-VQGAN . Unlike prior works, we introduce a non-autoregressive transformer with higher efficiency and flexibility.

Non-autoregressive transformers.

Concurrently, a few methods use non-autoregressive transformers for image synthesis [12, 75, 42, sohn2022visual]. Section 2 reviews a state-of-the-art model called MaskGIT . Compared with these approaches , we present an embedding mask to model multi-task video conditions with better quality.

Diffusion models.

Diffusion models have recently received much attention for image synthesis. For example, the state-of-the-art video diffusion model extends the image denoising diffusion model by incorporating 3D U-Net architectures and joint training on both images and videos. Despite its high-quality, sampling speed is a bottleneck hindering the application of diffusion models in video synthesis. We show a different solution to train a highly-efficient model that offers compelling quality.

Multi-task video synthesis.

Multi-task video synthesis is yet to be well-studied. Transframer is the closest to our work, which adopts an image-level representation for autoregressive modeling of tasks based on frame prediction. We present an efficient non-autoregressive multi-task transformer, and verify the quality and efficiency on ten video generation tasks.

Text-to-video.

All of our models are trained only on public benchmarks, except the Web video model. We leave the text-to-video task as future work. As shown in recent works , training such models requires large, and sometimes non-public, datasets of paired texts and images.

Conclusion

In this paper, we propose MAGVIT, a generic and efficient mask-based video generation model. We introduce a high-quality 3D-VQ tokenizer to quantize a video and design COMMIT for multi-task conditional masked token modeling. We conduct extensive experiments to demonstrate the video generation quality, efficiency, and flexibility for multi-task generation. Notably, MAGVIT establishes a new state-of-the-art quality for class conditional generation on UCF-101 and frame prediction on BAIR Robot Pushing and Kinetics-600 datasets.

References