PaperlayerView exact revision on arXiv

Source-linked AI summary

Greedy Hierarchical Variational Autoencoders for Large-Scale Video Prediction

Bohan Wu, Suraj Nair, Roberto Martin-Martin, Li Fei-Fei, Chelsea Finn

arXiv:2103.04174v3cs.CVcs.AIcs.LGcs.RO

TL;DR

Large, diverse video datasets expose underfitting, while memory limits and end-to-end optimization challenges constrain larger hierarchical video prediction models. GHVAEs greedily train hierarchical VAE modules, outperform state-of-the-art models across video datasets and real-robot tasks, and improve monotonically as modules are added.

  • Problem

    Existing video prediction models underfit on large, diverse datasets, while training larger models is constrained by GPU/TPU memory and hierarchical latent models are difficult to optimize end-to-end.

    Source: §1 · p. 1

  • Method

    GHVAEs greedily train sequentially stacked hierarchical VAE modules, freezing previously trained modules to address memory and optimization challenges.

    Source: §1 · p. 2

  • Results

    17–55% gains in FVD score across four video datasets and 35–40% gains in real-robot task success rate over state-of-the-art methods.

    Source: §5 · p. 8

  • Takeaways & Limitations

    GHVAE performance improves monotonically as modules are added, and the method is intended to support learning from large-scale offline video for complex visuomotor tasks.

    Source: §1 · p. 2, §5 · p. 8

    Evaluation beyond six modules and application to harder, longer-horizon manipulation and navigation tasks remain future directions.

    Source: §5 · p. 8

Abstract

from arXiv · show

A video prediction model that generalizes to diverse scenes would enable intelligent agents such as robots to perform a variety of tasks via planning with the model. However, while existing video prediction models have produced promising results on small datasets, they suffer from severe underfitting when trained on large and diverse datasets. To address this underfitting challenge, we first observe that the ability to train larger video prediction models is often bottlenecked by the memory constraints of GPUs or TPUs. In parallel, deep hierarchical latent variable models can produce higher quality predictions by capturing the multi-level stochasticity of future observations, but end-to-end optimization of such models is notably difficult. Our key insight is that greedy and modular optimization of hierarchical autoencoders can simultaneously address both the memory constraints and the optimization challenges of large-scale video prediction. We introduce Greedy Hierarchical Variational Autoencoders (GHVAEs), a method that learns high-fidelity video predictions by greedily training each level of a hierarchical autoencoder. In comparison to state-of-the-art models, GHVAEs provide 17-55% gains in prediction performance on four video datasets, a 35-40% higher success rate on real robot tasks, and can improve performance monotonically by simply adding more modules.

1. Introduction

Large, diverse video datasets expose severe underfitting, while larger models face GPU/TPU memory constraints and hierarchical VAEs are difficult to optimize end-to-end. GHVAEs greedily train hierarchical modules and report gains over state-of-the-art models on video prediction and real-robot tasks.

Motivation

  • Accurate video prediction from raw pixels could let robots plan actions toward task objectives, and broad unlabeled datasets offer a potential training source.The paper frames learning from large, diverse offline data as a challenge.Source: §1 · p. 1
  • Existing stochastic video prediction methods underfit on large datasets, while hierarchical VAEs can model multiple levels of future stochasticity but are difficult to optimize end-to-end.Future observations combine uncertainty in low-level features and higher-level attributes.Source: §1 · p. 2

Approach

  • Greedy, modular optimization is proposed to address both GPU/TPU memory constraints and the optimization challenges of large-scale video prediction.Sequentially training submodules avoids end-to-end training and breaks bidirectional dependencies among latent variables.Source: §1 · p. 2

Results

  • GHVAEs outperform state-of-the-art models by 17–55% in FVD score across four video datasets and by 35–40% in real-robot task success rate.The paper also reports monotonically improving performance as modules are added.Source: §1 · p. 2

2. Related Work

Prior work addresses stochasticity and hierarchical modeling in video prediction, but large-scale underfitting and optimization remain challenges. GHVAE is positioned as a greedy alternative that scales without additional inductive biases.

The Underfitting Challenge of Large-Scale Video Prediction

  • Stochastic latent-variable methods improve prediction on smaller datasets, but existing methods underfit on large, multi-domain, or multi-robot datasets.Dense connections improved hierarchical VAEs on smaller datasets, but deep optimization problems limited scaling.Source: §2 · p. 2
  • Alternative scaling approaches change pixel distributions or predict object-centric and goal-centric representations; GHVAE scales without requiring additional inductive biases.Source: §2 · p. 2

Greedy Machine Learning

  • Earlier greedy learning focused on local feature extraction and time-invariant information, whereas this work uses hierarchical latent variables to model time-varying information about the future.Prior greedy training also supported stacked autoencoders and high-resolution image generation.Source: §2 · p. 3

Hierarchical Variational Autoencoders

  • End-to-end hierarchical VAEs face optimization challenges from bidirectional dependencies among latent variables; prior approaches add structures or regularizers to alleviate them.Those approaches have largely succeeded in image generation, while this paper studies video prediction.Source: §2 · p. 3

3. Greedy Hierarchical VAEs (GHVAEs)

GHVAEs stack local VAE modules and train them sequentially, freezing earlier modules to add capacity without fitting the whole model in memory. The design uses compressed spatial representations and deepest-layer stochastic prediction, with theoretical guarantees for its training objective and module-wise improvement.

Overview

  • Sequentially trained modules can be stacked to incrementally add model capacity without fitting the entire model into memory.The method locally optimizes each module using the frozen weights of preceding modules.Source: §3 · pp. 3–4

Encoder

  • Encoders compress spatial dimensions while expanding channels, allowing the deepest representation to retain information with less spatial correlation.The encoder design addresses the mean-field approximation used for stochastic latent dimensions.Source: §3 · p. 4

Decoder

  • Decoders combine hidden representations with stochastic latent information to predict the next image, incrementally expanding spatial dimensions while compressing channels.The deepest decoder also takes the current hidden state so it can use information directly from the past.Source: §3 · p. 4

Prior Network

  • All K stochastic latent variables are used during training, but only the deepest latent is predicted from the prior and used at test time.This design lets decoders propagate uncertainty from the deepest layer toward pixel space.Source: §3 · p. 4

Posterior Inference Network

  • The posterior inference network maps the current module’s future hidden representation to a diagonal Gaussian latent distribution and is used only during training.Source: §3 · p. 5

Optimization

  • A K-module GHVAE is trained in K phases, optimizing only module k in phase k while previously trained module weights remain frozen.Each module comprises an encoder, decoder, prior network, and posterior inference network.Source: §3 · pp. 4–5

Theoretical Guarantees

  • Sequential module optimization maximizes a lower bound on the end-to-end ELBO, and adding modules can only raise the GHVAE ELBO.A fixed posterior standard deviation is used in the KL term to improve training stability.Source: §3 · p. 5

4. Experimental Evaluation and Analysis

The evaluation compares GHVAE video prediction with prior models across video datasets and tests downstream robot tasks, then examines scalability, optimization, overfitting, and learned priors. Results favor GHVAEs across the reported video benchmarks and robot tasks, while ablations characterize how module count and training choices affect performance.

Video Prediction Performance

  • GHVAEs improve FVD by 17–55% and LPIPS by 13–45% over SVG’ across three datasets, with human preferences favoring GHVAE predictions more than 85% of the time.The reported comparisons cover all metrics across the three datasets; Table 1 summarizes test performance.Source: Table 1 · p. 6, §4 · p. 6
  • On Cityscapes, six-module GHVAEs outperform Hier-VRNN by 26% in FVD, 18% in SSIM, and 27% in LPIPS.The six-module GHVAE’s convolutional layers were downsized by 40% for the 16GB-memory comparison.Source: Table 2 · p. 6, §4 · p. 6

Ablation 1: Monotonic Improvement and Scalability of GHVAEs

  • Greedily adding modules from two to four to six improves GHVAE performance across all metrics, supporting monotonic improvement in practice.The authors interpret this as evidence that greedy module addition can scale to large datasets.Source: Table 3 · p. 6, §4 · p. 7

Ablation 2: Greedy vs. End-to-End Optimization of GHVAEs

  • On RoboNet, greedy optimization performs better than end-to-end training, which failed to reach good performance in any of five runs.End-to-end fine-tuning after greedy pretraining produces minor performance gains, but end-to-end training from scratch is less stable.Source: Table 4 · p. 7, §4 · p. 7

Ablation 3: Train-Test Comparison for GHVAEs

  • A six-module GHVAE shows similar train and test performance on RoboNet, but some overfitting on KITTI, Human3.6M, and Cityscapes.The authors hypothesize that the latter datasets’ smaller sizes contribute to overfitting, and that Human3.6M’s test subjects are unseen.Source: Table 5 · p. 7, §4 · p. 7

Ablation 4: Performance Contribution of Learned Prior

  • A learned prior significantly outperforms a uniform diagonal Gaussian prior, particularly on action-conditioned datasets.The authors hypothesize that the learned prior contains action information that the uniform prior lacks.Source: Table 6 · p. 7, §4 · p. 7

Real Robot Performance

  • A six-module GHVAE outperforms SVG’ by 40% in Pick&Wipe success rate and 35% in Pick&Sweep success rate, reaching 85% success on Pick&Sweep.The robot was pretrained on RoboNet and fine-tuned on 5,000 interaction videos; test objects, tools, and containers were unseen during training.Source: Table 7 · p. 8, §4 · p. 8, Figure 3 · p. 8

5. Conclusion

GHVAEs stack local VAE modules trained greedily to form a stably optimized hierarchical video-prediction model. The conclusion reports gains across video datasets and robot tasks, while the appendix describes memory savings from freezing previously trained modules.

5. Conclusion

  • GHVAEs stack local VAE modules sequentially and optimize them greedily to build an expressive, stably optimized hierarchical video-prediction model.Source: §5 · p. 8
  • GHVAEs outperform state-of-the-art methods by 17–55% in FVD across four video datasets and by 35–40% in real-robot task success rate.Performance can also improve monotonically by stacking more modules.Source: §5 · p. 8
  • Experiments beyond six modules remain a future direction for assessing the method’s full potential.Source: §5 · p. 8

A.1.1 GHVAE

  • Greedy training avoids storing frozen-module gradients during back-propagation, and reported memory savings increase with the number of modules.The memory table uses batch size 1 per GPU, a rollout horizon of 10, two context frames, and 64 × 64 × 3 observations.Source: §A.1.1 · p. 12, Table 8 · p. 12

A.1.2 Other Methods

The appendix contrasts GHVAEs with other ways to address memory constraints and offers intuition for modeling multiple levels of uncertainty in video prediction. It also outlines the supplementary evaluation and visualization materials.

A.1.2 Other Methods

  • Increasing accelerator count, distributing examples, or partitioning weights are complementary to GHVAEs; GHVAEs can still enable larger models when per-GPU batch size falls to 1.Source: §A.1.2 · p. 13
  • End-to-end training from scratch failed to converge to good performance in any reported run, whereas greedy training showed higher optimization stability.Source: §A.1.2 · p. 13

A.2. Intuition

  • RoboNet predictions must account for uncertainty in both the robot’s end-effector position and the textures of objects revealed as its arm moves.The passage presents these as high- and low-level sources of uncertainty, respectively.Source: §A.2 · p. 13

B.1. Video Prediction

  • The video-prediction appendix covers qualitative results, metric calculations, human evaluation, and additional ablations.Source: §B.1 · p. 14

B.1.1 Visualizations

  • Supplementary rollouts show video-prediction examples and real-robot experiments for Pick&Sweep and Pick&Wipe.Source: §B.1.1 · p. 14

B.1.2 Performance Evaluation

The appendix states that its metric-calculation methodologies match those used in prior work and describes the GPU-memory reporting scope and a model-size adjustment for Cityscapes.

B.1.2 Performance Evaluation

  • Metric-calculation methodologies match prior work to support fair comparisons.Source: §B.1.2 · p. 14
  • Table 9 reports GPU memory usage for experiments in Tables 1 and 2, with all Cityscapes 6-module GHVAE convolutional layers downsized by 40% to fit 16GB GPU memory.Source: Table 9 · p. 14
  • Table 9’s memory comparisons are scoped to the experiments in Tables 1 and 2, rather than all possible model configurations.Source: Table 9 · p. 14

B.1.3 Human Evaluation

Human evaluation compares blinded GHVAE and SVG’ predictions against a known ground-truth video. A separate architecture ablation reports that spatial downsampling outperformed an alternative that retained or expanded spatial dimensions.

B.1.3 Human Evaluation

  • Workers completed 300 tasks, choosing which of two generated videos was more realistically similar to a ground-truth video they could identify.Selections were averaged and reported in Tables 1 and 2.Source: §B.1.3 · p. 14

B.1.4 Ablation for Encoder-Decoder Architectures

  • An encoder-decoder that reduced spatial dimensions iteratively and expanded channels significantly outperformed an alternative that retained or expanded spatial dimensions while reducing channels.The authors hypothesize that spatial reduction improves mean-field approximation in the deepest latent space.Source: §B.1.4 · p. 14

B.1.5 Ablation for Single vs. Multiple Latents

The single-deepest-latent design avoids noisy signals from earlier latents while allowing stochasticity to propagate through deterministic layers. The appendix also describes real-robot task setup, training data, and example executions.

B.1.5 Ablation for Single vs. Multiple Latents

  • Using all K latent variables produced suboptimal, noisy signals from earlier latents when spatial dimensions were large, motivating use of only the deepest latent.A single latent can represent multi-level stochasticity when decoders nonlinearly project it to preceding layers, which remain deterministic.Source: §B.1.5 · pp. 14, 16, 17
  • GHVAE visibly outperforms SVG′ in RoboNet, KITTI, Human3.6M, and Cityscapes prediction examples.The RoboNet comparison uses the same two context images; KITTI and Human3.6M use five, and Cityscapes compares against Hier-VRNN.Source: Figure 5 · p. 15, Figure 6 · p. 16, Figure 7 · p. 16

B.2.1 Setup

  • Pick&Wipe trials randomize plates, objects, tools, and initial locations; Pick&Sweep trials use unseen objects, dustpans, and tools with randomized placements.Each task has 20 trials and a 50-timestep success limit; Pick&Sweep requires using a present sweeping tool rather than the gripper.Source: §B.2.1 · pp. 17–18

B.2.2 Training Data

  • The robot models were pretrained on RoboNet and fine-tuned on 5,000 self-collected Franka videos because either dataset alone had a stated generalization drawback.RoboNet-only training had lower task success under unseen lighting and camera viewpoints, while training only on the 5,000-video dataset overfit to novel objects and tools; data collection was self-supervised.Source: §B.2.2 · p. 18

B.2.3 Task Execution

  • Figures 11 and 12 show example Pick&Sweep and Pick&Wipe executions by GHVAE and SVG′, respectively.The appendix describes visualizations of the real-robot task executions.Source: §B.2.3 · p. 18, Figure 11 · p. 18

B.2.4 Task Diversity

The appendix presents diverse real-robot evaluation scenes, describes the random-shooting planner, and states the ELBO-validity theorem used in the proof section.

B.2.4 Task Diversity

  • Evaluation uses objects unseen during training, and the examples illustrate varied real-robot environments and tools.The training-environment caption says evaluation objects are excluded from training and that Pick&Sweep and Pick&Wipe share the same 5,000-video training data.Source: §B.2.4 · p. 18, Figure 13 · p. 19, Figure 14 · p. 19

B.2.5 Planning

  • The real-robot planner samples 140 trajectories, selects the action subsequence with the lowest L1 loss to the goal image, and replans after execution.The selected sequence length is optimized within the sampled horizon; for Franka, T = 10.Source: §B.2.5 · pp. 19–20

C.1. Proof of Theorem 1

  • Theorem 1 is labeled ELBO Validity and is stated for any positive integer k and frozen weights of preceding modules.The proof setup considers optimal parameters for the last module of a k-module GHVAE.Source: §C.1 · p. 20

C.2. Proof of Theorem 2

The proof excerpts annotate derivation steps and state how the argument extends from action-free to action-conditioned prediction. A separate clarification specifies how actions enter Equation 2 in practice.

C.2. Proof of Theorem 2

  • The derivation labels steps as applying a variation lower bound, Jensen’s inequality, Bayes’ rule, conditional-independence removal, and algebra.The annotations distinguish unchanged derivation steps from annotations.Source: §C.2 · p. 21
  • A six-module GHVAE failure example misses the blue bowl’s movement; the authors describe this as slight RoboNet underfitting and hypothesize that more modules may resolve it.The proposed module counts are 8, 10, or 12.Source: Figure 15 · p. 22
  • The proof assumes action-free prediction; the authors state that the action-conditioned proof replaces conditioning on x_t with joint conditioning on x_t and a_t.The example changes p_k(x_t+1 | x_t, z_k) to p_k(x_t+1 | x_t, a_t, z_k).Source: §C.2 · p. 22

C.3. Clarification for Equation 2

  • In practice, actions enter the prior network; the clarification gives a more general Equation 2 form that conditions the likelihood on the action as well.The text notes that Equation 2 omits a_t from the likelihood term used in practice.Source: §C.3 · p. 22

D. Failure Case Analysis

The reported RoboNet failure analysis identifies residual underfitting and a possible partial-observability issue with monocular images.

D. Failure Case Analysis

  • Despite outperforming SVG′ and Hier-VRNN, the six-module model remains slightly underfit on RoboNet, with an example of inaccurate blue-bowl tracking.The authors hypothesize that 8, 10, or 12 modules could potentially address this failure.Source: D. Failure Case Analysis · p. 22
  • The authors hypothesize that monocular images can make robot grasping and future object motion partially observable; they suggest end-effector position or a second camera view as possible additions.Without precise 3D locations of the robot and bowl, the model may not determine whether the bowl has been grasped.Source: D. Failure Case Analysis · p. 22

Related papers

Loading 2103.04174v3…