Recurrent Independent Mechanisms

Learning modular structures which reflect the dynamics of the environment can lead to better generalization and robustness to changes which only affect a few of the underlying causes . We propose Recurrent Independent Mechanisms ( RIMs ) , a new recurrent architecture in which multiple groups of recurrent cells operate with nearly independent transition dynamics , communicate only sparingly through the bottleneck of attention , and are only updated at time steps where they are most relevant . We show that this leads to specialization amongst the RIMs , which in turn allows for dramatically improved generalization on tasks where some factors of variation differ systematically between training and evaluation . 1 INTRODUCTION . Physical processes in the world often have a modular structure , with complexity emerging through combinations of simpler subsystems . Machine learning seeks to uncover and use regularities in the physical world . Although these regularities manifest themselves as statistical dependencies , they are ultimately due to dynamic processes governed by physics . These processes are often independent and only interact sparsely . For instance , we can model the motion of two balls as separate independent mechanisms even though they are both gravitationally coupled to Earth as well as ( weakly ) to each other . They may , however , occasionally strongly interact via collisions . The notion of independent or autonomous mechanisms has been influential in the field of causal inference , where it is applied not only to dynamic processes but also to time independent datasets . For instance , it has been argued that the conditional distribution of the average annual temperature given the altitude of a place is an abstraction of a causal mechanism ( subsuming complex physical processes involving air pressure , etc . ) that is independent of the distribution of the altitudes of settlements ( Peters et al. , 2017 ) , and will thus apply invariantly for , say , different countries in the same climate zone with different altitude distributions . A complex generative model , temporal or not , can be thought of as the composition of independent mechanisms or “ causal ” modules . In the causality community , this is often considered a prerequisite of being able to perform localized interventions upon variables determined by such models ( Pearl , 2009 ) . It has been argued that the individual modules tend to remain robust or invariant even as other modules change , e.g. , in the case of distribution shift ( Schölkopf et al. , 2012 ; Peters et al. , 2017 ) . One may hypothesize that if a brain is able to solve multiple problems beyond a single i.i.d . ( independent and identically distributed ) task , it would be economical to learn structures aligned with this , by learning independent mechanisms that can flexibly be reused , composed and re-purposed . In the dynamic setting , we think of an overall system being assayed as composed of a number of fairly independent subsystems that evolve over time , responding to forces and interventions . A learning agent then need not devote equal attention to all subsystems at all times : only those aspects that significantly interact need to be considered jointly when taking a decision or forming a plan ( Bengio , 2017 ) . Such sparse interactions can reduce the difficulty of learning since few interactions need to be considered at a time , reducing unnecessary interference when a subsystem is adapted . Models learned this way may be more likely to capture the compositional generative ( or causal ) structure of the world , and thus better generalize across tasks where a ( small ) subset of mechanisms change while most of them remain invariant ( Simon , 1991 ; Peters et al. , 2017 ; Parascandolo et al. , 2018 ) . The central question motivating our work is how a machine learning approach can learn independent but sparsely interacting recurrent mechanisms in order to benefit from such modularity . 2 RECURRENT INDEPENDENT MECHANISMS WITH SPARSE INTERACTIONS . Our approach to modelling a dynamical system of interest divides the overall model into k small subsystems ( or modules ) , each of which is recurrent in order to be able to capture dynamics . We refer to these subsystems as Recurrent Independent Mechanisms ( RIMs ) , where each RIM has distinct functions that are learned automatically from data1 . We refer to RIM k at time step t as having state ht , k , where t = 1 , . . . , T . Each RIM has parameters θk , which are shared across all time steps . At a high level ( see . Fig . 1 ) , we want each RIM to have its own independent dynamics operating by default , and occasionally to interact with other relevant RIMs and with selected elements of the encoded input . The total number of parameters can be kept small since RIMs can specialize on simple sub-problems , similar to Parascandolo et al . ( 2018 ) . This specialization and modularization not only has computational and statistical advantages ( Baum & Haussler , 1989 ; Bengio et al. , 2019 ) , but also prevents individual RIMs from dominating and modelling complex , composite mechanisms . We expect this to lead to more robust systems than training one big homogeneous neural network ( Schmidhuber , 2018 ) . Moreover , modularity also has the desirable implication that a RIM should maintain its own independent functionality even as other RIMs are changed . A more detailed account of the desiderata for the model is given in Appendix A . 2.1 INDEPENDENT RIM DYNAMICS . Now , consider the default transition dynamics which we apply for each RIM independently and during which no information passes between RIMs . We use h̃ for the hidden state after the independent dynamics are applied ( and before attention is applied ) . First , for the RIMs which are not activated ( we refer to the activated set as St ) , the hidden state remains unchanged : h̃t+1 , k = ht , k ∀k /∈ St. ( 1 ) Note that the gradient still flows through a RIM on a step where it is not activated . For the RIMs that are activated , we run a per-RIM independent transition dynamics . The form of this is somewhat flexible , but in this work we opted to use either a GRU ( Chung et al. , 2015 ) or an LSTM ( Hochreiter & Schmidhuber , 1997 ) . We generically refer to these independent transition dynamics as Dk , and we emphasize that each RIM has its own separate parameters . Aside from being RIM-specific , the internal operation of the LSTM and GRU remain unchanged , and the active RIMs are updated by h̃t+1 , k = Dk ( ht , k ) = LSTM ( ht , k , A ( in ) k ; θ ( D ) k ) ∀k ∈ St ( 2 ) as a function of the attention mechanism A ( in ) k applied on the current input , described in the next two subsections below , after explaining the key-value mechanism used to select arguments for this update . 1Note that we are using the term mechanism both for the mechanisms that make up the world ’ s dynamics as well as for the computational modules that we learn to model those mechanisms . 2.2 KEY-VALUE ATTENTION TO PROCESS SETS OF NAMED INTERCHANGEABLE VARIABLES . Each RIM should be activated and updated when the input is relevant to it . We thus utilize competition to allocate representational and computational resources . As argued by Parascandolo et al . ( 2018 ) , this tends to produce independence among learned mechanisms , provided the training data has been generated by a set of independent physical mechanisms . In contrast to Parascandolo et al . ( 2018 ) , we use an attention mechanism for this purpose . In doing so , we are inspired by findings from experimental psychology in the study of the interplay of top-down attention and bottom-up information flow , conceptualized in the biased competition theory of selective attention ( Desimone & Duncan , 1995 ) : A brain ’ s capacity for parallel processing of complex entities is limited , and many brain systems representing visual information use competition ( operating in parallel across the visual field ) to allocate resources , often biased by feedback from higher brain areas . The introduction of content-based soft-attention mechanisms ( Bahdanau et al. , 2014 ) has opened the door to neural networks which operate on sets of typed interchangeable objects . This idea has been remarkably successful and widely applied to most recent Transformer-style multi-head dot product self attention models ( Vaswani et al. , 2017 ; Santoro et al. , 2018 ) , achieving new state-of-the-art results in many tasks . Soft-attention uses the product of a query ( or read key ) Q of dimensionality Nr × d matrix Q , and d dimension of each key ) to a set of No objects each associated with a key ( or write-key ) matrix KT ( No × d ) , and after normalization with a softmax yields outputs in the convex hull of the values ( or write-values ) Vi ( row i of matrix V ) . Its result is computed as Attention ( Q , K , V ) = softmax ( QKT√ d ) V , where the softmax is applied to each row of its argument matrix , yielding a set of convex weights . As a result , one obtains a convex combination of the values V . If the attention is focused on one element for a particular row ( i.e. , the softmax is saturated ) , this simply selects one of the objects and copies its value to row j of the result . Note that the d dimensions in the key can be split into heads which then have their attention matrix and write values computed separately . When the inputs and outputs of each RIM are a set of objects or entities ( each associated with a key and value vector ) , the RIM processing becomes a generic object-processing machine which can operate on “ variables ” in a sense analogous to variables in a programming language : as interchangeable arguments of functions . Because each object has a key embedding ( which one can understand both as a name and as a type ) , the same RIM processing can be applied to any variable which fits an expected `` distributed type '' ( specified by a query vector ) . Each attention head then corresponds to a typed argument of the function computed by the RIM . When the key of an object matches the query , it can be used as input for the RIM . Whereas in regular neural networks ( without attention ) neurons operate on fixed variables ( the neurons which are feeding them from the previous layer ) , the key-value attention mechanisms make it possible to select on the fly which variable instance ( i.e . which entity or object ) is going to be used as input for each of the arguments of the RIM dynamics , with a different set of query embeddings for each RIM . These inputs can come from the external input or from the output of other RIMs . So , if the individual RIMs can represent these “ functions with typed arguments , ” then they can “ bind ” to whatever input is currently available and best suited according to its attention score : the “ input attention ” mechanism would look at the candidate input object ’ s key and evaluate if its “ type ” matches with what this RIM expects ( specified in the query ) .

Summary: The authors argue that the world consists of largely independent causal mechanisms that sparsely interact. The authors propose a new kind of recurrent network (RIM) that presumably distills this world view into inductive biases. RIMs consist of largely independent recurrent modules that are sparsely activated and interact through soft attention. They tested RIMs on a number of supervised and reinforcement learning tasks, and showed that RIMs perform better than LSTMs and several other more recently proposed networks (e.g., Differential Neural Computers, Relational Memory Core, etc.) in these tasks. In particular, RIMs can generalize more naturally than other networks to out-of-distribution test set in presumably modular tasks.