Visual Representation Learning with Stochastic Frame Prediction

TL;DR

This paper tackles the challenge of learning robust visual representations from videos by addressing the under-determined nature of future-frame prediction with a stochastic frame-prediction model. The proposed RSP framework combines a time-dependent prior/posterior over latent variables with a shared decoder that also supports a masked autoencoding objective, enabling efficient joint learning of temporal relations and dense frame information. Empirical results across vision-based robot learning and video label propagation show that RSP consistently outperforms strong SSL baselines, with notable gains over SiamMAE, and ablations highlight the benefits of stochastic framing, discrete latent variables, and a shared decoder. The work demonstrates that modeling uncertainty in future frames yields representations that transfer well to downstream tasks requiring temporal understanding, and it suggests future extensions to diffusion-based video generation and larger-scale models.

Abstract

Self-supervised learning of image representations by predicting future frames is a promising direction but still remains a challenge. This is because of the under-determined nature of frame prediction; multiple potential futures can arise from a single current frame. To tackle this challenge, in this paper, we revisit the idea of stochastic video generation that learns to capture uncertainty in frame prediction and explore its effectiveness for representation learning. Specifically, we design a framework that trains a stochastic frame prediction model to learn temporal information between frames. Moreover, to learn dense information within each frame, we introduce an auxiliary masked image modeling objective along with a shared decoder architecture. We find this architecture allows for combining both objectives in a synergistic and compute-efficient manner. We demonstrate the effectiveness of our framework on a variety of tasks from video label propagation and vision-based robot learning domains, such as video segmentation, pose tracking, vision-based robotic locomotion, and manipulation tasks. Code is available on the project webpage: https://sites.google.com/view/2024rsp.

Figures (6)

• Figure 1: Representation learning with stochastic frame prediction. (a) We train a stochastic frame prediction model, which is built upon stochastic video generation model denton2018stochastic, which consists of an encoder that extracts representations, a posterior model with access to both current and future frames, a prior model with only access to the current frame, and a decoder that generates frame conditioned on features from the current frame and a sample from either posterior or prior distributions. We train the model to accurately generate the future frame while enforcing the posterior and prior distributions to be close to each other, i.e., encourage the posterior distribution to be more predictable and the prior distribution to predict the future. (b) We introduce an auxiliary masked autoencoding objective he2022masked with a shared decoder architecture. Our decoder makes the [MASK] tokens attend to different inputs via the cross-attention layer, enabling us to share the decoder parameters for different objectives.
• Figure 2: Examples of visual observations from CortexBench majumdar2023we, RLBench james2020rlbench, and FrankaKitchen gupta2019relay, which we used for training imitation learning agents that learn a mapping from observations to expert actions. Learning such agents requires representations that can understand both temporal and dense information.
• Figure 3: Aggregate results on vision-based robot learning. We report the interquartile mean agarwal2021deep over 20 vision-based robot learning tasks from CortexBench majumdar2023we, RLBench james2020rlbench, and Frana Kitchen gupta2019relay.
• Figure 4: Qualitative results. We provide examples of predicted propagation from RSP on video object segmentation pont20172017, video part segmentation zhou2018adaptive, and pose tracking jhuang2013towards benchmarks. "ref" indicates the ground-truth annotations, and 25, 50, and 100% refers to the propagated ratio of the videos. We provide additional qualitative results in \ref{['appendix:additional_qualitative_results']}.
• Figure 5: Effect of KL loss scale. We report the learning curves of models trained with different KL loss scales ($\beta$).