summaries: The authors propose a single stage pruning method (DAM) that jointly prunes and refines weights during training. The method uses a monotonically increasing gate function for the neurons/channels in each layer with one trainable parameter. The gate function only discriminates neurons based on the position of them in the layer. The proposed method achieves reasonable results without any fine-tuning.

summaries: This paper proposes a novel structured pruning method called DiscriminAtive Masking (DAM) for neuron-level sparsity. A relu-tanh gating function is proposed to apply on the neuron activations, and gradually zeroing out the neurons at lower order by introducing a L0 norm. The proposed method is a single stage method, and does not require further after-prune finetuning. Extensive experiments demonstrate the effectiveness of the proposed method

summaries: This paper proposes a gradual structured pruning method, DAM, that can achieve good performance with various applications. The paper conducts extensive experiments to support the effectiveness of the proposed methods. The method part is reasonable and makes sense.

summaries: This paper presents a simple and novel in-training model pruning algorithm to learn sparse feature representation with reduced computational load w/o compromising accuracy. To do this, authors introduce a single learnable parameter to simplify the optimization involved in enforcing L0 sparsity. Authors further provide theoretic analysis on how reducing the single learnable parameter would enforce L0 sparsity. Finally, authors demonstrate the remarkable performance over various applications comparing with the established strong baselines.

abstractText: A central goal in deep learning is to learn compact representations of features at every layer of a neural network, which is useful for both unsupervised representation learning and structured network pruning. While there is a growing body of work in structured pruning, current state-of-the-art methods suffer from two key limitations: (i) instability during training, and (ii) need for an additional step of fine-tuning, which is resource-intensive. At the core of these limitations is the lack of a systematic approach that jointly prunes and refines weights during training in a single stage, and does not require any fine-tuning upon convergence to achieve state-of-the-art performance. We present a novel single-stage structured pruning method termed DiscriminAtive Masking (DAM). The key intuition behind DAM is to discriminatively prefer some of the neurons to be refined during the training process, while gradually masking out other neurons. We show that our proposed DAM approach has remarkably good performance over a diverse range of applications in representation learning and structured pruning, including dimensionality reduction, recommendation system, graph representation learning, and structured pruning for image classification. We also theoretically show that the learning objective of DAM is directly related to minimizing the L0 norm of the masking layer. All of our codes and datasets are available https://github.com/jayroxis/ dam-pytorch.

authors:
Sanjeev Arora
Rong Ge
Behnam Neyshabur
Yi Zhang

title: Stronger generalization bounds for deep nets via a compression approach

venue: In International Conference on Machine Learning,

authors:
Haoli Bai
Jiaxiang Wu
Irwin King
Michael Lyu

title: Few shot network compression via cross distillation

venue: In AAAI,

authors:
Yoshua Bengio
Aaron Courville
Pascal Vincent

title: Representation learning: A review and new perspectives

venue: IEEE transactions on pattern analysis and machine intelligence,

authors:
Davis Blalock
Jose Javier Gonzalez Ortiz
Jonathan Frankle
John Guttag

title: What is the state of neural network pruning

venue: arXiv preprint arXiv:2003.03033,

authors:
Jie Bu
Anuj Karpatne

title: Quadratic Residual Networks: A New Class of Neural Networks for Solving Forward and Inverse Problems in Physics Involving PDEs, pages 675–683

venue: URL https://epubs.siam.org/doi/abs/10

authors:
Tim R Davidson
Luca Falorsi
Nicola De Cao
Thomas Kipf
Jakub M Tomczak

title: Hyperspherical variational auto-encoders

venue: arXiv preprint arXiv:1804.00891,

authors:
Misha Denil
Babak Shakibi
Laurent Dinh
Marc’Aurelio Ranzato
Nando de Freitas

title: Predicting parameters in deep learning

venue: In Proceedings of the 26th International Conference on Neural Information Processing Systems-Volume

authors:
Xinhan Di
Pengqian Yu
Rui Bu
Mingchao Sun

title: Mutual information maximization in graph neural networks

venue: In 2020 International Joint Conference on Neural Networks (IJCNN),

authors:
Utku Evci
Trevor Gale
Jacob Menick
Pablo Samuel Castro
Erich Elsen

title: Rigging the lottery: Making all tickets winners

venue: Proceedings of the 37th International Conference on Machine Learning,

authors:
Jonathan Frankle
Michael Carbin

title: The lottery ticket hypothesis: Finding sparse, trainable neural networks

venue: In ICLR,

authors:
Jonathan Frankle
Gintare Karolina Dziugaite
Daniel M Roy
Michael Carbin

title: Stabilizing the lottery ticket hypothesis

venue: arXiv preprint arXiv:1903.01611,

authors:
Jonathan Frankle
Gintare Karolina Dziugaite
Daniel M. Roy
Michael Carbin

title: Linear mode connectivity and the lottery ticket hypothesis, 2020

authors:
Trevor Gale
Erich Elsen
Sara Hooker

title: The state of sparsity in deep neural networks, 2019

authors:
Ariel Gordon
Elad Eban
Ofir Nachum
Bo Chen
Hao Wu
Tien-Ju Yang
Edward Choi

title: Morphnet: Fast & simple resource-constrained structure learning of deep networks

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Kathrin Grosse
Michael Backes

title: Do winning tickets exist before DNN training?, pages 549–557

venue: URL https://epubs.siam

authors:
Song Han
Jeff Pool
John Tran
William Dally

title: Learning both weights and connections for efficient neural network

venue: Advances in Neural Information Processing Systems,

authors:
I. Higgins
Loïc Matthey
A. Pal
Christopher P. Burgess
Xavier Glorot
M. Botvinick
S. Mohamed
Alexander Lerchner

title: beta-vae: Learning basic visual concepts with a constrained variational framework

authors:
G.E. Hinton
R.R. Salakhutdinov

title: Reducing the dimensionality of data with neural networks

authors:
Diederik P Kingma
Max Welling

title: Auto-encoding variational bayes

venue: arXiv preprint arXiv:1312.6114,

authors:
Thomas N Kipf
Max Welling

title: Semi-supervised classification with graph convolutional networks

venue: arXiv preprint arXiv:1609.02907,

title: Variational graph auto-encoders

venue: arXiv preprint arXiv:1611.07308,

authors:
Simon Kornblith
Mohammad Norouzi
Honglak Lee
Geoffrey Hinton

title: Similarity of neural network representations revisited

authors:
Duong Hoang Le
Binh-Son Hua

title: Network pruning that matters: A case study on retraining variants

venue: In International Conference on Learning Representations,

authors:
Quoc V Le
Jiquan Ngiam
Adam Coates
Ahbik Lahiri
Bobby Prochnow
Andrew Y Ng

title: On optimization methods for deep learning

authors:
Namhoon Lee
Thalaiyasingam Ajanthan
Philip Torr

title: Snip: Single-shot network pruning based on connection sensitivity

authors:
Hao Li
Asim Kadav
Igor Durdanovic
Hanan Samet
Hans Peter Graf

title: Pruning filters for efficient convnets

venue: arXiv preprint arXiv:1608.08710,

authors:
Shiwei Liu
Lu Yin
Decebal Constantin Mocanu
Mykola Pechenizkiy

title: Do we actually need dense over-parameterization? in-time over-parameterization in sparse training, 2021

authors:
Zhuang Liu
Jianguo Li
Zhiqiang Shen
Gao Huang
Shoumeng Yan
Changshui Zhang

title: Learning efficient convolutional networks through network slimming

authors:
Zhuang Liu
Mingjie Sun
Tinghui Zhou
Gao Huang
Trevor Darrell

title: Rethinking the value of network pruning

authors:
Christos Louizos
Max Welling
Diederik P Kingma

title: Learning sparse neural networks through l_0 regularization

authors:
Costas Mavromatis
George Karypis

title: Graph infoclust: Leveraging cluster-level node information for unsupervised graph representation learning

venue: arXiv preprint arXiv:2009.06946,

authors:
Decebal Constantin Mocanu
Elena Mocanu
Peter Stone
Phuong H. Nguyen
Madeleine Gibescu
Antonio Liotta

title: Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science

venue: Nature Communications,

authors:
Bruno A Olshausen
David J Field

title: Emergence of simple-cell receptive field properties by learning a sparse code for natural images

authors:
Shirui Pan
Ruiqi Hu
Guodong Long
Jing Jiang
Lina Yao
Chengqi Zhang

title: Adversarially regularized graph autoencoder for graph embedding

venue: arXiv preprint arXiv:1802.04407,

authors:
Marc’Aurelio Ranzato
Y-Lan Boureau
Yann LeCun

title: Sparse feature learning for deep belief networks

venue: Advances in neural information processing systems,

authors:
Alex Renda
Jonathan Frankle
Michael Carbin

title: Comparing rewinding and fine-tuning in neural network pruning

title: Stochastic complexity and modeling

venue: The annals of statistics,

authors:
Pedro Savarese
Hugo Silva
Michael Maire

title: Winning the lottery with continuous sparsification

authors:
Skyler Seto
Martin T Wells
Wenyu Zhang

title: Halo: Learning to prune neural networks with shrinkage

venue: In Proceedings of the 2021 SIAM International Conference on Data Mining (SDM),

authors:
Nitish Srivastava
Geoffrey Hinton
Alex Krizhevsky
Ilya Sutskever
Ruslan Salakhutdinov

title: Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning

authors:
Ilya Sutskever
James Martens
George Dahl
Geoffrey Hinton

title: On the importance of initialization and momentum in deep learning

venue: In International conference on machine learning,

title: Regression shrinkage and selection via the lasso

venue: Journal of the Royal Statistical Society: Series B (Methodological),

authors:
Rishabh Tiwari
Udbhav Bamba
Arnav Chavan
Deepak Gupta

title: Chipnet: Budgetaware pruning with heaviside continuous approximations

authors:
Pascal Vincent
Hugo Larochelle
Yoshua Bengio
Pierre-Antoine Manzagol

title: Extracting and composing robust features with denoising autoencoders

venue: In Proceedings of the 25th International Conference on Machine Learning,

authors:
Jiayun Wang
Yubei Chen
Rudrasis Chakraborty
Stella X Yu

title: Orthogonal convolutional neural networks

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

authors:
Hong Yang
Shirui Pan
Peng Zhang
Ling Chen
Defu Lian
Chengqi Zhang

title: Binarized attributed network embedding

venue: IEEE International Conference on Data Mining (ICDM),

authors:
Jihun Yun
Peng Zheng
Eunho Yang
Aurelie Lozano
Aleksandr Aravkin

title: Trimming the l1 regularizer: Statistical analysis, optimization, and applications to deep learning

authors:
Muhan Zhang
Yixin Chen

title: Inductive matrix completion based on graph neural networks

authors:
Wenda Zhou
Victor Veitch
Morgane Austern
Ryan P. Adams
Peter Orbanz

title: Compressibility and generalization in large-scale deep learning

authors:
Michael Zhu
Suyog Gupta

title: To prune, or not to prune: exploring the efficacy of pruning for model compression, 2017

authors:
Xiaotian Zhu
Wengang Zhou
Houqiang Li

title: Improving deep neural network sparsity through decorrelation regularization

venue: In IJCAI,

heading: 1 Introduction

text:
A central goal in deep learning is to learn compact (or sparse) representations of features at every layer of a neural network that are useful in a variety of machine learning tasks. For example, in unsupervised representation learning problems [3], there is a long-standing goal to learn lowdimensional embeddings of input features that are capable of reconstructing the original data [44, 35, 24, 33]. Similarly, in supervised learning problems, there is a growing body of work in the area of network pruning [4], where the goal is to reduce the size of modern-day neural networks (that are known to be heavily over-parameterized [7, 51, 22, 45]) so that they can be deployed in resource-constrained environments (e.g., over mobile devices) without compromising on their accuracy. From this unified view of representation learning and network pruning, the generic problem of “learning compact representations” has applications in several machine learning use-cases such as dimensionality reduction, graph representation learning, matrix factorization, and image classification.
*These authors contributed equally to this work.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
A theoretically appealing approach for learning compact representations is to introduce regularization penalties in the learning objective of deep learning that enforce L0 sparsity of the network parameters, θ. However, directly minimizing the L0 norm requires performing a combinatorial search over all possible subsets of weights in θ, which is computationally intractable. In practice, a common approach for enforcing sparsity is to use a continuous approximation of the L0 penalty in the learning objective, e.g., L1-based regularization (or Lasso [42]) and its variants [39, 47]. While such techniques are capable of pruning individual weights and thus reducing storage requirements, they do not offer any direct gains in inference speed since the number of features generated at the hidden layers can still be large even though the network connectivity is sparse [4]. Instead, we are interested in the area of structured network pruning for learning compact representations, where the sparsity is induced at the level of neurons by pruning features (or channels) instead of individual weights.
While there is a growing body of work in structured network pruning [26, 28, 43, 14], the basic structure of most state-of-theart (SOTA) methods in this area (e.g., ChipNet [43] and NetSlim [28]) can be described as training a learnable vector of mask parameters, g ∈ Rn,g = [g1, g2, ..., gn]⊺, which when multiplied with the features extracted at a hidden layer, h ∈ Rn, results in the pruned outputs of this layer, o = g ⊙ h. Sparsity in the mask parameters is generally enforced using different approximations of the L0 norm of g (e.g., use of Lasso in NetSlim and use of “crispness” loss in ChipNet). Despite recent progress in this area, current SOTA in structured pruning suffer from two key limitations. First, since most methods do not explicitly minimize the L0 norm of g during pruning, they often suffer from training instabilities. In particular, most SOTA methods [28, 43] involve thresholding techniques to set small non-zero weights to zero leading to large drops in accuracy during the training process, as
evidenced by our results in this paper. Second, once the pruning process is complete and we have converged at a compact network, most SOTA methods still need an additional step of fine-tuning the network in order to achieve reasonable accuracy. This is not only resource-intensive but there is also an on-going discussion on whether and how we should fine-tune [36] or rewind to initial weights [23] or train from scratch [29], making it difficult to prefer one approach over another.
At the core of these limitations is the lack of a systematic approach for structured network pruning that jointly prunes and refines weights during training in a single stage, and does not require any fine-tuning upon convergence to achieve SOTA performance. Notice that in existing methods for structured pruning, allowing every neuron j to be pruned differently using an independent mask parameter gj only increases the number of learnable (or free) parameters in the learning objective, increasing the complexity of the problem. Instead, we ask the question: “Can we leverage the intrinsic symmetry of neural network initialization to design a pruning mask using the least number of free parameters?”
We present a simple solution to this question by proposing a new single-stage structured pruning method that learns compact representations while training and does not require fine-tuning, termed DiscriminAtive Masking (DAM). The basic idea of DAM is to use a monotonically increasing gate function G for masking every neuron in a layer that only depends on the index j = 1, 2, ..., n (or position) of the neuron in the layer and a scalar parameter β to be learned during training (see Figure 1). At the start of training, the gate function admits non-zero values for all neurons in the layer, allowing all neurons to be unmasked (or active). As the training progresses, the gate function gradually shifts from “left” to “right” as a result of updating β such that upon convergence, only a subset of neurons (on the extreme right) are active while all others are masked out. The key intuition behind DAM is to discriminatively prefer some of the neurons (on the right) to be refined (or re-adapted) during the training process for capturing useful features, while gradually masking out (or pruning) neurons on the left. This preferential pruning of neurons using a very simple gate function helps in regulating the number of features transmitted to the next layer.*
Contributions: We show that our proposed DAM approach has remarkably good performance over various applications, including dimensionality reduction, recommendation system, graph represen-
*This is somewhat analogous to how a physical dam regulates the flow of water by shifting a movable gate.
tation learning, and structured pruning for image classification, achieving SOTA performance for structured pruning. This shows the versatility of DAM in learning compact representations on diverse problems and network choices, in contrast to baseline methods for structured pruning that are only developed and tested for the problem of image classification [28, 43]. Our approach is single-stage, does not require fine-tuning (and hence has lower training time), and does not suffer from training instabilities. We also theoretically show that the learning objective of DAM is directly related to minimizing the L0 norm of the discriminative mask, providing a new differentiable approximation for enforcing L0 sparsity. Our approach is also easy to implement and can be applied to any layer in a network architecture. The simplicity and effectiveness of DAM provides unique insights into the problem of learning compact representations and its relationship with preferential treatment of neurons for pruning and refining, opening novel avenues of systematic research in this rapidly growing area.

heading: 2 Related Work

text:
Learning with Sparsity: There is a long history of methods for enforcing compactness (or sparsity) in learned parameters [33, 42], where it is understood that compact hypotheses can not only result in more efficient data representations but also promote better generalizability [37]. This connection between compactness and generalization has also been recently explored in the context of deep learning [49, 1], motivating our problem of learning compact representations.
Unstructured Pruning: There are several methods that have been developed for unstructured network pruning [30, 10, 12, 38], where the goal is to prune individual weights of the network. A widely used technique in this field is referred to as the Lottery Ticket Hypothesis (LTH) [10], which suggests that a certain subset of weights in a network may be initialized in such a way that they are likely “winning tickets”, and the goal of network pruning is then to simply uncover such winning tickets from the initialization of an over-parameterized network using magnitude-based pruning. Similar to our work, there exists a line of work in the area of dynamic sparse training for unstructured pruning [50, 13, 32, 9, 27] that gradually prunes the model to the target sparsity during training. However, while unstructured pruning methods can show very large gains in terms of pruning ratio, they are not directly useful for learning compact representations, since the pruned weights, even though sparse, may not be arranged in a fashion conducive to the goal of reducing the number of features.
Structured Network Pruning: The goal of structured pruning is to enforce sparsity over the network substructures (e.g., neurons or convolutional channels) that is directly relevant to the goal of learning compact representations. Most structured pruning SOTA methods need a three-stage process [16], i.e., training, pruning and finetuning, in order to get highly compressed models. Early exploration on the L1 norm based filter pruning can be traced back to [26], followed by Network Slimming (NetSlim) [28], which uses a sparsity-inducing regularization based on L1 norm. Following this line of work, a recent development in the area of network pruning is ChipNet [43], which introduces an additional crispness loss to the L1 regularization, and achieves SOTA performance on benchmark tasks. However, network slimming can only work if the underlying architecture has batch normalization layers. It also has training instabilities due to soft-thresholding inherent to Lasso. ChipNet also iterates between soft- and hard-pruning, which is done after the training phase. Further, both these approaches require finetuning or retraining, which can be resource intensive.
Faster pruning schemes have also been explored by previous works. Bai et al [2] proposed a few-shot structured pruning algorithm utilizing filter cross distillation. The SNIP [25] took one step ahead by performing single-stage pruning. However, SNIP falls in the category of unstructured pruning. To the best of our knowledge, no existing structured pruning method has been demonstrated to achieve SOTA performance without finetuning. For example, while the method proposed in [30] can be potentially used as a single-stage structured pruning method with the help of group sparsity, the empirical analysis of the paper only focused on unstructured pruning.
Deep Representation Learning: Leveraging the power of deep neural networks, high-dimensional data can be encoded into a low-dimensional representation [18]. Such process is known as deep representation learning. A good representation extracts useful information by capturing the underlying explanatory factors of the observed input. The existence of noises, spurious patterns and complicated correlations among features make it a challenging problem. Previous work [44] show that an
Autoencoder architecture can be used for denoising images. To disentangle correlated features, VAE [19] and β-VAE [17] encourages the posterior distribution over the generative factors q(z|x) to be closer to the isotropic Gaussian N (0, I) that explicitly promotes disentanglement of the latent generative factors. Driven by the idea of the autoencoder, the deep representation learning has been found remarkably successful in various of applications, e.g., graph representation learning [21, 6, 8, 31, 46, 34] and recommendation system [48].
However, most existing works in representation learning treat the embedding dimension as a hyperparameter, which can be crucial and difficult to choose properly. Small embedding dimensions cannot sufficiently represent the important information in the data that leads to bad representations. On the other hand, large embedding dimensions allow some level of redundancy in the learned representations, even picking up spurious patterns that can degrade model performances.

heading: 3 Proposed Approach

text:
Problem Statement: Let us represent the general architecture of a neural network with l layers as F(x) = Ψl ◦ Ψl−1 ◦ ...Ψ2 ◦ Ψ1(x), where x are the input features to the network, Ψi is the learnable function at layer i that generates hidden features hi = Ψi(hi−1), and f ◦ g represents the composite function f(g(·)). Let us also consider the “masked” version of this network where the features learned at every hidden layer, hi are multiplied with a vector of mask values gi using a mask function Mi(hi) = gi ⊙ hi. The masked neural network can then be described as Fmasked(x) = Ψl ◦Ml−1 ◦Ψl−1 ◦ ...M2 ◦Ψ2 ◦M1 ◦Ψ1(x). The goal of learning compact representations can be framed as minimizing the following learning objective:
min θ Ex∼Dtr
[ L (Fmasked(x)) + λ
l − 1 l−1∑ i=1 ∥gi∥0
] , (1)
where θ are the learnable parameters of the network, Dtr is the training set, L(.) is the loss function, and λ is the trade-off parameter for enforcing L0 sparsity in the mask values gi. Note that this formulation provides a unified view of both representation learning and structured network pruning, which differ in the choice of loss function (typically, reconstruction error for representation learning and cross-entropy for structured pruning), and the number of layers that are pruned (in representation learning, we only aim to prune the bottleneck layer while in structured pruning, we aim to prune all of the layers in the network, see Figure 2).
DiscriminAtive Masking (DAM): Our proposed approach presents a novel perspective of structured pruning by introducing a new strategy of “ordering-and-masking.” To understand the basic idea of this strategy, let us assume that we are given the ideal “ordering” of neurons in a layer that reflect their importance (i.e., neurons with higher order number are more useful for the learning task). Given this ordering, we can perform structured pruning by incrementally “masking” out neurons with lower order (i.e., setting their mask values to 0) before going to neurons with higher order during the training process. In other words, we discriminatively (or preferentially) focus on pruning neurons lower in the order, while retaining and refining the features learned at neurons higher in the order. This approach for masking can be realized using a monotonic gate function G that maps neuron order to mask values, i.e., lower order neurons are assigned smaller mask values than higher order neurons. By dynamically regulating the gate function during the course of training, we can control the percentage of neurons pruned at every training epoch.
Specifically, let us assign every neuron j at layer i with an “order number,” µij . Before we describe how we arrive at these neuron orders, we first present the formulation of our proposed gate function G, whose value at neuron j monotonically increases with the order number of the neuron, µij , as follows:
gij = G(µij , αi, βi) = ReLU[tanh (αi(µij + βi))], (2) where αi is a constant scalar parameter (termed as the steepness parameter) that is not optimized during training, while βi is a ‘learnable’ scalar parameter (termed the offset parameter) that is the only parameter optimized during training to control (or regulate) the process of masking in DAM. Figures 2(c) and (d) show examples of gate function values (in red) at two different values of βi. Note that using a single learnable parameter βi to implement the masking process introduces significant simplicity in the design of the learnable mask, in contrast to SOTA methods in structured
pruning [28, 43] where a different masking parameter is trained for every neuron in the network. We hypothesize this simplicity of the gate function, in conjunction with the preferential pruning property of DAM, to result in remarkable gains in learning compact representations, as evidenced later in our experiments. Further, note that while we can use other monotonically increasing function to implement our gates, we found ReLU-tanh to perform best in practice.
Assigning Neuron Orders µij : We now discuss our approach for assigning neuron order numbers µij . Note that the order numbers only have to be assigned during the initialization step, and do not change during the training process. Further, note that every neuron is initialized with weights that are independent and identically distributed (because of the i.i.d nature of commonly used initialization methods in deep learning [41]). Hence, all neurons are symmetrically invariant to each other at initialization and any neuron has an equally good chance of refining themselves to capture useful features later in the training process as any other neuron. Hence, any random ordering of the neurons in the DAM approach would result in a similar pruning process (we empirically demonstrate the permutation invariance of µij in Section 5). For simplicity, we consider a trivial choice of µij = kj/ni, where k is a constant parameter that determines the domain size of µij , and ni is the total number of neurons at layer i.
How Does Reducing βi Enforce L0 Sparsity? Let us first understand how mask values gi respond to the only learnable parameter βi. Figure 2 (c) shows an example of gi when βi = −2. At this stage, we can see that only a small number of neurons have zero gate values (white). As we reduce βi to −6 in Figure 2 (d), we can see that the gate moves towards the ’right’, resulting in more excessive pruning of neurons. In general, the number of non-zero values of the gate function (and hence, its L0 norm) is directly related to βi as follows:
∥gi∥0 = ni∑ j=1 1(gij > 0) = ni∑ j=1 (1− 1(µij ≤ βi)) = ⌈ni(1 + βi/k)⌉, for βi ≥ −k, (3)
where 1 indicates the identity operator and ⌈·⌉ is the ceiling operator. In order to back-propagate gradients, we use a continuous approximation of Equation (3) (by dropping the ceiling operator) for regularization as ∥gi∥0 ≈ ni(1 + βi/k). Hence, minimizing L0 norm of gi is directly proportional to minimizing βi with a scaling factor of ni/k.
Learning Objective of DAM: While minimizing βi with scaling factors proportional to ni can explicitly minimize L0 norm of the gate values, we found in practice that layers with smaller number of neurons ni still require adequate pruning of features for learning compact representations proportional to layers with large ni. Hence, in our learning objective, we drop the scaling factor and directly minimize the sum of βi across all layers in the following objective function.
min θ,β Ex∼Dtr
[ L (Fmasked(x)) + λ
l − 1 l−1∑ i=1 βi
] , (4)
where, β is the set of all βi’s in the network, {β1, β2, . . . , βl−1}. Theoretical Analysis of DAM: Analysis of the dynamics of the gate functions and additional specifications for choosing αi and the gate function are provided in Appendix A.
Parameter Specification and Implementation Details: In all of our experiments, we used αi = 1 and k = 5. In our structured network pruning experiments, we used a cold-start of 20 epochs (i.e., the βi’s were frozen for the duration of cold-start), so as to allow the leftmost neurons to undergo
some epochs of refinement before beginning the pruning process. We also set the initial value of βi to 1, which can also be thought of as another form of cold-starting (since pruning of a layer only starts when βi becomes less than zero).

heading: 4 Experiments on Representation Learning Problems

heading: 4.1 Results on Dimensionality Reduction (DR) Problems

text:
Problem Setup: We evaluate the effectiveness of DAM in recovering the embeddings of synthetically generated dimensionality reduction problems. The general form of the problem is expressed by (5). Suppose Ω ∈ Rr and all of its r dimensions are independent from each other. A d-dimensional data X can be expressed using a transformation Ψ : Rr −→ Rd defined on Hilbert space, and X = Ψ(Ω), d > r. Let M be the learnable mask, we formulated the dimensionality reduction problem as:
min Θ
{ ∥FD ◦M ◦ FE(X)−X∥2F + λβ } , Θ = {θE , θD, β} (5)
The FE ,M,FD are trained end-to-end using gradient descent. In our experiments, Ω is sampled from a isotropic normal distribution, i.e., Ω ∼ Nr(0, Ir). We let Ψ to use the same structure as the decoder FD with randomly generated parameters θΨ, i.e., Ψ = FD|θD=θΨ , so that the minimum number of dimensions needed in the encoded representation for the decoder to reconstruct X is r.
Linear DR: We test DAM for removing linear correlations in a given rank-deficient matrix. In this case, Ψ is a matrix representing a linear projection from Rr to Rd. We use two real matrices as the encoder and decoder and train the DAM to find the full-rank representation of X .
Nonlinear DR: To test the ability of DAM for disentangling nonlinearly correlated dimensions, we present two cases: (i) Ψ is a polynomial kernel of degree 2 (we use a QRes layer from [5]); (ii) Ψ is a nonlinear transformation expressed by a neural network. We use a deep neural network as the encoder with sufficient complexity.
Observation: Figure 3 shows the convergence of bottleneck dimensions over the training epochs for different synthetic data with varying sizes of underlying factors (r). The curves and shades show the mean and standard deviation of five different runs. We can see that DAM consistently uncovers the exact dimension of the underlying factors r. We provides details of in-depth sensitivity analysis of DAM results based on its hyper-parameters (e.g., learning rate and λ) in Appendix C.
Theoretical Analysis of DAM for linear DR: We theoretically show that DAM converges to the optimal solution for the linear DR case (see Appendix B for details).

heading: 4.2 Recommendation System Results

text:
Problem Setup: We consider the state-of-the-art method for recommendation system problems (IGMC [48]), which transforms the rating matrix completion task to a link prediction problem, where the ratings are interpreted as links between users and items.
In particular, IGMC generates 128-dimensional embeddings to represent the enclosing subgraph and further avoids overfitting using a dropout layer with a dropout rate of 0.5. We replaced the dropout
layer with a DAM layer in the IGMC to reduce the dimension of the learned representations. We train our IGMC-DAM model for 100 epochs under the same training configurations.
Observation: Table 1 shows that DAM successfully reduces the dimensions of the learned representations without any increase in errors, demonstrating that DAM is able to boost the performance of IGMC by learning compact representations over the target graph.

heading: 4.3 Graph Representation Learning Results

text:
Problem Setup: Following the previous experiment, we further explore the effectiveness of DAM in learning compact graph representations. Our goal is to learn low-dimensional embeddings for each node in a graph that captures the structure of interaction among nodes. A simple graph autoencoder, e.g., GAE [21], uses a graph convolutional network (GCN) [20] as an encoder and an inner product as a decoder. The encoder calculates the embedding matrix Z from the node feature matrix X with the adjacency matrix A, and the decoder reconstructs a adjacency matrix Â such that: Â = σ(ZZ⊤), with Z = GCN(X,A). The reconstruction loss, BCE(A, Â) is backpropagated to train the GAE model. To reduce the dimension of the learned representation, we add a DAM layer in a GAE after the encoder (GAE-DAM).
Observation: Figure 4 shows that DAM based GAE method is able to learn meaningful structural information in compact latent embeddings. DAM improves the link prediction performance of the simple GAE for Cora, CiteSeer and PubMed dataset by learning compact representations.

heading: 4.4 Representation Learning Results for MNIST Dataset

text:
Problem Setup: We further evaluate DAM on the problem of dimensionality reduction using simple auto-encoders on the MNIST dataset, and also compare the effectiveness of DAM (that directly enforces L0 sparsity) with an ablation of our approach that instead minimizes the L1 norm (similar to Lasso). We vary the trade-off parameter λ to compare the sparsity of the learned representations between DAM and L1-norm based ablation method.
Observation: Figure 5 shows that DAM is able to achieve lower reconstruction error as well as higher F1 scores over the same bottleneck dimensions than L1 based method. In Figure 5 (c), we can see that DAM shows a near-linear descending trend of bottleneck dimensions as λ increases (further evidence presented in Appendix E1). In contrast, the L1 based method has a saturating effect at large λ values since the weights come close to zero but require some thresholding to be pruned. This shows that DAM is amenable to learning highly compact representations in contrast to L1 based methods.

heading: 5 Experiments on Structured Network Pruning

text:
Evaluation Setup Here we demonstrate the effectiveness of DAM on the task of structured network pruning. We compared the performance of our proposed DAM with two SOTA structured pruning algorithms, Network Slimming (Net-Slim) [28] and ChipNet [43]. Net-Slim is well-established SOTA that is widely used as a strong baseline, while ChipNet is the latest SOTA representing a very recent development in the field. ChipNet and Net-Slim are both pretrained for 200 epochs, while the DAM performs single-stage training and pruning in the same number of epochs. Since, our DAM approach does not require additional fine-tuning, we imposed a limited budget of 50 epochs for fine-tuning the SOTA methods. Note that ChipNet also has a pruning stage which involves 20 additional epochs. We evaluate the performance of these methods on the PreResnet-164 architecture on benchmark computer vision datasets, CIFAR-10 and CIFAR-100. Additional evaluation setup details are in Appendix E2.
Performance and Running time comparison: Figures 6a and 6b demonstrate the performance of different network pruning methods for various pruning ratios. We observe that for both datasets, DAM is able to outperform Net-Slim by a significant margin specially when the models are sparse. Additionally, DAM is able to achieve similar or slightly better performance than ChipNet with no additional fine-tuning. We also compare the total running time for pruning across the different models (divided into the three categories) in Figure 6c. We can see that the training time of the three methods are almost the same, with ChipNet being slightly small. However, the running time for the pruning and fine-tuning stages for DAM are both zero. Net-Slim also does not involve additional pruning after training. However, ChipNet involves 20 epochs of pruning which is significant and is almost comparable to its pretraining time. Finally, comparing the total running time taken by each of the structured pruning methods, we observe that DAM is significantly faster than the current SOTA counterparts owing to its single-stage pruning approach.
Stability Analysis: We further analyse the stability of our proposed DAM approach for network pruning through visualization of the training dynamics. We observe that the training cross-entropy (CE) loss and the validation CE loss are very similar to what we expect from a Preresnet model trained on benchmark vision datasets with learning rate changes at 100 and 150 epochs, respectively. We further notice that the total training loss for the DAM is also very stable and does not show any sudden variations, leading to steady convergence. Finally, from the convergence plot for E[βi] (green-dashed line), we can see that our pruning does not involve any sudden steps and happens in a continuous and steady rate throughout the training stage.
The training dynamics of the ChipNet (blue shaded region in the middle panel) is exactly the same as training vanilla vision models since it does not involve sparsity constraints. However, as we move to the pruning stage, we notice a sharp rise in both training and validation losses (white region). This is due to the use of a different optimizer AdamW as opposed to SGD which was used in the training stage. Further, a very interesting phenomenon happens in the fine-tuning stage (red shaded region) where the training and validation losses increases for the first 25 epochs and then takes a sharp decrease once the learning rate is adjusted. This suggests that after the pruning stage, finetuning the model with large learning rate forces the model to make large updates in the initial epochs (akin to training from scratch), which further drives the model out of the local minima. Once the learning rate is adjusted, the model is able to slowly progress towards a new local optima and ultimately converges to the fine-tuned minima.
Finally, both the training and fine-tuning dynamics of the Net-Slim (rightmost panel) appears to be highly unstable. This suggests that the pre-training with sparsity loss is not stable and can lead to large fluctuations in the training loss. Also, note that Net-Slim does not have a validation set as it instead uses the test set to choose the best model. We have refrained from using the best model by looking at the performance on the test set, and instead evaluate the model from the last epoch.
Permutation Invariance: To test the effect of permuting the neuron order on the performance of DAM, Table 2 provides results across five randomly permuted values of µij across every neuron in the network, while using the same initialization as previous experiments. We can see that DAM is invariant to permutations of neuron order during initialization, validating our simple choice of assigning neuron ordering based on their index values.
Additional Results: To illustrate generalizibility we provide additional network pruning results for more networks such as VGG-19, PreResnet-20, PreResNet-56 and PreResnet-110 on CIFAR datasets and LeNet-5 on MNIST in Appendix E3. To demonstrate decorrelated kernels obtained after pruning using DAMs we use CKA similarity [22]. The results can be found in Appendix E3. For extreme values of λ, we demonstrate that DAM is able to pruning entire bottleneck layers for ResNet architectures by visualizing the pruned architectures in Appendix E4.

heading: 6 Discussion, Open Questions, and Future Works

text:
The effectiveness of DAM in learning compact representations provides a novel perspective into structured network pruning and raises several open questions about its relationships with other concepts in this area. In the following, we list some of these connections and describe the limitations and future extensions of our work.
Connections with Lottery Ticket Hypothesis: Our work builds on the intuition of discriminative masking, i.e., we preferentially favor some neurons to be refined while favoring some other neurons to be pruned. This has interesting connections with the lottery ticket hypothesis (LTH) [10], where some subsets of initialized weights in the network are hypothesized to contain “winning tickets” that are uncovered after a few epochs of training. While there is some open discussion on whether the winning tickets exist even before initialization [11, 15], there are several studies supporting this hypothesis [38, 12]. Our results attempts to throw light on the question: “can we find winning tickets if we discriminatively search for it in a certain subregion of the network?” Further, we are able to show that the results of DAM are invariant to random permutation of the neuron indices at initialization, since all neurons receive identically distributed weights. Along these lines, we can also explore if there are certain ordering of neurons (e.g., in accordance with their likelihood of containing winning tickets revealed through LTH) that can perform better than random ordering in DAM.
Connections with Dropout: Another interesting connection of DAM is with the popular regularization technique of Dropout [40] that randomly drops out some neurons during training to break “co-adaptation” patterns in the learned features so as to avoid overfitting. Essentially, by randomly dropping a neuron at some epoch of training, the other neurons are forced to pick up the learned features of the dropped neuron, thus resulting in the learning of robust features. This is similar to the “re-adaptation” of weights at the active neurons at some epoch of DAM training, while the mask value of the neurons in the transitioning zone (when 0 < gj < 1) are gradually dropped by the shifting of the gate function. This motivates us to postulate the weight re-adaptation hypothesis as a potential reason behind the effectiveness of DAM in learning compact representations, which requires further theoretical justifications.
Budget-aware Variants of DAM: One promising future extension of our current DAM formulation would be to make it budget-aware, i.e., to stop the training process once we arrive at a target level of sparsity. Note that there is a direct correspondence between the value of the learnable gate offset parameter βi (that is minimized at every epoch) and the resulting L0 sparsity (see Equation 3). Hence, a simple budget-aware extension of our current DAM formulation would be to start from a large value of lambda (to provide sufficient scope for aggressive pruning) and keep monitoring the level of sparsity at every layer during the training process. As soon as the target sparsity mark is achieved, we can freeze βi’s at every layer thus essentially halting the pruning process (note that the network architecture becomes immutable if βi’s is fixed).
DAM Variants for Pre-trained Networks: One of the fundamental assumptions of our current DAM formulation is that all neurons are symmetrically invariant to each at initialization, and hence we can simply use random neuron ordering for discriminative masking. While this assumption is valid for “training from scratch”, this may not hold for neurons in a pre-trained network. Future extensions of DAM can include advanced ways of ordering neurons that do not rely on the above assumption and hence can even be used with pre-trained networks. For example, we can order neurons based on the magnitudes of weights of every neuron or the value of mutual information between neuron activations and outputs. This would open novel avenues of research in structured network pruning by finding effective neuron orderings for different initializations of network weights.

heading: 7 Acknowledgement

text: This work was supported by the NSF Eager Grant #2026710.

title: Learning Compact Representations of Neural Networks using DiscriminAtive Masking (DAM)


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.

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

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

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

abstractText: 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.

authors:
Bang An
Jie Lyu
Zhenyi Wang
Chunyuan Li
Changwei Hu
Fei Tan
Ruiyi Zhang
Yifan Hu
Changyou Chen

title: Repulsive attention: Rethinking multi-head attention as bayesian inference

authors:
Jacob Andreas
Marcus Rohrbach
Trevor Darrell
Dan Klein

title: Neural module networks

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Ivan Anokhin
Kirill Demochkin
Taras Khakhulin
Gleb Sterkin
Victor Lempitsky
Denis Korzhenkov

title: Image generators with conditionally-independent pixel synthesis, 2020

authors:
Dzmitry Bahdanau
Shikhar Murty
Michael Noukhovitch
Thien Huu Nguyen
Harm de Vries
Aaron Courville

title: Systematic generalization: What is required and can it be learned

authors:
David G.T. Barrett
Felix Hill
Adam Santoro
Ari S. Morcos
Timothy Lillicrap

title: Measuring abstract reasoning in neural networks, 2018

authors:
Yoshua Bengio
Tristan Deleu
Nasim Rahaman
Rosemary Ke
Sébastien Lachapelle
Olexa Bilaniuk
Anirudh Goyal
Christopher Pal

title: A meta-transfer objective for learning to disentangle causal mechanisms, 2019

authors:
Michel Besserve
Rémy Sun
Dominik Janzing
Bernhard Schölkopf

title: A theory of independent mechanisms for extrapolation in generative models

venue: arXiv preprint arXiv:2004.00184,

authors:
Shyamal Buch
Li Fei-Fei
Noah D. Goodman

title: Neural event semantics for grounded language understanding

venue: In Transactions of the Association for Computational Linguistics (TACL),

authors:
Michael B Chang
Abhishek Gupta
Sergey Levine
Thomas L Griffiths

title: Automatically composing representation transformations as a means for generalization

venue: arXiv preprint arXiv:1807.04640,

authors:
Xinlei Chen
Saining Xie
Kaiming He

title: An empirical study of training self-supervised vision

authors:
Anna Choromanska
Mikael Henaff
Michael Mathieu
Gérard Ben Arous
Yann LeCun

title: The loss surfaces of multilayer networks

venue: In Artificial intelligence and statistics,

authors:
Tarin Clanuwat
Mikel Bober-Irizar
Asanobu Kitamoto
Alex Lamb
Kazuaki Yamamoto
David Ha

title: Deep learning for classical japanese literature, 2018

authors:
Jean-Baptiste Cordonnier
Andreas Loukas
Martin Jaggi

title: On the relationship between selfattention and convolutional layers

venue: arXiv preprint arXiv:1911.03584,

authors:
Ekin D. Cubuk
Barret Zoph
Jonathon Shlens
Quoc V. Le

title: Randaugment: Practical automated data augmentation with a reduced search space, 2019

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: Bert: Pre-training of deep bidirectional transformers for language understanding, 2019

authors:
Alexey Dosovitskiy
Lucas Beyer
Alexander Kolesnikov
Dirk Weissenborn
Xiaohua Zhai
Thomas Unterthiner
Mostafa Dehghani
Matthias Minderer
Georg Heigold
Sylvain Gelly

title: An image is worth 16x16 words: Transformers for image recognition at scale

venue: arXiv preprint arXiv:2010.11929,

authors:
William Fedus
Barret Zoph
Noam Shazeer

title: Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity

venue: arXiv preprint arXiv:2101.03961,

authors:
Chrisantha Fernando
Dylan Banarse
Charles Blundell
Yori Zwols
David Ha
Andrei A Rusu
Alexander Pritzel
Daan Wierstra

title: Pathnet: Evolution channels gradient descent in super neural networks

venue: arXiv preprint arXiv:1701.08734,

authors:
Yaroslav Ganin
Evgeniya Ustinova
Hana Ajakan
Pascal Germain
Hugo Larochelle
François Laviolette
Mario Marchand
Victor Lempitsky

title: Domain-adversarial training of neural networks. The journal of machine learning

authors:
Anirudh Goyal
Yoshua Bengio

title: Inductive biases for deep learning of higher-level cognition

venue: arXiv preprint arXiv:2011.15091,

authors:
Anirudh Goyal
Aniket Didolkar
Nan Rosemary Ke
Charles Blundell
Philippe Beaudoin
Nicolas Heess
Michael Mozer
Yoshua Bengio

title: Neural production systems

venue: arXiv preprint arXiv:2103.01937,

authors:
Anirudh Goyal
Alex Lamb
Phanideep Gampa
Philippe Beaudoin
Sergey Levine
Charles Blundell
Yoshua Bengio
Michael Mozer

title: Object files and schemata: Factorizing declarative and procedural knowledge in dynamical systems

authors:
Anirudh Goyal
Alex Lamb
Jordan Hoffmann
Shagun Sodhani
Sergey Levine
Yoshua Bengio
Bernhard Schölkopf

title: Recurrent independent mechanisms

authors:
Vineet Gupta
Tomer Koren
Yoram Singer

title: Shampoo: Preconditioned stochastic tensor optimization, 2018

authors:
Dan Hendrycks
Kevin Gimpel

title: Gaussian error linear units (gelus), 2020

authors:
Drew A Hudson
Christopher D Manning

title: Compositional attention networks for machine reasoning

venue: arXiv preprint arXiv:1803.03067,

authors:
Andrew Jaegle
Felix Gimeno
Andrew Brock
Andrew Zisserman
Oriol Vinyals
Joao Carreira

title: Perceiver: General perception with iterative attention, 2021

authors:
Jason Jo
Yoshua Bengio

title: Measuring the tendency of cnns to learn surface statistical regularities, 2017

authors:
Simon L Kendal
Malcolm Creen

title: An introduction to knowledge

authors:
Daniel Keysers
Nathanael Schärli
Nathan Scales
Hylke Buisman
Daniel Furrer
Sergii Kashubin
Nikola Momchev
Danila Sinopalnikov
Lukasz Stafiniak
Tibor Tihon

title: Measuring compositional generalization: A comprehensive method on realistic data

venue: arXiv preprint arXiv:1912.09713,

authors:
Louis Kirsch
Julius Kunze
David Barber

title: Modular networks: Learning to decompose neural computation

venue: arXiv preprint arXiv:1811.05249,

authors:
Brenden Lake
Marco Baroni

title: Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks

venue: In International Conference on Machine Learning,

authors:
Alex Lamb
Di He
Anirudh Goyal
Guolin Ke
Chien-Feng Liao
Mirco Ravanelli
Yoshua Bengio

title: Transformers with competitive ensembles of independent mechanisms

venue: arXiv preprint arXiv:2103.00336,

authors:
Yann LeCun
Corinna Cortes
CJ Burges

title: Mnist handwritten digit database

venue: ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist,

authors:
Liyuan Liu
Haoming Jiang
Pengcheng He
Weizhu Chen
Xiaodong Liu
Jianfeng Gao
Jiawei Han

title: On the variance of the adaptive learning rate and beyond

authors:
Francesco Locatello
Dirk Weissenborn
Thomas Unterthiner
Aravindh Mahendran
Georg Heigold
Jakob Uszkoreit
Alexey Dosovitskiy
Thomas Kipf

title: Object-centric learning with slot attention, 2020

authors:
Lajanugen Logeswaran
Ann Lee
Myle Ott
Honglak Lee
Marc’Aurelio Ranzato
Arthur Szlam

title: Few-shot sequence learning with transformers

venue: arXiv preprint arXiv:2012.09543,

authors:
Joao Loula
Marco Baroni
Brenden M Lake

title: Rearranging the familiar: Testing compositional generalization in recurrent networks

venue: arXiv preprint arXiv:1807.07545,

authors:
Jiayuan Mao
Chuang Gan
Pushmeet Kohli
Joshua B Tenenbaum
Jiajun Wu

title: The neurosymbolic concept learner: Interpreting scenes, words, and sentences from natural supervision

authors:
Yuval Netzer
Tao Wang
Adam Coates
Alessandro Bissacco
Bo Wu
Andrew Y Ng

title: Reading digits in natural images with unsupervised feature learning

authors:
G. Parascandolo
N. Kilbertus
M. Rojas-Carulla
B. Schölkopf

title: Learning independent causal mechanisms

venue: In Proceedings of the 35th International Conference on Machine Learning (ICML),

authors:
Jonas Peters
Dominik Janzing
Bernhard Schölkopf

title: Elements of causal inference: foundations and learning algorithms

authors:
Alec Radford
Jeffrey Wu
Rewon Child
David Luan
Dario Amodei
Ilya Sutskever

title: Language models are unsupervised multitask learners

venue: OpenAI blog,

authors:
Nasim Rahaman
Aristide Baratin
Devansh Arpit
Felix Draxler
Min Lin
Fred A. Hamprecht
Yoshua Bengio
Aaron Courville

title: On the spectral bias of neural networks, 2019

authors:
Nasim Rahaman
Anirudh Goyal
Muhammad Waleed Gondal
Manuel Wuthrich
Stefan Bauer
Yash Sharma
Yoshua Bengio
Bernhard Schölkopf

title: S2rms: Spatially structured recurrent modules, 2020

authors:
John C Raven
John Hugh Court

title: Raven’s progressive matrices and vocabulary scales, volume 759

authors:
Clemens Rosenbaum
Ignacio Cases
Matthew Riemer
Tim Klinger

title: Routing networks and the challenges of modular and compositional computation

authors:
Clemens Rosenbaum
Tim Klinger
Matthew Riemer

title: Routing networks: Adaptive selection of non-linear functions for multi-task learning

venue: arXiv preprint arXiv:1711.01239,

authors:
Sara Sabour
Nicholas Frosst
Geoffrey E Hinton

title: Dynamic routing between capsules

venue: arXiv preprint arXiv:1710.09829,

authors:
B. Schölkopf
F. Locatello
S. Bauer
R. Nan Ke
N. Kalchbrenner
A. Goyal
Y. Bengio

title: Towards causal representation learning

venue: Proceedings of the IEEE,

authors:
Bernhard Schölkopf
Dominik Janzing
Jonas Peters
Eleni Sgouritsa
Kun Zhang
Joris Mooij

title: On causal and anticausal learning

venue: arXiv preprint arXiv:1206.6471,

authors:
Noam Shazeer
Azalia Mirhoseini
Krzysztof Maziarz
Andy Davis
Quoc Le
Geoffrey Hinton
Jeff Dean

title: Outrageously large neural networks: The sparsely-gated mixture-of-experts layer

venue: arXiv preprint arXiv:1701.06538,

authors:
Xander Steenbrugge
Sam Leroux
Tim Verbelen
Bart Dhoedt

title: Improving generalization for abstract reasoning tasks using disentangled feature representations

venue: arXiv preprint arXiv:1811.04784,

authors:
Hugo Touvron
Matthieu Cord
Matthijs Douze
Francisco Massa
Alexandre Sablayrolles
Hervé Jégou

title: Training data-efficient image transformers and distillation through attention, 2021

authors:
Sjoerd Van Steenkiste
Michael Chang
Klaus Greff
Jürgen Schmidhuber

title: Relational neural expectation maximization: Unsupervised discovery of objects and their interactions

venue: arXiv preprint arXiv:1802.10353,

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Lukasz Kaiser
Illia Polosukhin

title: Attention is all you need

venue: arXiv preprint arXiv:1706.03762,

authors:
Duo Wang
Mateja Jamnik
Pietro Lio

title: Abstract diagrammatic reasoning with multiplex graph networks

venue: arXiv preprint arXiv:2006.11197,

title: Pytorch image models

authors:
Yuhuai Wu
Honghua Dong
Roger Grosse
Jimmy Ba

title: The scattering compositional learner: Discovering objects, attributes, relationships in analogical reasoning

venue: arXiv preprint arXiv:2007.04212,

authors:
Kexin Yi
Chuang Gan
Yunzhu Li
Pushmeet Kohli
Jiajun Wu
Antonio Torralba
Joshua B Tenenbaum

title: Clevrer: Collision events for video representation and reasoning

title: Using anytime algorithms in intelligent systems

venue: AI magazine,

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.

title: Dynamic Inference with Neural Interpreters