PRUNING CONVOLUTIONAL NEURAL NETWORKS FOR RESOURCE EFFICIENT INFERENCE
We propose a new formulation for pruning convolutional kernels in neural networks to enable efficient inference. We interleave greedy criteria-based pruning with finetuning by backpropagation—a computationally efficient procedure that maintains good generalization in the pruned network. We propose a new criterion based on Taylor expansion that approximates the change in the cost function induced by pruning network parameters. We focus on transfer learning, where large pretrained networks are adapted to specialized tasks. The proposed criterion demonstrates superior performance compared to other criteria, e.g. the norm of kernel weights or feature map activation, for pruning large CNNs after adaptation to fine-grained classification tasks (Birds-200 and Flowers-102) relaying only on the first order gradient information. We also show that pruning can lead to more than 10× theoretical reduction in adapted 3D-convolutional filters with a small drop in accuracy in a recurrent gesture classifier. Finally, we show results for the largescale ImageNet dataset to emphasize the flexibility of our approach.
1 INTRODUCTION
Convolutional neural networks (CNN) are used extensively in computer vision applications, including object classification and localization, pedestrian and car detection, and video classification. Many problems like these focus on specialized domains for which there are only small amounts of carefully curated training data. In these cases, accuracy may be improved by fine-tuning an existing deep network previously trained on a much larger labeled vision dataset, such as images from ImageNet (Russakovsky et al., 2015) or videos from Sports-1M (Karpathy et al., 2014). While transfer learning of this form supports state of the art accuracy, inference is expensive due to the time, power, and memory demanded by the heavyweight architecture of the fine-tuned network.
While modern deep CNNs are composed of a variety of layer types, runtime during prediction is dominated by the evaluation of convolutional layers. With the goal of speeding up inference, we prune entire feature maps so the resulting networks may be run efficiently even on embedded devices. We interleave greedy criteria-based pruning with fine-tuning by backpropagation, a computationally efficient procedure that maintains good generalization in the pruned network.
Neural network pruning was pioneered in the early development of neural networks (Reed, 1993). Optimal Brain Damage (LeCun et al., 1990) and Optimal Brain Surgeon (Hassibi & Stork, 1993) leverage a second-order Taylor expansion to select parameters for deletion, using pruning as regularization to improve training and generalization. This method requires computation of the Hessian matrix partially or completely, which adds memory and computation costs to standard fine-tuning.
In line with our work, Anwar et al. (2015) describe structured pruning in convolutional layers at the level of feature maps and kernels, as well as strided sparsity to prune with regularity within kernels. Pruning is accomplished by particle filtering wherein configurations are weighted by misclassification rate. The method demonstrates good results on small CNNs, but larger CNNs are not addressed.
Han et al. (2015) introduce a simpler approach by fine-tuning with a strong `2 regularization term and dropping parameters with values below a predefined threshold. Such unstructured pruning is very effective for network compression, and this approach demonstrates good performance for intra-kernel pruning. But compression may not translate directly to faster inference since modern hardware
exploits regularities in computation for high throughput. So specialized hardware may be needed for efficient inference of a network with intra-kernel sparsity (Han et al., 2016). This approach also requires long fine-tuning times that may exceed the original network training by a factor of 3 or larger. Group sparsity based regularization of network parameters was proposed to penalize unimportant parameters (Wen et al., 2016; Zhou et al., 2016; Alvarez & Salzmann, 2016; Lebedev & Lempitsky, 2016). Regularization-based pruning techniques require per layer sensitivity analysis which adds extra computations. In contrast, our approach relies on global rescaling of criteria for all layers and does not require sensitivity estimation. Moreover, our approach is faster as we directly prune unimportant parameters instead of waiting for their values to be made sufficiently small by optimization under regularization.
Other approaches include combining parameters with correlated weights (Srinivas & Babu, 2015), reducing precision (Gupta et al., 2015; Rastegari et al., 2016) or tensor decomposition (Kim et al., 2015). These approaches usually require a separate training procedure or significant fine-tuning, but potentially may be combined with our method for additional speedups.
2 METHOD
The proposed method for pruning consists of the following steps: 1) Fine-tune the network until convergence on the target task; 2) Alternate iterations of pruning and further fine-tuning; 3) Stop pruning after reaching the target trade-off between accuracy and pruning objective, e.g. floating point operations (FLOPs) or memory utilization.
The procedure is simple, but its success hinges on employing the right pruning criterion. In this section, we introduce several efficient pruning criteria and related technical considerations.
Consider a set of training examples D = { X =
{x0,x1, ...,xN},Y = {y0, y1, ..., yN} }
, where x and y represent an input and a target output, respectively. The network’s parameters1W = {(w11, b11), (w21, b21), ...(w C` L , b C` L )} are optimized to minimize a cost value C(D|W). The most common choice for a cost function C(·) is a negative log-likelihood function. A cost function is selected independently of pruning and depends only on the task to be solved by the original network. In the case of transfer learning, we adapt a large network initialized with parametersW0 pretrained on a related but distinct dataset.
During pruning, we refine a subset of parameters which preserves the accuracy of the adapted network, C(D|W ′) ≈ C(D|W). This corresponds to a combinatorial optimization:
min W′ ∣∣∣∣C(D|W ′)− C(D|W)∣∣∣∣ s.t. ||W ′||0 ≤ B, (1) where the `0 norm in ||W ′||0 bounds the number of non-zero parameters B in W ′. Intuitively, if W ′ = W we reach the global minimum of the error function, however ||W ′||0 will also have its maximum.
Finding a good subset of parameters while maintaining a cost value as close as possible to the original is a combinatorial problem. It will require 2|W| evaluations of the cost function for a selected subset of data. For current networks it would be impossible to compute: for example, VGG-16 has |W| = 4224 convolutional feature maps. While it is impossible to solve this optimization exactly for networks of any reasonable size, in this work we investigate a class of greedy methods. Starting with a full set of parametersW , we iteratively identify and remove the least important parameters, as illustrated in Figure 1. By removing parameters at each iteration, we ensure the eventual satisfaction of the `0 bound onW ′.
1A “parameter” (w, b) ∈ W might represent an individual weight, a convolutional kernel, or the entire set of kernels that compute a feature map; our experiments operate at the level of feature maps.
Since we focus our analysis on pruning feature maps from convolutional layers, let us denote a set of image feature maps by z` ∈ RH`×W`×C` with dimensionality H` ×W` and C` individual maps (or channels).2 The feature maps can either be the input to the network, z0, or the output from a convolutional layer, z` with ` ∈ [1, 2, ..., L]. Individual feature maps are denoted z(k)` for k ∈ [1, 2, ..., C`]. A convolutional layer ` applies the convolution operation (∗) to a set of input feature maps z`−1 with kernels parameterized by w (k) ` ∈ RC`−1×p×p:
z (k) ` = g (k) ` R ( z`−1 ∗w(k)` + b (k) ` ) , (2)
where z(k)` ∈ RH`×W` is the result of convolving each ofC`−1 kernels of size p×p with its respective input feature map and adding bias b(k)` . We introduce a pruning gate gl ∈ {0, 1}Cl , an external switch which determines if a particular feature map is included or pruned during feed-forward propagation, such that when g is vectorized: W ′ = gW .
2.1 ORACLE PRUNING
Minimizing the difference in accuracy between the full and pruned models depends on the criterion for identifying the “least important” parameters, called saliency, at each step. The best criterion would be an exact empirical evaluation of each parameter, which we denote the oracle criterion, accomplished by ablating each non-zero parameter w ∈ W ′ in turn and recording the cost’s difference. We distinguish two ways of using this oracle estimation of importance: 1) oracle-loss quantifies importance as the signed change in loss, C(D|W ′)− C(D|W), and 2) oracle-abs adopts the absolute difference, |C(D|W ′) − C(D|W)|. While both discourage pruning which increases the loss, the oracle-loss version encourages pruning which may decrease the loss, while oracle-abs penalizes any pruning in proportion to its change in loss, regardless of the direction of change.
While the oracle is optimal for this greedy procedure, it is prohibitively costly to compute, requiring ||W ′||0 evaluations on a training dataset, one evaluation for each remaining non-zero parameter. Since estimation of parameter importance is key to both the accuracy and the efficiency of this pruning approach, we propose and evaluate several criteria in terms of performance and estimation cost.
2.2 CRITERIA FOR PRUNING
There are many heuristic criteria which are much more computationally efficient than the oracle. For the specific case of evaluating the importance of a feature map (and implicitly the set of convolutional kernels from which it is computed), reasonable criteria include: the combined `2-norm of the kernel weights, the mean, standard deviation or percentage of the feature map’s activation, and mutual information between activations and predictions. We describe these criteria in the following paragraphs and propose a new criterion which is based on the Taylor expansion.
Minimum weight. Pruning by magnitude of kernel weights is perhaps the simplest possible criterion, and it does not require any additional computation during the fine-tuning process. In case of pruning according to the norm of a set of weights, the criterion is evaluated as: ΘMW : RC`−1×p×p → R, with ΘMW (w) = 1|w| ∑ i w 2 i , where |w| is dimensionality of the set of weights after vectorization. The motivation to apply this type of pruning is that a convolutional kernel with low `2 norm detects less important features than those with a high norm. This can be aided during training by applying `1 or `2 regularization, which will push unimportant kernels to have smaller values.
Activation. One of the reasons for the popularity of the ReLU activation is the sparsity in activation that is induced, allowing convolutional layers to act as feature detectors. Therefore it is reasonable to assume that if an activation value (an output feature map) is small then this feature detector is not important for prediction task at hand. We may evaluate this by mean activation, ΘMA : RHl×W`×C` → R, with ΘMA(a) = 1|a| ∑ i ai for activation a = z (k) l , or by the standard deviation
of the activation, ΘMA_std(a) = √ 1 |a| ∑ i(ai − µa)2.
2While our notation is at times specific to 2D convolutions, the methods are applicable to 3D convolutions, as well as fully connected layers.
Mutual information. Mutual information (MI) is a measure of how much information is present in one variable about another variable. We apply MI as a criterion for pruning, ΘMI : RHl×W`×C` → R, with ΘMI(a) = MI(a, y), where y is the target of neural network. MI is defined for continuous variables, so to simplify computation, we exchange it with information gain (IG), which is defined for quantized variables IG(y|x) = H(x) +H(y)−H(x, y), where H(x) is the entropy of variable x. We accumulate statistics on activations and ground truth for a number of updates, then quantize the values and compute IG.
Taylor expansion. We phrase pruning as an optimization problem, trying to findW ′ with bounded number of non-zero elements that minimize ∣∣∆C(hi)∣∣ = |C(D|W ′)− C(D|W)|. With this approach based on the Taylor expansion, we directly approximate change in the loss function from removing a particular parameter. Let hi be the output produced from parameter i. In the case of feature maps, h = {z(1)0 , z (2) 0 , ..., z (C`) L }. For notational convenience, we consider the cost function equally dependent on parameters and outputs computed from parameters: C(D|hi) = C(D|(w, b)i). Assuming independence of parameters, we have:∣∣∆C(hi)∣∣ = ∣∣C(D, hi = 0)− C(D, hi)∣∣, (3) where C(D, hi = 0) is a cost value if output hi is pruned, while C(D, hi) is the cost if it is not pruned. While parameters are in reality inter-dependent, we already make an independence assumption at each gradient step during training.
To approximate ∆C(hi), we use the first-degree Taylor polynomial. For a function f(x), the Taylor expansion at point x = a is
f(x) = P∑ p=0 f (p)(a) p! (x− a)p +Rp(x), (4)
where f (p)(a) is the p-th derivative of f evaluated at point a, and Rp(x) is the p-th order remainder. Approximating C(D, hi = 0) with a first-order Taylor polynomial near hi = 0, we have:
C(D, hi = 0) = C(D, hi)− δC δhi hi +R1(hi = 0). (5)
The remainder R1(hi = 0) can be calculated through the Lagrange form:
R1(hi = 0) = δ2C
δ(h2i = ξ)
h2i 2 , (6)
where ξ is a real number between 0 and hi. However, we neglect this first-order remainder, largely due to the significant calculation required, but also in part because the widely-used ReLU activation function encourages a smaller second order term. Finally, by substituting Eq. (5) into Eq. (3) and ignoring the remainder, we have ΘTE : RHl×Wl×Cl → R+, with
ΘTE(hi) = ∣∣∆C(hi)∣∣ = ∣∣C(D, hi)− δC
δhi hi − C(D, hi) ∣∣ = ∣∣∣∣ δCδhihi ∣∣∣∣. (7)
Intuitively, this criterion prunes parameters that have an almost flat gradient of the cost function w.r.t. feature map hi. This approach requires accumulation of the product of the activation and the gradient of the cost function w.r.t. to the activation, which is easily computed from the same computations for back-propagation. ΘTE is computed for a multi-variate output, such as a feature map, by
ΘTE(z (k) l ) = ∣∣∣∣ 1M ∑ m δC δz (k) l,m z (k) l,m ∣∣∣∣, (8) where M is length of vectorized feature map. For a minibatch with T > 1 examples, the criterion is computed for each example separately and averaged over T .
Independently of our work, Figurnov et al. (2016) came up with similar metric based on the Taylor expansion, called impact, to evaluate importance of spatial cells in a convolutional layer. It shows that the same metric can be applied to evaluate importance of different groups of parameters.
Relation to Optimal Brain Damage. The Taylor criterion proposed above relies on approximating the change in loss caused by removing a feature map. The core idea is the same as in Optimal Brain Damage (OBD) (LeCun et al., 1990). Here we consider the differences more carefully.
The primary difference is the treatment of the first-order term of the Taylor expansion, in our notation y = δCδhh for cost function C and hidden layer activation h. After sufficient training epochs, the gradient term tends to zero: δCδh → 0 and E(y) = 0. At face value y offers little useful information, hence OBD regards the term as zero and focuses on the second-order term.
However, the variance of y is non-zero and correlates with the stability of the local function w.r.t. activation h. By considering the absolute change in the cost3 induced by pruning (as in Eq. 3), we use the absolute value of the first-order term, |y|. Under assumption that samples come from independent and identical distribution, E(|y|) = σ √ 2/ √ π where σ is the standard deviation of y, known as the expected value of the half-normal distribution. So, while y tends to zero, the expectation of |y| is proportional to the variance of y, a value which is empirically more informative as a pruning criterion.
As an additional benefit, we avoid the computation of the second-order Taylor expansion term, or its simplification - diagonal of the Hessian, as required in OBD.
We found important to compare proposed Taylor criteria to OBD. As described in the original papers (LeCun et al., 1990; 1998), OBD can be efficiently implemented similarly to standard back propagation algorithm doubling backward propagation time and memory usage when used together with standard fine-tuning. Efficient implementation of the original OBD algorithm might require significant changes to the framework based on automatic differentiation like Theano to efficiently compute only diagonal of the Hessian instead of the full matrix. Several researchers tried to tackle this problem with approximation techniques (Martens, 2010; Martens et al., 2012). In our implementation, we use efficient way of computing Hessian-vector product (Pearlmutter, 1994) and matrix diagonal approximation proposed by (Bekas et al., 2007), please refer to more details in appendix. With current implementation, OBD is 30 times slower than Taylor technique for saliency estimation, and 3 times slower for iterative pruning, however with different implementation can only be 50% slower as mentioned in the original paper.
Average Percentage of Zeros (APoZ). Hu et al. (2016) proposed to explore sparsity in activations for network pruning. ReLU activation function imposes sparsity during inference, and average percentage of positive activations at the output can determine importance of the neuron. Intuitively, it is a good criteria, however feature maps at the first layers have similar APoZ regardless of the network’s target as they learn to be Gabor like filters. We will use APoZ to estimate saliency of feature maps.
2.3 NORMALIZATION
Some criteria return “raw” values, whose scale varies with the depth of the parameter’s layer in the network. A simple layer-wise `2-normalization can achieve adequate rescaling across layers:
Θ̂(z (k) l )= Θ(z (k) l )√∑
j
( Θ(z (j) l ) )2 .
2.4 FLOPS REGULARIZED PRUNING
One of the main reasons to apply pruning is to reduce number of operations in the network. Feature maps from different layers require different amounts of computation due the number and sizes of input feature maps and convolution kernels. To take this into account we introduce FLOPs regularization:
Θ(z (k) l ) = Θ(z (k) l )− λΘ flops l , (9)
where λ controls the amount of regularization. For our experiments, we use λ = 10−3. Θflops is computed under the assumption that convolution is implemented as a sliding window (see Appendix). Other regularization conditions may be applied, e.g. storage size, kernel sizes, or memory footprint.
3OBD approximates the signed difference in loss, while our method approximates absolute difference in loss. We find in our results that pruning based on absolute difference yields better accuracy.
3 RESULTS
We empirically study the pruning criteria and procedure detailed in the previous section for a variety of problems. We focus many experiments on transfer learning problems, a setting where pruning seems to excel. We also present results for pruning large networks on their original tasks for more direct comparison with the existing pruning literature. Experiments are performed within Theano (Theano Development Team, 2016). Training and pruning are performed on the respective training sets for each problem, while results are reported on appropriate holdout sets, unless otherwise indicated. For all experiments we prune a single feature map at every pruning iteration, allowing fine-tuning and re-evaluation of the criterion to account for dependency between parameters.
3.1 CHARACTERIZING THE ORACLE RANKING
We begin by explicitly computing the oracle for a single pruning iteration of a visual transfer learning problem. We fine-tune the VGG-16 network (Simonyan & Zisserman, 2014) for classification of bird species using the Caltech-UCSD Birds 200-2011 dataset (Wah et al., 2011). The dataset consists of nearly 6000 training images and 5700 test images, covering 200 species. We fine-tune VGG-16 for 60 epochs with learning rate 0.0001 to achieve a test accuracy of 72.2% using uncropped images.
To compute the oracle, we evaluate the change in loss caused by removing each individual feature map from the fine-tuned VGG-16 network. (See Appendix A.3 for additional analysis.) We rank feature maps by their contributions to the loss, where rank 1 indicates the most important feature map—removing it results in the highest increase in loss—and rank 4224 indicates the least important. Statistics of global ranks are shown in Fig. 2 grouped by convolutional layer. We observe: (1) Median global importance tends to decrease with depth. (2) Layers with max-pooling tend to be more important than those without. (VGG-16 has pooling after layers 2, 4, 7, 10, and 13.) However, (3) maximum and minimum ranks show that every layer has some feature maps that are globally important and others that are globally less important. Taken together with the results of subsequent experiments, we opt for encouraging a balanced pruning that distributes selection across all layers.
Next, we iteratively prune the network using pre-computed oracle ranking. In this experiment, we do not update the parameters of the network or the oracle ranking between iterations. Training accuracy is illustrated in Fig. 3 over many pruning iterations. Surprisingly, pruning by smallest absolute change in loss (Oracle-abs) yields higher accuracy than pruning by the net effect on loss (Oracle-loss). Even though the oracle indicates that removing some feature maps individually may decrease loss, instability accumulates due the large absolute changes that are induced. These results support pruning by absolute difference in cost, as constructed in Eq. 1.
3.2 EVALUATING PROPOSED CRITERIA VERSUS THE ORACLE
To evaluate computationally efficient criteria as substitutes for the oracle, we compute Spearman’s rank correlation, an estimate of how well two predictors provide monotonically related outputs,
even if their relationship is not linear. Given the difference between oracle4 and criterion ranks di = rank(Θoracle(i))−rank(Θcriterion(i)) for each parameter i, the rank correlation is computed:
S = 1− 6 N(N2 − 1) N∑ i=1 di 2, (10)
where N is the number of parameters (and the highest rank). This correlation coefficient takes values in [−1, 1], where −1 implies full negative correlation, 0 no correlation, and 1 full positive correlation. We show Spearman’s correlation in Table 1 to compare the oracle-abs ranking to rankings by different criteria on a set of networks/datasets some of which are going to be introduced later. Data-dependent criteria (all except weight magnitude) are computed on training data during the fine-tuning before or between pruning iterations. As a sanity check, we evaluate random ranking and observe 0.0 correlation across all layers. “Per layer” analysis shows ranking within each convolutional layer, while “All layers” describes ranking across layers. While several criteria do not scale well across layers with raw values, a layer-wise `2-normalization significantly improves performance. The Taylor criterion has the highest correlation among the criteria, both within layers and across layers (with `2 normalization). OBD shows the best correlation across layers when no normalization used; it also shows best results for correlation on ImageNet dataset. (See Appendix A.2 for further analysis.)
3.3 PRUNING FINE-TUNED IMAGENET NETWORKS
We now evaluate the full iterative pruning procedure on two transfer learning problems. We focus on reducing the number of convolutional feature maps and the total estimated floating point operations (FLOPs). Fine-grained recognition is difficult for relatively small datasets without relying on transfer
4We use Oracle-abs because of better performance in previous experiment
learning. Branson et al. (2014) show that training CNN from scratch on the Birds-200 dataset achieves test accuracy of only 10.9%. We compare results to training a randomly initialized CNN with half the number of parameters per layer, denoted "from scratch".
Fig. 4 shows pruning of VGG-16 after fine-tuning on the Birds-200 dataset (as described previously). At each pruning iteration, we remove a single feature map and then perform 30 minibatch SGD updates with batch-size 32, momentum 0.9, learning rate 10−4, and weight decay 10−4. The figure depicts accuracy relative to the pruning rate (left) and estimated GFLOPs (right). The Taylor criterion shows the highest accuracy for nearly the entire range of pruning ratios, and with FLOPs regularization demonstrates the best performance relative to the number of operations. OBD shows slightly worse performance of pruning in terms of parameters, however significantly worse in terms of FLOPs.
In Fig. 5, we show pruning of the CaffeNet implementation of AlexNet (Krizhevsky et al., 2012) after adapting to the Oxford Flowers 102 dataset (Nilsback & Zisserman, 2008), with 2040 training and 6129 test images from 102 species of flowers. Criteria correlation with oracle-abs is summarized in Table 1. We initially fine-tune the network for 20 epochs using a learning rate of 0.001, achieving a final test accuracy of 80.1%. Then pruning procedes as previously described for Birds-200, except with only 10 mini-batch updates between pruning iterations. We observe the superior performance of the Taylor and OBD criteria in both number of parameters and GFLOPs.
We observed that Taylor criterion shows the best performance which is closely followed by OBD with a bit lower Spearman’s rank correlation coefficient. Implementing OBD takes more effort because of computation of diagonal of the Hessian and it is 50% to 300% slower than Taylor criteria that relies on first order gradient only.
Fig. 6 shows pruning with the Taylor technique and a varying number of fine-tuning updates between pruning iterations. Increasing the number of updates results in higher accuracy, but at the cost of additional runtime of the pruning procedure.
During pruning we observe a small drop in accuracy. One of the reasons is fine-tuning between pruning iterations. Accuracy of the initial network can be improved with longer fine tunning and search of better optimization parameters. For example accuracy of unpruned VGG16 network on Birds-200 goes up to 75% after extra 128k updates. And AlexNet on Flowers-102 goes up to 82.9% after 130k updates. It should be noted that with farther fine-tuning of pruned networks we can achieve higher accuracy as well, therefore the one-to-one comparison of accuracies is rough.
3.4 PRUNING A RECURRENT 3D-CNN NETWORK FOR HAND GESTURE RECOGNITION
Molchanov et al. (2016) learn to recognize 25 dynamic hand gestures in streaming video with a large recurrent neural network. The network is constructed by adding recurrent connections to a 3D-CNN pretrained on the Sports-1M video dataset (Karpathy et al., 2014) and fine tuning on a gesture dataset. The full network achieves an accuracy of 80.7% when trained on the depth modality, but a single inference requires an estimated 37.8 GFLOPs, too much for deployment on an embedded GPU. After several iterations of pruning with the Taylor criterion with learning rate 0.0003, momentum 0.9, FLOPs regularization 10−3, we reduce inference to 3.0 GFLOPs, as shown in Fig. 7. While pruning
increases classification error by nearly 6%, additional fine-tuning restores much of the lost accuracy, yielding a final pruned network with a 12.6× reduction in GFLOPs and only a 2.5% loss in accuracy.
3.5 PRUNING NETWORKS FOR IMAGENET
11.5 and 8 GFLOPs.
We also test our pruning scheme on the largescale ImageNet classification task. In the first experiment, we begin with a trained CaffeNet implementation of AlexNet with 79.2% top-5 validation accuracy. Between pruning iterations, we fine-tune with learning rate 10−4, momentum 0.9, weight decay 10−4, batch size 32, and drop-out 50%. Using a subset of 5000 training images, we compute oracle-abs and Spearman’s rank correlation with the criteria, as shown in Table 1. Pruning traces are illustrated in Fig. 8. We observe: 1) Taylor performs better than random or minimum weight pruning when 100 updates are used between pruning iterations. When results are displayed w.r.t. FLOPs, the difference with random pruning is only 0%−4%, but the difference is higher, 1%−10%, when plotted with the number of feature maps pruned. 2) Increasing the number of updates from 100 to 1000 improves performance of pruning significantly for both the Taylor criterion and random pruning.
Hardware Batch Accuracy Time, ms Accuracy Time (speed up) Accuracy Time (speed up)
For a second experiment, we prune a trained VGG-16 network with the same parameters as before, except enabling FLOPs regularization. We stop pruning at two points, 11.5 and 8.0 GFLOPs, and fine-tune both models for an additional five epochs with learning rate 10−4. Fine-tuning after pruning significantly improves results: the network pruned to 11.5 GFLOPs improves from 83% to 87% top-5 validation accuracy, and the network pruned to 8.0 GFLOPs improves from 77.8% to 84.5%.
3.6 SPEED UP MEASUREMENTS
During pruning we were measuring reduction in computations by FLOPs, which is a common practice (Han et al., 2015; Lavin, 2015a;b). Improvements in FLOPs result in monotonically decreasing inference time of the networks because of removing entire feature map from the layer. However, time consumed by inference dependents on particular implementation of convolution operator, parallelization algorithm, hardware, scheduling, memory transfer rate etc. Therefore we measure improvement in the inference time for selected networks to see real speed up compared to unpruned networks in Table 2. We observe significant speed ups by proposed pruning scheme.
4 CONCLUSIONS
We propose a new scheme for iteratively pruning deep convolutional neural networks. We find: 1) CNNs may be successfully pruned by iteratively removing the least important parameters—feature maps in this case—according to heuristic selection criteria; 2) a Taylor expansion-based criterion demonstrates significant improvement over other criteria; 3) per-layer normalization of the criterion is important to obtain global scaling.
A APPENDIX
A.1 FLOPS COMPUTATION
To compute the number of floating-point operations (FLOPs), we assume convolution is implemented as a sliding window and that the nonlinearity function is computed for free. For convolutional kernels we have:
FLOPs = 2HW (CinK2 + 1)Cout, (11)
where H , W and Cin are height, width and number of channels of the input feature map, K is the kernel width (assumed to be symmetric), and Cout is the number of output channels.
For fully connected layers we compute FLOPs as:
FLOPs = (2I − 1)O, (12)
where I is the input dimensionality and O is the output dimensionality.
We apply FLOPs regularization during pruning to prune neurons with higher FLOPs first. FLOPs per convolutional neuron in every layer:
VGG16: Θflops = [3.1, 57.8, 14.1, 28.9, 7.0, 14.5, 14.5, 3.5, 7.2, 7.2, 1.8, 1.8, 1.8, 1.8]
AlexNet: Θflops = [2.3, 1.7, 0.8, 0.6, 0.6]
R3DCNN: Θflops = [5.6, 86.9, 21.7, 43.4, 5.4, 10.8, 1.4, 1.4]
A.2 NORMALIZATION ACROSS LAYERS
Scaling a criterion across layers is very important for pruning. If the criterion is not properly scaled, then a hand-tuned multiplier would need to be selected for each layer. Statistics of feature map ranking by different criteria are shown in Fig. 10. Without normalization (Fig. 14a–14d), the weight magnitude criterion tends to rank feature maps from the first layers more important than last layers; the activation criterion ranks middle layers more important; and Taylor ranks first layers higher. After `2 normalization (Fig. 10d–10f), all criteria have a shape more similar to the oracle, where each layer has some feature maps which are highly important and others which are unimportant.
A.3 ORACLE COMPUTATION FOR VGG-16 ON BIRDS-200
We compute the change in the loss caused by removing individual feature maps from the VGG-16 network, after fine-tuning on the Birds-200 dataset. Results are illustrated in Fig. 11a-11b for each feature map in layers 1 and 13, respectively. To compute the oracle estimate for a feature map, we remove the feature map and compute the network prediction for each image in the training set using the central crop with no data augmentation or dropout. We draw the following conclusions:
• The contribution of feature maps range from positive (above the red line) to slightly negative (below the red line), implying the existence of some feature maps which decrease the training cost when removed.
• There are many feature maps with little contribution to the network output, indicated by almost zero change in loss when removed.
• Both layers contain a small number of feature maps which induce a significant increase in the loss when removed.
Table 3 contains a layer-by-layer listing of Spearman’s rank correlation of several criteria with the ranking of oracle-abs. In this more detailed comparison, we see the Taylor criterion shows higher correlation for all individual layers. For several methods including Taylor, the worst correlations are observed for the middle of the network, layers 5-10. We also evaluate several techniques for normalization of the raw criteria values for comparison across layers. The table shows the best performance is obtained by `2 normalization, hence we select it for our method.
A.4 COMPARISON WITH WEIGHT REGULARIZATION
Han et al. (2015) find that fine-tuning with high `1 or `2 regularization causes unimportant connections to be suppressed. Connections with energy lower than some threshold can be removed on the assumption that they do not contribute much to subsequent layers. The same work also finds that thresholds must be set separately for each layer depending on its sensitivity to pruning. The procedure to evaluate sensitivity is time-consuming as it requires pruning layers independently during evaluation.
The idea of pruning with high regularization can be extended to removing the kernels for an entire feature map if the `2 norm of those kernels is below a predefined threshold. We compare our approach with this regularization-based pruning for the task of pruning the last convolutional layer of VGG-16 fine-tuned for Birds-200. By considering only a single layer, we avoid the need to compute layerwise sensitivity. Parameters for optimization during fine-tuning are the same as other experiments with the Birds-200 dataset. For the regularization technique, the pruning threshold is set to σ = 10−5 while we vary the regularization coefficient γ of the `2 norm on each feature map kernel.5 We prune only kernel weights, while keeping the bias to maintain the same expected output.
A comparison between pruning based on regularization and our greedy scheme is illustrated in Fig. 12. We observe that our approach has higher test accuracy for the same number of remaining unpruned feature maps, when pruning 85% or more of the feature maps. We observe that with high regularization all weights tend to zero, not only unimportant weights as Han et al. (2015) observe in the case of ImageNet networks. The intuition here is that with regularization we push all weights down and potentially can affect important connections for transfer learning, whereas in our iterative procedure we only remove unimportant parameters leaving others untouched.
A.5 COMBINATION OF CRITERIA
One of the possibilities to improve saliency estimation is to combine several criteria together. One of the straight forward combinations is Taylor and mean activation of the neuron. We compute the joint criteria as Θjoint(z (k) l ) = (1− λ)Θ̂Taylor(z (k) l ) + λΘ̂Activation(z (k) l ) and perform a grid search of parameter λ in Fig.13. The highest correlation value for each dataset is marked with with vertical bar with λ and gain. We observe that the gain of linearly combining criteria is negligibly small (see ∆’s in the figure).
5In our implementation, the regularization coefficient is multiplied by the learning rate equal to 10−4.
A.6 OPTIMAL BRAIN DAMAGE IMPLEMENTATION
OBD computes saliency of a parameter by computing a product of the squared magnitude of the parameter and the corresponding element on the diagonal of the Hessian. For many deep learning frameworks, an efficient implementation of the diagonal evaluation is not straightforward and approximation techniques must be applied. Our implementation of Hessian diagonal computation was inspired by Dauphin et al. (2015) work, where the technique proposed by Bekas et al. (2007) was used to evaluate SGD preconditioned with the Jacobi preconditioner. It was shown that diagonal of the Hessian can be approximated as:
diag(H) = E[v Hv] = E[v ∇(∇C · v)], (13)
where is the element-wise product, v are random vectors with entries ±1, and ∇ is the gradient operator. To compute saliency with OBD, we randomly draw v and compute the diagonal over 10 iterations for a single minibatch for 1000 mini batches. We found that this number of mini batches is required to compute close approximation of the Hessian’s diagonal (which we verified). Computing saliency this way is computationally expensive for iterative pruning, and we use a slightly different but more efficient procedure. Before the first pruning iteration, saliency is initialized from values computed off-line with 1000 minibatches and 10 iterations, as described above. Then, at every minibatch we compute the OBD criteria with only one iteration and apply an exponential moving averaging with a coefficient of 0.99. We verified that this computes a close approximation to the Hessian’s diagonal.
A.7 CORRELATION OF TAYLOR CRITERION WITH GRADIENT AND ACTIVATION
The Taylor criterion is composed of both an activation term and a gradient term. In Figure 14, we depict the correlation between the Taylor criterion and each constituent part. We consider expected absolute value of the gradient instead of the mean, because otherwise it tends to zero. The plots are computed from pruning criteria for an unpruned VGG network fine-tuned for the Birds-200 dataset. (Values are shown after layer-wise normalization). Figure 14(a-b) depict the Taylor criterion in the y-axis for all neurons w.r.t. the gradient and activation components, respectively. The bottom 10% of neurons (lowest Taylor criterion, most likely to be pruned) are depicted in red, while the top 10% are shown in green. Considering all neurons, both gradient and activation components demonstrate a linear trend with the Taylor criterion. However, for the bottom 10% of neurons, as shown in Figure 14(c-d), the activation criterion shows much stronger correlation, with lower activations indicating lower Taylor scores.

Authors propose a strategy for pruning weights with the eventual goal of reducing GFLOP computations. The pruning strategy is well motivated using the taylor expansion of the neural network function with respect to the feature activations. The obtained strategy removes feature maps that have both a small activation and a small gradient (eqn 7).
(A) Ideally the gradient of the output with respect to the activation functions should be 0 at the optimal, but as a result of stochastic gradient evaluations this would practically never be zero. Small variance in the gradient across mini-batches indicates that irrespective of input data the specific network parameter is unlikely to change - intuitively these are parameters that are closer to convergence. Parameters/weights that are close to convergence and also result in a small activation are intuitively good candidates for pruning. This is essentially what eqn 7 conveys and is likely to be reason why just removing weights that result in small activations is not as good of a pruning strategy (as shown by results in the paper). There are two kind of differences in weights that are removed by activation v/s taylor expansion:
1. Weights with high-activations but very low gradients will be removed by taylor expansion, but not by activation alone.
2. Weights with low-activation but high gradients will be removed by activation criterion, but not by taylor expansion.
It will be interesting to analyze which of (1) or (2) contribute more to the differences in weights that are removed by the taylor expansion v/s activation criterion. Intuitively it seems that weight that satisfy (1) are important because they are converged and contribute significantly to network's activation. It is possible that a modified criterion - eqn (7) + \lambda feature activation, (where \lambda needs to be found by cross-validation) may lead to even better results at the cost of more parameter tuning.

(B) Another interesting comparison is with the with the optimal damage framework - where the first order gradients are assumed to be zero and pruning is performed using the second-order information (also discussed by authors in the appendix). Critically, only the diagonal of the Hessian is computed. There is no comparison with optimal damage as authors claim it is memory and computation inefficient. Back of envelope calculations suggest that this would result only in 50% increase in memory and computation during pruning, but no loss in efficiency during testing. Therefore from a standpoint of deployment, I don't think this missing comparison is justified.
(C) The eventual goal of the authors is to reduce GFLOPs. Some recent papers have proposed using lower precision computation for this. A comparison in GFLOPs with lower precision v/s pruning would be a great. While both these approaches are complementary and it is expected that combining both of them can lead to superior performance than either of the two - it is unclear when we are operating in the low-precision regime how much pruning can be performed. Any analysis on this tradeoff would be great (but not necessary).
(D) On finetuning, authors report results of AlexNet and VGG on two different datasets - Flowers and Birds respectively. Why is this the case? It would be great to see the results of both the networks on both the datasets.
(E) Authors report there is only a small drop in performance after pruning. Suppose the network was originally trained with N iterations, and then M finetuning iterations were performed during pruning. This means that pruned networks were trained for N + M iterations. The correct comparison in accuracies would be if we the original network was also trained for N + M iterations. In figure 4, does the performance at 100% parameters reports accuracy after N+M iterations or after N iterations?
Overall I think the paper is technically and empirically sound, it proposes a new strategy for pruning:
(1) Based on taylor expansion
(2) Feature normalization to reduce parameter tuning efforts.
(3) Iterative finetuning.
However, I would like to see some comparisons mentioned in my comments above. If those comparisons are made I would change my ratings to an accept.


NEURAL FUNCTIONAL PROGRAMMING
We discuss a range of modeling choices that arise when constructing an end-to-end differentiable programming language suitable for learning programs from inputoutput examples. Taking cues from programming languages research, we study the effect of memory allocation schemes, immutable data, type systems, and builtin control-flow structures on the success rate of learning algorithms. We build a range of models leading up to a simple differentiable functional programming language. Our empirical evaluation shows that this language allows to learn far more programs than existing baselines.
1 INTRODUCTION
Inductive Program Synthesis (IPS), i.e., the task of learning a program from input/output examples, is a fundamental problem in computer science. It is at the core of empowering non-experts to use computers for repeated tasks, and recent advances such as the FlashFill extension of Microsoft Excel (Gulwani, 2011) have started to deliver on this promise.
A related line of research is the extension of neural network architectures with components that correspond to hardware primitives (Giles et al., 1989; Graves et al., 2014; Weston et al., 2015; Joulin & Mikolov, 2015; Grefenstette et al., 2015; Kurach et al., 2016; Kaiser & Sutskever, 2016; Reed & de Freitas, 2016; Andrychowicz & Kurach, 2016; Zaremba et al., 2016; Graves et al., 2016), enabling them to learn program-like behavior. However, these models are usually tightly coupled to the idea of a differentiable interpretation of computer hardware, as names such as Neural Turing Machine (Graves et al., 2014), Neural Random-Access Machine (Kurach et al., 2016), and Neural GPU (Kaiser & Sutskever, 2016) indicate. We observe that while such architectures form the basis modern computing, they are usually not the models that are used to program computers. Instead, decades of programming languages research have lead to ever higher programming languages that aim to make programming simpler and less error-prone. Indeed, as recent comparisons show (Gaunt et al., 2016b), program synthesis methods from the programming languages community that actively exploit such constructs, e.g. by leveraging known semantics of loops, are currently achieving considerably better results than comparable neural architectures. Still, neural IPS techniques are clearly at an advantage when extending the problem setting from simple integer input/output examples to more complex cases, such as IPS problems with perceptual data (Gaunt et al., 2016a), imprecise examples, or leveraging additional cues such as a natural language description of the desired program.
Hence, we propose to adapt features of modern high-level programming languages to the differentiable setting. In this paper, we develop an end-to-end differentiable programming language operating on integers and lists, taking cues from functional programming. In our empirical evaluation, we show the effects on learning performance of our four modeling recommendations, namely automatic memory management, the use of combinators and if-then-else constructs to structure program control flow, immutability of data, and an application of a simple type system. Our experiments show that each of these features crucially improves program learning over existing baselines.
2 BACKGROUND
The basic building block of functional programs is the function, and programs are built by composing functions together. In the following, we highlight some common features in functional programs before discussing how to integrate them into an end-to-end differentiable model in Sect. 3.
Memory Management Most modern programming languages eschew manual memory management and pointer manipulation where possible. Instead, creation of heap objects automatically generates an appropriate pointer to fresh memory. Similarly, built-in constructs allow access to fields of objects, instead of requiring pointer arithmetic. Both of these choices move program complexity into the fixed implementation of a programming language, making it easier to write correct programs.
Immutable Data Functions are expected to behave like their mathematical counterparts, avoiding mutable data and side effects. This helps programmers reason about their code, as it eliminates the possibility that a variable might be left uninitialized or accessed in an inconsistent state. Moreover, no data is ever “lost” by being overwritten or mutated.
Types Expressive type systems are used to protect programmers from writing programs that will fail. Practically, a type checker is able to rule out many syntactically correct programs that are certain to fail at runtime, and thus restricts the space of valid programs. Access to types helps programmers to reason about the behavior of their code. In particular, the type system tells the programmer what kinds of data they can expect each variable to contain.
Structured Control Flow A key difference between hardware-level assembly languages and higher-level programming languages is that higher-level languages structure control flow using loops, conditional statements, and procedures, as raw gotos are famously considered harmful (Dijkstra, 1968). Functional languages go a step further and leverage higher-order functions to abstract over common control flow patterns such as iteration over a recursive data structure. In an imperative language, such specialized control flow is often repeated and mixed with other code.
3 OUR MODELS
In the following, we will discuss a range of models, starting with a simple assembly-like language and progressing to a differentiable version of a simple functional programming language. We make four modeling recommendations whose effect we demonstrate in our experiments in Sect. 4.
We first discuss the general format of our programs and program states, which we will refine step by step. Our programs operate on states consisting of an instruction pointer indicating the next instruction to execute, a number of registers holding inputs and intermediate results of executed instructions, and a heap containing memory allocated by the program. We focus on list-manipulating programs, so we create a heap consisting of standard cons cells, which are data and pointer value pairs where the pointer points to another cons cell or the special nil value. To represent a linked list, each cell points to the next cell in the list, except for the last cell, which points to nil.
3.1 PROGRAM AND DATA REPRESENTATION
We define our models by lifting simple instructions to the differentiable setting. To do so, we bound the domain of all values and parameters, following earlier work (e.g. (Graves et al., 2014; Kurach et al., 2016; Gaunt et al., 2016b)). We represent a value v from a domain {d1 . . . dD} as a tuple RD, interpreted as a discrete probability distribution. We pick a maximal integer value M that bounds all values occurring in our programs, a number of instructions I , and a number of registers R. In this setting, the size of the heap memory H has to be equal to the maximal integer value M , but we will relax this later. We limit the length of programs to some value P , and can then encode programs as a sequence of tuples (o(p), i (p), a(p)1 , a (p) 2 ), where i
(p) ∈ [1, I ] identifies the p-th instruction and o(p), a
(p) 1 , a (p) 2 ∈ [1,R] its output and argument registers respectively. To “execute” such a program,
we unroll it for T timesteps and keep a program state s(t) = (p(t), r (t)1 . . . r (t) R , h (t) 1 . . . h (t) H ) for each timestep t , where p(t) ∈ [1,P ] is an instruction pointer indicating which instruction to execute next, r (t)∗ are the values of registers, and h (t) ∗ are the values of the cons cells in the heap.
All of our models share a basic instruction set, namely the cons cell constructor cons, the heap accessors (car & cdr) which return the data (resp. pointer) element of a cons cell, integer addition, increment and decrement (add, inc, dec), integer equality and greater-than comparison (eq & gt), Boolean conjunction and disjunction (and & or), common constants (zero & one), and finally a
noop instruction. These all have the usual semantics as transformers on the program state, and we will discuss the behavior of cons in detail later. For example, executing (o, add, a1, a2) on a state at timestep t yields the following registers at the next timestep, where the addition operation is lifted to operate on distributions over natural numbers.
r (t+1)u =
{ r (t) a1 + r (t) a2 mod M if u = o
r (t) u otherwise.
∀u ∈ [1,R]
As we allow all involved quantities to be distributions over all possible choices, computing the next state requires a case analysis for all allowed values. The new state is then obtained by averaging the results of all possible execution steps, weighted by the probabilities assigned to each choice. Thus, if η(s(t), (o, i , a1, a2)) computes the state obtained by executing the instruction (o, i , a1, a2), we can compute the next state s(t+1) as follows, where Jx = nK denotes the probability that a variable x encoding a discrete probability distribution assigns to the value n.
s(t+1) = ∑
p∈[1,P ],i∈[1,I ], o,a1,a2∈[1,R]
Jp(t) =pK · Jo(p)=oK · Ji (p)= iK · Ja(p)1 =a1K · Ja (p) 2 =a2K · η(s(t), (o, i , a1, a2)) (1)
In practice, we developed our models in TerpreT (Gaunt et al., 2016b), which hides these technicalities.
Training Objective Our aim is to learn the program parameters (o(p), i (p), a(p)1 , a (p) 2 ) such that program “evaluation” according to (1) starting on a state s(0) initialized to an example input yields the target output in s(T). For scalar outputs such as a sum of values, our objective is simply to minimize the cross-entropy between the distribution in the output register r(T)R and a point distribution with all probability mass on the correct output value.
Handling list outputs is more complex, as many valid outputs exist (depending on how list elements are placed in the heap memory). Intuitively, we traverse the heap from the returned heap address until reaching the end of a linked list, recording the list elements as we go. To formalize this intuition, let dh(T)k (resp. ph (T) k ) denote the data (resp. pointer) information in the heap cell at address k at the final state of the evaluation. We then compute the traversal sequences of list element values v1, . . . , vH and addresses a1, . . . , aH as follows.
ai = { r (T) R if i = 1∑ a∈[1,H ]Jai−1 = aK · ph (T) a otherwise
vi = ∑
a∈[1,H ]
Jai = aK · dh(T)a
The probability that the computed output list is equal to an expected output list [v̄1, . . . , v̄k] is then Jak+1 = 0K · ∑k i=1Jvi = v̄iK.
Memory Management As the programs we want to learn need to construct new lists, we need a memory allocation mechanism that provides fresh cells. We explored two options for this allocator.
First, we attempt to follow stack-allocation models in which a stack of memory cells is used with a stack pointer sp which always points to the next free memory cell. We fix a maximum stack size H . Whenever a memory cell is allocated (i.e., a cons instruction is executed), the stack pointer is incremented, guaranteeing that no cell is ever overwritten. However, uncertainty about whether an instruction is cons translates into uncertainty about the precise value of the stack pointer, as each call to cons changes sp. This uncertainty causes cells holding results from different instructions in the stack to blur together, despite the fact that cells are immutable once created. As an example, consider the execution of two instructions, where the first is cons 1 0 with probability 0.5 and noop otherwise, and the second is cons 2 0 with probability 0.5 and noop otherwise. After executing starting with sp = 1 and an empty stack, the value of sp will be blurred across three values 1, 2 and 3 with probabilities 0.25, 0.5 and 0.25. Similarly, the value of the first heap cell will be 0 (the default) with probability 0.25, 1 with probability 0.5 and 2 with probability 0.25. This blurring effect becomes stronger with longer programs, and we found that it substantially impacted learning.
Both of these problems can be solved by transitioning to a fully immutable representation of the heap. In this variant, we allocate and initialize one heap cell per timestep, i.e., we set H = T . If the
current instruction is a cons, the appropriate values are filled in, otherwise both data and pointer value are set to a default value (in our case, 0). This eliminates the issue of blurring between outputs of different instructions. The values of a cons cell may still be uncertain as they inherit uncertainty about the executed instructions and the values of arguments, but depend only on the operations at one timestep. While this modification requires a larger domain to store pointers, we found not copying the stack significantly reduces memory usage during training of our models.
Recommendation (F): Use fixed heap memory allocation deterministically controlled by the model.
3.2 PROGRAM MODELS
1: out ← instr arg1 arg2 branch 2: out ← instr arg1 arg2 branch
. . . Our baseline program model corresponds closely to an assembly language as used in earlier work (Bunel et al., 2016), resulting in a program model as shown on the right, where boxes correspond to learnable parameters. We extend our instruction set with jump-if-zero (jz), jump-if-not-zero (jnz) and return instructions. Our assembly program representation also includes a “branch” parameter b specifying the new value of the instruction pointer for a successful conditional branch. To learn programs in this language, the model must learn how to create the control flow that it needs using these simple conditional jumps. Note that the instruction pointer suffers from the same problems as the stack pointer above, i.e., uncertainty about its value blurs together the effects of many possible program executions.
pre1: out ← instr arg1 arg2 cond
pre2: . . . foreach ele in reg
loop1: out ← instr arg1 arg2 cond
loop2: . . .
suf 1: out ← instr arg1 arg2 cond
suf 2: . . . Structured Control Flow We see structured control flow as a way to reduce the “bleeding” of uncertainty about the value of the instruction pointer into the values of registers and cells on the heap. To introduce structured control flow, we replace raw jumps with an if-then-else instruction and an explicit foreach loop that is suited for processing lists. We restrict our model to a prefix of instructions, a loop which iterates over a list, and a suffix of instructions. The parameters for instructions in the loop can access an additional register that contains the value of the current list element. To implement this behavior in practice, we unroll the loop for a fixed number of iterations derived from the bound on the size of the input, which ensures that every input list can be processed. After unrolling, the instruction executed at each timestep becomes deterministic, removing uncertainty about the value of the instruction pointer.
For the if-then-else instruction, we extend the instruction representation with a “condition” parameter c ∈ [1,R] and let the evaluation of if-then-else yield its first argument when the register c is non-zero and the second argument otherwise. An overview of the structure of such programs is displayed above. foldli:
acc ← reg
idx ← 0 foreach ele in reg
out ← instr arg1 arg2 cond
. . .
acc ← reg
idx ← idx + 1 out
← acc
We note that while fixing the iteration over the list elements is already helpful, learning most list-processing programs requires the model to repeatedly infer the concepts of creating a new list, aggregating results and keeping track of the current list index. In functional programming languages, such regular patterns are encapsulated in combinators. Thus, in a second model, we replace the simple foreach loop with three combinators: mapi creates a new list by applying a function to each element of the input list, zipWithi creates a new list by iterating over two lists in parallel and applying a function to both elements, and foldli computes a value by iterating over all list elements
and applying a function to the current list element and the value computed so far. A program model using the foldli combinator is shown on the left. The i suffix indicates that these combinators additionally provide the index of the current list element (the precise semantics of the combinators are presented in Sect. A.1).
Recommendation (L): Instead of raw jumps, use loop and if-then-else templates.
Immutable Data In training our models, we observed that many random initializations of the program parameters would overwrite input data or important intermediate results, and later steps would not be able to recover this data. In models with combinators that provide a way to accumulate result values, we can sidestep this issue by making registers immutable. To do so, we create one register per timestep, and fix the output of each instruction to the register for its timestep. Parameters for arguments then range over all registers initialized in prior timesteps, with an exception for the closures executed by a combinator. Here, each instruction only gets access to the inputs to the closure, values computed in the prefix, and registers initialized by preceding instructions in the same loop iteration. As in the heap allocation case, we can avoid keeping a copy of all registers for every timestep, and instead share these values over all steps, reducing memory usage.
Recommendation (I): Use immutable registers by deterministically choosing where to store outputs.
Types When training our models, we found that for many initializations, training would fall into local minima corresponding to ill-typed programs, e.g., where references to the heap would be used in integer additions. We expect the learned program to be well-typed, so we introduce a simple type system. We explored two approaches to adding a type system.
A first attempt integrated the well-typedness of the program into our objective function. In our programs, we use three simple types of data—integers, pointers and booleans—as well as a special type,⊥, which represents type errors. We extended the program state to contain an additional element tr for each register, encoding its type. Each instruction then not only computes a value that is assigned to the target register, but also a type for the target register. Most significantly, if one of the arguments has an unsuitable type (e.g., an integer in place of a pointer), the resulting type is⊥. We then extended our objective function to add a penalty for values with type ⊥. Unfortunately, this changed objective function had neither a positive nor negative effect on our experiments, so it seems that optimizing for the correct type is redundant when we are already optimizing for the correct return value.
In our second attempt, rather than penalizing ill-typed programs, we prevent programs from accessing ill-typed data by construction. We augment our register representation by adding an integer, pointer, and Boolean slot to each register, so each register can hold a separate value of each type. Instructions which read from registers now read from the slot corresponding to the type of the argument. When writing to a register, we write to the slot corresponding to the instruction’s return type, and set the other slots to a default value 0. This prevents any ill-typed sequence of instructions, i.e., it is now impossible to, for example, increment a pointer value or to construct a cons cell with a non-pointer value. Furthermore, this modification allows us to set the heap size H to a value different from the maximal integer M . Our experiments in Sect. 4.3 show that separating differently-typed values simplifies the learning of programs that operate on lists and integers at the same time.
Recommendation (T): Use different storage for data of different types.
4 EXPERIMENTS
We have empirically evaluated our modeling recommendations on a selection of program induction tasks of increasing complexity, ranging from straight-line programs to problems with loops and conditional expressions. All of our models are implemented in TerpreT (Gaunt et al., 2016b) and we learn using TerpreT’s TENSORFLOW (Abadi et al., 2015) backend. We aim to release TerpreT, together with these models, under an open source license in the near future.
For all tasks, three groups of five input/output example pairs were sampled as training data and another 25 input/output pairs as test data. For each group of five examples, training was started from 100 random initializations of the model parameters. After training for 3500 epochs (tests with longer training runs showed no significant changes in the outcomes), the learned programs were tested by discretizing all parameters and comparing program outputs on test inputs with the expected values. We perform 300 runs per model and task, and report only the ratio of successful runs. A run is successful if the discretized program returns the correct result on all five training and 25 test examples.1 The ratio of runs converging to zero loss on the training examples was within 1% of the number of successful runs, i.e., very few found solutions failed to generalize.
1We inspected samples of the obtained programs as well and verified that they were indeed correct solutions. See Sect. A.2 for some of the learned programs.
We performed a cursory exploration of hyperparameter choices. We varied the choice of optimization algorithm (Momentum, Adam, RMSProp), the learning rate (from 0.001 to 5), gradient noise (testing the recommended choices from Neelakantan et al. (2016b)), a decaying entropy bonus (starting from 0.001 to 20), and gradient clipping (to values between 0.1 and 10). We sampled 100 hyperparameter settings from this space and tested their effect on two simple tasks. We ran the remaining experiments with the best configuration obtained by this process: the RMSProp optimization algorithm, a learning rate of 0.1, clipped gradients at 1, and no gradient noise.
We consider the ratio of successful runs as earlier work has identified this as a significant problem. For example, Neelakantan et al. (2016b) reports that even after a (task-specific) “large grid search” of hyperparameters, the Neural Random Access Machine converged only in 5%, 7% and 22% of random restarts. Similar observations were made in Kaiser & Sutskever (2016); Bunel et al. (2016); Gaunt et al. (2016b) for related program learning models.
In our experiments we evaluate the effect of the choices discussed in Sect. 3, comparing seven model variants in total. We call our initial assembly model A and its variation with a fixed memory allocation scheme A+F. All other models use the same fixed memory allocation scheme. The extension of the assembly model with a built-in foreach loop is called A+L. The A+L model also allows a foreachZip loop structure that allows parallel iteration over two lists, similar to the zipWith combinator. The model including predefined combinators is called C, where C+I (resp. C+T) are its extensions with immutable registers (resp. typed registers). Finally, C+T+I combines all of these and is, in effect, a simple end-to-end differentiable functional programming language.
Additionally, we show results for λ2 (Feser et al., 2015), a strong program synthesis baseline from programming languages research, because of its built-in support for list-processing programs. As λ2 is deterministic, we only report a success rate of either 1 or 0.
4.1 STRAIGHT-LINE PROGRAMS
In our first experiment, we consider two families of simple problems—solvable with straight-line programs—to study the interaction of our modeling choices with program length. Our first benchmark task is to duplicate a scalar input a fixed number k times to create a list of length k. Our second benchmark is to retrieve the k-th element of a list, again fixing k beforehand (we will consider a generalization of this task where k is a program input later). We set the hyperparameters for all models to allow 11 statements, i.e., for A and A+F we have set the program length to 11, and for the A+L and C * models we have set the prefix and loop length to 0 and the suffix length to 11. For models where the number of registers does not depend on the number of timesteps, we set the number of registers to 3, with one initialized to the input. This allows for ∼ 1039 programs in the A, A+F, C+I, and C+T+I models and for ∼ 1028 programs in the remaining models. These parameters were chosen to be slightly larger than required by the largest program to be learned. For all of our experiments, the maximal integer M was set to 20 for models where possible (i.e., for A, C+T+I, C+T), and to H (derived from T , coming to 22) for the others.2
2We also experimented with varying the value of M . Choices over 20 showed no significant differences to smaller values.
We evaluated all of our models following the regime discussed above and present the results in Fig. 1 for k values from 1 to 9. The difference between A and A+F on the dupK task illustrates the significance of Recommendation (F) to fix the memory allocation scheme. Following Recommendation (T) to separate values of different types improves the results on both tasks, as the differences between C+T+I (resp. C+T) and C+I (resp. C) illustrate.
4.2 SIMPLE LOOP PROGRAMS
In our second experiment, we compare our models on three simple list algorithms: computing the length of a list, reversing a list and summing a list. Model parameters have been set to allow 6 statements for the A and A+F models, and empty prefixes, empty suffixes, and 2 instructions in the loop for the other models. For models where the number of registers does not depend on the number of timesteps, we set the number of registers to 4, with one initialized to the input.
The results of our evaluation are displayed in Tab. 1, starkly illustrating Recommendation (L) to use predefined loop structures. We speculate that learning explicit jump targets is extremely challenging because changes to the parameters controlling jump target instructions have outsized effects on all computed (intermediate and output) values. On the other hand, models that could choose between different list iteration primitives were able to find programs for all tasks. We again note the effect of Recommendation (T) to separate values of different types on the success rates for the len and sum examples, and the effect of Recommendation (I) to avoid mutable data on results for len and rev.
4.3 LOOP PROGRAMS
In our main experiment, we consider a larger set of common list-manipulating tasks (such as checking if all/one element of a list is greater than a bound, retrieving a list element by index, finding the index of a value, and computing the maximum value). Descriptions of all tasks are shown in Fig. 2 in the appendix. We do not show results for the A and A+F models, which always fail. We set the parameters for the remaining models to M = 32 where possible (M = H = 34 for the others), the length of the prefix to 1, the length of the closure / loop body to 3 and the length of the suffix to 2. Again, these parameters are slightly larger than required by the largest program to be learned.
The results for our experiments on these tasks are shown in Tab. 2. Note the changed results of the examples from Sect. 4.2, as the change in model parameters has increased the size of the program space from ∼ 107 to ∼ 1020. The relative results for the A+L model show the value of built-in iteration and aggregation patterns. The choice between immutable and mutable registers is less clear here, seemingly dampened by other influences. An inspection of the generated programs (eg. Fig. 8 in
the appendix) reveals that mutability of registers can sometimes be exploited to find elegant solutions. Overall, it may be effective to combine both approaches, using a small number of (mutable) “scratch value” registers and immutable default output registers for each statement.
5 RELATED WORK
Inductive Program Synthesis There has been significant recent interest in synthesizing functional programs from input-output examples in the programming languages community. Synthesis systems generally operate by searching for a program which is correct on the examples, using types or custom deduction rules to eliminate parts of the search space. Among the notable systems: MYTH (Osera & Zdancewic, 2015; Frankle et al., 2016) synthesizes recursive functional programs from examples using types to guide the search for a correct program, λ2 (Feser et al., 2015) synthesizes data structure manipulating programs structured using combinators using types and deduction rules in its search, ESCHER (Albarghouthi et al., 2013) synthesizes recursive programs using search and a specialized method for learning conditional expressions, and FlashFill (Gulwani, 2011) structures programs as compositions of functions and uses custom deduction rules to prune candidate programs. Our decision to learn functional programs was strongly inspired by this previous work. In particular, the use of combinators to structure control flow was drawn from Feser et al. (2015). However, our end-to-end differentiable setting is fundamentally different from discrete search employed in the programming languages community, and thus concrete techniques are largely incomparable.
Neural Networks Learning Algorithms A number of recent models aim to learn algorithms from input/output data. Many of these augment standard recurrent neural network architectures with differentiable memory and simple computation components (e.g. Graves et al. (2014); Kurach et al. (2016); Joulin & Mikolov (2015); Neelakantan et al. (2016a); Reed & de Freitas (2016); Zaremba et al. (2016); Graves et al. (2016)). The use of an RNN can be seen as fixed looping structure, and the use of fixed output registers for the modules in Neural Random Access Machines (Kurach et al., 2016) is similar to our modeling of immutable registers.
However, none of these works focus on producing source code. Gaunt et al. (2016b) show that this is an extremely challenging task for assembly-like program models. More recently, Bunel et al. (2016) and Riedel et al. (2016) have used program models similar to assembly (resp. Forth) source code to initialize solutions, and either optimize or complete them.
6 DISCUSSION AND FUTURE WORK
We have discussed a range of modeling choices for end-to-end differentiable programming languages and made four design recommendations. Empirically, we have shown these recommendations to significantly improve the success ratio of learning programs from input/output examples, and we expect these results to generalize to similar models attempting to learn programs.
In this paper, we only consider list manipulating programs, but are interested in supporting more data structures, such as arrays (which should be a straightforward extension) and associative maps. We also only support loops over lists at this time, but are interested in extending our models to also have built-in support for loops counting up to (and down from) integer values. A generalization of this concept would be an extension allowing the learning and use of recursive functions. Recursion is still more structured than raw goto calls, but more flexible than the combinators that we currently employ. An efficient implementation of recursion is a challenging research problem, but it could allow significantly more complex programs to be learned. Modeling recursion in an end-to-end differentiable language could allow us to build libraries of (learned) differentiable functions that can be used in later synthesis problems.
However, we note that with few exceptions on long straight-line code, λ2 performs better than all of our considered models, and is able to synthesize programs in milliseconds. We see the future of differentiable programming languages in areas in which deterministic tools are known to perform poorly, such as the integration of perceptual data, priors and “soft” side information such as natural language hints about the desired functionality. Gaunt et al. (2016a) was developed in parallel to this work and builds on many of our results to learn programs that can process perceptual data (in the current example, images).
A APPENDIX
A.2 SELECTED SOLUTIONS
We show example results of our training in Figs. 4-16. Note that these are the actual results produced by our system, and have only been slightly edited for typesetting. Finally, we have colored statements that a simple program analysis can identify as not contributing to the result in gray.

This paper presents design decisions of TerpreT [1] and experiments about learning simple loop programs and list manipulation tasks. The TerpreT line of work (is one of those which) bridges the gap between the programming languages (PL) and machine learning (ML) communities. Contrasted to the recent interest of the ML community for program induction, the focus here is on using the design of the programming language to reduce the search space. Namely, here, they used the structure of the control flow (if-then-else, foreach, zipWithi, and foldli "templates"), immutable data (no reuse of a "neural" memory), and types (they tried penalizing ill-typedness, and restricting the search only to well-typed programs, which works better). My bird eye view would be that this stands in between "make everything continuous and perform gradient descent" (ML) and "discretize all the things and perform structured and heuristics-guided combinatorial search" (PL).
I liked that they have a relevant baseline (\lambda^2), but I wished that they also included a fully neural network program synthesis baseline. Admittedly, it would not succeed except on the simplest tasks, but I think some of their experimental tasks are simple enough for "non-generating code " NNs to succeed on.
I wished that TerpreT was available, and the code to reproduce these experiments too.
I wonder if/how the (otherwise very interesting!) recommendations for the design of programming languages to perform gradient descent based-inductive programming would hold/perform on harder task than these loops. Even though these tasks are already interesting and challenging, I wonder how much of these tasks biased the search for good subset of constraints (e.g. those for structuring the control flow).
Overall, I think that the paper is good enough to appear at ICLR, but I am no expert in program induction / synthesis.
Writing:
 - The paper is at times hard to follow. For instance, the naming scheme of the model variants could be summarized in a table (with boolean information about the features it embeds).
 - Introduction: "basis modern computing" -> of
 - Page 3, training objective: "minimize the cross-entropy between the distribution in the output register r_R^{(T)} and a point distribution with all probability mass on the correct output value" -> if you want to cater to the ML community at large, I think that it is better to say that you treat the output of r_R^{(T)} as a classification problem with the correct output value (you can give details and say exactly which type of criterion/loss, cross-entropy, you use).
[1] "TerpreT: A Probabilistic Programming Language for Program Induction", Gaunt et al. 2016