SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Brown University, Google
Calvin Luo
Xingyi Zhou
Anurag Arnab
Cordelia Schmid
We aim to investigate whether end-to-end learning of visual reasoning can be achieved with general-purpose neural networks, with the help of visual pretraining. A positive result would refute the common belief that explicit visual abstraction (e.g. object detection) is essential for compositional generalization on visual reasoning, and confirm the feasibility of a neural network "generalist" to solve visual recognition and reasoning tasks. We propose a simple and general self-supervised framework which "compresses" each video frame into a small set of tokens with a transformer network, and reconstructs the remaining frames based on the compressed temporal context. To minimize the reconstruction loss, the network must learn a compact representation for each image, as well as capture temporal dynamics and object permanence from temporal context. We perform evaluation on two visual reasoning benchmarks, CATER and ACRE. We observe that pretraining is essential to achieve compositional generalization for end-to-end visual reasoning. Our proposed framework outperforms traditional supervised pretraining, including image classification and explicit object detection, by large margins.
This paper investigates if an end-to-end trained neural network is able to solve challenging visual reasoning tasks [23, 65, 66] that involve inferring causal relationships, discovering object relations, and capturing temporal dynamics. A prominent approach [49] for visual reasoning is to construct a structured, interpretable representation from the visual inputs, and then apply symbolic programs [42] or neural networks [17] to solve the reasoning task. Despite their appealing properties, such as being interpretable and enabling the injection of expert knowledge, it is practically challenging to determine what kinds of symbolic representations to use and how to detect them reliably from visual data. In fact, the most suitable symbolic representation may differ significantly across different tasks: the representation for modeling a single human's kinematics (e.g. with body parts and joints) is unlikely to be the same as that for modeling group social behaviors (e.g. each pedestrian can be viewed as its own complete, independent entity). With the success of unified neural frameworks for multi-task learning [7], it is desirable to utilize a unified input interface (e.g. raw pixels) and to have the neural network learn to dynamically extract suitable representations for different tasks. However, how to learn a distributed representation with a deep neural network that generalizes similarly to learning methods based on symbolic representation [66] for visual reasoning remains an open problem.
The key hypothesis we make in this paper is that a general-purpose neural network, such as a Transformer [55], can be turned into an implicit visual concept learner with self-supervised pretraining. An implicit visual concept refers to a vector-based representation in an end-to-end neural network, which can be "finetuned" directly on the downstream tasks. Some of the learned implicit representations may be discretized into human-interpretable symbols for the purposes of human understanding of and feedback to the model. Others may correspond to part of, or a combination of
(a) Neuro-Symbolic Concept Learning
(b) Attention over Object Embeddings
(c) Implicit Visual Concept Learning (ours)
Figure 1: Comparison between a neuro-symbolic approach, a hybrid approach with learned object embeddings [17], and our proposed approach for visual reasoning. The illustration of each model family flows upwards, where visual inputs are encoded by neural networks (stage 1), and then processed by symbolic programs or another neural network to generate reasoning predictions (stage 2). Compared to (a) and (b), our approach does not require a separate "preprocessing" stage to extract the symbolic representation from visual inputs, and the self-supervised pretrained neural network can be end-to-end "finetuned" to the downstream visual reasoning tasks.
human-interpretable symbols. As opposed to explicit symbolic representation (e.g. object detection), implicit visual concepts do not require pre-defining a concept vocabulary or constructing concept classifiers, and also do not suffer from the early commitment or loss of information issues which may happen when visual inputs are converted into explicit symbols or frozen descriptors (e.g. via object detection and classification). A comparison between our approach and those that utilize explicit symbols under a pipeline-styled framework is visualized in Figure 1.
Our proposed representation learning framework, implicit visual concept learner (IV-CL) consists of two main components: first, a single image is compressed into a small set of tokens with a neural network. This is achieved by a vision transformer (ViT) network [19] with multiple "slot" tokens (e.g. the [CLS] token in ViT) that attend to the image inputs. Second, the slot tokens are provided as context information via a temporal transformer network for other images in the same video, where the goal is to perform video reconstruction via the masked autoencoding [27] objective with the temporal context. Despite its simplicity, the reconstruction objective motivates the emergence of two desired properties in the pretrained network: first, to provide context useful for video reconstruction, the image encoder must learn a compact representation of the scene with its slot tokens. Second, to utilize the context cues, the temporal transformer must learn to associate objects and their implicit representation across time, and also capture the notion of object permanence – the existence of an object even when it is occluded from the visual observations.
We conduct extensive ablation experiments on the Compositional Actions and TEmporal Reasoning (CATER) [23] benchmark and the Abstract Causal REaasoning (ACRE) [66] benchmark. To better understand if and how end-to-end pretraining helps visual reasoning, we also consider the supervised pretraining paradigm, where the slot tokens in the Transformer network are pretrained to "decode" image-level labels or object locations and categories. Specifically, we adopt the Pix2Seq objective [13], which formulates object detection as an autoregressive "language" modeling task.
Our experimental results reveal the following observations: first, IV-CL learns powerful implicit representations that achieve competitive performance on CATER and ACRE, confirming that visual pretraining does help end-to-end reasoning. Second, the pretraining objective matters: networks pretrained on large-scale image classification benchmarks [15, 52] transfer poorly to the visual reasoning benchmarks, while object detection learns better representation for reasoning. However,
both are outperformed by IV-CL by large margins. Finally, we observe that the network inductive biases, such as the number of slot tokens per image, play an important role: on both datasets, we observe that learning a small number of slot tokens per image (1 for CATER and 4 for ACRE) lead to the best visual reasoning performance. To the best of our knowledge, our proposed framework is the first to achieve competitive performance on CATER and ACRE without the need to construct explicit symbolic representation from visual inputs.
In summary, our paper makes the following two main contributions: First, unlike common assumptions made by neuro-symbolic approaches, we demonstrate that compositional generalization for visual reasoning can be achieved with end-to-end neural networks and self-supervised visual pretraining. Second, we propose IV-CL, a self-supervised representation learning framework, and validate its effectiveness on the challenging CATER and ACRE visual reasoning benchmarks against supervised visual pretraining counterparts. Code and checkpoints will be released.
Neural Network Pretraining. Huge progress has been made towards building unified learning frameworks for a wide range of tasks, including natural language understanding [16, 48, 8, 40], visual recognition [36, 35, 63, 22], and multimodal perception [33, 50, 38, 24, 3]. Unfortunately, most of the "foundation models" [7] for visual data focus on perception tasks, such as object classification, detection, or image captioning. Despite improved empirical performance on the visual question answering task [32, 64], visual reasoning remains challenging when measured on more controlled benchmarks that require compositional generalization and causal learning [66, 23, 14]. It is commonly believed that symbolic or neurosymbolic methods [42, 62, 37, 4], as opposed to general-purpose neural networks, are required to achieve generalizable visual reasoning [61, 66, 65]. To our knowledge, our proposed framework is the first to demonstrate the effectiveness of a general-purpose end-to-end neural network on these visual reasoning benchmarks.
Self-supervised Learning from Images and Videos. Self-supervised learning methods aim to learn strong visual representations from unlabelled datasets using pre-text tasks. Pre-text tasks were initially hand-designed to incorporate visual priors [18, 69, 10]. Subsequent works used contrastive formulations which encourage different augmented views of the same input to map to the same feature representation, whilst preventing the model from collapsing to trivial solutions [45, 12, 28, 26, 2]. One challenge of the contrastive formulation is the construction of positive and negative views, which has been shown to critically impact the learned representation [12, 59, 51]. Whereas contrastively learned representations may not easily transfer across domains [46], our pretraining successfully generalizes to visually different datasets, such as from ACRE to RAVEN.
Our work is most related to masked self-supervised approaches. Early works in this area used stacked autoencoders [56] or inpainting tasks [47] with convolutional networks. These approaches have seen a resurgence recently, inspired by BERT [16] and vision transformers [19]. BEiT [6] encodes masked patches with discrete variational autoencoders and predicts these tokens. Masked Autoencoders (MAE) [27], on the other hand, simply regress the pixel values of these tokens. MAE has been extended to regress features [57] and to learn video representations [53, 20]. Our training objective is different, as it is predictive coding based on compressed video observations. We confirm empirically that the proposed method outperforms MAE and its video extension by large margins.
Object-centric Representation for Reasoning. Most of the existing neuro-symbolic [42, 62] and neural network [17] based visual reasoning frameworks require a "preprocessing" stage of symbolic representation construction, which often involves detecting and classifying objects and their attributes from image or video inputs. Our proposed framework aims to investigate the effectiveness of single-stage, end-to-end neural networks for visual reasoning, which is often more desirable than the two-stage frameworks for scenarios that require transfer learning or multi-task learning. In order to obtain the object-centric, or symbolic representation in the preprocessing stage, one can rely on a supervised object detector [42, 54], such as Mask R-CNN [29]. An alternative approach is to employ self-supervised objectives and learn low-level features that are correlated with objects, such as textures [21, 30, 44], or objects themselves [9, 41, 11]. In practice, supervised or self-supervised approaches for object detection and object-centric representation learning may suffer from the lack of supervised annotations, or from noisy object detection results. For example, it was previously observed that object-centric representations are beneficial for transfer learning to temporal event classification only when ground truth object detections are used for training and evaluation [68].
Figure 2: IV-CL self-supervised pretraining. We consider the video reconstruction objective via masked autoencoding: A ViT-B image encoder is tasked to (1) extract visual representations (orange) for the unmasked patches per image and (2) compress an image into a small set of slot tokens (blue). A temporal transformer then propagates the information from the slot representations and patch-level representations from neighboring frames, which are essential for successful reconstruction.
We now introduce the proposed implicit visual concept learning (IV-CL) framework. We follow the pretraining and transfer learning paradigm: during pretraining (Figure 2), we task a shared image encoder to output patch-level visual embeddings along with a small set of slot tokens that compress the image's information. The pretraining objective is inspired by masked autoencoding (MAE) for unlabeled video frames, where the aim is to reconstruct a subset of "masked" image patches given the "unmasked" image patches as context. Compared to the standard MAE for images [27], the image decoder has access to two additional types of context information: (1) The encoded patch embedding from the unmasked image patches of the neighboring frames; (2) The encoded slot tokens from a subset of context frames. The context information is encoded and propagated by a temporal transformer network. To successfully reconstruct a masked frame, the image encoder must learn a compact representation of the full image via the slot tokens, and the temporal transformer has to learn to capture object permanence and temporal dynamics.
After pretraining, the image decoder is discarded, and only the image encoder and temporal transformer are kept for downstream visual reasoning tasks. The inputs to the temporal transformer are the slot tokens encoded from individual, unmasked video frames. We use the full finetuning strategy where the weights of both the newly added task decoder (e.g. a linear classifier), and the pretrained image and temporal transformers are updated during transfer learning.
Image Encoder: We adopt the Vision Transformer (ViT) backbone to encode each image independently: an input image is broken into non-overlapping patches of  $16 \times 16$  pixels, which are then linearly projected into patch embeddings as inputs to the transformer encoder. Spatial information is preserved by sinusoidal positional encodings. We use the standard ViT-Base configuration which has 12 Transformer encoder layers. Each layer has hidden size of 768, MLP projection size of 3072, and 12 attention heads. During pretraining, a subset of video frames are spatially masked randomly given a masking ratio, only the unmasked image patches are fed into the ViT-B encoder. For context frames and during transfer learning, all image patches are provided as inputs to the image encoder.
Slot Tokens: In the seminal work by Locatello et al. [41], slot tokens are defined as soft cluster centroids that group image pixels, where the goal is unsupervised object detection. Each slot token repeatedly attends to the raw image inputs and is iteratively refined with a GRU network. We borrow their terminology, and use slots to denote the representational bottleneck in which we hope to encode implicit visual concepts, such as object-centric information. We generalize their slot update rules
by: (1) iteratively updating the visual representation with layers of the Transformer encoder (ViT); (2) replacing cross-attention with multi-headed self-attention; (3) using MLP layers with untied weights to update the intermediate slot representation as opposed to a shared GRU network. These two modifications allow us to implement "slot attention" directly with a Transformer encoder, simply by prepending slot tokens as additional inputs to the encoder (similar to [CLS] tokens). The initial slot embeddings at the input of the visual encoder are implemented as a learnable embedding lookup table. To compare the effectiveness of different methods to aggregate "slot" information, we also explore single-headed soft attention and Gumbel-max attention as used by [60].
Temporal Transformer: To propagate temporal information across frames, we use another transformer encoder (with fewer layers than the ViT-B image encoder) which takes the tokens encoded by the image encoder as its inputs. During pretraining, the slot tokens from context frames, along with the unmasked patch tokens from the query frames are concatenated together and fed into the temporal transformer. For each query image, the temporal transformer outputs its corresponding unmasked patch tokens contextualized from both the unmasked patches from neighboring query frames and the slot tokens from context frames. The contextualized patches are then fed into the image decoder to compute the reconstruction loss. To preserve temporal position information, we use learned positional embeddings (implemented with an embedding lookup table). When finetuned on a reasoning task, the temporal transformer takes the slot tokens encoded by the image encoder as its inputs. Putting the image encoder and the temporal transformer together, the overall video encoder used for finetuning can be viewed as a factorized space-time encoder proposed by [5]. It is more parameter-efficient than the vanilla video transformer used by [53].
Image Decoder for Pre-training: We use the same image decoder as in [27]. The query images are decoded independently given the contextualized unmasked patch tokens. The image decoder is implemented with another transformer, where masked patch tokens are appended to the contextualized unmasked patch tokens as inputs to the image decoder. Sinusoidal positional encodings are used to indicate the spatial locations of individual patch tokens. We use the same number of layers, hidden size, and other hyperparameters as recommended by [27]. During pre-training, we use mean squared error to measure the distance between the original query image patches and the reconstructed patches.
Transfer Learning: As the goal of pre-training is to learn the slot tokens which we hope to compress an input image into a compact set of implicit visual concept tokens, we only ask the image encoder to generate the slot tokens during finetuning, which are fed to the temporal transformer as its inputs. We then average pool the output tokens of the temporal transformer and add a task-specific decoder to make predictions. Both benchmarks used in our experiments can be formulated as multi-class classification: for CATER, the goal is to predict the final location of the golden snitch, where the location is quantized into one of the  $6 \times 6$  positions; in ACRE, the goal is to predict whether the platform is activated, unactivated, or undetermined given a query scenario. We use linear classifiers as the task-specific decoders with standard softmax cross-entropy for transfer learning.
Supervised Pretraining Baselines: To better understand if visual pretraining helps end-to-end reasoning, we consider two types of supervised pretraining baselines. The first is the "classical" image classification pretraining which often exhibits scaling laws [52] when transferred to other visual recognition benchmarks. The second is the object detection task, which intuitively may also encourage the emergence of object-centric representations (per task requirement) inside the neural network. Both pretraining objectives can be directly applied on the same Transformer architecture as utilized in IV-CL, with different designs on the task specific decoders (which are discarded for visual reasoning finetuning). For image classification, we directly treat the slot token as a [CLS] token and add a linear classifier on top of it. For object detection, to make minimal modification to our framework, we follow the design proposed by Pix2Seq [13], which parameterizes the bounding box annotations as discrete tokens, and formulates the training objective as an autoregressive sequence completion task. The inputs to the autoregressive decoder are the encoded slot tokens. We adopt the same sequence construction and augmentation strategies as in Pix2Seq.
Benchmarks: In the classic "shell game", a ball is placed under a cup and shuffled with other empty cups on a flat surface; then, the objective is to determine which cup contains the ball. Inspired by
this, CATER is a dataset composed of videos of moving and interacting CLEVR [34] objects. A special golden ball, called the "snitch", is present in each video, and the associated reasoning task is to determine the snitch's position at the final frame. Solving this task is complicated by the fact that larger objects can visually occlude smaller ones, and certain objects can be picked up and placed down to explicitly cover other objects; when an object is covered, it changes position in consistence with the larger object that covers it. In order to solve the task, a model must learn to reason not only about objects and movement, but also about object permanence, long-term occlusions, and recursive covering relationships. Each video has 300 frames, and we use the static camera split for evaluation.
The ACRE dataset tests a model's ability to understand and discover causal relationships. The construction of the dataset is motivated by the Bicket experiment in developmental psychology [25], where there is a platform as well as many distinct objects, some of which contain the "Bicketness" property. When at least one object with the "Bicketness" property is placed on the platform, music will be played; otherwise, the platform will remain silent. In ACRE, the platform is represented by a large pink block that either glows or remains dim depending on the combination of CLEVR objects placed on it. Given six evidence frames of objects placed on the platform, the objective of the reasoning task is to determine the effect a query frame, containing a potentially novel object combination, would have on the platform. Possible answers including activating it, keeping in inactive, or indeterminable.
Pretraining data: We use the unlabeled videos from the training and validation splits of the CATER dataset for pretraining. Both the static and moving camera splits are used, which contains 9,304 videos in total. In our experiments, we observe that ACRE requires higher resolution inputs during pretraining and finetuning. Our default preprocessing setup is to randomly sample 32 frames of size  $64 \times 64$  for pretraining the checkpoints that are transferred to CATER, and 16 frames of size  $224 \times 224$  for pretraining the checkpoints that are transferred to ACRE. The randomly sampled frames are sorted to preserve the arrow of time information. No additional data augmentations are performed.
Transfer learning: For CATER, we evaluate on the static split which has 3,065 training, 768 validation, and 1645 test examples. We select the hyperparameters based on the validation performance, then use both training and validation data to train the model to be evaluated on the test split. By default, we use 100 randomly sampled frames of size  $64 \times 64$  during training, and 100 uniformly sampled frames of stride 3 during evaluation. For ACRE, we explore all three splits, all of which contain 24,000 training, 8,000 validation, and 8,000 test examples. We use the validation set to select hyperparameters and use both training and validation to obtain the models evaluated on the test split.
Default hyperparameters: We use the Adam optimizer for pretraining with a learning rate of  $10^{-3}$ , and the AdamW optimizer for transfer learning with a learning rate of  $5 \times 10^{-5}$ . The pretraining checkpoints are trained from scratch for 1,000 epochs using a batch size of 256. For transfer learning, we finetune the pretrained checkpoints for 500 epochs using a batch size of 512. All experiments are performed on TPU with 32 cores. Below we study the impact of several key model hyperparameters.
4.2 IV-CL vs. Supervised Pretraining