Dynamic Inference with Neural Interpreters
Year: 2021

Abstract:
Modern neural network architectures can leverage large amounts of data to generalize well within the training distribution. However, they are less capable of systematic generalization to data drawn from unseen but related distributions, a feat that is hypothesized to require compositional reasoning and reuse of knowledge. In this work, we present Neural Interpreters, an architecture that factorizes inference in a self-attention network as a system of modules, which we call functions. Inputs to the model are routed through a sequence of functions in a way that is end-to-end learned. The proposed architecture can flexibly compose computation along width and depth, and lends itself well to capacity extension after training. To demonstrate the versatility of Neural Interpreters, we evaluate it in two distinct settings: image classification and visual abstract reasoning on Raven Progressive Matrices. In the former, we show that Neural Interpreters perform on par with the vision transformer using fewer parameters, while being transferrable to a new task in a sample efficient manner. In the latter, we find that Neural Interpreters are competitive with respect to the state-of-the-art in terms of systematic generalization.

heading: 1 Introduction
text: 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.

heading: 2 Neural Interpreters
text: 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).

heading: 3 Related Work
text: Modularity. Despite recent successes of deep neural networks, their ability to recombine and reuse meaningful units for systematic compositional generalization is shown to be limited [31, 33, 39]. To mitigate this issue, work has been done to modularize deep architectures to enable the reuse and recombination of learned modules [2, 9, 19, 27, 32, 49, 53], and connections have been drawn to the principle of Independent Causal Mechanisms [7, 22–24, 34, 42, 43, 46, 52]. Modular architectures vary in how the modules are learned and laid out during inference. For instance, Neural Module Networks [2] and Neural Event Semantics [8] dynamically compose the modules using a natural language parser, whereas neuro-symbolic architectures [21–24, 37, 40, 57, 61, 62] typically require entity-centric representations [37, 57, 62] but use neural components like attention for compositional reasoning. Unlike some of these methods, Neural Interpreters (NI) neither require domain knowledge nor objectcentric representations. To obtain a strong compositional inductive bias, NIs aim at factorizing the computational graph in terms of functional units which are dynamically recombined in order to solve a particular task. In comparison with SCOFF [23] and NPS [22] which also separate functions, arguments and values and enable a similar dynamic and composed application of functions, NIs introduce the notion of interpreter (which enables parameter sharing between functions and extending the set of functions easily) and sophisticated signature and typing mechanisms described in the previous section.
Dynamic Routing of Information. The systematic generalization of modular architectures is shown to depend on the layout of modules [4, 48]. There exist different frameworks to facilitate the joint learning of module layout and parameters [8, 19, 24, 46, 49, 50, 53]. For example, [49] uses collaborative multi-agent reinforcement learned routing, [53] uses a trainable gating network. In this work, we use kernel modulated dot product attention [46] between function signatures and the input types to assign inputs to functions at each step.
Transformers. Modularity in Transformers [58] is relatively less studied. Repulsive attention [1] is based on repelling transformer heads to encourage specialization, whereas TIM [34] draws on the principle of independent mechanisms by dividing computation over parallel transformers that interact via competitive top-k attention [24]. In contrast, functions in NIs iteratively acquire set elements with kernel dot-product attention, which, if need be, can uniformly activate all the functions for a given token, encouraging efficient use of computation. Closer to Neural Interpreters, Switch Transformers (STs) [18] employ a mixture of experts framework where experts specialize across tokens, guided by a routing scheme which activates only one expert per token. However, STs focus more on dynamically selecting experts, whereas Neural Interpreters focus on dynamically composing available experts. The starting point of this work is the vision transformer (ViT) [13, 17], which applies a standard Transformer directly to image patches. We use the same architectural scaffolding as ViTs, but implement additional inductive biases that facilitate fast adaptation and systematic generalization.

heading: 4 Experiments
text: In this section, we empirically evaluate Neural Interpreters in a variety of problem settings. The questions we seek to answer are the following. (a) Can Neural Interpreters learn reusable computational units, given a number of training tasks that can in principle be solved with a fixed set of primitives? (Section 4.1) (b) Do Neural Interpreters modularize information in a way that helps fast adaptation [6]? (Section 4.2) (c) Can Neural Interpreters gracefully handle a larger number of functions (computational units) than they were trained with? (Section 4.2) (d) Do the inductive biases help with systematic generalization required in abstract reasoning problems? (Section 4.3).

heading: 4.1 Learning Fuzzy Boolean Expressions
text: In this section, we construct a toy problem to investigate whether neural interpreters can learn reusable functions when trained on a large number of tasks that share the same underlying building blocks.
Task Definition. Consider a set of N scalars {x1, ..., xN}, where xi ∈ [0, 1] is a real number between 0 and 1 (inclusive). We now define the following operations on the elements of this set:
and(xi, xj) = xixj ; not(xi) = x̄i = 1− xi; or(xi, xj) = xi ⊕ xj = x̄ix̄j (15)
Note that the above operations map from [0, 1]2 to [0, 1]; if xi, xj ∈ {0, 1}, they reduce to their Boolean namesakes. By combining these primitives, it is possible to construct and sample from a family of 22N fuzzy Boolean functions mapping from [0, 1]N to [0, 1] (Appendix B). The problem is now set-up as follows. We sample 30 random fuzzy Boolean functions of N = 5 variables, {fi}30i=1 where fi : [0, 1]5 → [0, 1], of which we use 20 for the pre-training phase and reserve 10 for the later adaptation phase. For each function, we sample 163840 points from a uniform distribution over [0, 1]5, of which we use 80% for training, and the remaining for validation. This yields two multi-task regression datasets, one for pre-training and the other for adaptation. A sample from the pre-training dataset comprises a 5D input vector x ∈ [0, 1]5 and the scalar regression targets f1(x), ..., f20(x).
Method. We pre-train a Neural Interpreter to regress all 20 functions simultaneously. To do this, we feed it a set of 25 elements as input: the first 5 are the components of x (with learnable positional embeddings [44] added in) and the remaining 20 are learned CLS tokens (vectors), one for each function. We attach a (shared) regression head to the output elements corresponding to the CLS tokens, and train their outputs to regress the respective function. Having pre-trained the model for 20 epochs, we finetune it for an additional 3 epochs on the 10 reserved functions {f21, ..., f30}. For the latter, we always instantiate and train 10 new CLS tokens, corresponding to the new functions; however, we
investigate three settings in which three different sets of parameters remain frozen. In the first setting, all parameters remain frozen, implying that the CLS tokens are the only trainable parameters. In the
second setting, we unfreeze the parameters of the type matching mechanism (function signatures and parameters of the type inference MLP, in addition the CLS tokens). In the third setting, we unfreeze all parameters and finetune the entire model. Additional details in Appendix B.
Hypothesis. By finetuning just the type-matching parameters, we only permit adaptation in how information is routed through the network. In other words, we only allow the computational units to rewire themselves in order to adapt to the new task at hand, while preserving their functionality. Now if the computational primitives that are learned during pre-training are recomposable, one would expect the performance of Neural Interpreters having finetuned just the type-matching parameters to approach that obtained by finetuning all the parameters, the latter including ones that determine the functionality of the computational primitives.
Results. Table 1 compares the coefficients of determination1 (R2) obtained in each of the investigated finetuning settings. We find that relative to finetuning just the CLS tokens, the performance difference between finetuning all parameters and just the type matching parameters is small. This is in line with expectation, suggesting that Neural Interpreters are indeed capable of learning recomposable computational primitives.

heading: 4.2 Multi-Task Image Classification
text: In this section, we evaluate Neural Interpreters in a multi-task image classification setting. Our goals are (a) to determine whether the inductive bias helps with fast adaptation, (b) to investigate whether the interpreter can indeed function as a general instruction executor, as intuited in Section 2, and (c) to demonstrate that the proposed architecture produces modules that can function autonomously (without additional training objectives that encourage this behaviour). Additional results in Appendix C analyze the effects of varying hyper-parameters.
Task Definition. We consider three related datasets sharing the same label semantics, namely SVHN [41], MNISTM [20] and MNIST [35]. The images therein are upsampled to shape 32× 32 (if required), and the resulting image is augmented with RandAugment [14]. Subsequently, the augmented images are split [13] to 64 patches of shape 4× 4. In addition, we also use unaugmented samples from the K-MNIST dataset [12] of Hiragana characters to probe fast-adaptation to new data.
1R2 = 1 implies perfect fit; a model that regresses to the mean has R2 = 0.
Method. We train Neural Interpreters for 100 epochs on the combined digits dataset described above. The input set contains 67 vector valued elements – the first 64 corresponding to linear embeddings of the 4× 4 patches, and the remaining 3 to learnable CLS tokens (vectors), one for each dataset. For each CLS token, a linear classification head is attached to the corresponding output; given an input sample from a certain dataset, the respective classification head is trained to predict the correct class. To inject position information into the model, we use a variant of the relative positional encoding scheme described in [13], but applied only to the first 64 input elements (corresponding to the image patches). Having pre-trained on the digits dataset, we finetune the model on K-MNIST with varying numbers of samples for 10 epochs. For additional details, please refer to Appendix C.
Baseline. Vision Transformers (ViT) [13, 17] make the natural baseline for Neural Interpreters, given the fact that the former is a special case of the latter. The set-up with CLS tokens and classification heads is identical to that of Neural Interpreters, as is the training protocol. We use a light-weight 8-layer deep model with approximately 1.8M parameters, but ensure that the considered Neural Interpreter model is even lighter, with roughly 0.6M parameters.
Hypotheses. (1) It has been suggested that a model that appropriately modularizes the learning problem should excel at fast-adaptation to a new but related data distribution [6]. If Neural Interpreters obtain a useful modularization of the problem, we should expect it to be able to adapt faster to the K-MNIST dataset than a non-modular baseline like the ViT. This is because some of the visual primitives present in the digit datasets (e.g., motifs, strokes or shapes) might reoccur in the K-MNIST dataset, which Neural Interpreters should be able to leverage. (2) Recall that in Section 2, we positioned the interpreter as a general purpose instruction executor, capable of being programmed by a code vector. If this is indeed the case, and if the interpreter does not overfit to the functions it was trained with, we should expect the capacity (and therefore performance) to increase as we add and train new functions (i.e., tuples of code and signatures), while keeping all other parameters fixed. (3) If functions (modules) are indeed autonomous, we should expect that as some of them are removed at test time, the others are still able to maintain performance. Consequently, the accuracy of the Neural Interpreter (trained with all functions) should
only degrade gracefully as functions are (randomly) removed at test time.
Results. Figure 5a compares the fast-adaptation performance of a Vision Transformer to that of a Neural Interpreter. While both models achieve almost identical performance on the validation set for all datasets (the difference being less than 0.2% in favor of Vision Transformers), we find that Neural Interpreters significantly outperform Vision Transformers at fast-adaptation in the low-data
regime, and the performance gap is only gradually closed as more data becomes available. This observation supports hypothesis (1). Further, in Figure 6, we establish a baseline where the functions pretrained on the digits dataset are destroyed prior to adaptation. We find that keeping the functions acquired from pretraining intact leads to significantly improved adaptation performance, suggesting that knowledge acquired during pretraining is indeed being reused during adaptation. This lends further support to hypothesis (1). Figure 5b shows the validation performance of a Neural Interpreter that was pre-trained with 5 functions on the digits dataset, but tested on the K-MNIST dataset with varying numbers of functions, having finetuned just the function signatures and codes. We find that the performance improves as new functions are added at adaptation-time, supporting hypothesis (2). Figure 5c shows the effect of randomly removing functions on the model accuracy. We observe that the performance on all datasets degrades gracefully as more functions are removed, thereby supporting hypothesis (3). In addition, Figure 3 visualizes the computational path taken by input set elements as they progress through the network, verifying that there is diversity in how samples are routed through the network, and shows that the routing mechanism discriminates between samples. Figure 4 shows that the input set elements gradually develop a type as they progress through the model. Additional figures in Appendix C analyze the effect of varying the number of scripts, number of function iterations, number of functions, kernel truncation parameter, dimension of type space T and freezing the function codes / signatures on the validation performance. The key finding is that, on the one hand, a wide range of hyper-parameter settings leads to performant models; on the other hand, there are patterns in what hyper-parameters perform consistently well.

heading: 4.3 Abstract Reasoning
text: In this section, our goal is to (a) use visual abstract reasoning tasks to evaluate whether Neural Interpreters are capable of systematic (compositional) generalization, and (b) characterize how Neural Interpreters maintain performance when the amount of compute is reduced at test time.
Task Definition. Progressively Generated Matrices (PGM) [5] is a procedurally generated series of datasets based on Raven Progressive Matrices (RPMs) [47], a popular IQ test designed for humans. The datasets, each comprising around 1.2M samples, aim to assess abstract and analogical reasoning capabilities of machine learning models beyond surface statistical regularities [29, 45]. Each sample in a dataset (Figure 7) comprises 8 context panels arranged in an incomplete 3× 3 matrix, and 8 candidate answer panels (a panel being an 80× 80 image). Given access to the context panels, the model must predict which of the candidate images is the most appropriate choice to complete the matrix. The content of the panels in a matrix are related by one or more triples, where a triple comprises a logical rule (∈ progression, XOR, OR, AND, consistent union) applied to an attribute (∈ size, type, color, position, number) of an object of a type (∈ shapes, lines). There are 8 datasets in the series, of which 7 measure systematic generalization in different ways, i.e., the test datasets contain samples where the panels are related in ways that are not encountered during training. In this work, we consider 6 datasets that probe the compositional reasoning ability of the model, namely: Interpolation, Extrapolation, Held-out (H.O.) triples, H.O. pairs of triples, H.O. Attribute Pairs, and Neutral. We omit H.O. line-type and H.O. shape-type datasets, since they stress the convolutional feature extraction instead of compositional reasoning component of the model. Please refer to Appendix D and [5] for additional details.
Method. Each panel (context and choice) is embedded with a shallow convolutional neural network to obtain an embedding vector. The input to the model is a set of 10 elements, comprising the embeddings of 8 context panels, that of a single candidate panel, and a CLS token (learnable vector). The model output corresponding to the CLS token is fed as input to the prediction head, which is trained to output a score measuring the compatibility of the candidate panel to the context panels. The final prediction is obtained by applying a softmax over the scores of all candidate panels. We note that this set-up resembles the one proposed in [5] for WReNs, and refer to Appendix D for details.
Baselines. We train Vision Transformers [13, 17] and Neural Interpreters with the same architectural scaffolding (described above). Additional baselines include Wild Relational Networks (WReN) [5],
Multiplex Graph Networks (MXGNet) [59] and WReNs atop disentangled representations (VAEWReN) [54]. Regarding the latter, we note that disentangled representations only improves on the feature extraction component of the WReN model; as such, it is entirely complementary to our contribution, which is the abstract reasoning component. Future work may explore replacing the convolutional embeddings with disentangled representations.
Hypothesis. Recall that the test sets of (all but one) PGM datasets contain samples where the panels are not related in a way that is encountered during training. A model that is able to infer and reason about the underlying compositional structures is therefore more likely to be successful. If the inductive biases do indeed enable Neural Interpreters to factorize knowledge in to recomposable primitives, we might expect them to be able to dynamically recombine these primitives in order to generalize to problems it has not encountered during training and thereby excel at this task.
Results. Table 2 tabulates the validation (in-distribution) and test (out-of-distribution) accuracies obtained by all models for the various considered datasets. We make two observations: (1) Neural Interpreters are competitive in terms of test accuracy, outperforming both Vision Transformers and the prior state-of-the art in 4 of 6 datasets. (2) Vision Trans-
formers are competitive in-distribution: they outperform the state of the art (excluding Neural Interpreters) in 5 of 6 datasets in terms of validation accuracy. However, they are outperformed in terms of test accuracy by either Neural Interpreters or other baselines in all datasets. In addition, Figure 8 shows how the test-set performance varies with the number of function iterations, for a model that was trained with a fixed number of function iterations (8). We note that the number of function iterations is proportional to the amount of computation used by the model. In all cases, the model maintains at least 70% of its original performance while using only half the amount of compute; in most cases, at least 80% of the original performance is maintained. This suggests that Neural Interpreters can function as anytime predictors [63], which are models that can trade-off performance with compute.

heading: 5 Conclusion
text: We have presented Neural Interpreters, a self-attention based architecture capable of learning a system of recomposable computational primitives. Neural Interpreters relax the rigidity in how computation is reused in current neural architectures, and our experiments show that the modular inductive biases it incorporates facilitate systematic generalization and sample-efficient adaptation to new tasks.
There are multiple exciting avenues of future research. One line of work could leverage the capacity extension capability of Neural Interpreters in a continual learning setting, thereby enabling continuous integration of knowledge in a neural model. Another promising direction entails using Neural Interpreters as a backbone for learning world models, where systematic generalization and out-ofdistribution robustness are of paramount importance.

heading: Acknowledgements
text: The authors would like to thank Sharath Chandra Raparthy and Sarthak Mittal for providing valuable feedback on the draft. This work was in part supported by the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039B). YB acknowedges support from the Canada CIFAR AI Chair Program, and Samsung Electronics Co., Ltd.

Review summaries:
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. 

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.

This paper proposes a method for modularizing neural networks inspired by programming functions and compilation. Building on the success of attention and transformer architectures, the neural interpreter (NI) proposes a new architecture built on layered functions, each of which is composed of stacked MLP and attention layers. Functions and input variables both have embeddings, called “signature and type” respectively, that allow NIs to learn specialized functions that operate over a subset of inputs. The NI is applied to a suite of image datasets as well as to Raven matrices, where it is shown to outperform vision transformers.

This paper focuses on improving the generalization ability of deep models to the unseen but related distributions to training data, and thus enhancing interpretability. In specific, a Neural Interpreter is proposed, which factorizes the inference of self-attention networks as the dynamic routing within functions. The Neural Interpreter can flexibly compose computation on a per-sample basis, and it is end-to-end trainable. Some interesting experiments including image classiﬁcation and visual abstract reasoning are conducted to validate the effectiveness of the proposed method.