MAGVIT: Masked Generative Video Transformer

In plain words

MAGVIT generates short videos without working on raw pixels. A first model squeezes each clip into a small grid of codes, drawn from a fixed vocabulary of about a thousand entries. The squeezing looks at space and time together, so a single code describes a small patch of picture and the motion across a few frames.

A second model, a transformer, learns a fill-in-the-blanks game on these codes. Some codes are hidden and the transformer guesses them. Many video jobs fit this one game: predicting the next frame, filling in a missing region, or making a clip of a given class. Each job just changes which codes are hidden and which are shown. The known parts of the video are written straight into the hidden positions, so the model sees its clues in place.

To generate, the model starts with the clues and fills in the rest over a few rounds. Each round it keeps the guesses it is most sure of and redoes the others. That takes far fewer passes than producing one code at a time, which is where the speed comes from. One trained model covers several video tasks, and the authors compare it against diffusion and token-by-token generators.

TL;DR

• Two stages: a 3D VQ autoencoder maps a 16-frame 128x128 clip to 4x16x16 discrete tokens (codebook 1024), then a BERT-style transformer is trained with masked token modeling §3.1 §3.2 §4.1.
• One model handles ten tasks (frame prediction and interpolation, in- and outpainting, class-conditional generation). The condition pixels are tokenized and placed into the mask, so every task uses the same loss §3.2 app. B.1.
• Generation runs in 12 non-autoregressive steps with a conditional-transition rule alg. 1 app. B.3.
• Author-reported headline: FVD 76 on UCF-101 class-conditional (L), 62 on BAIR frame prediction (L), 9.9 on Kinetics-600 frame prediction (L) tab. 1 tab. 2.
• Released code is Apache-2.0, but the repo has no trained checkpoints; configs point to gs:// paths and only cover UCF-101, BAIR and BAIR 8-task (see Paper vs code).

Problem

• Diffusion video models need 256 to 1000 denoising steps with a 3D U-Net, which the authors give as the reason they are slow §4.3.
• Autoregressive token models emit one token at a time. Against 12 decoding steps, the authors' comparison uses 1024 steps for the autoregressive model §4.5.
• Earlier masked video models used a 2D tokenizer applied frame by frame, which gives 4096 tokens for a 16-frame 128x128 clip. The authors' 3D tokenizer gives 1024 §3.1 §4.5.
• Conditioning is the hard part. Prefix conditions give variable-length sequences that are hard to train across tasks §4.5. Directly unmasking the condition region leaks ground-truth information, because each tokenizer code depends on pixels outside its own patch §3.1 §3.2.
• The tokenizer sets an upper bound on generation quality and also sets the sequence length §3.1.

Method

Stage 1: 3D VQ tokenizer §3.1 app. A.1

• The tokenizer builds on the image VQGAN. Every 2D convolution is replaced by a 3D one §3.1.
• Down-sampling uses average pooling. Up-sampling uses nearest resizing followed by a convolution. Spatial-only layers sit near the latent, spatial-temporal layers near the pixels, and the decoder mirrors the encoder app. A.1.
• Weights are initialized by central inflation from a 2D VQ trained on ImageNet §3.1 app. B.2.
• Training uses a per-frame perceptual loss, a GAN loss with a StyleGAN discriminator inflated to 3D, and LeCAM regularization §3.1.
• The authors also say they replace zero "same" padding with reflect padding in the convolutions §3.1.
• Two sizes: B (41M, base channels 64) and L (158M, base channels 128), both without the discriminator app. A.1 app. B.2.
• Output: f_T maps a clip V (T x H x W x 3) to N discrete tokens with a 1024-entry codebook. For 16x128x128 input, N = 4x16x16 = 1024 §3.1 §4.1.

Stage 2: multi-task masked token modeling §3.2

Input sequence: one task prompt token, one class token, then N = 1024 visual tokens, 1026 in total §4.1 app. B.2.

Condition tokens. For a task, the known pixels are padded to the full clip shape and passed through the same frozen tokenizer. A position gets its condition token if its supervoxel contains condition pixels, and [MASK] otherwise eq. 2. Positions outside the condition are left as the target token eq. 2.

Mask ratio. The paper draws r uniformly and uses a cosine schedule, which it takes from MaskGIT app. B.2 §2. The released code uses ratio = cos(pi/2 * r) mask_schedule.py.

Objective. Eq. 3 is the sum over positions of the negative log-likelihood of each target token, given the task prompt, class token and corrupted sequence. Eq. 4 splits it into a refine term for condition-region positions, a mask term for [MASK] positions, and a reconstruction term for the rest eq. 3 eq. 4. The paper treats the reconstruction term as a regulariser that is optional in NLP, but keeps it here §3.2.

Tasks. Ten tasks, each with its own prompt token and condition shape. Settings from the appendix:

Training step:

# Masked token training for one batch  [§3.2] [eq. 3] [eq. 4]
# Inputs:  video (B, T, H, W, 3), task with padding rule
# Output:  scalar loss

def training_step(video, task):
    # Tokenize the clip into 4 x 16 x 16 targets with the frozen 3D-VQ  [§3.1]
    targets = tokenizer(video)                           # (B, N)
    # Pad the task's condition pixels and tokenize them the same way  [app. B.1]
    condition_tokens = tokenizer(task.pad(video))        # (B, N)
    # One ratio per step, shared across the batch: r ~ U(0,1)  [§2]
    # [code: videogvt/trainers/maskgvt_trainer.py#L206]
    ratio = cos(pi / 2 * uniform(0, 1))
    # Give each token a uniform score; mask the lowest floor(N * ratio)  [§3.2]
    # [code: videogvt/train_lib/mask_utils.py#L166-L175]
    scores = uniform(0, 1, shape=(B, N))
    chosen = lowest_k(scores, floor(N * ratio))          # bool (B, N)
    # Chosen positions inside the condition region get the condition token;
    # other chosen positions get [MASK]  [eq. 2]
    inputs = where(chosen, where(in_condition, condition_tokens, MASK),
                   targets)
    logits = transformer(task_prompt, class_token, inputs)  # (B, N, 1024)
    # Cross-entropy at every position, label smoothing 1e-4  [eq. 3]
    # [code: videogvt/trainers/maskgvt_trainer.py#L226-L243]
    token_loss = cross_entropy(logits, targets, smoothing=1e-4)
    return mean(token_loss)                              # batch and positions

Inference: non-autoregressive decoding alg. 1 §3.2

Decoding starts from the interior-condition mask rather than a blank canvas, and each step swaps some of the sampled tokens back to the mask. The step count is K = 12 by default §3.2 app. B.3.

# Non-autoregressive decoding with interior conditions  [alg. 1]
# Inputs:  prompt (task and class tokens), condition tokens (B, N), steps K
# Output:  predicted visual tokens (B, N)

def decode_video(prompt, condition_tokens, steps):
    # Start from a fully masked canvas; condition-region positions take
    # their condition token instead of [MASK]  [§3.2] [eq. 2]
    tokens = mask_with_condition(condition_tokens, ratio=1.0)  # (B, N)
    unknown = all_positions(tokens)                          # (B, N) bool
    for step in range(steps):
        # Predict logits at every position from the current sequence  [alg. 1]
        logits = transformer(prompt, tokens)                 # (B, N, 1024)
        # Sample a candidate for each unknown position  [alg. 1]
        # Released code samples at temperature 1.0
        # [code: videogvt/train_lib/parallel_decode.py#L78]
        sampled = categorical(logits)
        tokens = where(unknown, sampled, tokens)
        # Confidence is the log-probability of the candidate; fixed positions
        # get +inf so they are never re-masked  [alg. 1]
        # [code: videogvt/train_lib/parallel_decode.py#L147-L161]
        confidence = where(unknown, log_prob(logits, sampled), inf)
        # Add Gumbel noise scaled by the mask temperature and (1 - progress)
        # [alg. 1] [code: videogvt/train_lib/parallel_decode.py#L154-L161]
        progress = (step + 1) / steps
        confidence = confidence + mask_temperature * (1 - progress) * gumbel()
        # Re-mask the lowest-confidence floor(N * cos(pi/2 * progress)) tokens;
        # re-masked condition-region positions go back to their condition token
        # [§2] [code: parallel_decode.py#L169]
        ratio = cos(pi / 2 * progress)
        tokens, unknown = remask_lowest(
            tokens, confidence, ratio, condition_tokens)
    return tokens

Hyperparameters

Results

Author-reported. Each number below is the paper's own measurement with the setup shown.

Other author-reported results:

• Multi-task, 8 tasks on BAIR: average debiased FVD 32.8 (B) and 22.8 (L), against single-task models of 150.6 (unconditional) and 201.1 (frame prediction) §4.4.
• Multi-task, 10 tasks on SSv2: average FVD 43.4 (B) and 27.3 (L), against 258.8 and 402.9 for the single-task models §4.4.
• Speed, per frame at 128x128: B runs at 37 fps on one V100 §4.3. On a single TPUv4i, B runs at 190 fps and L at 65 fps §4.3. The paper reports the autoregressive TATS as about 60 times slower and MaskViT as 4 to 16 times slower §4.3.
• Conditioning ablation on BAIR (debiased FVD, FP and 8-task average): latent masking as in MaskGIT 74 and 151; prefix condition 55 and not converged; loss without refine or reconstruction terms 388 and 143; with reconstruction 51 and 53; full objective 48 and 33 §4.5.
• Decoding ablation on BAIR frame prediction (FVD, debiased in brackets): MaskGIT with 2D-VQ, 4096 tokens, 12 steps 222 (177); MaskGIT with 3D-VQ, 1024 tokens, 12 steps 122 (74); autoregressive with 3D-VQ, 1024 steps 91 (56); this paper, 1024 tokens, 12 steps 76 (48) §4.5.
• Tokenizer reconstruction on UCF-101 (FVD / IS): TATS 3D-VQ 162 / 80.6; MaskGIT 2D-VQ 240 / 80.9; this paper's B 127 / 82.1, with central inflation 58 / 87.0; L with central inflation 25 / 88.9 §4.5.

Paper vs code

Official code: https://github.com/MAGVIT/magvit, linked from the project page and described in its README as the official code for the CVPR 2023 paper. Commit 05e8cfd6559c47955793d70602d62a2f9b0bdef5, Apache-2.0 README.md LICENSE. The repo also points to a separate google-research/magvit repository for weights; that one was not compared.

Differences between the paper and the code:

1. Padding. The paper says reflect padding replaces zero padding in the convolutions §3.1. The released configs leave custom_conv_padding unset, which gives "SAME" zero padding enc_dec_3dcnn.py. The only mention of reflect padding is commented out in the BAIR config vqgan3d_bair_config.py.
2. LeCAM. The paper describes LeCAM regularization as part of training §3.1 app. A.1. The code adds it only when lecam_weight > 0 vqgan_trainer.py, and no released config sets that key. As configured, the term is off.
3. Entropy term. The tokenizer loss includes an entropy term (ratio 0.1, temperature 0.01) that the paper does not mention vqgan2d_ucf101_config.py vqvae.py.
4. Mask count. The paper's schedule and Algorithm 1 use ceil(gamma * N) §2 alg. 1. The code uses floor mask_utils.py. The gap is at most one token.
5. Decoding score. Algorithm 1 stores the probability p_theta. The code uses the log-probability plus Gumbel noise alg. 1 parallel_decode.py. This ranking is the Gumbel-max form, so implement the code's version.
6. Temperature. The paper calls 4.5 "temperature" app. B.3. The code uses it only as the masking-noise temperature, and sampling temperature is 1.0 parallel_decode.py maskgvt_train_lib.py.
7. BAIR decoding. The appendix gives exponential schedule and temperature 400 for BAIR app. B.3. In the released BAIR config, these values apply only to the single-clip evaluation option maskgvt_bair_config.py. The default BAIR evaluation inherits cosine and 4.5 from the UCF-101 config maskgvt_ucf101_config.py. Which setting produced Table 2 is not stated.
8. Task prompts. The paper defines ten tasks with their own settings §3.2 app. B.1. The code gives inpainting the same prompt token as outpainting unless a bug_patches flag is set task_registry.py.
9. Loss normalization. Eq. 3 is a sum over positions. The code takes a mean over positions for each example and then over the batch maskgvt_trainer.py. The loss is therefore 1/N of Eq. 3, which changes the scale that the gradient-norm clip at 1.0 sees.
10. Mask ratio sharing. The paper does not say whether the mask ratio is shared across a batch. The code draws one ratio per step maskgvt_trainer.py maskgvt_train_lib.py.
11. Weight decay. The paper gives one weight decay value app. B.2. The code applies it only to parameters whose path contains kernel maskgvt_train_lib.py.
12. Coverage. Released configs cover UCF-101 class-conditional generation, BAIR frame prediction and BAIR 8-task training. There are no configs for Kinetics-600, SSv2, nuScenes, Objectron or Web video, and so no reproduction path for those numbers from the repo . Checkpoints and ImageNet 2D initializations are gs:// paths that are not in the repo maskgvt_ucf101_config.py vqgan2d_ucf101_config.py.

Checked and matching: 4x16x16 latents with 1024 codebook entries, central inflation vqgan_train_lib.py, 12 decoding steps, cosine schedule, transformer sizes and dropout, Adam betas, label smoothing 1e-4, grad-norm clip 1.0, batch sizes, and the perceptual and adversarial weights.

Pitfalls

• Prompt sharing. Inpainting and outpainting share a prompt token in the released code, so the model cannot tell them apart by prompt. This follows the code, not the paper, and is the first thing to check if you adapt the task set task_registry.py.
• Missing tokenizer features. If you follow the paper and add reflect padding or LeCAM, you will not match the released tokenizer. Decide which one you are reproducing before you write code §3.1 enc_dec_3dcnn.py.
• Debiased FVD. The standard BAIR FVD (62) uses 256 evaluation clips with 100 predictions each. The debiased value (31) uses 3,840 uniformly sampled clips app. B.3. Compare debiased with debiased, and the multi-task and ablation numbers are debiased app. B.3.
• Kinetics-600 protocol. Videos are generated at 128x128 and evaluated at 64x64 after a central crop and bilinear resize app. B.3. Other papers' Kinetics-600 FVDs may use a different protocol.
• Speed numbers. The fps figures depend on hardware and per-frame timing. The paper does not state the batch size used for timing §4.3.
• Condition regions on the latent grid. The released code sets condition regions by integer division of the latent size (for example l_h // 4 to 3 * l_h // 4) rather than by checking pixels cond_utils.py. Off-by-one differences at other resolutions are possible.
• Minimum masking. The decoder re-masks at least one token and at most all but one unknown token at each step parallel_decode.py. The last step's re-masking does not change the output.
• No-class token. For tasks without a class label, the code uses a no-class token built from the number of classes plus the codebook offset maskgvt_train_lib.py. The paper only says there is one class token.
• Parameter counts. "B" and "L" refer to the transformer. The 128M and 464M totals in the text include the tokenizer §4.1 app. A.2.
• Tokenizer ceiling. Reconstruction quality bounds generation quality §3.1. The tokenizer's reconstruction FVD is 58 (B) and 25 (L) on UCF-101 §4.5, so a weak tokenizer will cap the final numbers.

Lineage

• Builds on: VQGAN for the image tokenizer and its losses §3.1; MaskGIT for masked token modeling and parallel decoding §2; BERT for the bidirectional transformer §2; 3D inflation of 2D weights §3.1.
• Changes from the image line: 3D tokenizer with average-pool down-sampling, nearest resize plus convolution up-sampling, codebook 1024 instead of 16,384 as in TATS app. A.1; interior-condition masks for multi-task training §3.2; a decoder that moves from denoising to conditional transition alg. 1.
• Compared against: CCVS, TATS and Video Diffusion §4.2; RaMViD tab. 2; MaskViT and autoregressive decoding §4.5; Make-A-Video, CogVideo, Phenaki, MCVD and FitVid in the tables tab. 1 tab. 2.
• Related concepts: [[vector-quantization]], [[masked-image-modeling]], [[multi-task-learning]].