Reasoning-Modulated Representations

1 INTRODUCTION . Neural networks are able to learn policies in environments without access to their specifics ( Schrittwieser et al. , 2020 ) , generate large quantities of text ( Brown et al. , 2020 ) , or automatically fold proteins to high accuracy ( Senior et al. , 2020 ) . However , such “ tabula rasa ” approaches hinge on having access to substantial quantities of data , from which robust representations can be learned . Without a large training set that spans the data distribution , representation learning is difficult ( Belkin et al. , 2019 ; LeCun , 2018 ; Sun et al. , 2017 ) . Here , we study ways to construct neural networks with representations that are robust , while retaining a data-driven approach . We rely on a simple observation : very often , we have some ( partial ) knowledge of 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 properties of the data—and such properties can not always be robustly extracted from natural observations . Motivation Consider the task of predicting the future state of a system of n bouncing balls , from a pixel input x ( Figure 1 ) . Reliably estimating future pixel observations , y , is a challenging reconstruction task . However , the generative properties of this system are simple . Assuming knowledge of simple abstract inputs ( radius , rc , position , xc , and velocity , vc ) for every ball , x̄c , the future movements 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 . While this gives us a potentially simpler path from pixel inputs to pixel outputs , via abstract inputs to abstract outputs ( x → x̄ ȳ → y ) , it still places potentially unrealistic demands on our task setup , every step of the way : x→ x̄ : Necessitates either upfront knowledge of how to abstract away x̄ from x , or a massive dataset of paired ( x , x̄ ) to learn such a mapping from ; x̄ ȳ : Implies that the algorithm perfectly simulates all aspects of the output . In reality , an algorithm 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 x̄ is incorrectly predicted , this will negatively compound in ȳ , hence y ; ȳ→ y : Necessitates a renderer that generates y from ȳ , or a dataset of paired ( ȳ , 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 use to our predictor , but we do not assume a trivial mapping or a paired dataset which would allow us to convert directly from x to x̄ or from ȳ to y . Our only remaining assumption is that the algorithm x̄ ȳ can be efficiently computed , allowing us to generate massive quantities of paired abstract input-output pairs , ( x̄ , ȳ ) . Present work In this setting , we propose Reasoning-Modulated Representations ( RMR ) , an approach that first learns a latent-space processor of abstract data ; i.e . a mapping x̄ f−→ z P−→ z′ g−→ ȳ , where z ∈ Rk are high-dimensional latent vectors . f and g 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 ( Veličković & Blundell , 2021 ) , that such a processor network can be used as a drop-in differentiable component for any task where the x̄ ȳ kind of reasoning may be applicable . Hence , we then learn a pipeline x f̃−→ z P−→ z′ g̃−→ y , which modulates the representations z obtained from x , forcing them to pass through the pretrained processor network . By doing so , we have ameliorated the original requirement for a massive natural dataset of ( x , y ) pairs . Instead , we inject knowledge from a massive abstract dataset of ( x̄ , ȳ ) pairs , directly through the pre-trained parameters of P . This has the potential to relieve the pressure on encoders and decoders f̃ and g̃ , 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 , on physics 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 , we recover a dataefficient 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 reusability : 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 to construct robust latent-space models that imitate abstract reasoning procedures . This makes it well aligned with neural algorithmic reasoning ( Cappart et al. , 2021 ) , 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 ( Veličković et al. , 2019 ; Tang et al. , 2020 ) , linearithmic algorithms ( Freivalds et al. , 2019 ) , and data structures ( Strathmann et al. , 2021 ; Veličković et al. , 2020 ) . Further , the XLVIN model ( Deac et al. , 2021 ) demonstrates how such primitives can be re-used for data-efficient planning , paving the way for a blueprint ( Veličković & Blundell , 2021 ) that we leverage in RMR as well . 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 algorithmic computation and simulations , as the latter are typically realised using the former . Within this space , abstract GNN models of physics have been proposed by works such as interaction networks ( Battaglia et al. , 2016 ) and NPE ( Chang et al. , 2016 ) , and extended to pixel-based inputs by visual interaction networks ( Watters et al. , 2017 ) . The generalisation power of these models has increased drastically in recent years , with effective models of systems of 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-based physical simulators , for example by using Hamiltonian ODE integrators in conjunction with GNN models ( SanchezGonzalez et al. , 2019 ) , or by coupling a ( 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 a form of object-centric neural network , in which visual objects in an image or video are represented as separate latent variables in the model and their temporal dynamics and pairwise interactions are modeled via GNNs or self-attention mechanisms . There is a rich literature on discovering objects and learning their dynamics from raw visual data without supervision , with objectcentric models such as R-NEM ( Van Steenkiste et al. , 2018 ) , SQAIR ( Kosiorek et 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 ( François-Lavet et al. , 2019 ) . Related approaches discover and use keypoints ( Kulkarni et 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 a keypoint discovery method . A related line of works integrate attention-mechanisms in modular and object-centric models to interface latent variables with visual input , including models such as RMC ( Santoro et al. , 2018 ) , RIM ( Goyal et al. , 2019 ) , Slot Attention ( Locatello et al. , 2020 ) , SCOFF ( Goyal et al. , 2021b ) , and NPS ( Goyal et al. , 2021a ) . 3 RMR ARCHITECTURE Having provided a high-level overview of RMR and surveyed the relevant related work , we proceed to carefully detail the blueprint of RMR ’ s various components . This will allow us to ground any subsequent RMR experiments on 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 a set of natural inputs , X , and a set of natural outputs , Y . These sets represent the possible inputs and outputs of a target function , Φ : X → Y , which we would like to learn based on a ( potentially small ) dataset of input-output pairs , ( x , y ) , where y = Φ ( x ) . We further assume that the inner workings of Φ can be related to an algorithm , A : X̄ → Ȳ . The algorithm operates over a set of abstract inputs , X̄ , and produces outputs from an abstract output set Ȳ . Typically , it will be the case that dim X̄ dimX ; 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 we do assume is that A is either known or can be trivially computed , giving rise to a massive dataset of abstract input-output pairs , ( x̄ , ȳ ) , where ȳ = A ( x̄ ) . Lastly , we assume a latent space , Z , and that we can construct neural network components to both encode and decode from it . Typically , Z will be a real-valued vector space ( Z = Rk ) 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 Z = Rn×k , where n is the assumed maximal number of objects ( typically a 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 a neural network approximator g ( P ( f ( x̄ ) ) ) ≈ A ( x̄ ) , which follows the encode-process-decode paradigm ( Hamrick et al. , 2018 ) . It consists of the following three building blocks : Encoder , f : X̄ → Z , tasked with projecting the abstract inputs into the latent space ; Processor , P : Z → Z , simulating individual steps of the algorithm in the latent space ; Decoder , g : Z → Ȳ , tasked with projecting latents back into the abstract output space . Such a pipeline is now widely used both in neural algorithmic reasoning ( Veličković et al. , 2019 ) and learning physical simulations ( Sanchez-Gonzalez et al. , 2020 ) , and can be trained end-to-end with gradient descent . For reasons that will become apparent , it is favourable for most of the computational effort to be performed by P . Accordingly , encoders and decoders are often designed to be simple learnable linear projections , while processors tend to be either deep MLPs or graph neural networks—depending on whether the latent space is factorised into nodes . Natural pipeline Once an appropriate processor , P , has been learned , it may be observed that it corresponds to a highly favourable component in our setting . Namely , we 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. , 2021 ) , where it yielded a generic implicit planner . As an ablation , we have also rediscovered the bottleneck effect in our settings ; see Appendix C. We now leverage similar insights for general representation learning tasks . On a high level , what we need to do is simple and elegant : swap out f and g for natural encoders and decoders , f̃ : X → Z and g̃ : Z → Y , respectively . We are then able to learn a function g̃ ( P ( f̃ ( x ) ) ) ≈ Φ ( x ) , which is 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 Φ . While we rely on a skip connection in our 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 . 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 g̃ . For example , our natural inputs will often be pixel-based , necessitating a convolutional neural network for f̃ . • Further , since the parameters of P are kept frozen , f̃ is left with a challenging task of mapping natural inputs into an 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 a challenging feat ( Greff et al. , 2020 ) , but can be successfully performed . Sometimes arbitrarily factorising the feature maps of a 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 f̃ is that it ’ s very difficult to enforce their slot representations to line up in the same way as in the abstract inputs . This implies that P should be permutation equivariant ( and hence motivates using a GNN for it ) .

Summary: 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.