METACONTROL FOR ADAPTIVE IMAGINATION-BASED OPTIMIZATION
Many machine learning systems are built to solve the hardest examples of a particular task, which often makes them large and expensive to run—especially with respect to the easier examples, which might require much less computation. For an agent with a limited computational budget, this “one-size-fits-all” approach may result in the agent wasting valuable computation on easy examples, while not spending enough on hard examples. Rather than learning a single, fixed policy for solving all instances of a task, we introduce a metacontroller which learns to optimize a sequence of “imagined” internal simulations over predictive models of the world in order to construct a more informed, and more economical, solution. The metacontroller component is a model-free reinforcement learning agent, which decides both how many iterations of the optimization procedure to run, as well as which model to consult on each iteration. The models (which we call “experts”) can be state transition models, action-value functions, or any other mechanism that provides information useful for solving the task, and can be learned on-policy or off-policy in parallel with the metacontroller. When the metacontroller, controller, and experts were trained with “interaction networks” (Battaglia et al., 2016) as expert models, our approach was able to solve a challenging decision-making problem under complex non-linear dynamics. The metacontroller learned to adapt the amount of computation it performed to the difficulty of the task, and learned how to choose which experts to consult by factoring in both their reliability and individual computational resource costs. This allowed the metacontroller to achieve a lower overall cost (task loss plus computational cost) than more traditional fixed policy approaches. These results demonstrate that our approach is a powerful framework for using rich forward models for efficient model-based reinforcement learning.
Many machine learning systems are built to solve the hardest examples of a particular task, which often makes them large and expensive to run—especially with respect to the easier examples, which might require much less computation. For an agent with a limited computational budget, this “one-size-fits-all” approach may result in the agent wasting valuable computation on easy examples, while not spending enough on hard examples. Rather than learning a single, fixed policy for solving all instances of a task, we introduce a metacontroller which learns to optimize a sequence of “imagined” internal simulations over predictive models of the world in order to construct a more informed, and more economical, solution. The metacontroller component is a model-free reinforcement learning agent, which decides both how many iterations of the optimization procedure to run, as well as which model to consult on each iteration. The models (which we call “experts”) can be state transition models, action-value functions, or any other mechanism that provides information useful for solving the task, and can be learned on-policy or off-policy in parallel with the metacontroller. When the metacontroller, controller, and experts were trained with “interaction networks” (Battaglia et al., 2016) as expert models, our approach was able to solve a challenging decision-making problem under complex non-linear dynamics. The metacontroller learned to adapt the amount of computation it performed to the difficulty of the task, and learned how to choose which experts to consult by factoring in both their reliability and individual computational resource costs. This allowed the metacontroller to achieve a lower overall cost (task loss plus computational cost) than more traditional fixed policy approaches. These results demonstrate that our approach is a powerful framework for using rich forward models for efficient model-based reinforcement learning.
1 INTRODUCTION
While there have been significant recent advances in deep reinforcement learning (Mnih et al., 2015; Silver et al., 2016) and control (Lillicrap et al., 2015; Levine et al., 2016), most efforts train a network that performs a fixed sequence of computations. Here we introduce an alternative in which an agent uses a metacontroller to choose which, and how many, computations to perform. It “imagines” the consequences of potential actions proposed by an actor module, and refines them internally, before executing them in the world. The metacontroller adaptively decides which expert models to use to evaluate candidate actions, and when it is time to stop imagining and act. The learned experts may be state transition models, action-value functions, or any other function that is relevant to the task, and can vary in their accuracy and computational costs. Our metacontroller’s learned policy can exploit the diversity of its pool of experts by trading off between their costs and reliability, allowing it to automatically identify which expert is most worthwhile.
We draw inspiration from research in cognitive science and neuroscience which has studied how people use a meta-level of reasoning in order to control the use of their internal models and allocation of their computational resources. Evidence suggests that humans rely on rich generative models of the world for planning (Gläscher et al., 2010), control (Wolpert & Kawato, 1998), and reasoning (Hegarty, 2004; Johnson-Laird, 2010; Battaglia et al., 2013), that they adapt the amount of computation they perform with their model to the demands of the task (Hamrick et al., 2015), and that they trade off between multiple strategies of varying quality (Lee et al., 2014; Lieder et al., 2014; Lieder & Griffiths, in revision; Kool et al., in press).
Our imagination-based optimization approach is related to classic artificial intelligence research on bounded-rational metareasoning (Horvitz, 1988; Russell & Wefald, 1991; Hay et al., 2012), which formulates a meta-level MDP for selecting computations to perform, where the computations have a known cost. We also build on classic work by Schmidhuber (1990a;b), which used an RL controller with a recurrent neural network (RNN) world model to evaluate and improve upon candidate controls online.
Recently Andrychowicz et al. (2016) used a fully differentiable deep network to learn to perform gradient descent optimization, and Tamar et al. (2016) used a convolutional neural network for performing value iteration online in a deep learning setting. In other similar work, Fragkiadaki et al. (2015) made use of “visual imaginations” for action planning. Our work is also related to recent notions of “conditional computation” (Bengio, 2013; Bengio et al., 2015), which adaptively modifies network structure online, and “adaptive computation time” (Graves, 2016) which allows for variable numbers of internal “pondering” iterations to optimize computational cost.
Our work’s key contribution is a framework for learning to optimize via a metacontroller which manages an adaptive, imagination-based optimization loop. This represents a hybrid RL system where a model-free metacontroller constructs its decisions using an actor policy to manage model-free and model-based experts. Our experimental results demonstrate that a metacontroller can flexibly allocate its computational resources on a case-by-case basis to achieve greater performance than more rigid fixed policy approaches, using more computation when it is required by a more difficult task.
2 MODEL
We consider a class of fully observed, one-shot decision-making tasks (i.e., continuous, contextual bandits). The performance objective is to find a control c ∈ C which, given an initial state x ∈ X , minimizes some loss function L between a known future goal state x∗ and the result of a forward process, f(x, c). The performance loss LP is the (negative) utility of executing the control in the world, and is related to the optimal solution c∗ ∈ C as follows:
LP (x ∗, x, c) = L(x∗, f(x, c)), (1)
c∗ = argmin c LP (x ∗, x, c). (2)
However, (2) defines only the optimal solution—not how to achieve it.
2.1 OPTIMIZING PERFORMANCE
We consider an iterative optimization procedure that takes x∗ and x as input and returns an approximation of c∗ in order to minimize (1). The optimization procedure consists of a controller, which iteratively proposes controls, and an expert, which evaluates how good those controls are. On the nth iteration, the controller πC : X × X ×H → C takes as input, x∗, x, and information about the history of previously proposed controls and evaluations hn−1 ∈ H, and returns a proposed control cn that aims to improve on previously proposed controls. An expert E : X × X × C → E takes the proposed control and provides some information en ∈ E about the quality of the control, which we call an opinion. This opinion is added to the history, which is passed back to the controller, and the loop continues for N steps, after which a final control cN is proposed.
Standard optimization methods use principled heuristics for proposing controls. In gradient descent, for example, controls are proposed by adjusting cn in the direction of the gradient of the reward with respect to the control. In Bayesian optimization, controls are proposed based on selection criteria such as “probability of improvement”, or a meta-selection criterion for choosing among
several basic selection criteria Hoffman et al. (2011); Shahriari et al. (2014). Rather than choosing one of several controllers, our work learns a single controller and instead focuses on selecting from multiple experts (see Sec. 2.2). In some cases f is known and inexpensive to compute, and thus the optimization procedure sets E ≡ f . However, in many real-world settings, f is expensive or non-stationary and so it can be advantageous to use an approximation of f (e.g., a state transition model), LP (e.g., an action-value function), or any other quantity that gives some information about f or LP .
2.2 OPTIMIZING COMPUTATIONAL COST
Given a controller and one or more experts, there are two important decisions to be made. First, how many optimization iterations should be performed? The approximate solution usually improves
with more iterations, but each iteration costs computational resources. However, most traditional optimizers either ignore the cost of computation or select the number of iterations using simple heuristics. Because they do not balance the cost of computation against the performance loss, the overall effectiveness of these approaches is subject to the skill and preferences of the practitioners who use them. Second, which expert should be used on each step of the optimization? Some experts may be accurate but expensive to compute in terms of time, energy and/or money, while others may be crude, yet cheap. Moreover, the reliability of the experts may not be known a priori, further limiting the effectiveness of the optimization procedure. Our use of a metacontroller address these issues by jointly optimizing over the choices of how many steps to take and which experts to use.
We consider a family of optimizers which use the same controller, πC , but vary in their expert evaluators, {E1, . . . , EK}. Assuming that the controller and experts are deterministic functions, the number of iterationsN and the sequences of experts k = (k1, . . . , kN−1) exactly determine the final control and performance loss LP . This means we have transformed the performance optimization over c into an optimization over N and k: (N,k)∗ = argmink,n LP (x
∗, x, c(N,k, x, x∗)), where the notation c(N,k, x, x∗) is used to emphasize that the control is a function N , k, x, and x∗.
If each optimizer has an associated computational cost τk, then N and k also exactly determine the computational resource loss of the optimization run, LR(N,k) = ∑N−1 n=1 τkn . The total loss is then the sum of LP and LR, each of which are functions of N and k, LT (x ∗, x,N,k) = LP (x ∗, x, c(N,k, x, x∗)) + LR(N,k) (3)
= L(x∗, f(x, πC(x∗, x, hN−1))) + N−1∑ n=1 τkn , (4)
and the optimal solution is defined as (N,k)∗ = argminN,k LT (x ∗, x,N,k). Optimizing LT is difficult because of the recursive dependency on the history, hN−1, and because the discrete choices of N and k mean LT is not differentiable.
To optimize LT we recast it as an RL problem where the objective is to jointly optimize task performance and computational cost. As shown in Figure 1a, the metacontroller agent aM is comprised of a controller πC , a pool of experts {E1, . . . , EK}, a manager πM , and a memory µ. The manager is a meta-level policy (Russell & Wefald, 1991; Hay et al., 2012) over actions indexed by k, which determine whether to terminate the optimization procedure (k = 0) or to perform another iteration of the optimization procedure with the kth expert. Specifically, on the nth iteration the controller produces a new control cn based on the history of controls, experts, and evaluations. The manager, also relying on this history, independently decides whether to end the optimization procedure (i.e., to execute the control in the world) or to perform another iteration and evaluate the proposed control with the kthn expert (i.e., to ponder, after Graves (2016)). The memory then updates the history hn by concatenating k, cn, and en with the previous history hn−1. Coming back to the notion of imagination-based optimization, we suggest that this iterative optimization process is analogous to imagining what will happen (using one or more approximate world models) before actually executing that action in the world. For further details, see Appendix A, and for an algorithmic illustration of the metacontroller agent, see Algorithm 1 in the appendix.
We also define two special cases of the metacontroller for baseline comparisons. The iterative agent aI does not have a manager and uses only a single expert. Its number of iterations are pre-set to a single N . The reactive agent, a0, is a special case of the iterative agent, where the number of iterations is fixed to N = 0. This implies that proposed controls are executed immediately in the world, and are not evaluated by an expert. For algorithmic illustrations of the iterative and reactive agents, see Algorithms 2 and 3 in the appendix.
2.3 NEURAL NETWORK IMPLEMENTATION
We use standard deep learning building blocks, e.g., multi-layer perceptrons (MLPs), RNNs, etc., to implement the controller, experts, manager, and memory, because they are effective at approximating complex functions via gradient-based and reinforcement learning, but other approaches could be used as well. In particular, we constructed our implementation to be able to make control decisions in complex dynamical systems, such as controlling the movement of a spaceship (Figure 1b-c), though we note that our approach is not limited to such physical reasoning tasks. Here we used mean-squared error (MSE) for our L and Adam (Kingma & Ba, 2014) as the training optimizer.
Experts We implemented the experts as MLPs and “interaction networks” (INs) (Battaglia et al., 2016), which are well-suited to predicting complex dynamical systems like those in our experiments below. Each expert has parameters θEk , i.e. en = Ek(x∗, x, cn; θEk), and may be trained either on-policy using the outputs of the controller (as is the case in this paper), or off-policy by any data that pairs states and controls with future states or reward outcomes. The objective LEk for each expert may be different depending on what the expert outputs. For example, the objective could be the loss between the goal and future states, LEk = L ( f(x, c), Ek(x ∗, x, c; θEk) ) , which is what we use in our experiments. Or, it could be the loss between LP and an action-value function that predicts LP directly, LEk = L ( LP (x ∗, x, c), Ek(x ∗, x, c; θEk) ) . See Appendix B.1 for details.
Controller and Memory We implemented the controller as an MLP with parameters θC , i.e. cn = πC(x∗, x, hn−1; θ
C), and we implemented the memory as a Long Short-Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997) with parameters θµ. The memory embeds the history as a fixedlength vector, i.e. hn = µ(hn−1, kn, cn, Ekn(x ∗, x, cn); θ µ). The controller and memory were trained jointly to optimize (1). However, this objective includes f , which is often unknown or not differentiable. We overcame this by approximating LP with a differentiable critic analogous to those used in policy gradient methods (e.g. Silver et al., 2014; Lillicrap et al., 2015; Heess et al., 2015). See Appendices B.2 and B.3 for details.
Manager We implemented the manager as a stochastic policy that samples from a categorical distribution whose weights are produced by an MLP with parameters θM , i.e. kn ∼ Categorical(k;πM (x∗, x, hn−1; θ
M )). We trained the manager to minimize (3) using REINFORCE (Williams, 1992), but other deep RL algorithms could be used instead. See Appendix B.4 for details.
3 EXPERIMENTS
To evaluate our metacontroller agent, we measured its ability to learn to solve a class of physicsbased tasks that are surprisingly challenging. Each episode consisted of a scene which contained a spaceship and multiple planets (Figure 1b-c). The spaceship’s goal was to rendezvous with its mothership near the center of the system in exactly 11 time steps, but it only had enough fuel to fire its thrusters once. The planets were static but the gravitational force they exerted on the spacecraft induced complex non-linear dynamics on the motion over the 11 steps. The spacecraft’s action space was continuous, up to some maximum magnitude, and represented the instantaneous Cartesian velocity vector imparted by its thrusters. Further details are in Appendix C.
We trained the reactive, iterative, and metacontroller agents on five versions of the spaceship task involving different numbers of planets.1 The iterative agent was trained to take anywhere from zero (i.e., the reactive agent) to ten ponder steps. The metacontroller was allowed to take a maximum of ten ponder steps. We considered three different experts which were all differentiable: an MLP expert which used an MLP to predict the final location of the spaceship, an IN expert which used an interaction network (Battaglia et al., 2016) to predict the full trajectory of the spaceship, and a true simulation expert which was the same as the world model. In some conditions the metacontroller could use exactly one expert and in others it was allowed to select between the MLP and IN experts. For experiments with the true simulation expert, we used it to backpropagate gradients to the controller and memory. For experiments with an MLP as the only expert, we used a learned IN as the critic. For experiments with an IN as one of its experts, the critic was an IN with shared parameters. We trained the metacontroller on a range of different ponder costs, τk, for the different experts. Further details of the training procedure are available in Appendix D.
3.1 REACTIVE AND ITERATIVE AGENTS
Figure 2 shows the performance on the test set of the reactive and iterative agents for different numbers of ponder steps. The reactive agent performed poorly on the task, especially when the task was more difficult. With the five planets dataset, it was only able to achieve a performance loss of 0.583 on average (see Figure 1 for a depiction of the magnitude of the loss). In contrast, the iterative agent with the true simulation expert performed much better, reaching ceiling performance on the
1Available from: https://www.github.com/deepmind/spaceship dataset
datasets with one and two planets, and achieving a performance loss of 0.0683 on the five planets dataset. The IN and MLP experts also improve over the reactive agent, with a minimum performance loss of 0.117 and 0.375 on the five planets dataset, respectively.
Figure 2 also highlights how important the choice of expert is. When using the true simulation and IN experts, the iterative agent performs well. With the MLP expert, however, performance is substantially diminished. But despite the poor performance of the MLP expert, there is still some benefit of pondering with it. With even just a few steps, the MLP iterative agent outperforms its reactive counterpart. However comparing the reactive agent with the N = 1 iterative agent is somewhat unfair because the iterative agent has more parameters due to the expert and the memory. However, given that there tends to also be an increase in performance between one and two ponder steps (and beyond), it is clear that pondering—even with a highly inaccurate model—can still lead to better performance than a model-free reactive approach.
3.2 METACONTROLLER WITH ONE EXPERT
Though the iterative agents achieve impressive results, they expend more computation than necessary. For example, in the one and two planet conditions, the performances of the IN and true simulation iterative agents received little performance benefit from pondering more than two or three steps, while for the four and five planet conditions they required at least five to eight steps before their performance converged. When computational resources have no cost, the number of steps are of no concern, but when they have some cost it is important to be economical.
Because the metacontroller learns to choose its number of pondering steps, it can balance its performance loss against the cost of computation. Figure 3 (top row, middle and right subplots) shows that the IN and true simulation expert metacontroller take fewer ponder steps as τ increases, tracking closely the minimum of the iterative agent’s cost curve (i.e., the metacontroller points are always near the iterative agent curves’ minima). This adaptive behavior emerges automatically from the manager’s learned policy, and avoids the need to perform a hyperparameter search to find the best number of iterations for a given τ .
The metacontroller does not simply choose an average number of ponder steps to take per episode: it actually tailors this choice to the difficulty of each episode. Figure 4 shows how the number of ponder steps the IN metacontroller chooses in each episode depends on that episode’s difficulty, as measured by the episode’s loss under the reactive agent. For more difficult episodes, the metacontroller tends to take more ponder steps, as indicated by the positive slopes of the best fit lines, and this proportionality persists across the different levels of τ in each subplot.
The ability to adapt its choice of number of ponder steps on a per-episode basis is very valuable because it allows the metacontroller to spend additional computation only on those episodes which require it. The total costs of the IN and true simulation metacontrollers’ are 11% and 15% lower (median) than the best achievable costs of their corresponding iterative agents, respectively, across the range of τ values we tested (see Figure 7 in the Appendix for details).
There can even be a benefit to using a metacontroller when there are no computational resource costs. Consider the rightmost points in Figure 3 (bottom row, middle and right subplots), which show the performance loss for the IN and true simulation metacontrollers when τ is low. Remarkably, these points still outperform the best achievable iterative agents. This suggests that there can be an advantage to stopping pondering once a good solution is found, and more generally demonstrates that the metacontroller’s learning process can lead to strategies that are superior to those available to less flexible agents.
The metacontroller with the MLP expert had very poor average performance and high variance on the five planet condition (Figure 3, top left subplot), which is why we restricted our focus in this section to how the metacontrollers with IN and true simulation experts behaved. The MLP’s poor performance is crucial, however, for the following section (3.3) which analyzes how a multipleexpert metacontroller manages experts which vary greater in their reliability.
3.3 METACONTROLLER WITH TWO EXPERTS
When we allow the manager to additionally choose between two experts, rather than only relying on a single expert, we find a similar pattern of results in terms of the number of ponder steps (Figure 5, left). Additionally, the metacontroller is successfully able to identify the more reliable IN network and consequently uses it a majority of the time, except in a few cases where the cost of the IN network is extremely high relative to the cost of the MLP network (Figure 5, right). This pattern of results makes sense given the good performance (described in the previous section) of the metacontroller with the IN expert compared to the poor performance of the metacontroller with the MLP expert. The manager should not generally rely on the MLP expert because it is simply not a reliable source of information.
However, the metacontroller has more difficulty finding an optimal balance between the two experts on a step-by-step basis: the addition of a second expert did not yield much of an improvement over the single-expert metacontroller, with only 9% of the different versions (trained with different τ values for the two experts) achieving a lower loss than the best iterative controller. We believe the mixed performance of the metacontroller with multiple experts is partially due to an entropy term which we used to encourage the manager’s policy to be non-deterministic (see Appendix B.4). In particular, for high values of τ , the optimal thing to do is to always execute immediately without pondering. However, because of the entropy term, the manager is encourage to have a non-deterministic policy and therefore is likely to ponder more than it should—and to use experts that are more unreliable— even when this is suboptimal in terms of the total loss (3).
Despite the fact that the metacontroller with multiple experts does not result in a substantial improvement over that which uses a single expert, we emphasize that the manager is able to identify and use the more reliable expert the majority of the time. And, it is still able to choose a variable number of steps according to how difficult the task is (Figure 5, left). This, in and of itself, is an improvement over more traditional optimization methods which would require that the expert is hand-picked ahead of time and that the number of steps are determined heuristically.
4 DISCUSSION
In this paper, we have presented an approach to adaptive, imagination-based optimization in neural networks. Our approach is able to flexibly choose which computations to perform as well as how many computations need to be performed, approximately solving a speed-accuracy trade-off that depends on the difficulty of the task. In this way, our approach learns to rely on whatever source of information is most useful and most efficient. Additionally, by consulting the experts on-the-fly, our approach allows agents to test out actions to ensure that their consequences are not disastrous before actually executing them.
While the experiments in this paper involve a one-shot decision task, our approach lays a foundation that can be built upon to support more complex situations. For example, rather than applying a force only on the first time step, we could turn the problem into one of trajectory optimization for continuous control by asking the controller to produce a sequence of forces. In the case of planning, our approach could potentially be combined with methods like Monte Carlo Tree-Search (MCTS) (Coulom, 2006), where our experts would be akin to having several different rollout policies to choose from, and our controller would be akin to the tree policy. While most MCTS implementations will run rollouts until a fixed amount of time has passed, our approach would allow the manager to adaptively choose the number of rollouts to perform and which policies to perform the rollouts with. Our method could also be used to naturally augment existing model-free approaches such as DQN (Mnih et al., 2015) with online model-based optimization by using the model-free policy as a controller and adding additional experts in the form of state-transition models. An interesting extension would be to compare our metacontroller architecture with a naı̈ve model-based controller that performs gradient-based optimization to produce the final control. We expect our metacontroller architecture might require fewer model evaluations and to be more robust to model inaccuracies compared to the gradient-based method, because our method has access to the full history of proposed controls and evaluations whereas traditional gradient-based methods do not.
Although we rely on differentiable experts in our metacontroller architecture, we do not utilize the gradient information from these experts. An interesting extension to our work would be to pass this gradient information through to the manager and controller (as in Andrychowicz et al. (2016)), which would likely improve performance further, especially in the more complex situations discussed here. Another possibility is to train some or all of the experts inline with the controller and metacontroller, rather than independently, which could allow their learned functionality to be more tightly integrated with the rest of the optimization loop, at the expense of their generality and ability to be repurposed for other uses.
To conclude, we have demonstrated how neural network-based agents can use metareasoning to adaptively choose what to think about, how to think about it, and for how long to think for. Our
method is directly inspired by human cognition and suggests a way to make agents much more flexible and adaptive than they currently are, both in decision making tasks such as the one described here, as well as in planning and control settings more broadly.
ACKNOWLEDGMENTS
We would like to thank Matt Hoffman, Andrea Tacchetti, Tom Erez, Nando de Freitas, Guillaume Desjardins, Joseph Modayil, Hubert Soyer, Alex Graves, David Reichert, Theo Weber, Jon Scholz, Will Dabney, and others on the DeepMind team for helpful discussions and feedback.
A METACONTROLLER DETAILS
Here, we give the precise definitions of the metacontroller agent. As described in the main text, the iterative and reactive agents are special cases of the metacontroller agent, and are therefore not discussed here.
The metacontroller agent aM is comprised of the following components:
• A history-sensitive controller, πC : X ×X ×H → C, which is a policy that maps goal and initial states, and a history, h ∈ H, to controls, whose aim is to minimize (1).
• A pool of experts {E1, . . . , EK}. Each expert E : X ×X ×C → E maps goal states, input states, and actions to opinions. Opinions can be either states-only (E = X ), states and rewards (E = X × R), or rewards-only (E = R). The expert corresponds to the evaluator for the optimization routine, i.e., an approximation of the forward process f .
• A manager, πM : X ×X ×Hn → {0, . . . ,K}, which is a policy which decides whether to send a proposed control to the world (k = 0) or to the kth expert for evaluation, in order to minimize (3). This formulation is based on that used by metareasoning systems (Russell & Wefald, 1991; Hay et al., 2012). Details on the corresponding MDP are given in Appendix A.1.
• A memory, µ : Hn−1 ×Z → Hn, which is a function that maps the prior history hn−1 ∈ Hn−1, as well as the most recent manager choice, proposed control, and expert evaluation (k, c, e) ∈ {0, . . . ,K} × C × E = Z , to an updated history hn ∈ Hn, which is then made available to the manager and controller on subsequent iterations. The history at step n is a recursively defined tuple which is the concatenation of the prior history with the most recently proposed control, expert evaluation, and expert identity: hn = hn−1 ∩ ((kn, cn, Ekn(x∗, x, cn))) = ((k1, c1, Ek1(x ∗, x, c1)), . . . , (kn, cn, Ekn(x ∗, x, cn))) where h0 = () represents an empty ini-
tial history. Similarly, the finite set of histories up to step n is: Hn = Hn−1 × Z = Zn where H0 = {()}.
The metacontroller produces:
aM (x∗, x) = πC(x∗, x, hN−1) = cN (5)
where N = n s.t. kn = 0. This function is summarized in Algorithm 1. The other agents (iterative and reactive), as mentioned in the main text, are simpler versions of the metacontroller agent and are summarized in Algorithms 2 and 3.
A.1 META-LEVEL MDP
To implement the manager for the metacontroller agent, we draw inspiration from the metareasoning literature (Russell & Wefald, 1991; Hay et al., 2012) and formulate the problem as a finite-horizon Markov Decision Process (MDP) 〈S,A, P,R〉 over the decision of whether to perform another iteration of the optimization procedure or to execute a control in the world.
• The state space S consists of goal states, external states, and internal histories, S = X ×X ×H. • The action space A contains K + 1 discrete actions, {0, . . . ,K}, which correspond to execute
(k = 0) and ponder (k ∈ {1, . . . ,K}), where ponder (after Graves (2016)) refers to performing an iteration of the optimization procedure with the kth expert.
• The (deterministic) state transition model P : S × C × S → [0, 1] is, P (x′, hn|x∗, x, hn−1, k) = { P (x′|x∗, x, hn−1, k) if k = 0 P (hn|x∗, x, hn−1, k) otherwise
where x′ = f(x, c) and c = πC(x∗, x, hn−1) and, P (x′|x∗, x, hn−1, k) = { 1 if x′ = f(x, c) 0 otherwise
P (hn|x∗, x, hn−1, k) = { 1 if hn = hn−1 ∪ {(k, c, Ek(x∗, x, c))} 0 otherwise
Algorithm 1 Metacontroller agent. x is the scene and x∗ is the target. 1: function aM (x, x∗) 2: h0 ← () . Initial empty history 3: k0 ← πM (x, x∗, h0) . Get an action from the manager 4: c0 ← πC(x, x∗, h0) . Propose a control with the controller 5: n← 0 6: while kn 6= 0 do . When k 6= 0, ponder with an expert 7: en ← Ekn(x, x∗, cn) . Get an expert’s opinion 8: hn+1 ← µ(hn, kn, cn, en) . Update the history 9: n← n+ 1 10: kn ← πM (x, x∗, hn) . Choose the next action 11: cn ← πC(x, x∗, hn) . Propose the next control 12: end while 13: return cn 14: end function
Algorithm 2 Iterative agent. x is the scene, x∗ is the target, and N is the number of ponder steps. 1: function aI (x, x∗, N ) 2: h0 ← () . Initial empty history 3: c0 ← πC(x, x∗, h0) . Propose a control with the controller 4: n← 0 5: while n < N do . Ponder with an expert for N steps 6: en ← E(x, x∗, cn) . Get the expert’s opinion 7: hn+1 ← µ(hn, kn, cn, en) . Update the history 8: n← n+ 1 9: cn ← πC(x, x∗, hn) . Propose the next control 10: end while 11: return cn 12: end function
Algorithm 3 Reactive agent. x is the scene and x∗ is the target. 1: function a0(x, x∗) 2: c0 ← πC(x, x∗, ()) . Propose a control with the controller 3: return c0 4: end function
• The (deterministic) reward function R : S × A × S → R maps the current state, current action, and next state to real-valued loss:
R(x∗, x, hn−1, k, x ′) = { L(x∗, x′) if k = 0 (see Eq. 1) τk otherwise (see Eq. 3)
where x′ = f ( x, πC(x∗, x, hn−1) ) .
We approximate the solution to this MDP with a stochastic manager policy πM . The manager chooses actions proportional to the immediate reward for taking action k in state sn plus the expected sum of future rewards. This construction imposes a trade-off between accuracy and resources, incentivizing the agent to ponder longer and with more accurate (and potentially expensive) experts when the problem is harder.
B GRADIENTS
B.1 EXPERTS
Training the experts is a straightforward supervised learning problem (Figure 6c). The gradient is:
∂LEk ∂θEk = ∂LEk ∂Ek ∂Ek ∂θEk , (6)
where Ek is the kth expert and LEk is the loss function for the kth expert. For example, in the case of an action-value function expert, this loss function might be LEk(f,Ek) =∥∥L(x∗, f(x, c))− Ek(x∗, x, c; θEk)∥∥2. In the case of an expert that predicts the final state using a model of the system dynamics, it might be LEk(f,Ek) =
∥∥f(x, c)− Ek(x∗, x, c; θEk)∥∥2.
B.2 CRITIC
The critic, L̂P , is an approximate model of the performance loss, LP , (1), which is used to backpropagate gradients to the controller and memory. This means the critic can either be an action-value function, which approximates L̂P = E0 ≈ LP directly, or a model of the system dynamics composed with a known loss function between the goal and future states, L̂P = L◦E0 ≈ L◦f . We train the critic, E0 : X ×X ×C → R, using the same procedure as the experts are trained (Figure 6d). A good expert may even be used as the critic.
B.3 CONTROLLER AND MEMORY
As shown in Figure 6a, we trained the controller and memory using backpropagation through time (BPTT) with an actor-critic architecture. Specifically, rather than assuming f is known and differentiable, we use a critic and backpropagate through it (Heess et al., 2015):
∂L ∂θC = ∂L ∂E∗ ∂E∗ ∂πCn ∂πCn ∂µn ∂+µn ∂πCn−1 · · · ∂π C 0 ∂θC ,
∂L ∂θµ = ∂L ∂E∗ ∂E∗ ∂πCn ∂πCn ∂µn ∂+µn ∂µn−1 · · · ∂µ0 ∂θµ
(7)
where E∗ is the critic, n is the maximum number of iterations the controller can use, and: ∂+µn ∂πCn−1 = ∂µn ∂Ekn−1 ∂Ekn−1 ∂πCn−1 + ∂µn ∂πCn−1 , ∂+µn ∂µn−1 = ∂+µn ∂πCn−1 + ∂µn ∂µn−1
(8)
where we are using the ∂+ notation to indicate summed gradients, following Pascanu et al. (2013). Since kn has already been produced by the manager it can be treated as a constant and will produce an unbiased estimate of the gradient. This is convenient because it allows for training the controller and manager separately, or testing the controller’s behavior with arbitrary actions post-training.
B.4 MANAGER
As discussed in the main text, we used the REINFORCE algorithm Williams (1992) to train the manager (Figure 6b). One potential issue, however, is that when training the controller and manager simultaneously, the controller will result in high cost early on in training and thus the manager will learn to always choose the execute action. To discourage the manager from learning what is an essentially deterministic policy, we included a regularization term based on the entropy, LH (Williams & Peng, 1991; Mnih et al., 2016):
LH( · ; θM ) = λ EπM [log πM ( · ; θM )] ∂EπM [r] ∂θM = ( r − LH( · ; θM ) ) ∂ ∂θM log πM ( · ; θM ),
r is the full return given by (3) and λ is the strength of the regularization term.
C SPACESHIP TASK
C.1 DATASETS
We generated five datasets, each containing scenes with a different number of planets (ranging from a single planet to five planets). Each dataset consisted of 100,000 training scenes and 1,000 testing scenes. The target in each scene was always located at the origin, and each scene always had a sun with a mass of 100 units. The sun was located between 100 and 200 distance units away from the target, with this distance sampled uniformly at random. The other planets had a mass between 20 and 50 units, and were located 100 to 250 distance units away from the target, sampled uniformly at random. The spaceship had a mass between 1 and 9 units, and was located 150 to 250 distance units away from the target. The planets were always fixed (i.e., they could not move), and the spaceship always started at the beginning of each episode with zero velocity.
C.2 ENVIRONMENT
We simulated our scenes using a physical simulation of gravitational dynamics. The planets were always stationary (i.e., they were not acted upon by any of the objects in the scene) but acted upon the spaceship with a force of:
Fp = G mpms r3
(xp − xs), (9) where Fp is the force vector of the planet on the spaceship,G = 1000000 is a gravitational constant, mp is the mass of the planet, ms is the mass of the spaceship, r is the distance between the centers of masses of the planet and the spaceship, xp is the location of the planet, and xs is the location of the spaceship. We simulated this environment using the Euler method, i.e.:
as = ( ∑ pFp)− dvs + c
ms x′s = xs + vs v ′ s = vs + as (10)
where as, vs, and xs are the acceleration, velocity, and position of the spaceship, respectively; d = 0.1 is a damping constant; c is the control force applied to the spaceship; and is the step size. Note that we set c to zero for all timesteps except the first.
D IMPLEMENTATION DETAILS
We used TensorFlow (Abadi et al., 2015) to implement and train all versions of the model.
D.1 ARCHITECTURE
In our implementation of the controller, we used a two-layer MLP each with 100 units. The first layer used ReLU activations and the second layer used a multiplicative interaction similar to van den Oord et al. (2016), which we found to work better in practice. In our implementation of the memory, we used a single LSTM layer of size 100. In our implementation of the manager, we used a MLP of two fully connected layers of 100 units each, with ReLU nonlinearities.
We constructed three different experts to test the various controllers. The true simulation expert was the same as the world model, and consisted of a simulation for 11 timesteps with = 0.05 (see Appendix C). The IN expert was an interaction network (Battaglia et al., 2016), which has previously been shown to be able to learn to predict n-body dynamics accurately for simple systems. The IN consists of a relational module and an object module. In our case, the relational module was composed of 4 hidden layers of 150 nodes each, outputting “effects” encodings of size 100. These effects, together with the relational model input are then used as input to the object model, which contained a single hidden layer of 100 nodes. The object model outputs the velocity of the spaceship and we trained it to predict the velocity on every timestep of the spaceship’s trajectory. The MLP expert was a MLP that predicted the final location of the spaceship and had the same architecture as the controller.
As discussed in Appendix B, we used a critic to train the controller and memory. We always used the IN expert as the critic, except in the case when the true simulation expert was used, in which case we also used the true simulation as the critic.
D.2 TRAINING PROCEDURE
All weights were initialized uniformly at random between 0 and 0.01. An iteration of training consisted of gradient updates over a minibatch of size 1000; in total, we ran training for 100,000 iterations. We additionally used a waterfall schedule for each of the learning rates during training, such that after 1000 iterations, if the loss was not decreasing, we would decay the step size by 5%.
We trained the controller and memory together using the Adam optimizer (Kingma & Ba, 2014) with gradients clipped to a maximum global norm of 10 (Pascanu et al., 2013). The manager was trained simultaneously, but using a different learning rate than the controller and memory. The IN and MLP experts were also trained simultaneously, but again with different learning rates. Learning rates were determined using a grid search over a small number of values, and are given in Table 1 for the iterative agent, in Table 2 for the metacontroller with one expert, and in Table 3 for the metacontroller with two experts.
The iterative agent was trained to take a fixed number of ponder steps, ranging from 0 (i.e., the reactive agent) to 10. The metacontrollers were allowed to take a variable number of ponder steps up to a maximum of 10. For the metacontroller with a single expert, we trained the manager using τ = 0 and 20 additional values of τ spaced logarithmically between 0.00004 and 0.4 (inclusive). For the metacontroller with multiple experts, we trained the manager on a grid of pairs of τ values, where each expert could have τ = 0 or one of 6 values spaced logarithmically between 0.00004 and 0.2 (inclusive). In all cases, the entropy penalty for the metacontroller was λ = 0.2.
D.3 CONVERGENCE
Reactive agent. Training for the reactive agents was straightforward and converged reliably on all datasets.
Iterative agent. For the iterative agent with the interaction network or true simulation experts, convergence was also reliable for small numbers of ponder steps. Convergence was somewhat less reliable for larger numbers of ponder steps. We believe this is because for some scenes, a larger number of ponder steps was more than necessary to solve the task (as is evidenced by the plateauing performance in Figure 2). So, the iterative agent had to effectively “remember” what the best control was while it took the last few ponder steps, which is a more complicated and difficult task to perform.
For the iterative agent with the MLP expert, convergence was more variable especially when the task was harder, as can be seen in the variable performance on the five planets dataset in Figure 2 (left). We believe this is because the MLP agent was so poor, and that convergence would have been more reliable with a better agent.
Metacontroller with a single expert. The metacontroller agent with a single expert converged more reliably than the corresponding iterative agent (see the bottom row of Figure 3). As mentioned in the previous paragraph, the iterative agent had to take more steps than actually necessary, causing it to perform less well for larger numbers of ponder steps, whereas the metacontroller agent had the flexibility of stopping when it had found a good control. On the other hand, we found that the metacontroller agent sometimes performed too many ponder steps for large values of τ (see Figures 3 and 7). We believe this is due to the entropy term (λ) added to the REINFORCE loss. This is because when then ponder cost is very high, the optimal thing to do is to behave deterministically and always execute (never ponder); however, the entropy term encouraged the policy to be nondeterministic. We plan to explore different training regimes in future work to alleviate this problem, for example by annealing the entropy term to zero over the course of training.
Metacontroller with multiple experts. The metacontroller agent with multiple experts was somewhat more difficult to train, especially for high ponder cost of the interaction network expert. For example, note how the proportion of steps using the MLP expert does not decrease monotonically in Figure 5 (right) with increasing cost for the MLP expert. We believe this is also an unexpected result of using the entropy term: in all of these cases, the optimal thing to do actually is to rely on the MLP expert 100% of the time, yet the entropy term encourages the policy to be non-deterministic. Future work will explore these difficulties further by using experts that complement each other better (i.e., so there is not one that is wholly better than the other).
Experts. The experts themselves always converged quickly and reliably, and trained much faster than the rest of the network.

Pros (quality, clarity, originality, significance:):
This paper presents a novel metacontroller optimization system that learns the best action for a one-shot learning task, but as a framework has the potential for wider application. The metacontroller is a model-free reinforcement learning agent that selects how many optimization iterations and what function or “expert” to consult from a fixed set (such as an action-value or state transition function). Experimental results are presented from simulation experiments where a spacecraft must fire its thruster once to reach a target location, in the presence of between 1 and 5 heavy bodies.
The metacontroller system has a similar performance loss on the one-shot learning task as an iterative (standard) optimization procedure. However, by taking into account the computational complexity of running a classical, iterative optimization procedure as a second “resource loss” term, the metacontroller is shown to be more efficient. Moreover, the metacontroller agent successfully selects the optimal expert to consult, rather than relying on an informed choice by a domain-expert model designer. The experimental performance is a contribution that merits publication, and it also exhibits the use of an interaction network for learning the dynamics of a simulated physical system. The dataset that has been developed for this task also has the potential to act as a new baseline for future work on one-shot physical control systems. The dataset constitutes an ancillary contribution which could positively impact future research in this area.
Cons:
It's not clear how this approach could be applied more broadly to other types of optimization. Moreover, the REINFORCE gradient estimation method is known to suffer from very high variance, yielding poor estimates. I'm curious what methods were used to ameliorate these problems and if any other performance tricks were necessary to train well. Content of this type this could form a useful additional appendix.
A few critiques on the communication of results:
- The formal explication of the paper’s content is clear, but Fig.’s 1A and 3 could be improved. Fig. 1A is missing a clear visual demarcation of what exactly the metacontroller agent is. Have you considered a plate or bounding box around the corresponding components? This would likely speed the uptake of the formal description.
- Fig. 3 is generally clear, but the lack of x-axis tick marks on any subplots makes it more challenging than necessary to compare among the experts. Also, the overlap among the points and confidence intervals in the upper-left subplot interferes with the quantitative meaning of those symbols. Perhaps thinner bars of different colors would help here. Moreover, this figure lacks a legend and so the different lines are impossible to compare with each other.
- Lastly, the second sentence in Appendix B. 2 is a typo and terminates without completion.


RECURRENT NEURAL NETWORKS
Recurrent Neural Networks (RNN) are widely used to solve a variety of problems and as the quantity of data and the amount of available compute have increased, so have model sizes. The number of parameters in recent state-of-the-art networks makes them hard to deploy, especially on mobile phones and embedded devices. The challenge is due to both the size of the model and the time it takes to evaluate it. In order to deploy these RNNs efficiently, we propose a technique to reduce the parameters of a network by pruning weights during the initial training of the network. At the end of training, the parameters of the network are sparse while accuracy is still close to the original dense neural network. The network size is reduced by 8× and the time required to train the model remains constant. Additionally, we can prune a larger dense network to achieve better than baseline performance while still reducing the total number of parameters significantly. Pruning RNNs reduces the size of the model and can also help achieve significant inference time speed-up using sparse matrix multiply. Benchmarks show that using our technique model size can be reduced by 90% and speed-up is around 2× to 7×.
1 INTRODUCTION
Recent advances in multiple fields such as speech recognition (Graves & Jaitly, 2014; Amodei et al., 2015), language modeling (Józefowicz et al., 2016) and machine translation (Wu et al., 2016) can be at least partially attributed to larger training datasets, larger models and more compute that allows larger models to be trained on larger datasets.
For example, the deep neural network used for acoustic modeling in Hannun et al. (2014) had 11 million parameters which grew to approximately 67 million for bidirectional RNNs and further to 116 million for the latest forward only GRU models in Amodei et al. (2015). And in language modeling the size of the non-embedding parameters (mostly in the recurrent layers) have exploded even as various ways of hand engineering sparsity into the embeddings have been explored in Józefowicz et al. (2016) and Chen et al. (2015a).
These large models face two significant challenges in deployment. Mobile phones and embedded devices have limited memory and storage and in some cases network bandwidth is also a concern. In addition, the evaluation of these models requires a significant amount of computation. Even in cases when the networks can be evaluated fast enough, it will still have a significant impact on battery life in mobile devices (Han et al., 2015).
Inference performance of RNNs is dominated by the memory bandwidth of the hardware, since most of the work is simply reading in the parameters at every time step. Moving from a dense calculation to a sparse one comes with a penalty, but if the sparsity factor is large enough, then the smaller amount of data required by the sparse routines becomes a win. Furthermore, this suggests that if the parameter sizes can be reduced to fit in cache or other very fast memory, then large speedups could be realized, resulting in a super-linear increase in performance.
The more powerful server class GPUs used in data centers can generally perform inference quickly enough to serve one user, but in the data center performance per dollar is very important. Techniques
∗Now at Google Brain eriche@google.com
that allow models to be evaluated faster enable more users to be served per GPU increasing the effective performance per dollar.
We propose a method to reduce the number of weights in recurrent neural networks. While the network is training we progressively set more and more weights to zero using a monotonically increasing threshold. By controlling the shape of the function that maps iteration count to threshold value, we can control how sparse the final weight matrices become. We prune all the weights of a recurrent layer; other layer types with significantly fewer parameters are not pruned. Separate threshold functions can be used for each layer, although in practice we use one threshold function per layer type. With this approach, we can achieve sparsity of 90% with a small loss in accuracy. We show this technique works with Gated Recurrent Units (GRU) (Cho et al., 2014) as well as vanilla RNNs.
In addition to the benefits of less storage and faster inference, this technique can also improve the accuracy over a dense baseline. By starting with a larger dense matrix than the baseline and then pruning it down, we can achieve equal or better accuracy compared to the baseline but with a much smaller number of parameters.
This approach can be implemented easily in current training frameworks and is agnostic to the optimization algorithm. Furthermore, training time does not increase unlike previous approaches such as in Han et al. (2015). State of the art results in speech recognition generally require days to weeks of training time, so a further 3-4× increase in training time is undesirable.
2 RELATED WORK
There have been several proposals to reduce the memory footprint of weights and activations in neural networks. One method is to use a fixed point representation to quantize weights to signed bytes and activations to unsigned bytes (Vanhoucke et al., 2011). Another technique that has been tried in the past is to learn a low rank factorization of the weight matrices. One method is to carefully construct one of the factors and learn the other (Denil et al., 2013). Inspired by this technique, a low rank approximation for the convolution layers achieves twice the speed while staying within 1% of the original model in terms of accuracy (Denton et al., 2014). The convolution layer can also be approximated by a smaller set of basis filters (Jaderberg et al., 2014). By doing this they achieve a 2.5x speedup with no loss in accuracy. Quantization techniques like k-means clustering of weights can also reduce the storage size of the models by focusing only on the fully connected layers (Gong et al., 2014). A hash function can also reduce memory footprint by tying together weights that fall in the same hash bucket (Chen et al., 2015b). This reduces the model size by a factor of 8.
Yet another approach to reduce compute and network size is through network pruning. One method is to use several bias techniques to decay weights (Hanson & Pratt, 1989). Yet another approach is to use the diagonal terms of a Hessian matrix to construct a saliency threshold and used this to drop weights that fall below a given saliency threshold (LeCun et al., 1989). In this technique, once a weight has been set to 0, the network is retrained with these weights frozen at 0. Optimal Brain Surgeon is another work in the same vein that prunes weights using the inverse of a Hessian matrix with the additional advantage of no re-training after pruning (Hassibi et al., 1993).
Both pruning and quantization techniques can be combined to get impressive gains on AlexNet trained on the ImageNet dataset (Han et al., 2015). In this case, pruning, quantization and subsequent Huffman encoding results in a 35x reduction in model size without affecting accuracy. There has also been some recent work to shrink model size for recurrent and LSTM networks used in automatic speech recognition (ASR) (Lu et al., 2016). By using a hybrid strategy of using Toeplitz matrices for the bottom layer and shared low-rank factors on the top layers, they were able to reduce the parameters of a LSTM by 75% while incurring a 0.3% increase in word error rate (WER).
Our method is a pruning technique that is computationally efficient for large recurrent networks that have become the norm for automatic speech recognition. Unlike the methods that need to approximate a Hessian (LeCun et al., 1989; Hassibi et al., 1993) our method uses a simple heuristic to choose the threshold used to drop weights. Yet another advantage, when compared to methods that need re-training (Han et al., 2015), is that our pruning technique is part of training and needs no additional re-training. Even though our technique requires judicious choice of pruning hyperparameters, we feel that it is easier than choosing the structure of matrices to guide the sparsification
for recurrent networks (Lu et al., 2016). Another approach for pruning feed forward neural networks for speech recognition is using simple threshold to prune all weights (Yu et al., 2012) at a particular epoch. However, we find that gradual pruning produces better results than hard pruning.
3 IMPLEMENTATION
Our pruning approach involves maintaining a set of masks, a monotonically increasing threshold and a set of hyper parameters that are used to determine the threshold. During model initialization, we create a set of binary masks, one for each weight in the network that are all initially set to one. After every optimizer update step, each weight is multiplied with its corresponding mask. At regular intervals, the masks are updated by setting all parameters that are lower than the threshold to zero.
The threshold is computed using hyper-parameters shown in Table 1. The hyper-parameters control the duration, rate and frequency of pruning the parameters for each layer. We use a different set of hyper-parameters for each layer type resulting in a different threshold for each layer type. The threshold is updated at regular intervals using the hyper-parameters according to Algorithm 1. We don’t modify the gradients in the back-propagation step. It is possible for the updates of a pruned weight to be larger than the threshold of that layer. In this case, the weight will be involved in the forward pass again.
We provide heuristics to help determine start itr, ramp itr and end itr in table 1. After picking these hyper parameters and assuming that ramp slope(φ) is 1.5× start slope (θ), we calculate (θ) using equation 1.
θ = 2 ∗ q ∗ freq
2 ∗ (ramp itr − start itr) + 3 ∗ (end itr − ramp itr) (1)
In order to determine q in equation 1, we use an existing weight array from a previously trained model. The weights are sorted using absolute values and we pick the weight corresponding to the 90th percentile as q. This allows us to pick reasonable values for the hyper-parameters required for pruning. A validation set can be used to fine tune these parameters.
We only prune the weights of the recurrent and linear layers but not the biases or batch norm parameters since they are much fewer in number compared to the weights. For the recurrent layers, we prune both the input weight matrix and the recurrent weight matrix. Similarly, we prune all the weights in gated recurrent units including those of the reset and update gates.
4 EXPERIMENTS
We run all our experiments on a training set of 2100 hours of English speech data and a validation set of 3.5 hours of multi-speaker data. This is a small subset of the datasets that we use to train our
Algorithm 1 Pruning Algorithm current itr = 0 while training do
for all parameters do param = (param and mask ) if current itr > start itr and current itr < end itr then
if (current itr mod freq) == 0 then if current itr < ramp itr then = θ ∗ (current itr − start itr + 1)/freq
else = (θ ∗ (ramp itr − start itr + 1) + φ ∗ (current itr − ramp itr + 1))/freq end if mask = abs(param) <
end if end if
end for current itr += 1
end while
state-of-the-art automatic speech recognition models. We train the models using Nesterov SGD for 20 epochs. Besides the hyper-parameters for determining the threshold, all other hyper-parameters remain unchanged between the dense and sparse training runs. We find that our pruning approach works well for vanilla bidirectional recurrent layers and forward only gated recurrent units.
4.1 BIDIRECTIONAL RNNS
We use the Deep Speech 2 model for these experiments. As shown in Table 2, this model has 2 convolution layers, followed by 7 bidirectional recurrent layers and a CTC cost layer. Each recurrent linear layer has 1760 hidden units, creating a network of approximately 67 million parameters. For these experiments, we prune the linear layers that feed into the recurrent layers, the forward and backward recurrent layers and fully connected layer before the CTC layer. These experiments use clipped rectified-linear units (ReLU) σ(x) = min(max(x, 0), 20) as the activation function.
In the sparse run, the pruning begins shortly after the first epoch and continues until the 10th epoch. We chose these hyper-parameters so that the model has an overall sparsity of 88% at the end of pruning, which is 8x smaller than the original dense model. The character error rate (CER) on the devset is about 20% worse relative to the dense model as shown in Table 3.
An argument against this sparsity result might be that we are taking advantage of a large model that overfits our relatively small dataset. In order to test this hypothesis, we train a dense model with 704 hidden units in each layer, that has approximately the same number of parameters as the final sparse model. Table 3 shows that this model performs worse than the sparse models. Thus sparse model is a better approach to reduce parameters than using a dense model with fewer hidden units.
In order to recover the loss in accuracy, we train sparse models with larger recurrent layers with 2560 and 3072 hidden units. Figure 1a shows the training and dev curves for these sparse models compared to the dense baseline model. These experiments use the same hyper-parameters (except for small changes in the pruning hyper-parameters) and the same dataset as the baseline model. As we see in Table 3, the model with 2560 hidden units achieves a 0.75% relative improvement compared to the dense baseline model, while the model with 3072 hidden units has a 3.95% improvement. The dense 2560 model also improves the CER by 11.85% relative to the dense baseline model. The sparse 2560 model is about 12% worse than the corresponding dense model. Both these large models are pruned to achieve a final sparsity of around 92%. These sparse larger models have significantly fewer parameters than the baseline dense model.
We also compare our gradual pruning approach to the hard pruning approach proposed in Yu et al. (2012). In their approach, all parameters below a certain threshold are pruned at particular epoch. Table 4 shows the results of pruning the RNN dense baseline model at different epochs to achieve final parameter count ranging from 8 million to 11 million. The network is trained for the same
number of epochs as the gradual pruning experiments. These hard threshold results are compared with the RNN Sparse 1760 model in Table 3. For approximately same number of parameters, gradual pruning is 7% to 9% better than hard pruning.
We conclude that pruning models to achieve sparsity of around 90% reduces the relative accuracy of the model by 10% to 20%. However, for a given performance requirement, it is better to prune a larger model than to use a smaller dense model. Gradually pruning a model produces better results than hard pruning.
4.2 GATED RECURRENT UNITS
We also experimented with GRU models shown in Table 5, that have 2560 hidden units in the GRU layer and a total of 115 million parameters. For these experiments, we prune all layers except the convolution layers since they have relatively fewer parameters.
Figure 1b compares the training and dev curves of a sparse GRU model a dense GRU model. The sparse GRU model has a 13.8% drop in the accuracy relative to the dense model. As shown in Table 3, the sparse model has an overall sparsity of 88.6% with 13 million parameters. Similar to the RNN models, we train a sparse GRU model with 3568 hidden units. The dataset and the hyperparameters are not changed from the previous GRU experiments. This model has an overall sparsity of 91.82% with 17.8 million parameters. As shown in Table 3, the model with 3568 hidden units is only 2.2% worse than the baseline dense GRU model. We expect to match the performance of the GRU dense network by slightly lowering the sparsity of this network or by increasing the hidden units for the layers.
In addition, we experimented with pruning only the GRU layers and keeping all the parameters in fully connected layers. The accuracy for these experiments is around 7% worse than the baseline dense model. However, this model only achieves 50% compression due to the size of the fully connected layers.
5 PERFORMANCE
5.1 COMPUTE TIME
The success of deep learning in recent years have been driven by large models trained on large datasets. However this also increases the inference time after the models have been deployed. We can mitigate this effect by using sparse layers.
A General Matrix-Matrix Multiply (GEMM) is the most compute intensive operation in evaluating a neural network model. Table 6 compares times for GEMM for recurrent layers with different number
Table 6: GEMM times for recurrent layers with different sparsity
LAYER SIZE SPARSITY LAYER TYPE TIME (µsec) SPEEDUP
1760 0% RNN 56 1 1760 95% RNN 20 2.8 2560 95% RNN 29 1.93 3072 95% RNN 48 1.16 2560 0% GRU 313 1 2560 95% GRU 46 6.80 3568 95% GRU 89 3.5
of hidden units that are 95% sparse. The performance benchmark was run using NVIDIA’s CUDNN and cuSPARSE libraries on a TitanX Maxwell GPU and compiled using CUDA 7.5. All experiments are run on a minibatch of 1 and in this case, the operation is known as a sparse matrix-vector product (SpMV). We can achieve speed-ups ranging from 3x to 1.15x depending on the size of the recurrent layer. Similarly, for the GRU models, the speed-ups range from 7x to 3.5x. However, we notice that cuSPARSE performance is substantially lower than the approximately 20x speedup that we would expect by comparing the bandwidth requirements of the 95% sparse and dense networks. State of the art SpMV routines can achieve close to device memory bandwidth for a wide array of matrix shapes and sparsity patterns (see Baxter (2016) and Liu et al. (2013)). This means that the performance should improve by the factor that parameter counts are reduced. Additionally, we find that the cuSPARSE performance degrades with larger batch sizes. It should be possible for a better implementation to further exploit the significant reuse of the weight matrix provided by large batch sizes.
5.2 COMPRESSION
Pruning allows us to reduce the memory footprint of a model which allows them to be deployed on phones and other embedded devices. The Deep Speech 2 model can be compressed from 268 MB to around 32 MB (1760 hidden units) or 64 MB (3072 hidden units). The GRU model can be compressed from 460 MB to 50 MB. These pruned models can be further quantized down to float16 or other smaller datatypes to further reduce the memory requirements without impacting accuracy.
6 DISCUSSION
6.1 PRUNING CHARACTERISTICS
Figure 2a shows the sparsity of all the recurrent layers with the same hyper-parameters used to prune the layers. The layers are ordered such that layer 1 is closest to input and layer 14 is the final recurrent layer before the cost layer. We see that the initial layers are pruned more aggressively compared to the final layers. We also performed experiments where the hyper parameters are different for the recurrent layers resulting in equal sparsity for all the layers. However, we get higher CER for these experiments. We conclude that to get good accuracy, it is important to prune the final layers slightly less than the initial ones.
In Figure 2b, we plot the pruning schedule of a 95% sparse recurrent layer of the bidirectional model trained for 20 epochs (55000 iterations). We begin pruning the network at the start of the second epoch at 2700 iterations. We stop pruning a layer after 10 epochs (half the total epochs) are complete at 27000 iterations. We see that nearly 25000 weights are pruned before 5 epochs are complete at around 15000 iterations. In our experiments, we’ve noticed that pruning schedules that are a convex curve tend to outperform schedules with a linear slope.
6.2 PERSISTENT KERNELS
Persistent Recurrent Neural Networks (Diamos et al., 2016) is a technique that increases the computational intensity of evaluating an RNN by caching the weights in on-chip memory such as caches, block RAM, or register files across multiple timesteps. A high degree of sparsity allows significantly large Persistent RNNs to be stored in on-chip memory. When all the weights are stored in float16, a NVIDIA P100 GPU can support a vanilla RNN size of about 2600 hidden units. With the same datatype, at 90% sparsity, and 99% sparsity, a P100 can support RNNs with about 8000, and 24000 hidden units respectively. We expect these kernels to be bandwidth limited out of the memory that is used to store the parameters. This offers the potential of a 146x speedup compared to the TitanX GPU if the entire RNN layer can be stored in registers rather than the GPU DRAM of a TitanX.
Additionally, sparse matrix multiplication involves scheduling and load balancing phases to divide the work up evenly over thousands of threads and to route corresponding weights and activations to individual threads. Since the sparsity patterns for RNNs are fixed over many timesteps these scheduling and load balancing operations can be factored outside of the loop, performed once, and reused many times.
7 CONCLUSION AND FUTURE WORK
We have demonstrated that by pruning the weights of RNNs during training we can find sparse models that are more accurate than dense models while significantly reducing model size. These sparse models are especially suited for deployment on mobile devices and on back-end server farms due to their small size and increased computational efficiency. Even with existing sub-optimal sparse matrix-vector libraries we realize speed-ups with these models. This technique is orthogonal to quantization techniques which would allow for even further reductions in model size and corresponding increase in performance.
We wish to investigate whether these techniques can generalize to language modeling tasks and if they can effectively reduce the size of embedding layers. We also wish to compare the sparsity generated by our pruning technique to that obtained by L1 regularization.
We are investigating training techniques that don’t require maintaining dense matrices for a significant portion of the calculation. Further work remains to implement optimal small batch sparse matrix-dense vector routine for GPUs and ARM processors that would help in deployment.
ACKNOWLEDGMENTS
We would like to thank Bryan Catanzaro for helpful discussions related to this work.

The paper proposes a method for pruning weights in neural networks during training to obtain sparse solutions. The approach is applied to an RNN-based system which is trained and evaluated on a speech recognition dataset. The results indicate that large savings in test-time computations can be obtained without affecting the task performance too much. In some cases the method can actually improve the evaluation performance.
The experiments are done using a state-of-the-art RNN system and the methodology of those experiments seems sound. I like that the effect of the pruning is investigated for networks of very large sizes. The computational gains are clearly substantial. It is a bit unfortunate that all experiments are done using a private dataset. Even with private training data, it would have been nice to see an evaluation on a known test set like the HUB5 for conversational speech. It would also have been nice to see a comparison with some other pruning approaches given the similarity of the proposed method to the work by Han et al. [2] to verify the relative merit of the proposed pruning scheme. While single-stage training looks more elegant at first sight, it may not save much time if more experiments are needed to find good hyperparameter settings for the threshold adaptation scheme. Finally, the dense baseline would have been more convincing if it involved some model compression tricks like training on the soft targets provided by a bigger network.
Overall, the paper is easy to read. The table and figure captions could be a bit more detailed but they are still clear enough. The discussion of potential future speed-ups of sparse recurrent neural networks and memory savings is interesting but not specific to the proposed pruning algorithm. The paper doesn’t motivate the details of the method very well. It’s not clear to me why the threshold has to ramp up after a certain period time for example. If this is based on preliminary findings, the paper should mention that.
Sparse neural networks have been the subject of research for a long time and this includes recurrent neural networks (e.g., sparse recurrent weight matrices were standard for echo-state networks [1]). The proposed method is also very similar to the work by Han et al. [2], where a threshold is used to prune weights after training, followed by a retraining phase of the remaining weights. While I think that it is certainly more elegant to replace this three stage procedure with a single training phase, the proposed scheme still contains multiple regimes that resemble such a process by first training without pruning followed by pruning at two different rates and finally training without further pruning again. The main novelty of the work would be the application of such a scheme to RNNs, which are typically more tricky to train than feedforward nets.
Improving scalability is an important driving force of the progress in neural network research. While I don’t think the paper presents much novelty in ideas or scientific insight, it does show that weight pruning can be successfully applied to large practical RNN systems without sacrificing much in performance. The fact that this is possible with such a simple heuristic is a result worth sharing.
Pros:
The proposed method is successful at reducing the number of parameters in RNNs substantially without sacrificing too much in performance.
The experiments are done using a state-of-the-art system for a practical application.
Cons:
The proposed method is very similar to earlier work and barely novel.
There is no comparison with other pruning methods.
The data is private and this prevents others from replicating the results.
[1] Jaeger, H. (2001). The “echo state” approach to analyzing and training recurrent neural networks-with an erratum note. Bonn, Germany: German National Research Center for Information Technology GMD Technical Report, 148, 34.
[2] Han, Song, Pool, Jeff, Tran, John, and Dally, William J. Learning both weights and connections for efficient neural networks. In Advances in Neural Information Processing Systems, 2015.