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: This paper proposes Neural Interpreter (NI), a drop-in replacement for Transformer layers that can be applied on set-valued inputs. The motivation behind the architecture is that common computation units in deep learning such as convolution layer is rigid in how they process the input, so to overcome this, NI leverages self-attention to dynamically determine computation path using its modular functions depending on the given input. The paper shows how NI performs well in transfer-learning and systematic generalization tasks.