SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Aran Nayebi $^{1,*}$ , Rishi Rajalingham $^{1,4}$ , Mehrdad Jazayeri $^{1,2}$ , and Guangyu Robert Yang $^{1,2,3}$
<sup>1</sup>McGovern Institute for Brain Research, MIT; Cambridge, MA 02139
$^{2}$ Department of Brain and Cognitive Sciences, MIT; Cambridge, MA 02139
$^{3}$ Department of Electrical Engineering and Computer Science, MIT; Cambridge, MA 02139
$^{4}$ Reality Labs, Meta; 390 9th Ave, New York, NY 10001
Correspondence: anayebi@mit.edu
Humans and animals have a rich and flexible understanding of the physical world, which enables them to infer the underlying dynamical trajectories of objects and events, plausible future states, and use that to plan and anticipate the consequences of actions. However, the neural mechanisms underlying these computations are unclear. We combine a goal-driven modeling approach with dense neurophysiological data and high-throughput human behavioral readouts that contain thousands of comparisons to directly impinge on this question. Specifically, we construct and evaluate several classes of sensory-cognitive networks to predict the future state of rich, ethologically-relevant environments, ranging from self-supervised end-to-end models with pixel-wise or object-slot objectives, to models that future predict in the latent space of purely static image-pretrained or dynamic video-pretrained foundation models. We find that "scale is not all you need", and that many state-of-the-art machine learning models fail to perform well on our neural and behavioral benchmarks for future prediction. In fact, only one class of models matches these data well overall. We find that neural responses are currently best predicted by models trained to predict the future state of their environment in the latent space of pretrained foundation models optimized for dynamic scenes in a self-supervised manner. These models also approach the neurons' ability to predict the environmental state variables that are visually hidden from view, despite not being explicitly trained to do so. Finally, we find that not all foundation model latents are equal. Notably, models that future predict in the latent space of video foundation models that are optimized to support a diverse range of egocentric sensorimotor tasks, reasonably match both human behavioral error patterns and neural dynamics across all environmental scenarios that we were able to test. Overall, these findings suggest that the neural mechanisms and behaviors of primate mental simulation have strong inductive biases associated with them, and are thus far most consistent with being optimized to future predict on reusable visual representations that are useful for Embodied AI more generally.
Within the span of a couple seconds, we are able to draw rich inferences and make predictions about novel scenes [Smith and Vul, 2013, Battaglia et al., 2013]. A dominant cognitive theory has been that the brain builds mental models of the physical world, using those models to make inferences about the
future state of its environment [Craik, 1943]. In the past decade, this hypothesis has been supported by comparisons of human behavior to computational models which predict what will happen next in physical scenarios via forward simulations resembling those of game engines in modern video games [Battaglia et al., 2013, Hamrick et al., 2016, Ullman et al., 2017]. Both neuroimaging work in humans [Zacks, 2008, Fischer et al., 2016, Schwettmann et al., 2019, Pramod et al., 2022] and recent electrophysiological work in monkeys [Rajalingham et al., 2022a,b] has further provided evidence for the neurobiological basis of mental simulations in the frontoparietal network (FPN) of primates, a large-scale network consisting of several interacting brain regions. In this work, we make progress towards understanding the neural and behavioral mechanisms of mental simulation by constructing models that perform this behavior in rich, naturalistic environments. Specifically, we aim to determine what inductive biases (in the form of a loss function, architecture class, and pretraining environment) enable the brain to generally perform mental simulation across a range of environments and scenarios, from unstructured, continuous sensory inputs. In particular, we assess not only model generalization to novel, high-variation examples within the same environment, but also structural generalization to new environments and scenarios altogether.
Predicting the physical dynamics of environments is also critical to progress in Embodied AI. One common paradigm for learning these dynamics has been as a next frame prediction problem via a pixel-wise loss [Villegas et al., 2019, Wu et al., 2021, Babaeizadeh et al., 2021, Nash et al., 2022]. These losses emphasize prioritizing accurate prediction of every detail of a given scene's dynamics. However, fine-grained prediction of upcoming video frames would require near-perfect knowledge of the world's physical state (akin to Laplace's Demon), which may explain the observation why many of these models tend to underfit in high variation, naturalistic visual environments, with recent efforts aimed at scaling these methods up primarily by increasing their parameter count [Dasari et al., 2019, Babaeizadeh et al., 2021]. It is therefore unclear how much these resultant learned representations are able to successfully capture general physical understanding.
Another recent class of approaches involves the design of visual "foundation models" [Bommasani et al., 2021], trained on large amounts of webscale images and egocentric videos to develop an implicit representation of the world, that can then be deployed to downstream robotic manipulation tasks [Nair et al., 2022, Ma et al., 2023, Majumdar et al., 2023]. Of course, these models are not directly designed to do explicit physical simulation, but we equip them with a forward dynamics model that can be rolled out for an arbitrary number of timesteps. We ask whether such dynamically-equipped foundation models have learned physical knowledge by evaluating their generalization both to new scenarios and environments, and whether their representations bear any similarity to humans and non-human primates performing the same tasks?
In particular, we find strong constraints on primate mental simulation, especially when examining generalization within and across diverse environments. Our core result is that a small class of models best match primate frontal cortex neural dynamics while the animal plays a ball interception task in which the ball trajectory is partially occluded prior to hitting the paddle ("Mental-Pong"), developed previously by Rajalingham et al. [2022a]. Overall, one model class can match both neural response dynamics and human behavioral patterns reasonably well – namely, dynamics that are optimized to future predict in the latent space of VC-1, a video foundation model pretrained on the largest variety of egocentric sensorimotor settings overall. We therefore currently observe a tight correspondence between the ability to predict fine-grained neural and behavioral responses for the mental simulation phenomenon, and developing useful representations for Embodied AI more generally.
Mental simulations have been studied at the level of neural activity only very recently. Prior human neuroimaging studies [Zacks, 2008, Fischer et al., 2016, Pramod et al., 2022] in the past decades showed elevated levels of blood-oxygen-level-dependent (BOLD) signal to mental simulation, although they do not have the required resolution to verify that these dynamics are actually represented in underlying neural activity. Rajalingham et al. [2022b] was the first study to show that neural dynamics recorded from macaque dorsomedial frontal cortex (DMFC), track the occluded ball by comparing these dynamics to Recurrent Neural Networks (RNNs) that simulate the occluded ball's position in Mental-Pong, and finding that they better match these dynamics than RNNs that only perform ball endpoint prediction. However, monkeys can perform these tasks without substantial training, suggesting that they are already equipped with the necessary neural foundations for mental
simulation in this environment. Therefore, we aim to also build networks that are not explicitly trained on Mental-Pong itself, but are tasked to generalize to this novel setting as a test of their general understanding of physical scene dynamics – chiefly developed through three factors: their architecture, optimization objective, and pretraining on a naturalistic environment.
Additionally, we constrain our models by evaluating them against high-throughput human behavioral data (from Bear et al. [2021]) in more naturalistic, 3D environments than Mental-Pong alone, which goes beyond prior behavioral studies that either rely on a narrow range of physical scenarios [Shepard and Metzler, 1971, Cooperau and Shepard, 1973], such as block towers with several cubes of different colors [Groth et al., 2018, Li et al., 2016], or 2D environments that may not generalize to the real world [Bakhtin et al., 2019]. A key challenge to addressing these questions is a common standard to evaluating the everyday physical scene understanding and neural predictivity of these models, especially since they are usually trained on vastly different scenarios and input types. Towards this end, we require models to operate under similar constraints as the brain, namely (i) to take in unstructured visual inputs across a range of physical phenomena, (ii) to generate physical predictions for each scene (i.e. producing "behavioral outputs"), and (iii) to consist of internal units that can be compared to biological units (i.e. containing "artificial neurons").
Taken together, these three requirements encompass a large class of functionally reasonable hypotheses that we call "sensory-cognitive networks", and includes the two broad approaches mentioned in  $\S 1$ . However, they do exclude some approaches – for example, particle-based graph neural network dynamics predictors [Battaglia et al., 2016, 2018, Li et al., 2019, Sanchez-Gonzalez et al., 2020] that take the ground truth simulator state as input (which may not be readily available in real-world situations, failing to satisfy requirement (i)) or probabilistic programs [Battaglia et al., 2013, Hamrick et al., 2016, Ullman et al., 2017] (which fail to satisfy requirements (i) and (iii)).
Nonetheless, we believe these latter approaches from cognitive science are a useful guide for building improved models from pixels that satisfy the three requirements - especially in terms of assessing whether the prediction problem lies at the level of vision, or the dynamics that interfaces with it. For example, [Bear et al., 2021, Figure 5] demonstrated that particle-based graph neural networks (e.g. DPI-Net [Li et al., 2019]) with access to the ground truth simulator state approach human-level physical predictions on the OCP task that we consider in Section 4 across a diverse range of scenarios. This observation indicates that the onus mainly rests on learning a good visual representation, as these models assume perfect observability of physical dynamics. This further motivates our enforcement of requirement (i) as being a worthwhile challenge, as well as our latent future prediction approach of "factorizing" the mental simulation problem into separately pretrained vision and dynamics modules.
To tackle this question, we took a hypothesis-driven approach and built sensory-cognitive networks that performed mental simulations of their environment in a variety of different ways (schematized in Figure 1). Specifically, we tasked models to operate on the Physion dataset [Bear et al., 2021], a large-scale video dataset that focuses on everyday physical understanding, consisting of eight different scenarios in a simulated Unity3D-based environment (the ThreeDWorld simulator [Gan et al., 2021]) with roughly 2,000 scenes each. Models are pretrained on all eight scenarios of Dominoes, Support, Collide, Contain, Drop, Link, Roll, and Drape; which together cover many scenes involving rigid and soft bodies.
We additionally pretrain a subset of models on Kinetics-700 [Carreira et al., 2019], which consists of over 500,000 training videos from real-world (rather than simulated) scenes from 700 different action categories, in order to assess whether a different dataset of increased scale is beneficial or not, relative to Physion.
We consider models from several sensory-cognitive hypothesis classes:
1. End-to-end self-supervised future prediction:
(a) Pixel-wise: This class of models consists of three parts: encoder, dynamics, and decoder; which are altogether pretrained end-to-end with a pixel-wise loss to predict the next frame. We consider a current state-of-the-art model in this model family, FitVid [Babaeizadeh et al., 2021], and we pretrain it on Physion (with and without temporal augmentations such as RandAugment [Cubuk et al., 2020]); along with an
(A) Model Pretraining
(B) Model Evaluations
Figure 1: Model Pretraining and Evaluation Pipeline: We built and evaluate sensory-cognitive networks for mental simulation. (A) Models are pretrained to predict the future state of naturalistic environments in a variety of ways. Their dynamics are then compared to high-throughput human behavioral judgements and dense primate neurophysiological response dynamics. Model pretraining can be (1) two-stage, or (2) end-to-end. For the two-stage latent future prediction models, there is first visual module pretraining (blue), then dynamics module pretraining on Physion (peach). Across all models, the dynamics modules are pretrained on the same dataset (Physion), and then evaluated against neural and behavioral data with model weights fixed. (B) Evaluations are multi-fold: (1) Comparison to human behavior in OCP on held-out scenes across pretrained Physion scenarios. (2) Comparison to neural activity in Mental-Pong, an out-of-distribution environment that none of the models were pretrained on.
earlier variant, SVG [Denton and Fergus, 2018], that can be additionally trained at larger image scales (  $128 \times 128$  rather than  $64 \times 64$  pixels).
(b) Object-slot: This class of models is based on the Contrastive Learning of Structured World Models framework ("C-SWM" [Kipf et al., 2020]) which contrastively learns more structured object-slot latent states. It consists of an encoder (whose size we vary: "small", "medium", and "large") that outputs a fixed number of object slots  $N = 10$ , and a graph neural network forward dynamics module that operates on this representation. Since the dynamics module itself does not have temporal dependencies (only spatial), we enable temporal dependencies by passing frames in the input channel dimension of the encoder, as a sliding window of temporal dimension of  $T$  context frames to predict the object-slot latent at timestep  $T + 1$ . This model family is one instantiation of the cognitive theory that humans tend to reason about scenes with regards to objects and their relations [Baillargeon et al., 1985, Spelke, 1990, Spelke and Kinzler, 2007]. These Gestalt principles therefore may provide physical knowledge that are better-positioned to generalize than more fine-grained (pixel-wise) prediction - a hypothesis that we explicitly evaluate on high-throughput neural and behavioral data, in what follows.
2. Latent self-supervised future prediction: These models consist of two parts: an encoder  $\mathcal{E}$  that produces a latent state space  $h$  and a dynamics module  $\mathcal{D}$  that predicts the future state
purely in the latent space  $h$  of  $\mathcal{E}$ .  $\mathcal{E}$  is fixed and pretrained to perform a challenging vision task ("foundation model"), through which it learns a partial, implicit representation of the physical world. However, since  $\mathcal{E}$  is not pretrained to do explicit physical simulation, we train additional dynamics  $\mathcal{D}$  on Physion, that can later be "rolled out" an arbitrary number of steps. More concretely, given  $T$  context frames,  $\mathcal{E}$  produces the latent sequence  $h_{1:T}$ , from which  $\mathcal{D}$  is trained to predict  $h_{T+1}$ .  $\mathcal{D}$  is relatively simple, being either an LSTM [Hochreiter and Schmidhuber, 1997] or a continuous-time RNN (CTRNN) [Miller and Fumarola, 2012]; or "No Dynamics" as a control, which always outputs the last context frame latent  $h_T$  of  $\mathcal{E}$ . We consider a variety of foundation encoders, divided into two primary classes based on their type of large-scale pretraining dataset:
(a) Image Foundation Models (Static scenes):
- Standard Convolutional Neural Networks (CNNs) VGG16 [Simonyan and Zisserman, 2014] and ResNet-50 [He et al., 2016] pretrained on ImageNet with a supervised categorization objective.
- Vision Transformer (ViT) [Dosovitskiy et al., 2021] based architectures such as DeiT [Touvron et al., 2021] and DINO [Caron et al., 2021] pretrained on ImageNet with a self-supervised objective.
- DINOv2 [Oquab et al., 2023], which is a very recent ViT pretrained on a larger curated dataset called LVD-142M (142 million images).
- CLIP [Radford et al., 2021], which is a ViT pretrained on 400 million image-text pairs curated from the Internet (250 times larger than ImageNet).
(b) Video Foundation Models (Dynamic scenes):
- R3M [Nair et al., 2022], which is a ResNet-50 architecture pretrained with a temporally-contrastive video-language alignment objective on 5 million frames from a subset of the recent large-scale Ego4D human video dataset [Grauman et al., 2022].
- VIP [Ma et al., 2023], which is a ResNet-50 architecture pretrained with a goal-conditioned value function objective on 5 million frames from a subset of Ego4D.
- VC-1 [Majumdar et al., 2023], which is a very recent ViT pretrained on 7 different egocentric video sources (over 5.6 million frames, including Ego4D) relevant to sensorimotor skills, using a self-supervised masked autoencoding (MAE) [He et al., 2022] objective. While MAE is pretrained on individual frames, the statistics of the frames are informed by egocentric motion, unlike webscale images.
In total, these networks encompass both qualitatively distinct hypotheses (pixel-wise vs. object-slot vs. latent future prediction), alongside several variations within each hypothesis class. This combination of diverse networks allows us to potentially differentiate hypothesis classes across a range of functionally capable instantiations of each hypothesis.
4 Comparison to Human Physical Judgements
4.1 OCP task evaluation
In the object contact prediction (OCP) task [Bear et al., 2021], each evaluation scenario involves a red "agent" object and a yellow "patient" object (which did not appear during model pretraining). Both humans and models are tasked to predict the probability that they will come into contact. This prediction requires understanding of the relevant physical phenomenon in the given scenario, corresponding to a higher-order readout of the underlying scene dynamics.
4.2 End-to-end pixel-wise future predictors best predict human behavior in the same environment
The model comparison results are summarized in Figure 2. In Figure 2A, we examine held-out scene accuracy of models across all eight scenarios in the OCP task, where chance performance is  $50\%$ . We can see that humans are quite reliable on this task, attaining  $74.04\%$  average held-out accuracy across scenarios (grey horizontal line). Furthermore, the best models that approach human accuracy are models that are pretrained end-to-end with pixel-wise losses (FitVid and SVG). Despite
Figure 2: Model Comparisons to Human Physical Judgements: (A) Each model is compared to binary ("yes/no") accuracy on the OCP task (chance is  $50\%$ ). Grey horizontal bar represents the human accuracy on this task. Mean and s.e.m. across all eight Physion scenarios, weighted by number of scenes in each scenario. Hatched bars represent models pretrained on the Kinetics-700 dataset, rather than on Physion. (B) Each model is compared to human subject judgement probabilities for the OCP task, via Pearson's correlation from the Logistic Regression classifier trained on each of the model dynamics. Grey is the correlation of human judgement probabilities to each other. Mean and s.e.m. across all eight Physion scenarios, weighted by number of scenes in each scenario. (C) A model's match to human judgement probabilities is strongly correlated with its OCP task accuracy. (D) Per scenario model comparisons to human subject judgement probabilities for the OCP task. Mean and s.e.m. across all 41 models.
their differences in architecture, both FitVid and SVG are comparable to one another in their OCP test set accuracy (61.12%-63.31%), attaining non-significant differences in the distribution of their accuracies across scenarios (Paired  $t$ -test, minimum Bonferroni corrected  $p$ -value of 0.444 across pairwise comparisons). For a fixed architecture, pretraining with increased image size helps improve accuracy somewhat on the OCP prediction task (SVG  $128 \times 128$ , rightmost periwinkle bar, vs. SVG  $64 \times 64$ , leftmost one), along with temporal augmentations for a high capacity model (FitVid, rightmost dark green bar vs. leftmost one), but these differences are overall non-significant (Paired  $t$ -test  $p = 0.461$  and  $p = 0.196$  within the SVG and FitVid architectures, respectively). However, not all end-to-end models match human accuracy well. In particular, we see that the object-slotted C-SWM class of models matches human accuracies least well compared to other model classes, despite varying the encoder size. The class of latent future prediction models overall better matches these behavioral patterns than the more explicitly structured C-SWM models, but there does not appear to be strong differentiation across foundation models, and also not much differentiation between having a dynamics module vs. using the encoder latents directly (rightmost two bars vs. leftmost bar in each group). This suggests that at least for the OCP task, either Physion is not sufficiently high-variation enough to pretrain the dynamics module or most of the predictivity comes from the physical scene understanding encapsulated in the foundation model encoder latents. The latter appears to be more likely, since pretraining the dynamics module on a larger-scale dataset such as Kinetics-700 for a subset of the models (VIP+LSTM/CTRNN, VC-1+LSTM/CTRNN, and R3M+LSTM/CTRNN; hatched bars) does not improve OCP accuracy over the base encoder latents either, relative to pretraining them on Physion (Paired  $t$ -test, minimum Bonferroni corrected  $p$ -value of 0.066 across architectures).
Additionally, we can look at finer-grained error patterns of the probabilities reported by humans compared to those of models (rather than accuracies alone), summarized in Figure 2B. Here we see that despite the metric being more detailed, humans are quite consistent with each other, suggesting that this type of behavioral metric is quite reliable (grey horizontal line). In fact, all models appear to be further from the human-to-human consistency than the OCP accuracy. However, overall, similar trends appear to hold across models as with the OCP accuracy measure - where the end-to-end pixel-wise models (FitVid and SVG) match these consistency scores the best across models, and the object-slotted C-SWM models match them the least well. From an AI perspective, it is actually quite relevant to work towards matching human error patterns, as it is highly correlated with the primarily performance based measure of OCP accuracy, as seen in Figure 2C ( $R \approx 0.865$ ,  $p \ll 0.001$ ).