Visual Pre-Training on Unlabeled Images using Reinforcement Learning

TL;DR

This paper introduces a novel approach to image pre-training by framing it as a reinforcement learning problem, enabling the learning of visual features from unlabeled data through a value function that models image transformations.

Contribution

It proposes a reinforcement learning-based framework for self-supervised image pre-training, extending the analogy between RL value functions and feature learning from unlabeled images.

Findings

Improved visual representations on unlabeled image datasets.

Effective learning from diverse data sources like videos and web images.

Enhanced feature quality compared to traditional self-supervised methods.

Abstract

In reinforcement learning (RL), value-based algorithms learn to associate each observation with the states and rewards that are likely to be reached from it. We observe that many self-supervised image pre-training methods bear similarity to this formulation: learning features that associate crops of images with those of nearby views, e.g., by taking a different crop or color augmentation. In this paper, we complete this analogy and explore a method that directly casts pre-training on unlabeled image data like web crawls and video frames as an RL problem. We train a general value function in a dynamical system where an agent transforms an image by changing the view or adding image augmentations. Learning in this way resembles crop-consistency self-supervision, but through the reward function, offers a simple lever to shape feature learning using curated images or weakly labeled captions when they exist. Our experiments demonstrate improved representations when training on unlabeled images in the wild, including video data like EpicKitchens, scene data like COCO, and web-crawl data like CC12M.