Neural Attentive Circuits
Authors: Martin Weiss
Nasim Rahaman
Francesco Locatello
Christopher Pal
Yoshua Bengio
Bernhard Schölkopf
Li Erran Li
Nicolas Ballas
Venue: NeurIPS 2022 Accept (2022)
Keywords: Deep Learning
Low-Shot Adaptation
Attention Mechanisms
General Purpose Neural Architectures

Abstract: Recent work has seen the development of general purpose neural architectures that can be trained to perform tasks across diverse data modalities. General purpose models typically make few assumptions about the underlying data-structure and are known to perform well in the large-data regime. At the same time, there has been growing interest in modular neural architectures that represent the data using sparsely interacting modules. These models can be more robust out-of-distribution, computationally efficient, and capable of sample-efficient adaptation to new data. However, they tend to make domain-specific assumptions about the data, and present challenges in how module behavior (i.e., parameterization) and connectivity (i.e., their layout) can be jointly learned. In this work, we introduce a general purpose, yet modular neural architecture called Neural Attentive Circuits (NACs) that jointly learns the parameterization and a sparse connectivity of neural modules without using domain knowledge. NACs are best understood as the combination of two systems that are jointly trained end-to-end: one that determines the module configuration and the other that executes it on an input.  We demonstrate qualitatively that NACs learn diverse and meaningful module configurations on the Natural Language and Visual Reasoning for Real (NLVR2) dataset without additional supervision. Quantitatively, we show that by incorporating modularity in this way, NACs improve upon a strong non-modular baseline in terms of low-shot adaptation on CIFAR and Caltech-UCSD Birds dataset (CUB) by about 10 percent, and OOD robustness on Tiny ImageNet-R by about 2.5 percent. Further, we find that NACs can achieve an 8x speedup at inference time while losing less than 3 percent performance. Finally, we find NACs to yield competitive results on diverse data modalities spanning point-cloud classification, symbolic processing and text-classification from ASCII bytes, thereby confirming its general purpose nature. 

Neural Attentive Circuits
Martin Weiss 1.6 Nasim Rahaman 11 Francesco Locatello"
Chris Pal 1,077 Li Erran Lit.3
Yoshua Bengio 1,5,7
Bernhard Schölkopf² Nicolas Ballas f.4
* Equal contribution, random order.
Mila, Quebec AI Institute AWS AI
2 Max Planck Institute for Intelligent Systems, Tübingen
Meta AI
5 Université de Montréal Polytechnique Montréal
7 Canada CIFAR Al Chair
Abstract
Recent work has seen the development of general purpose neural architectures that can be trained to perform tasks across diverse data modalities. General purpose models typically make few assumptions about the underlying data-structure and known perform well in the large-data regime. At the same time, there has been growing interest modular neural architectures that represent the data using sparsely interacting modules. These models be more robust out-of- distribution, computationally efficient, and capable of sample-efficient adaptation new data. However, they tend make domain-specific assumptions about the data, and present challenges in how module behavior (i.e., parameterization) and connectivity (i.e., their layout) can be jointly learned. In this work, we introduce a general purpose, yet modular neural architecture called Neural Attentive Circuits (NACs) that jointly learns the parameterization and sparse connectivity of neural modules without using domain knowledge. NACs are best understood as the combination of two systems that are jointly trained end-to-end: one that determines the module configuration and the other that executes it an input. We demonstrate qualitatively that NACs learn diverse and meaningful module configurations on the Natural Language and Visual Reasoning for Real (NLVR2) dataset without additional supervision. Quantitatively, we show that by incorporating modularity this way, NACs improve upon a strong non-modular baseline in terms of low-shot adaptation on CIFAR and Caltech- UCCD Birds dataset (CUB) by about 10 percent, and OOD robustness Tiny ImageNet-R by about 2.5 percent. Further, we find that NACs can achieve an 8x speedup at inference time while losing less than 3 percent performance. Finally, we find NACs to yield competitive results diverse data modalities spanning point-cloud classification, symbolic processing and text- classification from ASCII bytes, thereby confirming its general purpose nature.
Introduction
General purpose neural models like Perceivers 29 do not make significant assumptions about the underlying data-structure of the input and tend perform well in the large-data regime. This enables the application of the same model on a variety of data modalities. including images, text, audio, point-clouds, and arbitrary combinations thereof (29 28). This is appealing from an case-of-use perspective, since the amount of domain- can function well out-of-the-bon in larger machine learning pipelines, e.g., AlphaStar (28).
components is minimized, and the resulting models
At the same time, natural data generating processes can often be well-represented by a system of sparsely interacting independent mechanisms [1 46). and the Sparse Mechanism Shift hypothesis
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

stipulates that real-world shifts are often sparse when decomposed, i.e., that most mechanisms may remain invariant 47. Modular architectures seek to leverage this structure by enabling the learning of systems of sparsely interacting neural modules 3
45 Such systems tend to excel in low-data regimes, systematic generalization, fast (sample-efficient adaptation to new data and can maintain larger degree of robustness
20
44
[52] 3 4 [1]. However, many of these models make domain-specific assumptions about the data distribution and modalities [4] - c.g., Mao et al. (37 relying on visual scene graphs, Andreas et al. [3] using hand-specified modules and behavioral cloning. As a result, such models are often less flexible and more difficult to deploy.
In this work, we propose Neural Attentive Circuits (NACs), an architecture that incorporates modular inductive biases while maintaining the versatility of general purpose models. NACs implement a system of sparsely and attentively interacting neural modules that pass messages to each other along connectivity graphs that are learned or dynamically inferred in the forward pass and subject to priors derived from network science [25] 6 [8]. NACs demonstrate strong low-shot adaptation performance and better out-of cuttof-ddisributioo robustness when compared to other general purpose models such as Perceiver IOs (28). Besides scaling linearly with input size (like Perceivers), we show that by pruning modules, the computational complexity NACs can be reduced at inference time while preserving performance. In addition, we propose conditional variant where the graph structure and module parameterization is conditioned on the input. This enables conditional computation 10] []. where the functions computed by neural modules and the sparse connectivity between modules is determined only at inference time. A qualitative analysis on Natural Language and Visual Reasoning for Real (NLVR2) shows that conditional NACs can learn connectivity graphs that are meaningfully diverse.
Figure D shows the schematics of the proposed architecture, with its two core components: the circuit generator and the circuit executor. The circuit generator produces the configuration over modules, which we call a circuit design, defining (a) the connectivity pattern between the neural modules, and (b) instructions that condition the computation performed by each module. The circuit design may either be conditioned on (part of) the sample (in the case conditional NACs), or it may simply be learned over the course of training via gradient descent (in the case of unconditional NACs). The circuit executor consumes the circuit design and an input sample to perform inference.
Circult Design
Processor Descriptora
Read Out Decor
"There cat outdoors the right image."
Circuit Generator
Sigrattunes
Codes
Preproses
Read-In
Propagator
Propagator
Read-Out
Circuit Executor
Figure 1: Circuit Generator and Executor. We show NAC applied to a natural language and visual reasoning task where text conditions the circuit generator and the program executor consumes images.
Our contributions are as follows. (a) We propose Neural Attentive Circuits (NACs), a general purpose neural architecture that jointly learns the module parameterization and connectivity end-to- end and gracefully supports more than one thousand sparsely interacting modules. (b) We demonstrate the general-purpose nature of NACs by training it on a selection of diverse data domains covering natural images, 3D point-clouds, symbolic processing, text understanding from raw ASCII bytes, and robustness and few-shot adaptation performance of NACs. We find that in the low-shot regime, NACs can outperform Perceiver IOs by approximately 10% on 8- 8-Way CIFAR and CUB-2011 1. Further, the proposed model yields roughly 2.5% improvement on TinylmageNet-R, an out-of-diisribution test-set for TinyImageNet based on ImageNet-R (24) (d) We explore the adaptive-come computation aspect of NACs, wherein the computational complexity of trained model can be reduced at inference time by roughly 8 times, at the cost of less than 3% loss in accuracy on Imay- (e) We qualitatively demonstrate that it is possible to condition the circuit design (i.e., configuration over modules) on the input. On NLVR2, we the text modality of the sample to condition the circuit design, which is
natural- Lanuual- visual reasoning. (c) We quantitatively evaluate the
2

---

then applied to the image modality of the same sample to produce a result. We find that connectivity graphs generated by sentences that involve similar reasoning skills have similar structures.
2 Neural Attentive Circuits
This section describes Neural Attentive Circuits (NACs), a neural architecture that can consume arbitrary (tokenized) set-valued inputs including images and text. NACs are system of neural modules that learn interact with each other (see Figure 2 Module connectivity learned end-to- end and is subject to regularization that encourages certain patterns (e.g., sparsity, scale-freeness and the following, we first describe two key building blocks of the proposed
formation of cliques). architecture. We then introduce the circuit executor that infers an output given some input and circuit design. We finally introduce the circuit generator which produces said circuit design.
Read-In
Propagator
X
Read-Out
ModFFN
ModFFN
ModFFN
ModFFN
ModFFN
ModFFN
ModFFN
ModFFN
ModFFN
ModFIN
Cross Attention
Steshorto Korna Hookdaaed
Soschastcc Kerna Modklted
Circuit Design
Processer Descriptors
Read-Outt Descriptors
S
%
f
Signatures
LayerNorm W, c.
Codes
X2
y
ModFC
Figure 2: Schematics of Neural Attentive Circuits. Top: a NAC with Up 4 processor modules and U, = read-out modules. Bottom right: each module has a unique signature and code (colored circles and boxes). The signatures determine which modules interact and the code given module conditions the computation it performs (via ModFFNs and ModFCs). Bottom left: visual description of ModFFNs and ModFC.
Circuit Design. Consider a NAC (Figure 2 with U modules, of which Up are processor modules (responsible for the bulk the computation) and U, are read-out modules (responsible for extracting model outputs). By circuit design, we refer to the set of U descriptors that condition the module connectivity and computation carried out at each module. Each descriptor comprises three vectors (Sa. Cu. 0 ): signature vector E Rddis code vector C. E Relodd and module initial state e E RRdmade! The set of signature vectors (su), determines the connectivity between modules while the code vector conditions the computation performed by module at index u. Further, the vector d!! chooode tracks the state of the module before the [-th layer, and by (()) we denote the set of all such states (over modules). Each module updates its state as the computation progresses in depth (see below). Finally, processor module signatures and codes are shared across depth.
ModFC and ModFFN. The Modulated Fully-Connected (ModFC) layer is the component that enables module condition its computation supports multiplicative conditioning of the input vector x by the code vector c. It can be expressed
code. It replaces the fully-connected layer and
ModFC(x; = ( @ (1 + a LayerNorm/ W.c)) + b.
(1)
where @ denotes element-wise multiplication, W E Rduu x Rd is a weight matrix, b E Bdout is a bias vector, E Redco the conditioning code vector, W. Rdi.. x Rdevde is the conditioning weight matrix, and a is learnable scalar parameter that controls the amount of conditioning. Setting = 0 removes all conditioning on e, and we generally initialize a = 0.1. Multiple such ModFC layers (sharing the same code vector e) and activation functions can be stacked to obtain a ModFFN
3

---

layer. Like in transformers, multiple copies of ModFCss and ModFFNs are executed in parallel. But unlike in transformers, each copy is conditioned by a unique learnable code vector and therefore performs different computation. Consequently, though the computation performed in each module different, the total number of parameters does not noticeably increase with the number of modules. See clarifying example in Appendix
SKMDPA. Stochastic Kernel Modulated Dot-Product Attention (SKMDPA) is a sparse-attention mechanism that allows the modules to communicate. Given two modules at index i and j. the distance between their signatures 8; and 8, determines the probability with which these modules can interact via dot- dot-product attention smaller distance implies higher likelihood that the modules are allowed to communicate. Like vanilla dot-product attention, SKMDPA operates on a set of query qı key kj. and value 2j vectors, where indexes queries and j indexes keys and values. But, in addition, each query qi and key k; is equipped with a signature embedding vector, which may be learned as parameters or dynamically inferred. These signatures are used to sample kernel matrix K [ via
d(s;, 8.)
Ki Concrete(PPy with P. exp
(2)
where di is some distance metric, is a bandwidth (a hyper-parameter). Concrete is the continuous relaxation of the Bernoulli distribution [36) with sampling temperature т, and the sampling operation is differentiable via the reparameterization trick. K2 is likely to be close 1 if the distance between s; and sy is small, and close to O if not. As 7 + 0. the kernel matrix Kyy becomes populated with either 1s or Os. Further, if P.j ensured to be sufficiently small on average, Kij is a sparse matrix. In practice, we set T to a small but non-zero value (e.g. 0.5) to allow for exploration. Now, where Art; the dot-produce attention score ks/va between the query q, and key ky, the net attention weight Wses and the output y: are
Wi+s= softmax, (Ac-s + log Ki]] where Rij Ku/[6+2i, K.)) and y Win (3)
Here, Wi+j denotes the weight of the message U, passed by element j to element i, and 5 is a small scalar for numerical stability. We note that if K, R O, we have Wi-j Rs 0, implying if K2 is sparse, so Wi-j- The proposed attention mechanism effectively allows the query q: to interact with the key ky with a probability that depends on the distance between their respective signatures. In this work, this distance is derived from the cosine similarity, i.e., note the connection to Lin et al. (34), where the attentive interactions between queries and keys are dropped at random, albeit with probability that is not differentiably learned.
= 1 - CosineSimilarityys: s,). We
2.1 Circuit Executor
The circuit executor takes the circuit design and executes it on an input to make a prediction. It has four components: (a) tokenizer convert the input (e.g., images or text) to a set of vectors, (b) a read-in layer which allows the processor modules to attentively read from the inputs, (c) a sequence of propagator layers which iteratively update the processor modules states through round of communication and computation, (d) read-out layers that enable the read-out modules to attentively poll the final outputs of the processor modules.
Tokenizer. The tokenizer is any component that converts the input to the overall model to a set of representation vectors (with positional encodings, where appropriate), called the input set X. can be a learned neural network (e.g., the first few layers of a ConvNet), or a hard-coded patch extractor. Role: The tokenizer standardizes the inputs to the model by converting the said inputs to set of vectors.
Read-In. The read-in component a read-in attention followed by a ModFFN. Read-in attention is a
cross-attention that uses the initial processor module states @ to generate a query vector and uses the inputs X to generate keys and values. Unlike typical cross-attention 32 29), we use ModFCs instead of FCs as query, key and value projectors. Where æ; is an element of X, we have:
Tu = (o cw))
k.jj = ModFC(2; c.)
Vuj = ModFC(z,, Cw)
(4)
in
Σ
qu .ua Uuy
yu ModFC
@ ModFFN(yu.. C.))
(5)
Vd

---

Each query 0 sees its own set of keys kuy and values Duj. owing to the conditioning on the code Cu. The read-in in attention is followed by Up copics of a two layer ModFFN each conditioned on a code Cu. The read-in outputs set of vectors 0!" E 6 Role: The read-in attends to the inputs via a cross-attention mechanism that scales linearly with the size of the input.
Propagator Layers. After the read-in operation, propagator layers sequentially update the state of each processor module L. times. Each propagator layer implements a round of module communication via SKMDPA, followed by the computation of Up copies of ModFFN. Both operations are condi- tioned on their processor module descriptor codes. The 1-th propagator (where / 1,..., L) ingests the set of states 0 and outputs another set 6% of the same size. Role: A propagator layer updates each processor module state vector by applying learned stochastic attention and module computation.
The read-out component is a SKMDPA-basCe cross-attention mechanism (read-out attention) followed by a ModFFN, and it is responsible for extracting the output from the processor module states after the final propagator layer. Read-out modules have their own signatures and codes parameters which are different from the processor modules. The read-out attention uses the initial states @ of these modules as queries, and the final states of the processor modules e to obtain keys and values. Finally, we note that in vanilla classification problems, we found it useful to use several read-out modules. Each read-out module produces a confidence score, which is used to weight their respective additive contribution the final output (e.g.. classification logits). Role: The read-out is responsible for attending to the final processor module states to produce an overall output.
Read-Out.
2.2 Circuit Generator
Recall that the circuit generator produces circuit design that specifies the connectivity between the modules (via signatures) and the computation performed by each module (via codes). There are two ways in which the circuit generator can be instantiated, as described below.
Unconditional Circuit Generator. In the first variant (unconditional), the circuit design is induced by signatures S. and codes C. that are freely learnable parameters (where u indexes the module descriptor). The initial state e! is obtained passing Cu through a two layer MLP. We refer to the overall model as an unconditional NAC, or simply a NAC. Role: The unconditional circuit generator implements a mechanism to learn the circuit design via free parameters with gradient descent.
Regularization via Graph Priors. In the unconditional setting, we observe a task-driver pressure that collapsed all signatures s, to the same value (see Appendix This results in graph that is fully connected, i.e., all modules interact with all other modules, thereby forgoing the benefits of sparsely interacting modules. To counteract this pressure, we impose structural prior on the learned graph via a regularization term in the objective, which we call the Graph Structure Regularizer. The graph prior is expressed as a target link probability matrix and the extent to which the prior is satisfied is obtained comparing the closest permutation of P. (the generated link probability matrix) to the target values. This regularizer encourages the graph to have certain chosen properties, while still preserving the flexibility afforded by learning the graph structure. We experiment with several such properties, including ccale-freeness 6 (leading to few well connected nodes or hubs, and a heavy tail of less connected nodes), and planted partition [15] one based on the stochastic block model (25 (encouraging modules to group themselves in cliques, where communication within 4 we will find that some graph properties do indeed yield better out-of distriibuitt performance, whereas others perform better in-distribution. Figure 3 shows shows the class of graphs we experiment with in this work, and Appendix [ contains the details. Role: The Graph Structure Regularizer enforces sparsity by preventing the module connectivity graph from collapsing to an all-to-all connected graph, while inducing graph properties that may unlock additional performance depending on the task.
cliques is less constrained than that between cliques). In Section
Conditional Circuit Generator. In the second variant (conditional). each sample is assigned a different circuit design. In this case, the circuit generator takes as input the sample X (or part thereof), and outputs a set of signatures and codes (X), cu (X))u. This setting is particularly interesting for multi-modal data, where the circuit generator may ingest one modality and the circuit executor may consume another. In this work, we implement the conditional circuit generator as simple cross attention mechanism, where learned parameters are used as queries (one per module) and elements of X as keys and values (32. We refer to the resulting model as conditional NACs,
5

---

Figure 3: Samples from Graph priors explored here. Node colors represent degree (# of edges). Left: Planted-
partition 15 25.. Center-left: Scale-free 7 6). Center-right: Ring-of-Cliques 38
Right: Benos 17
while noting that the circuit generator could in principle be implemented by other means (e.g.. auto-regressive models and GFLowNets (12) in future work. Role: The conditional circuit generator relies on the sample when producing a circuit design. This way. the connectivity between modules and the computation performed by each module can be dynamically changed for each sample.
In summary, we have introduced Neural Attentive Circuits (NACs) and its conditional variant, general purpose neural architecture of attentively interacting modules. We discussed the two key components of NACs: (a) circuit generator that generates circuit design determining which modules connect to which others (via signatures) and the function performed by each (via codes), and (b) circuit executor that executes the configuration on an input to infer an output.
3 Related work
General Purpose Models. A recent line of work develops general-purper models that make few assumptions about the data domain they are applied to [29 28 23). Of these models, Perceiver IO 28 is most similar to NACs, in that (a) they both rely on cross- aooss-attention to both read (tokenized) set- valued inputs into the model and extract set-valued output from the model, and (b) their computational complexity scales linearly with the input size. However, unlike Perceiver IO (PIO), (a) the connectivity between NAC modules is not all-to-all, but sparse in a learned way, and (b) the computation performed by NAC's equivalent of PIO latent is conditioned at all layers via the ModFC component.
Modular Inductive Biases. Modular neural architectures are known to yield models that are capable of systematic generalization and fast adaptation 5 2) Of these architectures, Neural Interpreters (NIs) are most similar to NACs, in that they both use the framework of signatures and codes, where signatures determine the computational pathway and codes condition individual modules. However, there are four key differences: (a) unlike in NIs, the routing of information in NACs is stochastic, (b) the mechanism to condition on codes (ModFC) is different, and (c) the computational complexity and memory requirements of NIs scales quadratically with the size of the input, whereas in NACs, this scaling is linear, and (d) NIs are only known work with up to 8 modules, whereas NACs can support more than a thousand modules on a single GPU (i.e., without model parallelism).
Neuro-symbolic Models. Neuro-symbolic models have shown strong systematic generalization on synthetic datasets such as CLEVR 30 and SQOOP [.. However, models of this class such as the Symro-Symbolic Concept Learner 337 and Neural Module Networks 3 are not general-purpose, using semantic parsers and Nomain-specific modules. In contrast, NACs are general-purpos and can jointly learn sparse layout and parameterization for more than one thousand homogeneous modules.
4 Experiments
In this section, we empirically investigate Neural Attentive Circuits (NACs) in five problem settings with different modalities: image classification, point cloud classification, symbolic processing (ListOps (49), text classification from raw bytes 49 and multi-modal reasoning over natural language and images [48). We compare against baselines from the literature, as well as those that we train under controlled settings. The questions we ask are the following:
(a) Does the inductive bias of sparsely interacting modules improve NACs -iittribuitt robustness and fast (few-shot) adaptation to new data? (b) How important is it that these modules be capable of modelling different functions (via ModFC layers), and that the connectivity between modules learned (via SKMDPA)? (c) Do the modules strongly co-adapt to each other, can the system still function with some modules

---

0.95
Neural Attentive Circuits
Neural Attentive Circults
Perceiver 10
0.85
Accurscy (R-Wagg)
0.7
0.75
0.6
0.70
0.60
0.55
0.4
26
4 Number of Samples per Class (Shots)
16
Number of Samples per Class (Shats)
(a) CUB 8-Way
(b) CIFAR S-y
Figure 4: Few-Shot Adaptation Results on CUB and CIFAR. On each dataset, we fine-ture ImageNet pretrained NAC and PIO models to classify images into 8 classes and vary the number of examples ("shots") per class from I 16. In the low-shot transfer regime (e.g., samples per class), we find that NAC compares favorably with the baseline PIO, suggesting that the modular inductive biases in the former can fast adaptation.
pruned inference time? (d) In the sample-conditional setting, there interesting patterns in the generated circuit designs? (e) Are NACs general-purpose?
4.1 Image Classification
In this section, conduct quantitative evaluation of NACs image classification problems. We use Imy--eegeNeet as our primary test-bed. Immy-Imageeee is a miniaturized version of ImageNet(T6), with smaller dataset size and lower image resolution. allows for compute efficient experimentatder but makes the overall classification task significantly more challenging, especially for larger models that easily overfit [33] To evaluate out-of-distribution (OOD) robustness in this regime, use Tiny--maaeNet-RR which is a down-sampled subset of ImageNet-R(endition
experimentation,
[24]. It contains roughly 12000 samples categorized in 64 classes (a subset of Tiny-ImageNet classes), spread across multiple visual domains such as art, cartoons, sculptures, origami, graffiti, embroidery, etc. For low-shot adaptation experiments, we also train NAC on ImageNet, in order compare with pretrained baseline. Additional details can be found in Appendix C
Baselines. For certain choice of hyper-parameters, NACs can simulate Perceiver IOs (PIOs) 28 making the latter the natural baseline for NACs. For all Finy-ImageNet runs for both NACs and PIOs, the same convolutional preprocessor that reduces the 64 64 input image to 8 feature map. For ImageNet training, we use the same convolutional preprocessors as we do for Tiny-ImageNet. Please refer to Appendix for additional details.
Pre-training. We pretrain all models on Tiny-ImageNer for 400 epochs, and perform model selection based accuracy the corresponding validation set. We evaluate the selected models Tiny- ImageNet-R. For few-shot adaptation experiments, we use the official ImageNet pretrained weights for 48-layer deep PIO (28), which yields 82.1% validation accuracy. To compare few-shot adaptation performance with this model, we train 8-layer deep NAC with a scale-free prior graph on full ImageNet for 110 epochs, which yields 77% validation accuracy with 1024 modules. Additional details and hyperparameters can be found in Appendix C
Few-Shot Adaptation. Bengio et al. [1] proposed that models that learn good modularization of the underlying data generative process should adapt in a sample-efficient manner to unseen but related data-distributions. Motivated by this reasoning, we measure NACs few-shot adaptation performance and compare with a non-modular PIO baseline. To this end, we fine-tune ImageNet pre-trained NACs and PIOs small numbers of samples between 8 and 128 in total from two different datasets: CUB-2011 (Birds) and CIFAR. The results shown in Figure support the hypothesis that the modular inductive biases in NACs help with sample-efficient adaptation, and we show the few-shot experiments with varying numbers of classes in the Appendix with similar results.
A Ablations. In this experiment (Table we start with a Perceiver IO and cumulatively add design elements to arrive at the NAC model. We evaluate the and OOD performances and make two key observations. (a) We find that adding ModFC layers does not significantly affect IID generalization, improves OOD generalization. Moreover, using a static Barabasi-Albert attention graph (via frozen signatures) reduces the IID performance while further improving OOD performance. Jointly
7

---

learning the connectivity graph and module codes yields the largest improvements. Learning the graph allows recover the IID performance, and the OOD performance further improved. (b) The choice of graph prior exposes a trade-off between IID) and OOI) performance. 'he scale-free prior yields the best IID performance, whereas the ring-of-cliques prior yields the best OOD performance.
III
0OD Imay-emageNet-R)
(Tiny-mmageNet Validation)
Acc@1
Acc@5
Acc@1 1
Acc@5
Perceiver I0
58.89
80.09
17.57
37.16
ModFC
58.91
79.43
18.01
37.05
Sparse graph (not learnable)
58.31
80.06
18.50
37.92
NAC with Scale-Free Prior
60.76
80.86
19.52
38.52
NAC with StantedPPartttin Prior
60.71
81.34
19.42
39.84
NAC with
Prior
60.54
81.18
20.03
40.44
NAC with Erdos-Renyi Prior
60.33
81.08
19.83
39.32
Table 1: Ablation and Graph Results. In this experiment, start with Perceiver IO and cumulatively add design elements arrive the NAC model. At each stage, train and observe the effect on generalization. find that adding ModFC layers alone does not have clear impact on either IID or OOD generalization. However, OOD is then clearly improved by adding a sparse kernel sampled from a frozen prior distribution. Below the midrule, we evaluate NACs performance with several different graph prior regularizers. We see that learning the graph via SKMDPA allows the model to improve both its IID and OOD performance.
Effect of Dropping Modules. In this experiment, seek to measure the extent to which the NACs modules depend on other modules to function. To this end, we drop modules order of their connectivity (i.e., least connected modules are dropped first) and measure the performance of the model on Tiny-ImageNet along the way. We perform this experiment for models trained with different types of graph priors, and show the result in Figure 5 We make two observations: (a) is possible to remove a large fraction (e.g., more than 80%) of the modules at little cost to accuracy, obtaining 8x speedup in the process, and (b) different graph priors offer different amounts of robustness towards dropped modules. The former observation is particularly interesting from computational efficiency perspective (Figure (5b), since dropping modules increases throughput and decreases memory requirements at inference-time time without additional implementation overhead. The latter observation suggests a way in which the choice of graph prior can affect the properties of the NACs model. We provide additional details in Appendix FF where Figure 15 compares the effect of sparsification on NACs (with different graph priors) and Perceiver I0.
700
Tihy- Inngeeest Vol. Acc
7-
GFLOPs Per Samole
6-
0.550
500
4-
0.500
Endos-Renyi Prior
3-
0.475
Scale
Ring of
0.450
Panted--Parthn Parttiiin Prior
200
Porcsiver
0.50
0.65
3.70
0.75
0.80
0.55
0.90
0.0
0.2
0.4
0.5
Fradton of hades (MAcdules Dropoed
Nodoc (Moduloo) Dropped
(a) Validation accuracy modules are dropped.
(b) Inference speed modules are dropped.
Figure 5: Adaptive Computation. We train models to convergence and then sparsify them at inference time and evaluate their validation accuracy. In both plots, the x-axis shows the number of modules dropped inference time, and the Y-axis shows the accuracy of the model on Tiny-ImageNet after sparsification inference time. We observe that the Neural Attentive Circuit (NAC) much more robust than Perceiver IO to sparsification at inference time, and can drop over 80% of circuits before its performance drops below Pereciver 10. As a result, NACs can be adapted at inference time to a wide range of computational budgets.

---

4.2 Point Cloud Classification
We evaluate the NAC architecture on point cloud classification benchmark without additional data augmentation, geometric inductive biases or self-supervised pre-training [14]. Comparably, Perceiver 10 obtained 77.4%, Hierarchical Perceivers achieved 75.9% without self-supervised pre-training. and when trained with self-supervise pre-training Hierarchical Perceivers obtained 80.6% (cf. Table 7 of 14)) We trained a NAC without specialized tokenization or pre-training to achieve a test accuracy of 83%, demonstrating that the sparse inductive prior was quite effective in this setting. The original Perceiver [29] relies on additional data augmentation beyond that used in [44 obtain 85% test accuracy. We do not know how well the original Perceiver would perform without the additional augmentations, but expect it to be similar to Perceiver IO.
Method
Test Accuracy
Hierarchical Perceiver (No MAE) [4]
75.9%
Perceiver IO [14]
77.4%
Hierarchical Perceiver (with MAE)
14
80.6
NAC (ours) (No MAE)
83.0% %
Table 2: Point Cloud Classification. In this experiment, we compare against results reported in Carreira et al. [14. In that work, the authors were able to improve over Perceiver IO results through a masked auto eitoeenn (MAE) pre-training step. We find that, without pre-training, NACs with Scale-Free sparse prior graph are able achieve superior test accuracy.
4.3 Symbolic Processing ListOps and Text Classification from Raw Bytes
In this set of experiments, we evaluate the ability of NACs to reason over raw, low-level inputs. We two tasks to this end: symbolic processing task (ListOps) and text-llasifiiattion task from raw ASCII bytes. The former (ListOps) is a challenging 10-way classification task on sequences of up 2,000 characters [39]. This commonly used benchmark (49) tests the ability of models to reason hierarchically while handling long contexts. The latter is a binary sentiment classification task defined sequences of up 4000 bytes. In partieular, bytes, and not tokenized text as common in most NI.P benchmarks [49] This tests the ability of NACs handle long streams of low-level and unprocessed data.
stress that the models ingest raw ASCII
The results are presented in Table 3 We find that NACs are competitive, both against general- purpose baseline (our implementation of Perceiver IO) and other third-party baselines, including full-attention transformers and the Linformer. This confirms that NACs are indeed general purpose, and that general-purpose architectures be effective for handling such data.
4.4 Sample-Conditional Circuit Design Generation
In the previous sections, we quantitatively analyzed the performance NAC model equipped with connectivity graph that is learned, but fixed at inference time. In this section, we investigate cumple-conditionss circuit generation without regularizing the graph structure. Concretely, we study the circuit designs generated by conditional NAC trained to perform multi-modal reasoning task. In our experiments with conditional NACs, the circuit generator was implemented by a simple cross attention layer followed by feed forward network (FFN), and self-attention layer followed by
Test Accuracy
Method
ListOps 39
Text Classification
Full Attention Transformers (49
37.13%
65.35%
Linformer 449
37.38%
56.
Perceiver 10 (our impl.)
39.70%
66.50%
NAC (ours)
41.40%
68.18%
Table 3: ListOps and Text Classification Results. In this experiment, we train a Perceiver 10 and a NAC from scratch on two tasks from the the Long Range Arena [49]: ListOps symbolic processing and text classification from ASCII bytes. We observe that NAC excels in processing long-range dependencies in this setting.

---

two FFNs run in parallel, see Appendix G for further details. We use Natural Language for Visual Reasoning for Real (NLVR2) dataset 48 which contains natural language sentences grounded in images. The task is to determine whether natural language statement holds true for given pair of photos. Solving this task requires reasoning about sets objects, making comparisons, and determining spatial relations. Accordingly, we would expect that the circuit generator creates circuit designs that are qualitatively different for different logical relations.
Training. We condition the circuit generator on natural language captions, and the executor outputs binary label conditioned both on the image pairs and on the output the circuit generator. Given that our goal of understanding how conditional NACs perform reasoning, we use (frozen) pre-trained CLIP [42 backbone to pre-process both texts and images. We select CLIP because it allows to measure zero-shot performance out-of-the-bol (53% validation accuracy, details in Appendix (G). and observe that conditional NACs improve on it (obtaining approximately 64% validation accuracy).
Analysis. We select a trained model and analyze the behaviour of its circuit generator. Figure 6 we visualize the 2D embeddings of the link probability matrix / (see Equation 2 well selection of inferred connectivity graphs for simple (non-compound) natural language statements that capture subset of the semantic phenomena found in Suhr et 48 (see Appendix G To obtain the former, we flatten the lower-triangular part of the link probability matrix P.y into a vector. We first reduce this di0k-dimeniionall vector to 64 dimensions via PCA B1] Subsequently, TSNE [1 is used to embed the 64 dimensional vector in 2 dimensions. To obtain the plots of module connectivity graphs, we use a force-based graph visualization subroutine in the network analysis library NetworkX (22). wherein the attractive force between two modules is inversely proportional to the distance of their respective signatures. We draw an edge between modules they are more than 50% likely to be connected with each other, and the module colors are consistent over all shown graphs. Key observations:
(a) there diversity in the generated graph the conditional setting, even though we do not use regularizing prior like we did in the image classification task; (b) different sentence types are assigned to different graph structures. Further, the generated graphs share similar structure with two large cliques, with nodes that bridge them.
"hhre purple laptop
There door.'
many
"hhre boats"
"The's arc four balls"
Type of Linguistic Reasoning
cheetah."
Eoewns (lirgari
Ensuenialónn merrill@whha Aductiial
Spatbl Raladonentp
TSNE of Conditionally Generated
Herd Carinallty
Link Probability Matrices Py
in frant above dog
Figure Link Probability Matrix Embeddings. This figure shows a TSNE (1 emhedding of the generated graphs provided different kinds of text input. We color-code text inputs based on the type of reasoning that it requires; here show hard cardinality which includes precise number of objects, seft cardinality which includes words like "many" or "few". existential which only indicates that an object is present, and spatial relationships which use prepositional phrase denote a physical relationship.
see by clustering that the low-dimensional embedding space of the generated graph captures aspects of the reasoning task. We also note that all Iearned graphs are similarly structured with two large cliques, but the specific arrangement varies widely.
In conclusion, have shown framework in Neural Attentive Circuits (NACs) for constructing a general-purposs purposs modular neural architecture that can jointly learn meaningful module parame- terizations and connectivity graphs. We have demonstrated empirically that this modular approach yields improvements in fhw-shot adaptation, computational efficiency, and OOD robustness.
10

---

Acknowledgements. This work was partially conducted while Nasim Rahaman was interning at Meta AI. Chris Pal and Martin Weiss thank NSERC for support under the COHESA program and IVADO for support under their AI, Biodiversity and Climate Change Thematic Program. Nasim Rahaman and Bernhard Schölkopf was supported by the German Federal Ministry of Education and Research (BMBF): Tübingen AT Center, FK7: 01SS88339 and by the Machine Learning Cluster of Excellence, EXC number 2064/1 Project number 390727645.
References
[1] D. Achlioptas,, F. McSherry, and B. Schölkopf. Sampling techniques for kernel methods. In T. G. Dietterich, S. Becker, and Z. Ghahramani, editors, Advances in Neural Information Processing Systems 14, pages 335-342, Cambridge, MA, USA, 2002. MIT Press.
[2] Ferran Alet, Erica Weng, Tomás Lozano-Péree and Leslie Pack Kaelbling, Neural relational inference with fast modular meta-learning. H. Wallach, Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. /ttp neurips pdf
[3] Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Deep compositional question answering with neural module networks. CoRR, abs/1511.022999 2015. URL /tttp://arxiv org/abs 01511.02799
[4] Dzmitry Bahdanau, Shikhar Murty, Michael Noukhovitch, Thien Huu Nguyen, Harm de Vries, and Aaron C. Courville. Systematic generalization: What is required and can it be learned? CoRR, abs/181 1. 12889, 2018. /ttp://///arx http://wwwi.r..rraaa////// abs 1811. 12889
[5] Dzmitry Bahdanau, Harm de Vries, Timothy O' O'Donnell, Shikhar Murty, Philippe Beaudoin, Yoshua Bengio, and Aaron C. Courville. CLOSURE: assessing systematic generalization of CLEVR models. CoRR, abs/1912.057783 2019. URL htttt://arxiii org/abs/ /aabs//991..... 05783
[6] Albert-László Barabási and Réka Albert. Emergence of scaling in random networks. science, 286(5439):509 512, 1999.
[7] Albert-Lászíó Barabási and E. Bonabeau. Scale-free networks. Scientific American, 288 (60-69), 2003. Uttphtt::/ httpt:tttp::w..///////////: nd. hadu:/nnttxorks/PD///oo... http://pp..aaaaeerrreee//ss..aa. 20May03.pdF
[8] Robert G. Bartle. The Elements New York, 1995.
of Integration and Lebesgue Measure. John Wiley Sonn,,
[9] Emmanuel Bengio, Pierre-Lue Bacon, Joelle Pineau, and Doina Precup. Conditional computa- tion in neural networks for faster models. arXiv preprint arXiv:1511.064: 10299,, 2015.
[10] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.
[11] Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Nan Rosemary Ke, Sebastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. A meta-transfer objective for learning to disentangle causal mechanisms. International Conference on Learning Representations, 2020. htt.::ttt////ttti/ httpp:///.ppppeeeeeeeeeeeeeeerra
[12] Yoshua Bengio, Tristan Deleu, Edward J. Hu, Salem Lahlou, Mo Tiwari, and Emmanuel Bengio. Gflownet foundations. CoRR, abs/211 1.09266, 2021. URL http:://axxxi........a www.:/rgg/bbb////ee (92266
[13] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers, 2020. URL https //arxiv. org/ atsg//bb/2000.0.. 12872
[14] Joao Carreira, Skanda Koppula, Daniel Zoran, Adria Recasens, Catalin Ionescu, Olivier Henaff, Evan Shelhamer, Relja Arandjelovic, Matt Botvinick, Oriol Vinyals, Karen Simonyan, Andrew Zisserman, and Andrew Jaegle. Hierarchical perceiver, 2022. URL, htttp::////// :// ://arxiv...i.. wwwsxiiooorg org/
abs/ abs/2202. 2202. 10890
11

---

[15] Anne Condon and Richard M. Karp. Algorithms for graph partitioning on the planted partition model. 18:1 16-140, 1999.
[16] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pages 248-255. leee, 2009.
[17] Erdös and A. Rényi. On random graphs i. Publicationes Mathematicae Debrecen, 6:290, 1959.
[18] Paul Erdös, Alfréd Rényi, al. On the evolution of random graphs. Publ. Math. Inst. Hung. Acad. Sci, 5(1):17-60, 1960.
[19] Daniel Glasscock. What is a graphon? 2016. doi: 10.48550/ARXIV..111.007 URL
nttpps://arsiv.org/aaaa
00718
[20] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. CoRR, abs/1909.10893, 2019.
URL httpp:///aaxiivcoom org abs/ 1909. 10893
[21] Nitish Gupta, Kevin Lin, Dan Roth, Sameer Singh, and Matt Gardner. Neural module networks for reasoning over text. CoRR, abs/1911.00971, 2019. http::///aax haiggg/bbs/19111
04971
[22] Aric Hagberg, Pieter Swart, and Daniel S Chult. Exploring network structure, dynamics, and function using networkx. Technical report, Los Alamos National (LANL,,,), Los Alamos, NM (United States), 2008
[23] Curtis Hawthorne, Andrew Jaegle, Cätälina Cangea, Sebastian Borgeaud, Charlic Nash, Mateusz Malinowski, Sander Dieleman, Oriol Vinyals, Matthew Botvinick, Ian Simon, Hannah Sheahan, Neil Zeghidour, Jean-Baptiss Alayrac, João Carreira, and Jesse Engel. General-purpose, long- context autoregressive modeling with perceiver ar, 2022. URL httts:://arxxiiii org/ abs/ 2202.07765
[24] Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, Dawn Song, Jacob Steinhardt, and Justin Gilmer. The many faces of robustness: A critical analysis of out-of-distribttion generalization. CoRR, abs/2006.16244 2020. URL https:::///. htttpps://aaxi...o wttprvvv.rrg/ba/////0//// 16241
[25] Paul Holland, Kathryn Blackmond Laskey, and Samuel Leinhardt. Stochastic blockmodels: First steps. Social networks, 10): 1983.
[26] Ronghang Hu, Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Kate Saenko. Learning to reason: End-to-end module networks for visual question answering. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), 2017.
[27] Drew A Hudson and Christopher D Manning, Compositional attention networks for machine reasoning. In International Conference on Learning Representations, 2018.
[28] Andrew Jaegle, Sebastian Borgeaud, -aatistae Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Hénaff, Matthew M. Botvinick, Andrew Zisserman, Oriol Vinyals, and João Carreira. Perceiver io: A general architecture for structured inputs and outputs, 2021. URLhttps: http::ttp///// https://arxiv. www.ooggg/a........ 2107.14795 14795
[29] Andrew Jaegle. Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and João Carreira. Perceiver: General perception with iterative attention. CoRR, abs/2103.03206, 2021. URL
[30] Justin Johnson, Bharath Hariharan, Laurens van der Maaten, Li Fei-Fei, C. Lawrence Zitnick, and Ross B. Girshick. CLEVR: A diagnostic dataset for compositional language and elementary visual reasoning. CoRR, htd//162200890, 2016. http:/t////axi org/abs/ 1612. 06890
[31] I.T. Jolliffe. Principal Component Analysis. Springer Verlag, 1986.
12

---

[32] Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R. Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer. CoRR, abs/1810.000255, 2018. httt://///axx trrg/bs//8811 00825
[33] Seung Hoon Lee, Seunghyun Lee, and Byung Cheol Song. Vision transformer for small-size datasets, 2021. URL https://arxi:: org abs wttp:////.a......////22 13492
[34] Zehui Lin, Pengfei Liu, Luyao Huang, Junkun Chen, Xipeng Qiu, and Xuanjing Huang. Dropattention: A regularization method for fully-connected self-attention networks. CoRR, 11065
abs/1907. 1065, 2019, URL http://axxxvv...gabba/
[35] Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. arXiv preprint arX1v:2201.035545 2022.
[36] Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. CoRR, abs/1611.00712 2016. http://aarrxxi oxg/abs/1611.000
[37] Jiayuan Mao, Chuang Gan, Pushmeet Kohli, Joshua B. Tenenbaum, and Jiajun Wu. The neuro- symbolic concept learner: Interpreting scenes, words, and sentences from natural supervision. CoRR, abs/1904. 12584, 2019. UtL.
12584
[38] Atsushi Miyauchi and Yasushi Kawase. Ring of cliques network. 2016. doi: 10.137 I/journal. none.0147805.o URL http:::/plos...... www.lo...gghare...o com/ httpa/are...comrrrrrrcc/////geee _iin__inn____ of
[39] Nikita Nangia and Samuel R. Bowman. Listops: A diagnostic dataset for latent tree learning, 2018. URL httpt:ttt////iiiii orgooogg//ooo.o abs/1804 10tp//000000008.
[40] Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron C. Courville. Film: Visual reasoning with a general conditioning layer. CoRR, abs/1709.0778 2017. URL
http:::/aaxvv.orr///a/. abs/ 170. 07871
[41] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of Causal Inference: Foundations and Learning Algorithms. The MIT Press, 2017. ISBN 0262037319.
[42] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. CoRR, abs/2103.00020 2021. URL https://arxi:: org/ wwg.ab//////oooooooo..ggoo.aaeet 2103 .00200
[43] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang. Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. CoRR, abs/1910 10683, 2019.
10683
[44] Nasim Rahaman, Anirudh Goyal, Muhammad Waleed Gondal, Manuel Wuthrich, Stefan Bauer, Yash Sharma, Yoshua Bengio, and Bernhard Schölkopf. S2rms: Spatially structured recurrent modules. CoRR, abs/2007.0657 06533, 2020. URL /Rttsss:://xxx org/ abaaaba/////// 06533
[45] Nasim Rahaman, Muhammad Walced Gondal, Shruti Joshi, Peter V. Gehler, Yoshua Bengio, Francesco Locatello, and Bernhard Schölkopf. Dynamic inference with neural interpreters. CoRR, abs/2110.06399,, 2021. URL https:://// //ta///axviiv otgg/a//////aaraa 06399
[46] Bernhard Schölkopf. Causality for machine learning. CoRR, abs/1911.10500 2019. URL
http://arxiv.orge 19t./////9911..11000000 10500
[47] Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 109(5):612-634 2021.
[48] Alane Suhr, Stephanie Zhou, Iris Zhang, Huajun Bai, and Yoav Artzi. A corpus for reasoning about natural language grounded photographs. CoRR, abs/181 11,00991 2018. hRLp
//arxiv org.abs//////eoooooooo.oooee hbss/bb//1811...../1111111111/11 00491
13

---

[49] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representations, 2021. URL
https///peaneerieee.ee///orr wow.eorruuuudeeee-e-rre.rre
[50] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficiem image transformers & distillation through attention. In International Conference on Machine Learning, pages 10347-10357. PMLR, 2021.
[51] Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of Machine Learning Research, 9:2579-2605, 2008. URL http::/wwww jmlr wtg/papers/v9/
wandermaatenO8a html
[52] Kan Xu, Hamsa Bastani, and Osbert Bastani. Robust generalization of quadratic neural networks via function identification, 2022. URL hutpss//openrerrie..tttffooo Xx4MNjSmQQ9
14


Reviews:
[{"review_id": "xZ04J7X2nvf", "reviewer": "Reviewer_g8fa", "summary": "This paper presents a novel general purpose modular neural architecture; Neural Attentive Circuits. In NAC, information is routed through between a set of modules that perform conditioned computations. Information routing and computation conditioning is specified by the circuit generator and information is processed by the circuit executor. The paper proposes different instantiations for the circuit generator depending on the use case. NAC is compared to PerceiverIO, a general purpose architecture used as a basis for developing NAC, on image classification, few-shot adaptation and OOD generalization. Although NAC does not reach PIO’s performance level on the ImageNet validation set, it performs better on OOD generalization and Few-shot adaptation. Furthermore, pruning several modules significantly decreases the model’s computational cost without a significant loss in accuracy. Ablations of the model highlight the importance of design choices such the conditioned computations, modularity and sparsity priors. Finally, the usefulness of data conditioning in the circuit generator is demonstrated on a visual reasoning task where similar instructions generated similar circuit connectivity patterns.", "questions": "The paper lacks in clarity. It would be beneficial to explain in more detail the model design and experiments following the issues mentioned in weaknesses.\n\nSince signatures and codes are shared across layers, the connectivity and modulation are similar in across layers. It would be interesting to vary dynamic signatures and codes while maintaining the weights shared across layers. Was model design explored previously ? If so, how did it perform compared to the current design ?", "limitations": "The limitations are addressed in the supplementary material and no negative societal impact was discussed in the paper.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n- The paper proposes a novel method for modularizing neural networks.\n- This method is has a flexible and efficient architecture design; module connectivity and modulation can be reduce the computational cost in the model and could potentially be incorporated in any model to replace MLPs.\n- The model’s signature initialization can be used to embed specific priors in the model.\n- A large percentage of modules can be pruned without significantly decreasing performance.\n- Module connectivity provides insight on how the model solves tasks.\n- The architecture design improves learning in low data regimes, OOD robustness and few-shot adaptation on two benchmarks.\n\nWeaknesses\n\n- While the architecture is claimed to be general purpose (applicable to arbitrary data types), its usefulness is demonstrated only on visual data.\n- As the authors mentioned in the limitations, the method is not applied to large scale data compared to Perceiver and PerceiverIO. Furthermore, performance on the ImageNet validation set is significantly lower compared to PIO.\n- If I understood section 2 correctly, the state $\\theta^{i}_u$ of a module is only used to compute the query in layer $l+1$. If this is the case, the state does not represent a state of the model (as much as the weights $W$ and $W_c$). It would be less confusing to give $\\theta^{i}_u$ a different name.\n- Looking at the hyper-parameters in table 3, the number of layers and number of layers that do not share weights both equal 8 which means the module weights $W$ are not shared across layers. If this is the case, the sharing of signatures and codes across layers is an unnecessary constraint on the model (since modules in different layers perform different computations). Furthermore, this means the modules in different layers are completely different and the notion of a state at different layers does not make sense for the model.\n- Considering the model’s similarity with PerceiverIO, the input is a MxC array of data that is transformed into a NxD latent array using the processor. PIO performs transformer-based attention on the N vectors before performing transforming them with an MLP. In NAC, SKMDPA is performed over U modules. It’s unclear how vectors of the latent array are channeled through modules. Can several vectors be processed by the same module ? And can several modules process the same vector ? It would be better to include the latent array dimension N in the description of the model and to clarify in detail the changes made to PIO to arrive at NAC.\n- What is the architecture of the circuit generator ?\n- How are CLIP text embeddings fed to the circuit generator ? How are all embeddings fed to PIO in the comparison experiment ? If both PIO and NAC achieve 64% accuracy on NLVR2 then NAC does not improve over PIO in this task.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "dNqMXAIUDe", "reviewer": "Reviewer_myJN", "summary": "This work proposes Neural Attentive Circuits (NAC) which learns module parameterization and the sparse connectivity between modules trained jointly End-to-end. The architecture is based on the Perciever-IO arhictecture which serves as a baseline for their ablation experiments. They introduce a SKMDPA which is a sparse attention mechanism utilizing with the Concrete distribution (Gumbel Softmax Distribution) and ModFC and ModFFN which conditions their compution on an input code. These modules control the attention connectivity and their computation based on the output of the Circuit Generator which can either be input-conditional or unconditional. They demonstrate the input-uncondition version of their work on a few-shot image classification tasks on the Cifar and CUB datasets and achieve better results than their baseline Perciever network.  They show with ablation on tiny imagenet that simply adding their modules and learnable sparse connectivity does not significantly improve in distirbution results, but greatly improve out of distribution results. They further demonstrate adding different graph-connectivity priors to the learned connectivity also significantly improves performance. They experimented with the sample-conditioned circuit design generation trained on the Natural Language for Visual Reasoning for Real using the CLIP encoder to preprocess both text and images. The circuit generator is conditioned only on the input text and only the image is given to the network executor. They visualize a TSNE of the connectivity graphs output by the circuit generator and show that their is a strong grouping in graph connectivity based on type of question given.", "questions": "Can you compare your model in the same setting as the Perciever-IO paper? Do you have metrics for comparison on language tasks? It would make it more convincing to compare this model as a general perpose architecture if there comparisons in more modalities.\n\nIt is likely useful to compare this work to other works for sparse attention and dynamic routing to increase inference efficiency as well as NAS works which also use the concrete distribution and other hard selection mechanisms for architecture search. [1], [2], [3].\n\nThe uncondition circuit generator work is significantly similar to many works in NAS and sparsification and this work would benefit greatly from comparison.\n\nThe Conditional Circuit Generator is quite novel especially if applicable to multimodal data, but the performance and properties of the network could be explored in much more detail.\n\nCould you elaborate about the ModFNN properties?\n\"unlike in transformers, each copy is conditioned by a unique learnable code vector and therefore performs a different computation. Consequently, though the computation performed in each module is different, the total number of parameters does not noticeably increase with the number of modules\"\nI'm not sure I understand. Are they unique and increases parameters or are they shared? Some of it is elaborated on in section 2.2, but it still not entirely clear if they are unique to each layer or module.\n\nCan accuracy and performance be discussed in section 4.2 with baselines? It would allow stronger conclusions be ot made about the results.\n\nFor \"We observe that the Neural Attentive Circuit (NAC) is much more robust than Perceiver IO to sparsification at inference time\" do you have results with PercieverIO sparsification baselines?\n\nFor Figure 5b, do you have a baseline for the compute cost and inference speed of Perciever IO?\n\nMay be useful to compare to [5]. The results with the meta-network inference such as in section 6.3 is somewhat similar to the work with an Unconditional Circuit Generator.\n\nUnrelated to the paper quality, but could this architecture be effective for continual learning? Some methods in continual learning also try to dynamically change the network to partition off tasks to prevent catastrophic forgetting.\n\n[1] Cai, Shaofeng, Yao Shu, and Wei Wang. \"Dynamic routing networks.\" Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. 2021.\n[2] Xie, Sirui, et al. \"SNAS: stochastic neural architecture search.\" International Conference on Learning Representations. 2018.\n[3] Correia, Gonçalo M., Vlad Niculae, and André FT Martins. \"Adaptively Sparse Transformers.\" Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP). 2019.\n[4] Yoon, Jaehong, et al. \"Lifelong Learning with Dynamically Expandable Networks.\" International Conference on Learning Representations. 2018.\n[5] Shaw, Albert, et al. \"Meta architecture search.\" Advances in Neural Information Processing Systems 32 (2019).", "limitations": "I'm not sure the authors should claim this is a general-purpose modular neural network architecture without experiments in multiple modalities.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Proposes a novel architecture for jointly learning the sparse connectivity between modules and the model parameters.\n- Ran ablation which demonstrated their improvements with learned non sample-conditioned sparse connectivity can improve out of distribution accuracy on tiny-imagenet and further graph priors improve in and out of distribution accuracy.\n- Demonstrated improvements over Perciever-IO on few-shot adaptation.\n- Demonstrated with a visualization that the architecture can learn different paths for different types of problems on the Natural Language for Visual Reasoning for Real datset.\n- Demonstrated that their method can be used to sparsify the architecture and achieve the same performance as Perciever-IO with 80% of circuits dropped. They further showed that this leads to a real inference speedup.\n\nWeaknesses:\n- Does not demonstrate the architecture as a General purpose architecture after introducing it as such. The results with the text modality use clip embeddings.\n- Weak empirical results. No strong baselines for the non sample dependent results. Very similar to architecture search and sparsification, but no reference. No comparisons on the same tasks as the original Perciever-IO paper.\n- The sparsification results seem strong, but still there aren't comparisons to other sparsification methods or Neural Architecture search methods.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "TGW00zi0_xv", "reviewer": "Reviewer_p9jH", "summary": "The paper presents a novel neural architecture (NAC) which learns to jointly predict the optimal (sparse) connectivity (via signatures) and parameterization (via codes) of its general purpose modules (ModFC/FFN). The architecture appears parameter- and computation-efficient, compares favorably against Perceiver IO in terms of few-shot adaptation (CUB/CIFAR) and OOD generalization (Tiny-ImageNet), and meaningfully generates similar connectivity patterns when conditioned on similar input samples (NLVR2).", "questions": "- All experiments seem to use some form of preprocessing (convolutional or CLIP) which in a way renders the general purpose claim of the paper questionable. Does the proposed architecture not work with raw inputs (i.e. byte arrays) like Perceiver IO?\n- In Table 1, is (Perceiver IO + ModFC) equivalent to (NAC - sparse graph prior) (i.e. without minimizing Eq 9)? Shouldn’t a dense connectivity perform better under ID validation due to less constraints? Mechanistically, how does the spare connectivity in NACs help achieve better OOD generalization?\n", "limitations": "The authors have adequately addressed the limitations of their work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n+ (Originality, Major) The paper’s novel architectural features (ModFC, sparse graph priors) compared to Perceiver IO, and performance improvements (linear scaling w.r.t. input size, adaptive computation with large speed boost) compared to Neural Interpreters, are impressive.\n+ (Clarity) The writing is clear and easy to follow. The figures are well-made and very helpful.\n\nWeaknesses\n- (Quality) The empirical validation of the proposed architecture is insufficient in my opinion. To claim that NAC is a stronger general purpose architecture, the scope of the comparisons (now only few-shot adaptation and OOD generalization in CV datasets) and the number of baselines (now only Perceiver IO) both need to be reasonably improved. E.g. Perceiver IO was evaluated on 6 different modalities (Table 5 in their paper) and against non-general purpose baselines within each modality.\n- (Significance)\nDue to the insufficient empirical validation, it’s hard to conclude that the proposed architecture is significant and will inspire or lead to future breakthroughs.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept