VideoGPT: Video Generation using VQ-VAE and Transformers

In plain words

Generating realistic video is much harder than generating images, because a video is a stack of frames and each frame is itself a big grid of pixels. This paper takes a simple route. It first learns a compact, shorthand version of the video, where each small patch of space and time is replaced by one entry from a fixed dictionary of building blocks. A second model then learns to write sequences of these shorthand entries, one at a time, the same way language models write words.

To make a new video, the second model writes a shorthand sequence from scratch, and the first model's decoder turns that sequence back into pixels. Because the shorthand is much smaller than the raw video, the second model can be larger and still fit in memory. The authors stress that the design is deliberately plain, built from standard parts, so it is easy to understand and adapt.

The results are reported on a few well-known video datasets. The samples are described as realistic and varied, though the authors say the method is not the best on every benchmark.

TL;DR

VideoGPT is a two-stage video generator. Stage one trains a VQ-VAE with 3D convolutions and axial self-attention that maps a video to a grid of discrete codes and back §3. Stage two trains a GPT-style transformer, following the Image-GPT architecture with extra dropout, to model those codes autoregressively §3. Sampling draws codes from the transformer and decodes them with the VQ-VAE fig. 2. The paper reports FVD 103.3 on BAIR and Inception Score 24.69 on UCF-101 tab. 1, tab. 2. It says the method is not state of the art on BAIR but competitive with GAN baselines §4.3. Read this if you want a simple, reproducible latent autoregressive video baseline. The official code is short, but it departs from the paper in several hyperparameters and details (see Paper vs code).

Problem

Video has much higher input dimension than images, and its correlations run across space and time, so modelling it directly is costly §1. The authors choose likelihood-based models over GANs for easier optimisation and evaluation §1. Among likelihood models they choose autoregressive models, and they model a downsampled latent space rather than every pixel, because natural video has heavy spatio-temporal redundancy that compression exploits §1. The paper's stated assumption is that the latent representation keeps enough detail for a lossy decoder to produce realistic samples §1.

Method

VideoGPT has two stages, shown in fig. 2 §1. The prose overview is in §3.

1. Stage 1: VQ-VAE for latent codes §3 §2.1.
   • Encoder: a stack of 3D convolutions that downsample over space and time, followed by attention residual blocks. The blocks use LayerNorm and axial attention, following CCNet and axial attention work §3 fig. 4.
   • Quantisation: each latent vector is replaced by its nearest entry in a codebook §2.1.
   • Decoder: the mirror of the encoder, with attention residual blocks and 3D transposed convolutions that upsample over space and time §3.
   • Position encodings: learned spatio-temporal embeddings, shared by all axial attention layers in encoder and decoder §3.
   • Codebook learning: an exponential-moving-average (EMA) update replaces the codebook loss §2.1. The paper's training recipe adds random restarts and initialises codes by copying encoder latents §4.1.
   • Training is unconditional §3.
2. Stage 2: prior over codes §3.
   • The prior follows the Image-GPT architecture, with dropout after the attention and feed-forward layers §3.
   • It is trained on discrete latent code sequences produced by the frozen stage-1 VQ-VAE §3 fig. 2.
   • Sparse Transformers with local and strided attention across space and time are used for the transformer §4.1.
3. Conditioning §3.
   • Frame conditioning: the conditioning frames go through a 3D ResNet, and the prior cross-attends to its output §3.
   • Conditional norms: the gain and bias of the transformer LayerNorm layers are affine functions of a conditioning vector. The paper uses this for action and class conditioning §3.
4. Sampling fig. 2 §3.
   • Draw a latent code sequence from the prior, then decode it with the VQ-VAE decoder to get a video.
   • The paper gives no sampling temperature, top-k, or sampler settings (not stated).

The VQ-VAE objective as written in the paper, reconstruction plus codebook plus β-weighted commitment, is given as an unnumbered display equation §2.1:

\mathcal{L} = \lVert x - D(e)\rVert_2^2 + \lVert sg[E(x)] - e\rVert_2^2 + \beta\lVert sg[e] - E(x)\rVert_2^2 §2.1, unnumbered

With the EMA update in use, the codebook term is replaced as described above §2.1. The official code implements the reconstruction term as a normalised MSE and uses a fixed commitment weight of 0.25 (see Paper vs code).

Pseudocode for one stage-1 training step, following the official code's structure (the codebook itself is updated separately, below):

# VQ-VAE training step  [§2.1] [§3] [code: videogpt/vqvae.py#L50-L63]
# Input:  video (B, 3, T, H, W), values in [-0.5, 0.5]  [§4.1]
# Output: loss (scalar); codebook is updated by EMA separately

def vqvae_training_step(video):
    # Encode with 3D convolutions and attention residual blocks  [§3]
    hidden = encoder(video)                        # (B, n_hid, t, h, w)
    latents = pre_vq_conv(hidden)                  # (B, D, t, h, w)

    # Replace each latent vector by its nearest codebook entry  [§2.1]
    codes = nearest_code_index(latents, codebook)  # (B, t, h, w)
    quantized = codebook[codes]                    # (B, t, h, w, D)

    # Straight-through: forward pass uses the code, gradient reaches latents
    straight = latents + stop_gradient(quantized - latents)
    recon = decoder(post_vq_conv(straight))

    # Reconstruction term, divided by a constant 0.06  [code: vqvae.py#L54]
    recon_loss = mse(recon, video) / 0.06

    # Commitment term, weight beta = 0.25  [tab. 8] [code: vqvae.py#L174]
    commit_loss = 0.25 * mse(latents, stop_gradient(quantized))

    return recon_loss + commit_loss

The codebook update, which the paper specifies as EMA §2.1:

# EMA codebook update, decay 0.99  [§2.1] [code: vqvae.py#L177-L190]
# Input:  latents (B, t, h, w, D), codes (B, t, h, w)
# Output: updated codebook (K, D)

def ema_update(latents, codes):
    # Running counts and sums per code, decayed by 0.99  [code: vqvae.py#L184]
    counts = 0.99 * counts + 0.01 * count_per_code(codes)
    sums = 0.99 * sums + 0.01 * sum_latents_per_code(latents, codes)
    # Laplace-smoothed mean, rescaled so the counts sum to the total
    # [code: vqvae.py#L187-L189]
    codebook = sums / smoothed(counts)[:, None]
    # Codes with a recent count below 1 are replaced by random batch
    # latents, on every training step  [code: vqvae.py#L192-L198]
    codebook = replace_unused(codebook, counts, random_batch_latents)
    return codebook

Stage-2 training step. The prior sees the codes of the same video, shifted right so that each position is predicted only from earlier positions:

# Prior training step  [§3] [code: videogpt/gpt.py#L139-L155]
# Input:  video (B, 3, T, H, W); frozen VQ-VAE
# Output: cross-entropy loss over codebook indices

def prior_training_step(video):
    # Map the video to code indices with the frozen VQ-VAE  [§3]
    targets = vqvae_encode(video)                  # (B, t, h, w) integers
    embeddings = codebook[targets]                 # (B, t, h, w, D)

    # Project to hidden size, shift right by one learned start token
    # [code: videogpt/attention.py#L513-L525]
    inputs = right_shift(fc_in(embeddings))        # (B, N, hidden)
    inputs = inputs + positional_embedding         # (B, N, hidden)

    # Pre-norm transformer blocks with causal self-attention  [§3]
    # [code: videogpt/attention.py#L90-L118]
    for each block:
        h = layer_norm(inputs)
        inputs = inputs + dropout(causal_attention(h))
        h = layer_norm(inputs)
        inputs = inputs + dropout(feed_forward(h))  # width 4 x hidden

    logits = fc_out(layer_norm(inputs))            # (B, N, n_codes)
    return cross_entropy(logits, targets)

Sampling. The code draws one code at a time in raster order and decodes at the end:

# Ancestral sampling of latent codes, then decoding  [§3] [fig. 2]
# [code: videogpt/gpt.py#L87-L117]
# Output: videos (n, 3, T, H, W) in [0, 1]

def sample_video(n):
    # Start from all-zero codes; fill them in raster order
    codes = zeros(n, t, h, w)
    for position in raster_order(t, h, w):
        # Logits for this position, conditioned on the codes before it
        logits = prior(codes, position)            # (n, n_codes)
        probs = softmax(logits)
        # Sample at temperature 1, with no top-k  [code: gpt.py#L110-L111]
        codes[:, position] = multinomial(probs)
    video = vqvae_decode(codes)
    # Clamp to the training range, then shift to [0, 1]  [code: gpt.py#L115]
    return clamp(video, -0.5, 0.5) + 0.5

Hyperparameters

Values are from the paper's appendix tab. 8 tab. 9 app. A. Columns follow the paper's dataset groupings; rows cited to §4.1 or §4.5 come from the main text.

Hardware: at most 8 Quadro RTX 6000 GPUs (24 GB each) for all results §4.1.

Not stated in the paper: optimiser for the prior, learning-rate schedule, warmup, gradient clipping, weight decay, sampling temperature, number of samples used for FVD or Inception Score, training time, and the exact 128×128 UCF-101 settings.

Results

Author-reported. Setup: BAIR Robot Pushing at 64×64 with 16-frame sequences §4.1 §4.3; UCF-101 with unconditional models trained at 16 frames and 64×64 and 128×128 resolutions §4.5; Inception Score computed with the TGANv2 code §4.5 footnote. The paper does not state which UCF-101 resolution Table 2 uses.

Headline results tab. 1 tab. 2:

Qualitative results only: Moving MNIST samples conditioned on one frame fig. 3, BAIR reconstructions and single-frame-conditioned samples fig. 5 fig. 6, UCF-101 and TGIF unconditional samples fig. 7 fig. 8, and ViZDoom samples, unconditional and action-conditioned fig. 9 fig. 10. The paper reports no quantitative numbers for Moving MNIST, ViZDoom or TGIF samples; for TGIF it reports train and test losses of 2.87 and 2.86, and for UCF-101 train and test losses of 3.40 and 3.12 §4.5 §4.6.

Ablations (BAIR unless noted). Each table is reproduced in full from the paper.

Axial attention in the VQ-VAE tab. 3 §4.7:

Number of prior transformer layers tab. 4 §4.7:

Latent size, with rows grouped by the paper's partitions of equal total latent count tab. 5 §4.7:

Number of codebook codes tab. 6 §4.7:

Number of codebooks (multiple codebooks slice the encoder channels) tab. 7 §4.7:

Paper vs code

Official code: https://github.com/wilson1yan/VideoGPT, owned by the first author, at commit 856c282a8078ee8c3e9db7196799455a2b6ad084, MIT license. The project page links it [project page: wilsonyan.com/videogpt]. The README says the repo is "primarily designed for simplicity"; reproducing the paper's results uses a separate repo, which the README describes as "not as clean" [README]. This card follows the main repo, not that one.

Differences the implementer will hit:

• VQ-VAE learning rate. Paper: 7e-4 for all datasets tab. 8. Code: Adam with lr 3e-4 and betas (0.9, 0.999), no scheduler vqvae.py.
• Reconstruction normalisation. Paper: MSE divided by the variance of the dataset §4.1. Code: MSE divided by the constant 0.06 vqvae.py. The code never computes a dataset variance.
• Residual-block normalisation. Paper: LayerNorm in the attention residual blocks §3. Code: BatchNorm3d, with a 3×3×3 conv to half width, then a 1×1×1 conv back, then axial attention vqvae.py.
• Encoder filter size. Paper: "Encoder filter size 3" tab. 8. Code: the strided downsampling convolutions use kernel 4, and only the final convolution uses kernel 3 vqvae.py. The paper does not say which convolution "encoder filter size" refers to.
• Prior feed-forward width. Paper, Moving MNIST: 384 hidden features, feed-forward 384, embedding 192 §4.2 tab. 9. Code: the feed-forward width is always 4× the embedding size attention.py, so a 192-dim embedding gets a 768-wide feed-forward layer. The paper's Moving MNIST configuration cannot be reproduced exactly with this code.
• Frame-conditioning ResNet. Paper: ResNet-18 with 512 units for Moving MNIST, ResNet-34 with 512 units for BAIR tab. 9. Code: always ResNet-34 with 240 channels, for every dataset gpt.py.
• Action conditioning. Paper: action conditioning on BAIR and ViZDoom §1, item 4 §3. Code: class labels only, through the conditional LayerNorm attention.py gpt.py. The HDF5 dataset returns only the video, not the BAIR actions, so the action path from the paper is not in this repo data.py.
• Transformer type. Paper: Sparse Transformers with local and strided attention §4.1. Code: full attention by default; sparse attention is optional and needs DeepSpeed gpt.py [README].
• Prior defaults. Code defaults: hidden size 576, 8 layers, 4 heads, dropout 0.2, attention dropout 0.3 gpt.py. Paper BAIR: hidden 512, 16 layers, dropout 0.2 tab. 9. The paper does not mention attention-weight dropout, so 0.3 is a code default the paper does not state.
• Learning-rate schedule. Code: cosine annealing over max_steps, stepped per iteration gpt.py. Paper: not stated.
• Sampling. Code: multinomial sampling at temperature 1 and clamping of the decoded video gpt.py. Paper: not stated.
• Start token and output layer. Code: a learned start-of-sequence vector replaces the first position attention.py, and the output projection is zero-initialised gpt.py. Paper: neither is stated.
• Training length defaults. Code: train_vqvae.py sets max_steps=200000 train_vqvae.py. Paper: 20k for Moving MNIST and 100K for the others tab. 8. The code's single default is not dataset-specific.
• Batch size and resolution defaults. Code: prior script default batch 8 per GPU and resolution 128 train_videogpt.py. Paper: batch 32 and 64×64 for BAIR tab. 8 tab. 9 §4.1. The paper does not say whether 32 is per GPU or in total.
• Commitment weight. Code: fixed at 0.25 in the forward pass, not exposed as an argument vqvae.py. Paper: β = 0.25 tab. 8. These agree.
• Codebook update. Paper: EMA update, random restarts, initialisation by copying encoder latents §2.1 §4.1. Code: EMA with decay 0.99, initialisation from random encoder latents plus small noise vqvae.py, and replacement of unused codes on every training step, not only at initialisation vqvae.py. The paper's "random restarts" is consistent with this, but it happens throughout training.

No discrepancy was found in the data preprocessing: values are scaled to [-0.5, 0.5] and UCF-101 is resized to the shorter side and centre-cropped data.py.

Pitfalls

• Inconsistent headline numbers. The introduction gives FVD 103 and FVD* 94 on BAIR §1. Table 1 gives FVD 103.3 tab. 1. In the layer ablation, the 8-layer prior gives 103.3 tab. 4, while the 16-layer prior, which Table 9 lists for BAIR, gives 103.6 tab. 4 tab. 5. The ablation's FVD* for the chosen latent is 94.6 ± 1.5 tab. 5, which does not round to the introduction's 94. Decide which configuration you are reproducing before comparing numbers.
• Moving MNIST prior sizes. The text says 384 hidden features §4.2; Table 9 lists embedding 192 and feed-forward 384 tab. 9. The paper does not resolve which one the text means. The code's fixed 4× feed-forward width attention.py matches neither exactly for this configuration.
• UCF-101 at 128×128. Table 8 and Table 9 cover only the 64×64 input. The 128×128 settings are not stated, and Table 9's "200K / 600K" training steps for UCF-101 / TGIF is unexplained.
• Unexplained table entries. Table 9's second BAIR conditional size (64) and the ViZDoom values (60, 315) are not explained in the paper. Table 8's "Encoder filter size" is ambiguous (see Paper vs code).
• Cost of the prior. The prior attends over every latent position. From the latent shapes in tab. 8 (derived here), BAIR has 8×32×32 = 8192 positions, UCF-101 4×32×32 = 4096, and Moving MNIST 4×16×16 = 1024. The paper uses Sparse Transformers for its experiments §4.1; the code treats sparse attention as an optional install [README].
• Sampling cost. The code produces one latent position per step in a raster loop, so generation time grows with the number of latent positions gpt.py. The paper reports no sampling time.
• Latent size trade-off. Larger latents reconstruct better, but the largest latent (16×32×32) does not give the best samples, which the paper attributes to limited compute for the prior §4.7 tab. 5. The sweet spot it reports for BAIR is 8×32×32 §4.7.
• Axial attention details not stated. The number of heads in the VQ-VAE axial attention is not stated in the paper. The code uses 2 heads per axis and sums the three axial outputs vqvae.py vqvae.py.
• Learned positional embedding. The code builds one learned embedding per axis of size embedding/3, concatenates them, and adds the result to the input attention.py. The paper says only "learned spatio-temporal embeddings" §3.
• Conditional LayerNorm initialisation. The class-conditional gain is 1 plus a linear map of the one-hot label, initialised with weight 1/√(class count) attention.py. The paper does not describe this initialisation.
• Evaluation details missing. The number of generated samples for FVD and Inception Score is not stated. Inception Score is computed with the TGANv2 code §4.5 footnote.
• Dataset size claim. The paper says TGIF has 103,068 GIFs "totalling around 100,000 hours" §4.6. The paper gives no source for this figure; treat it as author-stated.
• Data preparation. The code expects HDF5 files or a train/test folder layout. The README gives preprocessing scripts for BAIR and UCF-101 [README].

Lineage

• Builds on: VQ-VAE (arXiv:1711.00937), using EMA codebook updates and the Image-GPT architecture for the prior §2.1 §2.2 §3. Also Sparse Transformers (arXiv:1904.10509) for the attention pattern §4.1.
• Compared against: SV2P, LVT, SAVP, DVD-GAN-FP, TrIVD-GAN-FP and Video Transformer on BAIR tab. 1; VGAN, TGAN, MoCoGAN, Progressive VGAN, TGAN-F, TGANv2 and DVD-GAN on UCF-101 tab. 2.