Reasoning-Modulated Representations
Authors: Petar Veličković
Matko Bošnjak
Thomas Kipf
Alexander Lerchner
raia hadsell
Razvan Pascanu
Charles Blundell
Venue: ICLR 2022 Submitted (2022)
Keywords: representation learning
algorithmic reasoning
graph neural networks
relational learning

Abstract: Neural networks leverage robust internal representations in order to generalise. Learning them is difficult, and often requires a large training set that covers the data distribution densely. We study a common setting where our task is not purely opaque. Indeed, very often we may have access to information about the underlying system (e.g. that observations must obey certain laws of physics) that any "tabula rasa" neural network would need to re-learn from scratch, penalising data efficiency. We incorporate this information into a pre-trained reasoning module, and investigate its role in shaping the discovered representations in diverse self-supervised learning settings from pixels. Our approach paves the way for a new class of data-efficient representation learning.

Under review as a conference paper at ICLR 2022
REASONING-MODULATEH REPEESSTTOTORREEEEEETTTTTTOOS
REPRESENTATIONS
Anonymous authors
Paper under double-blind review
ABSTRACT
Neural networks leverage robust internal representations in order generalise. Learning them is difficult, and often requires a large training set that covers the data distribution densely. We study common setting where our task is not purely opaque. Indeed, very often we may have access to information about the under- lying system (e.g. that observations must obey certain laws of physics) that any "tabula rasa" neural network would need to re-learn from scratch, penalising data efficiency. We incorporate this information into a pre-trained reasoning module, and investigate its role in shaping the discovered representations in diverse self- supervised learning settings from pixels. Our approach paves the way for new class of data-efficient representation learning.
INTRODUCTION
Neural networks able to learn policies in environments without access to their specifies (Schrit- twieser et al. 2020), generate large quantities of text (Brown et al. 2020), or automatically fold pro- teins high accuracy (Senior al. 2020). However, such "tabula rasa" approaches hinge having access to substantial quantities of data, from which robust representations can be learned. Without large training that spans the data distribution, representation learning is difficult (Belkin al., 2019
LeCun 2018 Sun et al. 2017).
Here, study ways construct neural networks with representations that robust, while re- taining a data-driven approach. We rely on a simple observation: very often, we have some (partial) knowledge the underlying dynamics of the data, which could help make stronger predictions from fewer observations. This knowledge, however, usually requires us to be mindful of abstract proper- ties of the data and such properties cannot always be robustly extracted from natural observations.
Nilu
Vbtu
Xorg
Y org
(r v₁)
Xred
Fred
8yed E vi(xes &yod)
Xyel
Xpur
Y pur
Xgrn
gra
Figure I: Bouncing balls example (re-printed, with permission, from Battaglia et al. (2016)). Natural inputs, x, correspond to pixel observations. Predicting future observations (natural outputs, y), can be simplified as follows: if we able to extract set of abstract inputs, X. (e.g. the radius, position and velocity for each ball), the movements in this space must obey the laws of physics.
Motivation Consider the task of predicting the future state of a system of n bouncing balls, from a pixel input x (Figure []. Reliably estimating future pixel observations, y, is a challenging recon- struction task. However, the generative properties of this system are simple. Assuming knowledge of simple abstract inputs (radius, re- position, X., and velocity, v.) for every ball, x.. the future move- ments in this abstract space are the result of applying the laws of physics to these low-dimensional quantities. Hence, future abstract states, y, can be computed via a simple algorithm that aggregates pair-wise forces between objects.

---

Under review as a conference paper at ICLR 2022
While this gives a potentially simpler path from pixel inputs to pixel outputs, via abstract inputs to abstract outputs (x - X ~4 y + y), it still places potentially unrealistic demands on our task setup, every step of the way:
x: Necessitates either upfront knowledge of how to abstract away from x, or massive dataset of paired (x, x) to learn such a mapping from;
~ y: Implies that the algorithm perfectly simulates all aspects of the output. In reality, an algo- rithm may often only give partial context about y. Further, algorithms often assume that x is provided without error, exposing an algorithmic bottleneck (Deac et al., 2021 if is incorrectly predicted, this will negatively compound in y, hence y:
y y Necessitates renderer that generates y from y, or a dataset of paired (R.y) to learn it.
We will assume a general setting where none of the above constraints hold: we know that the mapping x - y is likely of to our predictor, but we do not assume trivial mapping paired dataset which would allow us to convert directly from x to x or from to y. Our only remaining assumption is that the algorithm & y can be efficiently computed, allowing us to generate massive quantities of paired abstract onput-output pairs, (x. y).
Present work In this setting, we propose Geasoning-Modduland Representations (RMR), ap-
proach that first learns latent-space processor of abstract data; i.c. mapping x 4 - z' 2 y, where z E I are high-ding nigh-dimensional latent vectors. f and are an encoder and decoder, designed to take abstract representations to and from this latent space, and p is a processor network which simulates the algorithm X ~ y in the latent space.
We then observe, in the spirit of neural algorithmic reasoning (Veličković & Blundell 202T). that such a processor network can be used as drop-in differentiable component for any task where the
y kind of reasoning may be applicable. Hence, we then learn pipeline 1 - z' 2 y, which modulates the representations z obtained from x, forcing them to pass through the pre- trained processor network. By doing so, have ameliorated the original requirement for massive natural dataset of yx pairs. Instead, we inject knowledge from a massive abstract dataset of y) pairs, directly through the pre-trained parameters of P. This has the potential relieve the pressure on encoders and decoders f and ÿ, which we experimentally validate on several challenging representation learning domains.
Our contributions can be summarised as follows:
Verifying and extending prior work, we show that meaningful latent-space models can be learned from massive abstract datasets, physies simulations and Atari 2600 games;
We then show that these latent-space models can be used as differentiable components within neural pipelines that process raw observations. In doing so, recover a data- efficient neural network pipeline that relies solely on the existence of massive abstract datasets (which can often be automatically generated).
Finally, we demonstrate early signs of processor rensability: latent-space abstract models can be used in tasks which do not even directly align with their environment, so long as these tasks can benefit from their underlying reasoning procedure.
2 RELATED WORK
Neural algorithmic reasoning RMR relies on being able construct robust latent-space models that imitate abstract reasoning procedures. This makes it well aligned with neural algorithmic rea- soning (Cappart et al., 202T), which is concerned with constructing neuralised versions of classical algorithms (typically by learning to execute them in a manner that extrapolates). Leveraging the ideas of algorithmic alignment (Xu et al. 2019), several known algorithmic primitives have already been successfully neuralised. This includes iterative computation (Veličković et al.) 2019 Tang et al. 2020), linearithmic algorithuss (Freivalds et al.] 2019). and data structures (Strathmann et al.
202 Veličković et al.) (020). Further, the XLVIN model (Deac et al. (021)) demonstrates how such
primitives can be re- re-used for data--efficad planning, paving the way for a blueprint (Veličković & Blundell 2021) that we leverage in RMR as well.
2

---

Under review as a conference paper at ICLR 2022
Physical simulation with neural networks Our work also has contact points with prior art in using (graph) neural networks for physics simulations. In fact, there is a tight coupling between al- gorithmic computation and simulations, as the latter typically realised using the former. Within this space, abstract GNN models of physies have been proposed by works such as interaction net- works (Battaglia et al. 2016) and NPE (Chang al. 20T6), and extended pixel-based inputs by visual interaction networks (Watters et 2017). The generalisation power of these models has in- creased drastically recent years, with effective models of systems particles Sanchez- Gonzalez et al. (2020) as well as meshes (Pfaff et al., 2020) being proposed. Excitingly, it has also been demonstrated that rudimentary laws of physics can occasionally be recovered from the update rules of these GNNs (Cranmer et al. 2020), and that they can be used to uncover new physical knowledge
(Bapst et al. 2020).
Recent work has also explored placing additional constraints on learning- baaninggbaased physical simulators, for example by using Hamiltonian ODE integrators in conjunction with GNN models (Sanchez- Gonzalez al.) 2019), or by coupling (non-neural) differentiable physics engine directly to visual inputs and optimizing its parameters via backprop (Jaques et al. 2019; Jatavallabhula et al. 2021)
Object-centric and modular models for dynamic environments RMR with factored latents can be viewed as form of object-centric neural network, in which visual objects in image video are represented as separate latent variables in the model and their temporal dynamics and pairwise interactions modeled via GNNs self-attention mechanisms. There rich literature on dis- covering objects and learning their dynamics from raw visual data without supervision, with object- centrie models such R-NEM (Van Steenkiste al.] 2018). SQAIR (Kosiorek al. 2018), OP3
/Veerapaneni et al. 2020). SCALOR (Jiang et al. 2019). G-SWM (Lin et al.) 2020). Recent work has explored using contrastive losses (Kipf et al. 2019) in this context or other losses directly in
latent space |François- Lavet et al. 2019p. Related approaches discover and use keypoints (Kulkarni
Lavett
al. 2019) to describe objects and even discover causal relations from visual input (Li et al. 2020) using neural relational inference (Kipf et al. 2018) in conjunction with keypoint discovery method. related line of works integrate attention mechanisms modular and object-centric models to in- terface latent variables with visual input, including models such as RMC (Santoro et al. 20188). RIM
Goyal et al. 2019). Slot Attention (Locatello et al. 2020), SCOFF (Goyal et al. 2021b), and NPS
Goyal al. 2021a)
RMR ARCHITECTURE
P
Abtract impats
Latents
Abatrect cutpptt
Sutu
Vidu
Sorg
Yorg
%
x
x;
X
Spur
X
y
Natural inputs
Nutural outputs
X
y
Figure representation learner (RMR).
3

---

Under review as a conference paper at ICLR 2022
Having provided high-level overview of RMR and surveyed the relevant related work, proceed to carefully detail the blueprint of RMR's 's various components. This will allow us to ground any subsequent RMR experiments diverse domains directly in our blueprint. Throughout this section, it will be useful to refer to Figure 2 which presents a visual overview of this section.
Preliminaries We assume set of natural inputs, X, and a set of natural outputs, V. These sets represent the possible inputs and outputs of target function, ④b X - V, which we would like to learn based on (potentially small) dataset of input-output pairs, (x, where y = ((x).
We further assume that the inner workings of & can be related to an algorithm, A + The algorithm operates over a set of abstract inputs, X. and produces outputs from an abstract output sei V: Typically, it will be the case that dim X' < dim X'; that is, abstract inputs are assumed substantially lower-dimensional than natural inputs. We do not assume existence of any aligned input pairs (x, x), and we do not assume that A perfectly explains the computations of Φ. What do assume is that A is either known or can be trivially computed, giving rise to a massive dataset of abstract input-output pairs, y), where y A(x).
Lastly, we assume a latent space, 2, and that we can construct neural network components to both encode and decode from it. Typically, 2 will be real-valued vector space (Z R$) which is high-dimensional: that is, k > dim X'. This ensures that any neural networks operating over Z are not vulnerable to bottleneck effects.
Note that either the natural or abstract input set may be factorised, e.g.. into objects; in this case, we can accordingly factorise the latent space, enforcing = R where is the assumed maximal number of objects (typically hyperparameter of the models if not known upfront).
Abstract pipeline RMR training proceeds by first learning a model of the algorithm A, which is bound to pass through a latent- space representation. That is, we learn neural network approximator paradigm (Hamrick et al. 2018). It consists of the following three building blocks: Encoder, f X 2, tasked with projecting the abstract inputs into the latent space; Processor, P 2 + 2, simulating individual steps of the algorithm in the latent space; Decoder, 9 Z + 3r tasked with projecting latents back into the abstract output space. Such pipeline is now widely used both in neural algorithmic reasoning Veličković et al. 2019) and learning physical simulations (Sanchez-
g( ((((()))) v A(x), which follows the mncode-process-deee
et al. 2020), and can
be trained end-to-end with gradient descent.
For reasons that will become apparent, performed by P. Accordingly, encoders and decoders are often designed to be simple learnable lin- ear projections, while processors tend to be either deep MLPs or graph neural networks depending on whether the latent space is factorised into nodes.
favourable for most of the computational effort to be
Natural pipeline Once an appropriate processor, P. has been learned, it may be observed that it corresponds to highly favourable component in our setting. Namely, can relate its operations to the algorithm A, and since it stays high-dimensional, it is a differentiable component we can easily plug into other neural networks without incurring any bottleneck effects. This insight was originally recovered in XLVIN (Deac et al., 202T), where it yielded a generic implicit planner. As an ablation, we have also rediscovered the bottleneck effect in our settings; see Appendix [ We now leverage similar insights for general representation learning tasks.
On a high level, what we need do is simple and elegant: swap out f and g for natural encoders and decoders, f X + Z and ğ Z + y. respectively. We are then able to learn a function g (P(((x))) x Φ (xx), which once again to be optimised through gradient descent. We would like P to retain its semantics during training, and therefore it is typically kept frozen in the natural pipeline. Note that P might not perfectly represent A which in turn might not perfectly represent d. While we rely on a skip connection in implementation of P. it has no learnable parameters and does not offer the system the ability to learn a correction of P in the natural setting. Our choice is motivated by the desire to both maintain the semantics and interpretability of P and to make this processor a bottleneck forcing the model to rely on it. We show empirically that our pipeline is surprisingly robust to imperfect P models even with weak (linear) encoders/decoders.
4

---

Under review as a conference paper at ICLR 2022
It is worth noting several potential challenges that may arise while training the natural pipeline, especially if the training data for it is sparsely available. We also suggest remedies for each:
If x and/or y exhibit any nontrivial geometry, simple linear projections will rarely suffice for f and ÿ. For example, our natural inputs will often be pixel-based, necessitating a convolutional neural network for
Further, since the parameters of P are kept frozen, f is left with a challenging task of mapping natural inputs into appropriate manifold that P can meaningfully operate over. While we demonstrate clear empirical evidence that such meaningful mappings definitely occur, we remark that its success may hinge on carefully tuning the hyperparameters of f.
A very common setting assumes that the abstract inputs and latents are factorised into objects, but the natural inputs are not. In this case, f is tasked with predicting appropriate object representations from the natural inputs. This is known to be challenging feat (Greff et al.] 2220)), but can be successfully performed. Sometimes arbitrarily factorising the feature maps of CNN (Kipf et al. 2019) is sufficient, while at other times, models such as slot attention (Locatello et al. 2020) may be required.
One corollary of using automated object extractors for is that it's very difficult to enforce their slot representations to line up in the same way as the abstract inputs. This implies that P should be permutation equivariant (and hence motivates using a GNN for it).
RMR FOR BOUNCING BALLS
To evaluate the capability of the RMR pipeline for transfer from the abstract space to the pixel space, we apply it on the "bouncing balls" problem. The bouncing balls problem is an instance of a physics simulation problem, where the task is to predict the next state an environment in which multiple balls are bouncing between each other and a bounding box. Though this problem had been studied in the the context of physics simulation from (abstract) trajectories (Battaglia et al. (016) and from (natural) videos (Watters et al.) 2017; Van Steenkiste et al. 2018 Löwe et al. (2020), here we focus on the aptitude of RMR to transfer learned representations from trajectories to videos.
Our results affirm that strong abstract models can be trained such tasks, and that including them in video pipeline induces more robust representations. See Appendix for more details on hyper- parameters and experimental setup.
Preliminaries Here, trajectories are represented by 2D coordinates of 10 balls through time, defin- ing our abstract inputs and outputs X' y R R00x2 We slice these trajectories into series of moving windows containing the input, x" spanning history of three previous states, and the target, y. representing the next state. We obtain these trajectories from 3D simulator (MuJoCo (Todorov) et al. 2).2)), together with their short-vide renderings, which represent our natural input and out- put space X y R Our goal is to train an RMR abstract model on trajectories and transfer learned representations to improve a dynamics model trained on these videos.
Abstract pipeline So as to model the dynamics of trajectories, we closely follow the RMR desiderata for the abstract model. We to linear projection the input concatenation, P to a Message Passing Neural Network (MPNN), following previous work (Battaglia et al. 2016 Sanchez-Gonzalez al. 2020), and g a linear projection.
Our model learns transition function g (P(((X))) & y. supervised using Mean Squared Error (MSE) over ball positions in the next step. It achieves an MSE of 4.59 x 10-4, which, evaluated qualitatively, demonstrates the ability of the model to predict physically realistic behavior when unrolled for 10 steps (the model trained on 1-step dynamics only). Next, we take the processor P from the abstract pipeline and re-use it in the natural pipeline.
Natural pipeline Here we evaluate whether the pre-trained RMR processor can be reused for leamning the dynamics of the bouncing balls from videos. The pixel- encoder ff is concatenation of per-inuut-imag Slot Attention model (Locatello et al., 2020), passed through a linear layer, and Broadcast Decoder (Watters et al. 2019) for the pixel-based decoder g.
5

---

Under review as a conference paper at ICLR 2022
The full model is transition function ((P(f(x') = y, supervised by pixel reconstruction loss over the next step image. We compare the performance of the RMR model with a pre-trained processor P against baseline in which P is trained fully end-to-end. The RMR model achieves an MSE of 7.94++0.41 (x10-1), whereas the baseline achieves 9. 9.47++0.24 ( x x10-1). We take a qualitative look at the reconstruction rollout of the RMR model in Figure and expose the algorithmic bottleneck properties for this task in Appendix C
Figure 3: RMR for bouncing balls reconstruction rollout. States marked in green are the natural input, followed by the reconstructed output. The states below the reconstruction is the ground truth.
CONTRASTIVE RMR FOR ATARI
We evaluate the potential of our RMR pipeline for state representation learning on the Atari 2600 (Bellemare et al., 2013). We find the RMR applicable here because there is a potential wealth of by inspecting its RAM traces.
information that can be obtained about the Atari's s
Preliminaries Accordingly, we will define our set of abstract inputs and outputs as Atari RAM matrices. Given that the Atari has 128 bytes of memory, X' = y = B128888 (where = {0, 1} is the set of bits). We collect data about how the console modifies the RAM by acting in the environment and recording the trace of RAM arrays we observe. These traces will be of the form (x. a, y) which signify that the agent's initial RAM state was x. and that after performing action A. its RAM state was updated to y. We assume that a is encoded as an 18-way one-hot vector.
We would like to leverage any reasoning module obtained over RAM states to support representation learning from raw pixels. Accordingly, our natural inputs, X, are pixel arrays representing the Atari's framebuffer.
Mirroring prior work, we perform contrastive learning directly in the latent space, and set y = 2; that is, our natural outputs correspond an estimate of the "updated" latents after taking action. All our models use latent representations of 64 dimensions per slot, meaning Z = IR 128
We note that it is important to generate a diverse dataset of experiences in order to train robust RAM model. To simulate a dataset which might be gathered by human players of varying skill, we sample our data using the 32 policy heads of a pre-trained Agent57 (Badia al. (020). Each policy head collects data over three episodes in the studied games. Note that this implies substantially more challenging dataset than the one reported by Anand et al. (2099), wherein data was collected by purely random policy, which may well fail to explore many relevant regions of the games.
Abstract pipeline Firstly, we set out to verify that it is possible to train nontrivial Atari RAM transition models. The construction of this abstract experiment follows almost exactly the abstract RMR setup: ff and are appropriately sized linear projections, while P needs to take into account which action was taken when updating the latents. To simplify the implementation and allow further model re-use, we consider the action a part of the P's inputs. See Appendix
F for detailed equations.
This implies that transition model learns function g (P(f(x), a)) y. We supervise this model using binary cross-entrog cotsooy to predict each bit of the resulting RAM state. Since RAM transitions assumed deterministic, assume a fully Markovian setup and learn 1-step dynamics.
For brevity purposes, we detail our exact hyperparameters and results per each Atari game consid- ered in Appendix B Our results ascertain the message passing neural network (MPNN) (Gilmer

---

Under review as a conference paper at ICLR 2022
et al. 2017) highly potent processor network in Atari: it ranked most potent in 17 out of 24 games considered, compared to MLPs and Deep Sets (Zaheer et al. 2017). Accordingly, we will focus on leveraging pre-trained MPNN processors for the next phase of the RMR pipeline.
Natural pipeline We now set out to evaluate whether our pre-trained RMR processors can be meaningfully re-used by an encoder in pixel-based contrastive learning pipeline.
For our pixel-based encoder f. we use the same CNN trunk as Anand et al. require slot-level rather than flat embeddings, the final layers our encoder are different. Namely, we apply 1 x 1 convolution computing 128m feature maps (where m is the number of feature maps per-slot). We then flatten the spatial axes, giving every slot m X h x w features, which we finally linearly project to 64-dimensional features per-slot, aligning with our pre-trained P. Note that setting m = recovers exactly the style of object detection employed by C-SWM (Kipf et al. 2019). Since desired outputs are themselves latents, is single linear projection to 64 dimensions.
Overall, our pixel-based transition model learns function ((P(f(x)) a)) ss f(y). where y is the next state observed after applying action a in state x. To optimise it, we re-use exactly the same TransEEinsppred (Bordes et al. 2013) contrastive loss that C-SWM (Kipf et al., (209)) used.
Once state representation learning concludes, all components are typically thrown away except for the encoder, f. which is used for downstream tasks. As a proxy for evaluating the quality of the encoder, we train linear classifiers on the concatenation of all slot embeddings obtained from f to predict individual RAM bits, exactly as in the abstract model case. Note that we have not violated our assumption that paired (x, x) samples will not be provided while training the natural model-in this phase, the encoder f is frozen, and gradients can only flow into the linear probe.
Our comparisons for this experiment, evaluating our RMR pipeline against an identical architecture with an unfrozen P (equivalent to C-SWM (Kipf et al.) 2019)) is provided in Table For each of 20 random seeds, we feed identical batches to both models, hence we can perform a paired Wilcoxon test to assess the statistical significance of any observed differences on validation episodes. The results are in line with our hypothesis: representations learnt by RMR are significantly better (p < 0.05) on 15 out of 24 games, significantly worse only on Pitfall! and indistinguishable to C-SWM's on others. This is despite the fact their architectures are identical, indicating that the pre- trained abstract model induces stronger representations for predicting the underlying data factors. As was the case for bouncing balls, we expose the algorithmic bottleneck here too; see Appendix C
As relevant initial baseline-and and to emphasise the difficulty of our task-wa we also include in Table the performance of the latent embeddings extracted from a pre-trained Agent57 (Badia et al 2020). These embeddings are, perhaps unsurprisingly, substantially worse than both the RMR and C-SWM models. As Agent57 was trained on a reward reward-maximising maximising objective, its embeddings are likely to capture the controllable aspects of the input, while filtering out the environment's background.
Abstract model transfer While the results above indicate that endowing a self-supervised Atari representation learner with knowledge of the underlying game's RAM transitions yields stronger representations, one may still argue that this constitutes form of "privileged information". This is due the fact we knew upfront which game we were learning the representations for, and hence could leverage the specific RAM dynamics of this game.
In the final experiment, we study more general setting: we are given access RAM traces showing us how the Atari console manipulates its memory in response to player input, but we cannot guaran- tee these traces came from the same game that we are performing representation learning over. Can we still effectively leverage this knowledge?
Specifically, we test abstract model transfer in the following way: first, we take pre-trained abstract processor network P from one game (the "train game") and freeze it. Then, using this processor, we perform the aforementioned natural pipeline training and testing over frames from another game (the "test game") "). We then evaluate whether the recovered performance improves over the C-SWM model with unfrozen weights once again using a paired one-sided Wilcoxon test over 20 seeds to ascertain statistical significance of any differences observed. The final outcome of our experiment is hence a 24 x 24 matrix, indicating the quality of abstract model transfer from every game to every other game. Table T corresponds the "diagonal" entries of this matrix.
7

---

Under review as a conference paper at ICLR 2022
Table 1: Natural modelling results for Atari 2600. Bit-level F reported for slots with high entropy,
as in Anand et al. (2019). Results assumed significant < 0.05 (one-sided paired Wilcoxon test).
Game
Agent57
C-SWM
RMR
p-value
Asteroids
0.514=0.001
0.582#0.009
0.593=0.001
10-5
Battlezone
0.351 10.003
0.59210.005
0.589
0.056
Berzerk
0.4544
0,463 10.053
0.425
0.364
Bowling
0.054
1.944.00.06
0.946
0.071
Boxing
0.558=0.002
0.667
0.669=0.011
0.215
Breakout
0.657=0.001
0.836
0.852=0.008
10-6
Demon Attack
0.539=000004
0.653
0.658=0.0m
0.002
Freeway
0.424=0.052
0.912=0.025
0.919=0.035
0.032
Frostbite
0.405=0.000
0.580
0.594=0.016
0.035
H.E.R.O.
0.481
0.729 0.002
0.779 0.001
10-5
Montezuma's Revenge
0.743 10.000
0.82410.012
0.821 0.001
0.156
Ms. Pac-Man
0.506:20.00
0.599
0.602=0.006
0.038
Pitfall!
0.495=0.000
0.626
0.603=0.010
10-5
Pong
0.39210.001
0.750
0.762=0.010
0.001
Private Eye
0.594
00.000
0.867
0.045
Q*Bert
0.536=20.010
0.508
0.590=0.017
0.165
River Raid
0.68610.000
0.762
0.764 0.007
0.032
Seaquest
0.472
0.6341
0.653
10-5
Skiing
0.599
0.766
0.775
0.174
Space Invaders
0.0882
0.719
0.761 00.006
10-5
Tennis
0.533
0.724
0.729 = 0.005
0.007
Venture
0.567=0.001
0.632++0000
0.633=0.001
0.392
Video Pinball
0.375=0.011
0.724
0.745=0.00
10-5
Yars' Revenge
0.608=0.001
0.715
0.751=0.010
< 10-5
The results, presented in Figure 4 testify to the performance of RMR. Representations learned by RMR transfer better in 64.6% of the train/test game pairs, are indistinguishable from C-SWM in 33.7% of the game pairs and perform worse than C-SWM only 1.7% of game pairs. Therefore, in plentiful circumstances, the answer to original question is positive: discovering trace of Atari RAM transitions of unknown origin can often be of high significance for representation learning from Atari pixels, regardless of whether the underlying games match.
Qualitative analysis of transfer While we find this result interesting in and of itself, it also raises interesting follow-up questions. Is representation learning on certain games more prone to being improved just by knowing anything about the Atari console? Figure 4 certainly implies so: several games (such as Seaquest or Space Invaders) have "fully-green" columns, making them "universal recipients". Similarly,
may be interested about the "donor" properties of each game to what extent are their RAM models useful across a broad range of test games?
We study both of the above questions by performing hierarchical cemmpleee-linke clustering of the rows and columns of the 24 x 24 matrix of transfer performances, to identify clusters of related donor and recipient games. Both clusterings are marked in Figure (on the sides of the matrix).
The analysis reveals several well-formed clusters, from which we are able to make some preliminary observations, based on the properties of the various games. To name few examples:
The strongest recipients (e.g. Yars' Revenge, H.E.R.O., Seaquest, Space Invaders and As- teroids) tend to include elements of "shooting" in their gameplay.
Conversely, the weakest recipients (River Raid, Berzerk, Private Eye, Ms. Pac-Man and Pitfall!) are all games in which movement is generally unrestricted across most of the screen, indicating a larger range of possible coordinate values to model.
Pong, Breakout, Battlezone and Skiing cluster closely in terms of donor properties- and they are all games in which movement restricted axis only.
8

---

Under review as a conference paper at ICLR 2022
8:2
8.24
17 0P4
Freeway
9:31
0.40
Pitfall!
8:35
0.2h
8.2%
0.1
TA
Private Eye
PMN
914
Tennis
Domon
13
0.8
Bowling
Montezuma's Revenge
Yars' Revenge
0:31
0:51
8:4
8.59
1X
Q"Bert
Train Geme
837
Invaders
Venture
227
0.22
53
031
833
644
Boxing
Sking
924
8 Battlezone
Breakout
Video Pinball
11
Ms. Pa Man
H.E.R.O.
Berzerk
8.00
4.2
0.0 -00
0.50
844
29 4.2
River Rald
827
32
126
Seaquest
River Roid
Berzerk
Private Eyo
Ms. Pac Man-
Pitall!
Demon Attack-
Tonnis-
Breakout
Yors Revarge-
H.E.R.O. 0.
Seaquest
Space Invaders
Asteroids
Video Piribll:
Bowling
Frostbite
Eattlezone:
Freeway
Rovarg Revenge
Venture-
Pong
Boxing
Sking
better
difference
worse
Test Game
Figure 4: A hierarchically clustered heatmap depiction the transfer results, where the Train Game abstract models have been trained the Test Games. The results are summarised by Wilcoxon test, denoting where the performance of RMR is better, does not differ or worse than C-SWM's. In each cell, the mean relative % improvement is noted. The significance cutoff is p 0.05.
Lastly, strong donors (Montezuma's Revenge, Yars' Revenge, QeBert and Venture) all feature massive, abrupt, changes to the game state (as they feature multiple rooms, for example). This implies that they might generally have seen more diverse of transitions. Further, on Yars' Revenge, a massive laser features, which, conveniently, prints an RGB projection of the RAM itself on the frame.
CONCLUSIONS
We presented Bessoning-MModulan Representations (RMR), novel approach for leveraging back- ground algorithmic knowledge within a representation learning pipeline. By encoding the under- lying algorithmic priors weights of processor neural network, we alleviate requirements on alignment between abstract and natural inputs, protect our model against bottlenecks, and make more data-efficient in the process. We believe that RMR paves the way novel class of algorith- mic ceusoneng-inspired representations, with a high potential for transfer across tasks a feat that has largely eluded deep reinforcement learning research.

---

Under review as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
Our work proposes general pipeline for representation learning under algorithmic priors. We have made a careful effort to expose all of the key elements of our RMR pipeline, both within figures and in text. We also report the hyperparameters of our experiments in the Appendices.
We carefully described the manner in which our datasets have been collected. Further, we will ex- plore the options for open-sourcing the specific datasets we employed in this paper post-uublication. We note that Atari trajectory datasets similar to the one we explore here (albeit with a significantly weaker actor) are already publicly available e.g. through the work of Anand et al. (2019).
REFERENCES
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté. and R Devon Hjelm. Unsupervised state representation learning in atari. arXiv preprint arXiv: arXiv:1906.088 1906.0826.. 2019.
Adria Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi, Zhaohan Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. In International Conference on Machine Learning, pp. 507-517. PMLR, 2020.
Victor Bapst, Thomas Keck, A Grabska-Barwinsskee Craig Donner, Ekin Dogus Cubuk, Samuel Schoenholz, Annette Obika, Alexander WR Nelson, Trevor Back, Demis Hassabis, et al. Un- veiling the predictive power static structure glassy systems. Nature Physics, 16(4):448-454, 2020.
Peter W Battaglia, Razvan Pascanu, Matthew Lai, Danilo Rezende, and Koray Kavukcuoglu. Interaction networks for learning about objects, relations and physics. arXiv preprint 2016.
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine- learning practice and the classical biri-vari trade-off. Proceedings of the National Academy of Sciences, 116(32):5549--1588
2019.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Antoine Bordes, Nicolas Usumier, Alberto Garcia-Duran, Jason Weston, and Oksana Yakhnenko. Translating embeddings for modeling multi-relational data. In Neural Information Processing Systems (NIPS), pp. 1-9, 2013.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 2020.
Quentin Cappart, Didier Chételat, Elias Khalil, Andrea Lodi, Christopher Morris, and Petar Veličković. Combinatorial optimization and reasoning with graph neural networks. arXiv preprint wrXiv:2102.005544 2021.
Michael B Chang, Tomer Ullman, Antonio Torralba, and Joshua B Tenenbaum. A compositional object-based approach to learning physical dynamics. arXiv preprint arXiv:1612.0033: 2016.
Miles Cranmer, Alvaro
Peter Battaglia, Rui Xu, Kyle Cranmer, David Spergel, and Shirley Ho. Discovering symbolic models from deep learning with inductive biases, arXiv preprint arXiv:2006.1188/ 2020.
Andreea- Loanaa Deac, Petar Veličković. Ognjen Milinkovie, Pierre-Lue Bacon, Jian Tang, and Mladen Nikolic. Neural algorithmic reasoners are implicit planners. Advances in Neural In- formation Processing Systems, 34, 2021.
10

---

Under review as a conference paper at ICLR 2022
Vincent François-Lavet, Yoshua Bengio, Doina Precup, and Joelle Pineau. Combined reinforce- ment learning via abstract representations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 3582-3589 2019.
Kärlis Freivalds, Emīls Ozoling, and Agris Sostaks. Neural shuffle-exchane networks sequence processing in (n log n) time. arXiv preprint arXiv:1907.000 0779777777777 2019.
setuurks-sequence
Justin Gilmer, Samuel Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning, pp. 1263-1272. PMLR, 2017.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. arXiv preprint arXiv: 10893, 2019.
Anirudh Goyal, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nico- las Heess, Michael Mozer, and Yoshua Bengio. Neural production systems. arXiv preprint arXiv:2103.01997, 2021a.
Anirudh Goyal, Alex Lamb, Phanideep Gampa, Philippe Beaudoin, Sergey Levine, Charles Blun- dell, Yoshua Bengio, and Michael Mozer. Factorizing declarative and procedural knowledge in structured, dynamic environments. In International Conference on Learning Representations, 2021b.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. arXiv preprint arXiv:2012.05208 2020.
Jessica B Hamrick, Kelsey R Allen, Victor Bapst, Tina Zhu, Kevin R McKee, Joshua B Tenen- baum, and Peter W Battaglia. Relational inductive bias for physical construction in humans and machines. arXiv preprint arXiv: 01 0030002200220 2018
Miguel Jaques, Michael Burke, and Timothy Hospedales. Physiesass-invers--grpiiiss Unsuper- vised physical parameter estimation from vidco. arXiv preprint
2019.
Krishna Murthy Jatavallabhula, Miles Macklin, Florian Golemo, Vikram Voleti, Linda Petrini, Mar- tin Weiss, Breandan Considine, Jerome Parent-Levesque, Kevin Xie, Kenny Brleben, et al. grad- sim: Differentiable simulation for system identification and visuomotor control. arXiv preprint arXiv:2 104.06604 2021.
Jindong Jiang, Sepehr Janghorbani, Gerard De Melo, and Sungjin Ahn. Scalor: Generative world models with scalable object representations. arXiv preprint arXiv:1910.023384, 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6996 2014.
Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning, pp. 2688- 2697. PMLR, 2018.
Thomas Kipf, Elise van der Pol, and Max Welling, Contrastive learning of structured world models. arXiv preprini 1911::911.122
2019.
Adam R Kosiorek, Hyunjik Kim, Ingmar Posner, and Yee Whye Teh. Sequential attend, infer, repeat: Generative modelling of moving objects. arXiv preprint arXiv:1806.00 017800060774,, 2018.
Tejas D Kulkarni, Ankush Gupta, Catalin lonescu, Sebastian Borgeaud, Malcolm Reynolds, Andrew Zisserman, and Volodymyr Mnih. Unsupervised learning of object keypoints for perception and control. Advances in neural information processing systems, 12:107441 2019.
Yann LeCun. The power and limits of deep learning. Research-Technnlogy Management, 61(6): 22-27, 2018.
Yunzhu Li, Antonio Torralba, Animashree Anandkumar, Dieter Fox, and Animesh Garg. Causal discovery in physical systems from videos. arXiv preprint arXiv:2007.0063), 2020.
11

---

Under review as a conference paper at ICLR 2022
Zhixuan Lin, Yi-Fu Wu, Skand Peri, Bofeng Fu, Jindong Jiang. and Sungjin Ahn. Improving gen- erative imagination in object-centric world models. In International Conference on Machine Learning. pp. 6140-6149. PMLR, 2020.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot atten- tion. arXiv preprint arXiv:2006. 15055, 2020.
Sindy Löwe, Klaus Greff, Rico Jonschkowski, Alexey Dosovitskiy, and Thomas Kipf. Learning object-centric video models by contrasting sets. arXiv preprint arXiv:201 10287, 2020.
Tobias Pfaff, Meire Fortunato, Alvaro Sunchez-Gonzalez and Peter W Battaglia. Learning mesh- based simulation with graph networks. arXiv preprint arXiv:2010.03400 2020.
Alvaro Sanchez-Gonzalez, Victor Bapst, Kyle Cranmer, and Peter Battaglia. Hamiltonian graph networks with ode integrators. arXiv preprint arXiv: 12790,, 2019.
Alvaro Sanchez-Gonzalez Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovee, and Peter Battaglia. Learning to simulate complex physics with graph networks. In International Confer- ence on Machine Learning, pp. 8469- PMLR, 2020.
Adam Santoro, David Raposo, David GT Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. A simple neural network module for relational reasoning. arXiv preprint 1706.07776.01222 2017.
Adam Santoro, Ryan Faulkner, David Raposo, Jack Rac, Mike Chrzanowski, Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy Lillicrap. Relational recurrent neural networks. arXiv preprint arXiv: 806.01822, 2018.
Julian Schrittwieser, Toannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):004-669 2020.
Andrew W Senior, Richard Evans, John Jumper, James Kirkpatrick, Laurent Sifre, Tim Green, Chongli Qin, Augustin Židek, Alexander WR Nelson, Alex Bridgland, et al. Improved protein structure prediction using potentials from deep learning. Nature, 577(7792):706-710, 2020.
Heiko Strathmann, Mohammadamin Barekatain, Charles Blundell, and Petar Veličković. Persistent message passing. arXiv preprint 103.0:2103,0104,, 2021.
Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. Revisiting unreasonable ef- fectiveness of data deep learning era. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.
Hao Tang, Zhiao Huang, Jiayuan Gu, Bao-Liang Lu, and Hao Su. Towards scale-invariant graph-related problem solving by iterative homogeneous graph neural networks. arXiv preprint arXiv:2010. 13547, 2020.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033 IEEE, 2012.
Sjoerd Van Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural ex- pectation maximization: Unsupervised discovery of objects and their interactions. arXiv preprint arXiv: 10353, 2018.
Rishi Veerapaneni, John D Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine. Entity abstraction in visual model-based reinforcement learning. In Conference Robot Learning, pp. 1439-1456. PMLR, 2020.
Petar Veličković and Charles Blundell. 105.v22105.02761, 2021.
Neural algorithmic reasoning.
arXiv preprint
12

---

Under review as a conference paper at ICLR 2022
Petar Veličković, Rex Ying, Matilde Padovano, Raia Hadsell, and Charles Blundell. Neural execu- tion of graph algorithms. arXiv preprint TrXiv:1910.10590 2019.
Petar Veličkovié, Lars Buesing, Matthew C Overlan, Razvan Pascanu, Oriol Vinyals, and Charles Blundell. Pointer graph networks, arXiv preprint arXiv:2006.06388 2020.
Nicholas Watters, Daniel Zoran, Theophane Weber, Peter Battaglia, Razvan Pascanu, and Andrea Tacchetti. Visual interaction networks: Learning physics simulator from video. Advances in neural information processing systems, 30:4539 2017.
Nicholas Watters, Loic Matthey, Christopher P Burgess, and Alexander Lerchner. Spatial broadcast decoder: A simple architecture for learning disentangled representations in vaes. arXiv preprint arXiv:1901.07017 2019.
Keyulu Xu, Jingling Li, Mozhi Zhang, Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. What can neural networks reason about? arXiv preprint 1rXiv:1905.132111 2019.
Keyulu Xu, Mozhi Zhang. Jingling Li, Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. arXiv preprint arXiv:2009.1 1848, 2020.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets. arXiv preprint arXiv: M:xNi:/703.000 /70306/ 14, 2017.
A BOUNCING BALLS MODELLING SETUP
Abstract pipeline The f is linear projection over the concatenation inputs, outputting a 128- long representation for each of the 10 balls in the input. The P is MPNN over the fully connected graph of the ball representations. message function, without the final layer activation, projecting to the same /28-dimensional space, per object. Finally, g is a linear projection applied on each object representation of the output P.
a 2-pass MPNN with a 3-layered -eattvacti MIPP as
The model is MSE-supervised with ball position the next step. It is trained on a 8-core TPU for 10000 epochs, with a batch size 512, and the Adam optimizer with the initial learning rate of 0.0001.
Natural pipeline f is Slot Attention model (Locatello al. 2020) each input image, con- catenating the images and passing them through a linear layer, outputting a 228-dimensional vector for each of the objects. is a Broadcast Decoder (Watters et al. 2019) containing sequence of transposed convolutions and linear layer mapping before calculating the reconstructions and their masks.
The model is MSE-supervised by pixel reconstruction (per-pixel MSE) over the next step image. Both the RMR and the baseline trained on 8-core TPU for 1000 epochs, with a batch size of 512, with the Adam optimiser and the initial learning rate of 0.0001, all over random seeds.
B ATARI ABSTRACT MODELLING SETUP AND RESULTS
Our best processor network is MPNN (Gilmer al. 2017 over fully connected graph (Santoro et al. 2017) of RAM slots, which concatenates the action embedding to every node (as done in Kipl et al. (2019)) ). It uses three-layer MLPs as message functions, with the ReLU activation applied after each hidden layer. The entire model is trained for every game in isolation, over 48 distinct episodes of Agent57 experience.
use the Adam SGD optimiser (Kingma Ba 2014) with a batch size of 50 and learning rate of 0.001 across all Atari experiments. To evaluate the benefits of message passing, we also compare our model to Deep Sets (Zaheer et al., (017)). which equivalent to our MPNN model only it passes messages over the identity adjacency matrix. Lastly, evaluate the benefits of factorised latents by comparing our methods against a three-layer MLP applied on the flattened RAM state.
13

---

Under review as a conference paper at ICLR 2022
Table 2: Abstract modelling results for Atari 2600. Entire-slot accuracies and bit-level F₁ scores reported only for slots with high entropy, as per Anand et al. (2019).
Copy baseline
MLP
Deep Sets
MPNN
Game
Slot acc.
Bit F1
Slot acc.
Bit F1
Slot acc.
Bit F1
Slot acc.
Bit F
Asteroids
70.65%
0.856
71.28%
0.872
72.84%
0.879
80.69%
0.930
Battlezone
57.09%
0.841
61.19%
0.840
61.71%
0.867
71.06%
0.892
Berzerk
84.32%
0.905
86.17%
0.930
84.16%
0.923
86.67% %
0.933
Bowling
93.86%
0.972
97.43%
0.991
90.72%
0.966
98.41%
0.995
Boxing
59.78%
0.848
54.45%
0.834
59.79%
0.890
58.56%
0.877
Breakout
89.80%
0.949
92.77%
0.970
94.34%
0.979
96.45%
0.988
Demon Attack
67.90%
0.850
68.43%
0.864
66.70%
0.877
69.51%
0.879
Freeway
46.65%
0.787
75.93%
0.921
84.68%
0.959
89.17%
0.965
Frostbite
76.83%
0.904
79.09%
0.904
78.25%
0.946
76.52%
0.918
H.E.R.O.
76.71%
0.891
82.96%
0.932
80.17%
0.929
89.07%
0.956
Montezuma's Revenge
82.58%
0.907
87.30%
0.941
85.90%
0,951
85.44%
0,932
Ms. Pac-Man
83.80%
0.941
80.60%
0.935
81.50%
0.952
85.88%
0.966
Pitfall!
66.60%
0.862
78.28%
0.923
81.92%
0.947
80.40%
0.941
Pong
68.76%
0.873
73.58%
0.911
74.71%
0.920
83.23%
0.952
Private Eye
75.25%
0.889
81.95%
0.932
84.77%
0.954
86.41 %
0.955
Q*Bert
83.00%
0.915
90.07%
0.966
87.87%
0.943
89.26%
0.946
River Raid
76.95%
0.895
80.82%
0.927
69.96%
0.865
86.79%
0.954
Seaquest
71.23%
0.859
78.48%
0.898
75.53%
0.906
70.94%
0.798
Skiing
91.02%
0.966
93.42%
0.980
93.51%
0.983
96.37%
0.992
Space Invaders
81.67%
0.942
84.62%
0.957
89.38%
0.974
91.98%
0.985
Tennis
78.13%
0.890
82.13%
0.926
71.60%
0.856
80.20%
0.893
Venture
61.29%
0.858
63.16%
0.863
64.88%
0.886
76.56%
0.935
Video Pinball
76.71%
0.848
85.92%
0.912
78.64%
0.877
86.61% %
0.913
Yars' Revenge
69.25%
0.896
74.87%
0.929
72.69%
0.948
84.11%
0,969
One immediate observation that RAM updates in Atari are extremely sparse, with copy baseline already being very strong for many games. To prevent the model from having to repeatedly re-learn identity functions, we also make it predict masks of the shape M = 128 specifying which cells be overwritten by the model at this step. This strategy, coupled with teacher forcing (as done by Veličković et al. (2020); Strathmann et al. (202T)) yielded substantially stronger predictors. We also this observation to prevent ever-inflating our prediction scores: we only display prediction accuracy over RAM slots with label entropy larger than 0.6 (as done by An|an et al. (2019)).
The full results of training Atari RAM transition models, for the games studied in Anand et al. (019), are provided in Table We evaluate both bit-level Fj scores, as well as slot-level accuracy (for which all bits need be predicted correctly in order to count), over the remaining 48 Agent57 episodes as validation. To the best of our knowledge, this is the first comprehensive feasibility study for learning Atari RAM transition models.
C REDISCOVERING THE ALGORITHMIC BOTTLENECK
As mentioned in the main text body, one of the key reasons favour high-dimensional algo- rithmic component is to avoid the algorithmic bouleneck, first exposed by Deac et al. (202T)
In short, the performance guarantees of running classical algorithms rely on having the exactly cor- rect inputs for them. If there are any errrrs in the predictions of these (usually very -indeinnnioona) inputs, these errors may propagate to the algorithmic computations and yield suboptimal results. Further, there is room for any kind of fallback if such event occurs.
In contrast, the high-dimensional neural processors like the ones we study here are not vulnerable to bottleneck effects: if any dimensions of the latent state poorly predicted, the other components of it could step in and compensate for this. Further, we can easily support skip connections in the case where the algorithm not fully descriptive of the problem we're re solving.
14

---

Under review as a conference paper at ICLR 2022
In this section, reaffirm the bottleneck effect for both Atari and bouncing ball experiments by providing additional sets of ablations on the processor network's ' latent size.
We ablate various RMR processor network architectures, for dim z {2, 4, 8, 32, 64), but leav-
ing all other components and operations unchanged.
The results of this ablation for the Atari representation learning setting are provided in Figure It can be clearly observed that, as we reduce the size of the latents, this typically induces a performance regression in the downstream bit Fj scores. This indicates the algorithmic bottleneck effect.
On the bouncing balls task, we observe the algorithmic bottleneck as well, with more pronounced effects. Namely, for all latent sizes 32 and under, the validation MSE shoots up even though the train- ing MSE keeps improving. Ultimately, we also attempted using an Interaction Network (Battaglia al., 2016) bottlenecked RMR processor which requires projecting on rather than which exhibited exactly the same behaviour, at larger scale. The IN processor achieved final validation MSE of 341.40 5.60 (x10-4). This is roughly x worse than the non-bottlenecked RMR.
D ON THE WEAK ACCEPTOR PROPERTIES OF BATTLEZONE
Out of all the Atari games studied in our work, Battlezone could been singled out the least "recep- tive" to RMR to its pixel based encoder (Figure 4).
only four games successfully donating their RAM representations
We set out to study why this effect took place. Upon inspection of the game's dynamics we deter- mined that Battlezone has certain properties that make representation learning uniquely challenging and somewhat decoupled from the underlying console computation.
Namely, Battlezone is the only game in our dataset that is played in the "first-person". Therefore, from the point of view of the pixel inputs, it may seem if the player is always in the same place, and we would expect the mechanics and the way of thinking about the 'avatar' to be substantially different from all other Atari games considered.
E ON THE DATA DISTRIBUTION AND SETUP USED FOR TRAINING P
In many cases, while algorithmic prior may be known, generative aspects of the relevant abstract inputs for the natural task could be unknown. For example, it may be known that the natural task could benefit from sorting subroutine, but it may not known upfront how many objects will need to be sorted.
When such details are unknown, it may still be possible to fall back to abstract inputs sampled from some sensible generic random distribution- long as the processor network is trained in way that promotes extrapolation. For recent theoretical treatment of OOD generalisation algorithmic reasoners, we refer the reader to Xu et al. (2020). Further, as observed by Deac et al. (2021) in the case of implicit planning, even generic random abstract distributions can promote useful transfer to noisy pixel-based acting settings such as Atari.
Lastly, it is important to ensure that, in the abstract pipeline, the processor network P carries the brunt of the computational burden; otherwise, the reasoning task may be partially captured in either or g, and not left to P's weights. our work, we generally promote such behaviour by keeping and g to be only linear layers; but even in settings where this is not appropriate, recommend taking care to not overpower the abstract encoder and decoder.
F ATARI ABSTRACT MODEL EQUATIONS
In this appendix, we provide a "bird's eye" view of the modelling steps taken by our abstract Atari pipeline, aiming support future implementations of the RMR blueprint. We will readily re-use the notation from the main text for the various components.
httpps:////t htttp:://ww.youtuubtt
15

---

Under review as a conference paper at ICLR 2022
Latent size
Asterbies
Bettlezone
16
64
Dowling
Box ing
Breakout
Demon Atteck
freeway
Frostbite
H.E.R.D.
Montezuma's is Revenge
Ms. Pac- Man
Privte Bye
O'Bert
River
Skiing
Space Invaders
Tennis
Venture
Video
Revenge
0.0
0.2
0.4
0.6
0.8
Figure 5: Bit-level Fj scores for the Atari experiments while varying the latent size. A clear de- creasing trend is apparent a dim z is reduced, indicating the bottleneck effect.
16

---

Under review as a conference paper at ICLR 2022
Firstly, the abstract encoder f is applied on the relevant Atari RAM representations, x, augmented by a one-hot representation which aiming to provide a generic embedding of the semantics each RAM slot. f is implemented linear layer. hence:
fx]|o) w'x+ U'o by
where W Uf, b' are learnable weights, and 12 is one-hot encoding s.t. L128- In a separate pipeline (officially part of the processor network), the performed actions are also en- coded using linear action encoder:
a = fa(a) w b/a
where a is one-hot encoded action representation.
Then, the processor network GNN is called to update these latents, and a sum-based skip connection employed to promote model-free path:
= P(z + a) + +
(3)
Here, the processor network is implemented as standard message passing neural network (Gilmer al. 2017) fully connected graph. Equations of such P are commonly exposed in e.g. Velicković et al. (2019).
Lastly, the relevant decoder networks predict properties: a mask which signifies which RAM slots have been changed as result of applying a, and the updated states y.
f 9u (z') = (W"z' + b²)
(4)
y
g(z') I1 (W + +b²) Iu>00.5
(5)
Here, a is the logistic sigmoid activation, @ is the elementwise product, and is the indicator function, which thresholds the mask. While this thresholding is non-differentiable, we note that the mask values can be directly supervised from known trajectories, and at training time, teacher forcing can be applied-slottng ground-trun truth masks in place of the indicator function.
17


Reviews:
[{"review_id": "u-CCY8y5fHT", "reviewer": "Reviewer_HnMg", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths**:\n\\\nThe paper is well presented, and explores an important and exciting territory. The empirical assessment, though rudimentary, still demonstrates great potential of the proposed approach.\n\n**Weaknesses**:\n1. My major concern is that abstract inputs, $\\bar{\\bf{x}}$, is given, thus substantially simplifies the problem. While in most real-world problems, $\\bar{\\bf{x}}$ is unknown. (For example, the representation of sentences, speech, or images for object detection)\n\\\nThe key difficulty of representation learning is to learn internal representations for natural inputs (images here). However, the authors pre-define the abstract inputs based on their knowledge. In bouncing balls problem, they define $\\bar{\\bf{x}}$ to be coordinates of balls, and those are the key to predict future trajectories.\n\\\nGiven the pre-defined internal representations, it becomes much simpler to supervise a task. And $P$ is learned easily. It is not surprising that $P$ can be directly (re)-used in a natural setting, because the only thing left to do is to learn a mapping between abstract representations and natural representations, i.e., $\\tilde{f}$ and $\\tilde{g}$.\n\\\nI think a more meaningful setting is that the form of abstract features ($\\bar{\\bf{x}}$ and $\\bar{\\bf{y}}$) is unknown. And the goal is to use some external knowledge to learn them instead of defining them for tasks.\n\n2. I also found tasks being evaluated are still toy tasks. I’d suggest test on more realistic tasks, such as visual question answering.\n\n3. There isn’t significant improvement compared to an end-end model (The difference might not be even distinguishable in the reconstruction rollout..); while RMR involves much more complicated modules. I’m not fully convinced to use RMR at least for bouncing balls problem.\n\n4. Not weakness, but a question for style/typo: In Section 5 **natural pipeline**, is there a typo in the function? The right hand side of equation should be $\\bf{y}$, not $\\tilde{f}(\\bf y)$?", "summary_of_the_paper": "This paper presents a learning paradigm referred to as \"Reasoning-Modulated Representation\" (RMR). The main argument is that it first learns a latent-space processor for abstract simplification of a problem, then use the processor for any tasks that can be modeled with algorithmic reasoning. The authors then verify RMR in two scenarios, bouncing balls problem (physics simulation) and Atari 2600 games.\n\nFor real-world problems, the proposed method alleviates requirements on alignment between abstract and natural inputs, making the learning pipeline more data-efficient, because the pipeline could rely solely on automatically generated, abstract training data. In addition, it is possible that the latent-space processor can be re-used to tasks that do not even directly align with the abstract environment.", "summary_of_the_review": "The paper tries to tackle a meaningful and intriguing problem. However, I find the pre-assumption of the whole learning paradigm is questionable. And I do not recommend acceptance of the paper in its current form.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "CVrz2X_QoqK", "reviewer": "Reviewer_fio5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "First, this paper addresses a very important topic. The fusion of deep learning and classical algorithms is a frontier in AI research with significant potential going forwards.\n\nThis paper is a somewhat unusual submission. The methodological contribution is minimal (the approach is easily described in words). The work doesn’t follow the usual sorts of patterns e.g., explaining why something works, or shedding new light on some occurrence. Instead, the main contribution of this work is to spell out a philosophy on how to combine deep learning and classical algorithms. It is for this reason that I wouldn’t be surprised if the paper receives some negative reviews. However, I myself find the conceptual contribution to be valuable, with the potential to change the way others think.\n\nWith this in mind, here a few high-level strengths and weaknesses of this work:\n\n**Strengths:**\n- The paper addresses a fundamental conceptual question, a nice change from many papers I read.\n- The approach is simple and general, and could be replicated for other related problems.\n- Good empirical performance compared to C-SWM.\n- Excellent writing.\n- Proper treatment of statistical significance. All too uncommon in our community. \n\n**Weaknesses:**\n- Light on ablations. There is a study of different architectures for $P$, but nothing beyond that. I’ll try to offer some suggestions below, but the authors probably can think of many more. \n- As best I can tell the authors don’t plan on releasing code (please correct me if this is untrue). \n- I wouldn’t call this work reproducible in its current state. E.g., the paper mentions in passing a skip connection in $P$ but I was unable to find further information on this (please point to it if I missed it). Hyperparameters are given, but a lot of the implementation details are left the to imagination. \n\n**Questions and suggestions:**\n\n- One of the central ideas in this work (also present in a couple of prior works) is that of a dimensionality bottleneck if one were to fix the reasoning module $P$ to have the same low dimensionality as the abstract inputs (size, position etc.). While this is an intuitive point, and somewhat observed already, it is not yet a widely known thing. So I would suggest an additional ablation that shows that RMR performs favorably compared to an architecture with a dimensionality bottleneck. I fully expect RMR to do better, but I would be interested in seeing the delta so I can understand how much of the work is being done by the avoidance of a dimensionality bottleneck.\n- Any thoughts on why with Battlezone as the test game the performance is often “no different” from C-SWM (Fig. 4)? Perhaps the reasoning module plays a less important role for this task?\n- The fact that $P$ can be transferred effectively between tasks is somewhat surprising to me. It makes me wonder: what is $P$ learning that is helping performance? It seems it isn’t the precise steps of the algorithm in question… \n- The authors mention a greater need for careful hyperparameter tuning. It would be nice to quantify/demonstrate that point somehow. \n- The method doesn't require $(x, \\bar{x})$ pairs etc., as you mention. I am curious how close RMR comes to matching the performance of a model trained with this sort of privileged information. \n- One take home from your results is that training $P$ is best done with lower dimensional abstract representations than the high dimensional raw input. It would be interesting to see some further experimentation on the relation between the dimensionality of the data $P$ is trained on, and its performance when inserted into the final model.  ", "summary_of_the_paper": "This paper proposes a training method based on ideas on neural algorithmic reasoning. As well using a model architecture that is aligned with the underlying algorithmic process this method proposes to pre-train the reasoning module using data from an abstract output stimulator (the work considers problems where this simulator is readily available, and so can generate lots of abstract data cheaply). A key motivation is that the reasoning module $P$ should ultimately be a mapping between two medium dimensional latent spaces (lower dimension than input space, higher dimension than the abstract representation space). Finding this balance allows the module to avoid modeling high dimensional data, whilst also not introducing low-dimensional bottlenecks into the architecture.  This work finds that inserting the pre-trained $P$ into a model that processes raw input data leads to improved performance compared to end-to-end training of a model on raw input data. ", "summary_of_the_review": "In all, I opt for arguing for acceptance. Since the contribution of this work is somewhat subtle, I am especially open to further discussions. The basis for my current position is this: \n\n- This works adds further weight behind a school of thought on combining deep learning and classical algorithms: instead of exactly running classical algorithms as subroutines in network architectures, we should instead insert network modules that behave similarly to classical algorithms, without exact execution.\n- The key requirement in order to use this method, that there are problems where a cheap abstract simulator is available, is fairly reasonable. For example, any case in which we wish to insert a classical algorithm as a subroutine in a network architecture, we have available a simulator (i.e., the classical algorithm). \n- Given the previous two points, the takeaways from this work have to potential to be of use in various problem settings. \n\nI opt for a weak accept instead of higher since although I enjoyed reading this paper, I did feel that 1) this paper appears to be somewhat modest extension of existing ideas, and 2) I feel the paper is missing some ablations. \n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "76YJiVjE1lR", "reviewer": "Reviewer_o4To", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper is concerned with the important problem of incorporating prior knowledge about relevant knowledge for solving a task in neural networks. The paper is well written and fairly accessible and clear. The novelty of the approach is somewhat limited, since to a large extent the idea of pretraining a processor in this way was outlined previously in Velickovic & Blundell (2021). At the same time, the implementation of this idea is by no means trivial and overall the technical contribution is sufficient in my view.\n\nThe idea of training a neural processor in the way outlined in this paper is appealing, and the possibility of using it for pixel-level prediction (eg. for bouncing balls) as well as for representation learning (eg. atari) is appealing. At the same time, I don't find that the current experiments sufficiently validate its effectiveness at these tasks, and I would like to see some improvements in this regard. In particular:\n\n* No stronger baseline is provided for the bouncing balls experiment, which leaves it unclear how well RMR performs at this task to alternative methods that also incorporate prior knowledge about the the underlying state of the environment (eg. it being composed of objects, and object interacting mostly pair-wise). From Figure 3 it can be seen that the predictions are quite different after only 5 steps from comparing the position of the dark purple ball. A stronger baseline that also incorporates assumptions about the underlying state of the world is desirable. RIMs seems like a good candidate for this or perhaps VIN.\n\n* On bouncing balls there is a clear expectation for what the processor should learn (i.e. something like how interaction networks are designed for this task) and I would like to see a comparison to using a processor that can be understood in this way. In particular, a comparison to a pre-trained IN as a processor would reveal how much knowledge about the underlying task is distilled. A key difference between using IN and RMR is the lack of an intermediary high-dimensional space on which the processor operates (as opposed to abstract inputs). While it is intuitive that this can be desirable when both the processor and the natural encoder are relatively unstructured, it would be good to ablate this design choice.\n\n* While the results on Atari in Table 1 are promising, it is a little concerning that these were obtained by first pre-training an Agent57 to generate sufficiently diverse RAM state data. Part of the motivation for this work is to alleviate the need to densely sample data in the natural space by incorporating prior knowledge, yet arguably this is exactly what was needed to train agent57 in the first place. As such I would like to see an attempt made at performing this experiment on transitions collected in a way that does not necessitate recording billions of pixel-level transitions. Alternatively, one could compare to a stronger baseline that leverages the data used to train agent57 for representation purposes as well.\n\nFinally some minor comments\n\n* When generating simulation data from abstract inputs as for bouncing balls, how is it ensured that a sufficient range of values is seen during training the processor? In the environments considered, the abstract and natural data always comes from the same simulator, but what if this is not the case? I.e. some real-world phenomena are observed for which no simulator is available, and abstract data can only be provided via a separate simulator (for which ranges must be chosen). It could be useful to comment on this.\n\n* When training the processor for the contrastive learning task, how is it ensured that the transformation of interest is performed by the processor and not the encoder or decoder (which are discarded afterward)? \n\n* A good reference regarding the third bullet in the intro on the \"challenging feat\" of learning object representations is \"On the Binding Problem in Artificial Neural Networks\" by Greff et al. (2020), which could be helpful to include beyond the two specific methods cited there. \n", "summary_of_the_paper": "This paper contributes an implementation of the neural algorithmic reasoning blueprint outlined in Velickovic & Blundell (2021). The setting that is considered is one where, for a given complex reasoning task that one is interested in solving, there is access to information about the underlying system that should be incorporated to simplify learning. For example, in the case of making predictions from videos about the future state of a physical system composed of bouncing balls bouncing balls, given access to each ball's radius, position, and velocity it is well known how their future movements follow the laws of physics. \n\nIn order to incorporate this prior knowledge, Velickovic & Blundell (2021) proposed the idea of training a neural processor to capture this prior knowledge between abstract inputs (eg. the 3-dimensional state of each ball) under the assumption that sufficient data obeying the known input-output mapping in this space can be generated (eg. pairs of states at time t and time t+1). An encoder then maps abstract inputs to high dimensional space, the processor transforms it to another high-dimensional space, and a decoder maps it back onto the corresponding abstract output pair. In this way, the processor can distill knowledge about the transformation between abstract inputs (eg. governing the laws of physics in its weights) into its weights. By incorporating this processor in a network that is trained to solve the original task (eg. prediction in pixel spaces) it is expected that this extracted prior knowledge can be leveraged to improve learning and transfer to new tasks.\n\nThe implementation considered here (Reasoning-Modulated Representations), is evaluated on two environments: the bouncing balls environment, and various Atari games. On bouncing balls it is shown how RMR can transfer learned representations from trajectories to videos. However, the absence of a stronger baseline makes it difficult to interpret these results. On Atari, RMR is trained on transitions of RAM states to improve predicting latent space transitions from pixels. Compared to C-SWM (Kipf et al., 2019), it is shown how RMR leads to better representations in several games, and it is shown how processors may transfer across games. \n", "summary_of_the_review": "This paper contributes an implementation of an interesting approach to incorporating prior knowledge about a problem in neural networks that was outlined previously. While this constitutes a sufficient contribution, the experimental evaluation should be improved to validate the proposed implementation. I have suggested a number of concrete improvements that can be made, which would have me reconsider the current score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "3N2YHH3jt7b", "reviewer": "Reviewer_Yu3T", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think this approach could potentially be useful, but I think the current set of experiments don't adequately demonstrate that it poses much of an advantage or will be scalable in a practical manner, for a number of reasons:\n\n- The central assumption of the proposed approach is that we have access to a groundtruth algorithm that governs the interaction of entities in the target domain. This works well in the synthetic domains that are evaluated in this paper, but it's unclear to me whether this approach will be useful in real-world settings. One possibility is that perhaps a simulation engine could be made realistic enough to permit meaningful transfer to real-world data. Is that the sort of approach that the authors envision? If so, it would be helpful to include some kind of proof-of-principle demonstration for this strategy. \n- Given that one has access to such a groundtruth algorithm, wouldn't it be more straightforward to simply use supervised learning to learn representations of the naturalistic inputs of the form that the algorithm expects, and then directly use the algorithm, rather than a learned proxy for it? The paper states that this approach is infeasible because it would require a 'massive dataset of paired $(x,\\bar{x})$' for the purposes of training the encoder (and a similar dataset for training the decoder). But why should this dataset need to be any more massive than the one that is necessary to train the encoder/decoder in the context of the learned interaction model? The paper also suggests that this approach may be impractical because of imperfections in the groundtruth algorithm (or the information available to it). Both of these points would be far more convincing if empirically evaluated (e.g. by demonstrating that it requires more training to learn an effective encoder directly via supervised learning vs. using the current approach), and this would help to justify the motivation for the proposed approach.\n- The primary contributions emphasized in the conclusion are that the proposed approach protects against bottlenecks, and that it is more data-efficient, but as far as I can tell there is no systematic evaluation of either of these claims. The claim about protection against bottlenecks could be evaluated by systematically manipulating the size of the latent representations $z$ (smaller $z$ should impair performance). The claim about data efficiency could be evaluated by manipulating the size of the training set (the advantage of the proposed approach over the end-to-end model should be particularly pronounced with smaller training sets).\n- Most of the comparisons in the paper are between the proposed approach and an end-to-end analog of that approach. This is an important comparison, but the authors should also clarify how these results compare to the current state-of-the-art for the metrics that are evaluated. One exception is the comparison with embeddings from Agent57, but, as the authors note, it's unsurprising that these embeddings don't work very well given that they were learned for the purpose of performing a different task, so that particularly baseline seems a bit contrived.\n\nOther considerations:\n- It's not very easy to resolve what's going on in Figure 3. It requires a very close inspection to even detect that the balls are moving, and it isn't easy to compare the predictions to the groundtruth. It might be easier to use an example with fewer balls, or to skip frames to make the movement more salient. If possible, a link to an animation might be even better.\n- Can Figure 4 be made quantitative rather than qualitative, e.g. by representing the performance of RMR as a percentage of the baseline performance? It would be nice to see how much of an improvement is yielded for each pair of games.\n- The qualitative clustering analyses are kind of cool, but what do we learn about the proposed approach from these analyses? These results seem to pertain primarily to the games themselves, in which case this analysis might be better suited to the supplementary material. ", "summary_of_the_paper": "This paper proposes a method for using domain knowledge to learn an interaction model given a groundtruth algorithm that operates over abstract inputs, which can then be used to learn an encoder and decoder for interfacing between this interaction model and naturalistic inputs. The method is evaluated, with favorable results, on both the bouncing balls domain and Atari suite.", "summary_of_the_review": "The proposed approach is favorably evaluated in synthetic domains, but it is unclear from the present set of experiments whether and how this approach will be scalable to more real-world problems. There are also some additional experiments that need to be performed to evaluate some of the paper's primary claims.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes Reasoning-Modulated Representations (RMR). That is, it incorporates how to incorporate (structure) prior knowledge (such as a law in physics) into a pre-trained reasoning modules, and investigates how doing so shapes the discovered representations in a number of self-supervised learning settings from pixels. The reviews and (short) discussion have presented salient arguments about the suitability of the paper for publication at this stage. One review argues that the "methodological contribution is minimal," another one is asking for "systematic evaluation" of the main claims made. Moreover, while we all agree that the direction is interesting, the RMR approach presented is not shown to "scale well" (yet), as pointed out by one review. This, however, is important since the general idea that prior knowledge shapes the representation learned is common wisdom in the literature. Indeed, one may now argue that the paper is much more about "how best to combine pixel-based deep learning and neural algorithmic reasoning algorithms" as one reviewer puts it. From this perspective the ATARI experiments are more interesting but here the benefit compared to C-SWM seems to be marginal and one should compare to other deep baseline conditions on the RAM; the significance is not looking at the difference in score and degree in freedoms but just the number of wins. Additionally, there should be other baselines that directly make use of more structured models (structure = prior knowledge, e.g., HMMs or some other way to have bit of a memory), other datasets (where no access to RAM exists) as well as a discussion of other approaches that combines (combinatorial) reasoning with pixel-based deep learning. That is, while pushing for a more high-level contributions is fine, this also requires some more illustrations and discussion of the broader context. Therefore, my overall recommendation is reject at this stage of the paper.
Decision: Reject