Dynamic Inference with Neural Interpreters

1 Introduction . Rule-based programming is the basis of computer science , and builds the foundation of symbolic-AI based expert systems that attempt to emulate human decision-making to solve real-world problems . The process of inference entails channeling information through a chain of computational units ( e.g. , logical primitives ) and culminates in a conclusion that answers a given query . Such systems have the advantage that they permit efficient reuse of computational units and enable iterative reasoning over multiple cycles of computation . As a simple example , consider the relation parent_child ( A , B ) , which can be used to construct a new relation sibling ( U , V ) = parent_child ( A , U ) AND parent_child ( A , V ) , which in turn can be used to construct yet another relation cousin ( U , V ) = parent_child ( X , U ) AND parent_child ( Y , V ) AND sibling ( X , Y ) , and so on . However , such systems are known to suffer from the knowledge acquisition problem , which is the inability to leverage unstructured data to derive new computational units and improve existing ones [ 30 ] . In stark contrast to the symbolic paradigm , modern machine learning models excel at absorbing large amounts of unstructured data and yield strong performance in many challenging domains , ranging from large-scale image classification to language modelling . However , they are relatively rigid in how they share and reuse computation in order to process information : convolutional neural networks , for instance , process the content of an image at every location to arrive at the class label of a given image . In doing so , they only reuse computational units ( here , convolutional filters ) laterally at a constant depth i.e. , amongst information being processed simultaneously . In the same spirit , recurrent neural networks only reuse computational units ( here , RNN cells ) vertically , i.e. , in depth . Such rigidity ∗ , †Equal contribution , ordered alphabetically by last name . 1Max-Planck-Institute for Intelligent Systems , Tübingen . 2Mila , Québec . 3Amazon Web Services . ‡Work partially done during an internship at Amazon Web Services . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . in how computation is reused is believed to be one of the reasons current deep neural networks are less capable of systematically generalizing to problems not encountered during training [ 4 , 33 , 39 ] . In the present work , we draw inspiration from typed programming languages to develop a self-attention based neural architecture that relaxes this rigidity in computation reuse . The resulting class of models , which we call Neural Interpreters , learns to flexibly use and compose computational units directly from data without additional supervision . Neural Interpreters factorize a self-attention network [ 58 ] as a system of computational units that we call functions . The input to the network is set-valued and processed in the network by a dynamically inferred sequence of functions – potentially by the same function multiple times , enabling vertical sharing of computational units . This aligns with earlier ideas on independent mechanisms [ 21–24 , 42 , 43 , 46 , 51 ] : a set of mechanisms can be combined and reused in many different ways depending on context ( the current input or task ) , thus factorizing knowledge in independent pieces which can lead to better systematic generalization . Neural Interpreters have two key benefits over vanilla self-attention networks . First , the modular inductive biases facilitate generalization beyond the training distribution and adaptation to new tasks in a sample-efficient manner . Second , but consistent with the notion of factorizing knowledge into approximately independent and composable pieces , the proposed parameterization is ( by construction ) not only agnostic to the cardinality of the input set , but also to the number of functions . The latter implies that given a new task , additional functions can be non-disruptively added and fine-tuned . In other words , knowledge acquired from the prior training tasks can be effectively repurposed for new tasks . Our primary contributions are as follows . ( a ) We introduce the Neural Interpreter , an attentionbased architecture that can be applied on arbitrary set-valued inputs or representations . ( b ) We quantitatively evaluate the proposed architecture in two distinct problem settings : multi-task image classification and abstract reasoning . In the former setting , we show that Neural Interpreters are capable of sample-efficient adaptation to new tasks and can exploit additional capacity added after pre-training . In the latter setting , we demonstrate that Neural Interpreters are capable of out-of-distribution generalization . In particular , we find that Neural Interpreters outperform the Vision Transformer baseline [ 13 , 17 ] , which in turn is competitive in-distribution with the prior state-of-the-art . In both settings , we find that Neural Interpreters develop the ability to gracefully trade-off performance with compute at inference time [ 63 ] . In addition , we include results on a toy problem , where we explicitly probe the ability of Neural Interpreters to learn recomposable computational primitives . ( c ) We ablate over the architectural parameters of Neural Interpreters and qualitatively visualize what the model learns . We find patterns in how the input is routed through the network and that a wide range of hyperparameters yield promising results . 2 Neural Interpreters . In this section , we describe in detail the components of a Neural Interpreter ; see Figure 1 for an overview . This architecture can be used as a drop-in replacement for a self-attention network , e.g. , Transformers [ 58 ] . In line with prior work [ 13 , 17 , 55 ] , we focus on applications to visual data . Input and Output . The input to the Neural Interpreter is any set with vector-valued elements { xi } i , xi ∈ Rdin and the output is another set of vectors { yj } j , yj ∈ Rdout with the same cardinality as the input set . We assume in this work that the input set contains vector embeddings of image patches [ 13 ] or of entire images [ 5 ] . The input set additionally includes one or more learned vectors , called CLS Tokens [ 16 ] , for which the corresponding outputs interface with their respective classifiers [ 38 ] . Scripts . A Neural Interpreter can be expressed as a stack of ns scripts , where a script maps one set of vectors to another with the same number of elements : { yj } j = NeuralInterpreter ( { xi } i ) = [ Scriptns ◦ ... ◦ ( ns times ) ◦ ... ◦ Script1 ] ( { xi } i ) ( 1 ) A script has four components : a type inference module , a type matching mechanism , a set of functions and an interpreter , as explained below . The parameters of these components are not shared between scripts . Role : Scripts function as independent building blocks that can be dropped in any set-to-set architecture , and Neural Interpreters with a single script can perform well in practice . Functions . Functions make the computational units in Neural Interpreters ; they can be represented as a tuple of vector valued parameters , i.e. , fu = ( su , cu ) where u indexes functions . Here , su is referred to as the signature of the function fu ( with a meaning similar to that in programming languages ) , and it is a normalized vector of dtype dimensions . The signature vector specifies to the type matching mechanism ( see below ) what inputs are to be routed to fu . We refer to cu as the code vector of fu , as it instructs an interpreter module ( shared between functions , see below ) how to process the inputs to fu . Role : Functions are vector-valued instructions to other components in the script . They implement the computational units that can be reused in the computational graph . Type Matching and Inference . The type matching mechanism ( Figure 2 ) enables the learned routing of set elements through functions , and it proceeds in three steps . First , given a set element xi , it is processed by an MLP module that outputs its type vector ti . This module is called the type inference module , and the resulting type vector is an element of the same topological space as function signatures , i.e. , a dtype-dimensional hypersphere T . Next , given a function fu and its signature vector su ∈ T , the compatibility Cui between the function fu and a set element xi is determined by the cosine similarity between su and ti in T ( larger =⇒ more compatible ) . Finally , if this compatibility is larger than a threshold ( τ ) , fu is permitted access to xi . Formally , let { xi } i be a set of intermediate representation vectors ( indexed by i ) . With a learnable parameter σ and a hyper-parameter τ , we have : ti = TypeInference ( xi ) ∈ T ; dT ( su , ti ) = ( 1− su · ti ) ( 2 ) Cui = C̃ui + ∑ u C̃ui where C̃ui = exp [ −dT ( su , ti ) σ ] if dT ( su , ti ) > τ , else 0 . ( 3 ) τ is called the truncation parameter of the kernel and is a small scalar for numerical stability . The compatibility matrix Cui ∈ [ 0 , 1 ] will serve as a modulation mask [ 46 ] for the self-attention mechanism in the interpreter ( see below ) . Role : The type matching mechanism is responsible for routing information through functions . The truncation parameter controls the amount of sparsity in routing . ModLin Layers and ModMLP . The components described below make use of linear layers conditioned by some code c. Consider a linear layer with weight matrix W ∈ Rdout × Rdin and a bias vector b ∈ Rdout . Let x ∈ Rdin denote the input vector to the layer , and c ∈ Rdcond a condition vector , and Wc ∈ Rdin × Rdcond . a learnable matrix . The output y ∈ Rdout is given as : y = ModLin ( x ; c ) = W ( x⊗ LayerNorm ( Wcc ) ) + b ( 4 ) where ⊗ denotes element-wise product and we call the resulting layer a modulated linear layer , or a ModLin layer [ 3 ] . Further , one may stack ( say ) L such ModLin layers ( sharing the same condition or code vector c ) interspersed with an activation function ( we use GELUs [ 26 ] ) to obtain a ModMLP : y = ModMLP ( x ; c ) = ( ModLinL ( · ; c ) ◦ Activation ◦ ... ◦ModLin1 ( · ; c ) ) ( x ) ( 5 ) Role : ModLin layers and the ModMLP can be interpreted as programmable neural modules , where the program is specified by the condition or code vector c. ModAttn . ModAttn is a conditional variant of the kernel modulated dot product attention ( KMDPA ) [ 46 ] , where the key , query and value vectors are obtained from ModLin layers conditioned by a vector . In our case , this vector is the code vector cu of function fu , and the corresponding key , query and value vectors are computed as follows ( with { xi } i as input and h indexing attention heads ) : kuhi = ModLinhkey ( xi , cu ) quhi = ModLin h query ( xi , cu ) vuhi = ModLin h value ( xi , cu ) ( 6 ) Note that further below ( e.g. , in Equation 9 ) , we will encounter xui , where the extra u in the subscript denotes that the set element at index i is specific to the function u ; in this case , xi is substituted with xui in Equation 6 . Next , given the keys , queries and the function-variable compatibility matrix Cui , the modulated self-attention weights Wuhij are given by : Wuhij = W̃uhij + W̃uhij where W̃uhij = CuiCuj [ softmaxj ( quhi · kuhj√ dkey ) ] ( 7 ) Here , the quantity Wuhij denotes the attention weights in function fu between elements xi and xj at head h and the softmax operation normalizes along j ; intuitively , information about xi and xj is mixed by fu at head h only if Wuhij 6= 0 . Now , on the one hand , this can be the case only if both Cui and Cuj are non-zero , i.e. , fu is granted access to both variables xi and xj by the typing mechanism . But on the other hand , fu does not necessarily mix xi and xj even if both Cui and Cuj are non-zero , for the self-attention weights ( square brackets ) may still be close to zero depending on the context ( i.e. , the content of xi and xj ) . Next , the values are linearly mixed using the computed attention weights , which is then processed by a final ModLin layer to yield the output yui : yui = ModLin ( ỹui ; h ; cu ) where ỹuhi = ∑ j Wuhijvuhj ( 8 ) Here , ỹui ; h means the head-axis is folded into channels . Role : ModAttn enables interaction between the elements of its input set in multiple parallel streams , one for each function . The query , key , value , and output projectors of each stream are conditioned on the corresponding code vectors , and the interaction between elements in each stream is weighted by their compatibility with the said function . Line of Code ( LOC ) . An LOC layer is a ModAttn layer followed by a ModMLP layer ( Figure 1 , rightmost ) , where both layers share the same condition vector cu , and there are weighted residual connections between the layers . Assuming inputs { xui } u , i to the LOC , we have : ãui = ModAttn ( { LayerNorm ( xuj ) } j ; cu , { Cuj } j ) aui = xui + Cuiãui ( 9 ) b̃ui = ModMLP ( LayerNorm ( aui ) ; cu ) yui = aui + Cuib̃ui ( 10 ) This parameterization ensures that yui = xui if Cui = 0 . In words , if a function ( fu ) is not granted access to a variable ( xi ) by the typing mechanism , it acts as an identity function for this variable . Further , note that we allow the input set to be indexed only by i ; in this case , we assume xui = xi for all u . Role : A LOC can be thought of as multiple instances of a layer in the original transformer architecture ( comprising a self-attention and a MLP module with residual connections ) , applied in parallel streams , one per function . Computations therein are conditioned on the respective code and signature vectors . Interpreter . The interpreter layer is a stack of nl LOCs sharing the same function codes cu and function-variable compatibilities Cui . Assuming the input to the interpreter is a set { xi } i , we have : yi = xi + ∑ u Cui ( LOCnl ◦ ... ( nl times ) ... ◦ LOC1 ) ( { xj } j ; cu , { Cuj } j ) ( 11 ) In words : the interpreter broadcasts a given set element to multiple parallel computational streams , one for each function . After the streams have processed their copy of the input element , the results are aggregated by a weighted sum over the streams , where the weights correspond to the compatibility of the input element with the respective function . Equation 11 can be justified by making two observations . First , if a variable xi is not matched with any function by the typing mechanism , it is left unmodified by the interpreter ; i.e. , if Cui = 0 for all u , then yi = xi . This allows signals to be propagated through the interpreter without interference from existing functions , if so determined by the type inference module . Second , the additive aggregation over the function index u implies that the overall parameterization of Neural Interpreters does not depend on the number of functions . This allows one to add a new function fv simply by including its signature and code ( sv , cv ) as learnable parameters and finetuning these on ( say ) a new problem . Role : The interpreter serves as a general-purpose instruction executor ( one that is shared between functions ) . Given a set of inputs and an instruction ( here , the function code ) , it executes said instruction to compute the output . Function Iterations in Script . Having expressed the overall model as a stack of multiple scripts , we are now equipped to describe the computational graph of a single script . A script can be expressed as a recurrent application of the type matching mechanism and the interpreter , where we refer to the composition of the latter two as a Function Iteration ( FnIter ) : { yi } i = FnIter ( { xj } j ) = Interpreter ( { xj } j , { cu } u , { Cuj } u , j ) ( 12 ) where Cuj = TypeMatching ( su , xj ) ( 13 ) Here , the TypeMatching component encapsulates both type inference and kernel matching , as detailed in Equation 2 . A script ( cf . Equation 1 ) can now be expressed as a recurrent application of FnIter : { yj } j = ( FnIter ◦ ... ◦ ( ni times ) ◦ ... ◦ FnIter ) ( { xi } i ) ( 14 ) Role : Inside a script , function iterations enable sharing of computational units in depth . Increasing the number of function iterations can increase depth without increasing the number of parameters . Preventing function signatures and variable types from collapsing on each other . One might obtain a scenario where the signatures of all functions and the types of all possible set elements collapse to a single point in type-space . This causes all set elements to be routed to all functions , thereby undermining the inductive bias of modularity . One effective way of preventing this degeneracy is to keep the function signatures fixed at initialization ( i.e . a high-entropy distribution ) . This effectively encourages the ( learnable ) type-inference module to produce diverse types , in order to use all available functions . In summary , we observe that any input element xi may trace a large number of computational paths as it progresses through the model , where a computational path is a partially ordered set ( poset ) of functions ( refer to Figure 3 for a visualization of such computational paths ) . In particular , owing to the fact that we recurrently apply function iterations in scripts , this poset may contain repeated functions , enabling weight sharing in depth . Further , the use of an interpreter that is shared between functions but explicitly conditioned ( or programmed ) by their code vector allows us to remain independent of the number of functions and retain the ability to add more functions after the model has been trained . Future work may investigate manipulating the code vectors themselves , thereby enabling higher-order functions ( i.e. , functions that manipulate other functions ) .

Summary: The paper proposes a model based on self-attention, factorised in different modules. It processes a sequence of tokens and could be viewed as a generalisation of a Transformer. The key idea is to have different modules (associated with a function) that behave differently but are general and actually share parameters. The core module of ModAttn can be viewed as a self-attention module with dynamic weights that depend on a context / code c. Using n_f learnable codes, each assigned to a ‘function’ leads to n_f modules, each with different functionality. An important aspect is the fact that each function only affects a subset of tokens, just the ones that are close to a type (learnable parameters) of the function.