In this paper, we propose and investigate a novel memory architecture for neural networks called Hierarchical Attentive Memory (HAM). It is based on a binary tree with leaves corresponding to memory cells. This allows HAM to perform memory access in Θ(log n) complexity, which is a significant improvement over the standard attention mechanism that requires Θ(n) operations, where n is the size of the memory. We show that an LSTM network augmented with HAM can learn algorithms for problems like merging, sorting or binary searching from pure input-output examples. In particular, it learns to sort n numbers in time Θ(n log n) and generalizes well to input sequences much longer than the ones seen during the training. We also show that HAM can be trained to act like classic data structures: a stack, a FIFO queue and a priority queue.
1 INTRO
Deep Recurrent Neural Networks (RNNs) have recently proven to be very successful in real-word tasks, e.g. machine translation (Sutskever et al., 2014) and computer vision (Vinyals et al., 2014). However, the success has been achieved only on tasks which do not require a large memory to solve the problem, e.g. we can translate sentences using RNNs, but we cannot produce reasonable translations of really long pieces of text, like books.
A high-capacity memory is a crucial component necessary to deal with large-scale problems that contain plenty of long-range dependencies. Currently used RNNs do not scale well to larger memories, e.g. the number of parameters in an LSTM (Hochreiter & Schmidhuber, 1997) grows quadratically with the size of the network’s memory. In practice, this limits the number of used memory cells to few thousands.
It would be desirable for the size of the memory to be independent of the number of model parameters. The first versatile and highly successful architecture with this property was Neural Turing Machine (NTM) proposed by Graves et al. (2014). The main idea behind the NTM is to split the network into a trainable “controller” and an “external” variable-size memory. It caused an outbreak of other neural network architectures with external memories (see Sec. 2).
However, one aspect which has been usually neglected so far is the efficiency of the memory access. Most of the proposed memory architectures have the Θ(n) access complexity, where n is the size of the memory. It means that, for instance, copying a sequence of length n requires performing Θ(n2) operations, which is clearly unsatisfactory.
1.1 OUR CONTRIBUTION
We propose a novel memory module for neural networks, called Hierarchical Attentive Memory (HAM). The HAM module is generic and can be used as a building block of larger neural architectures. Its crucial property is that it scales well with the memory size — the memory access requires only Θ(log n) operations, where n is the size of the memory. This complexity is achieved by using a new attention mechanism based on a binary tree with leaves corresponding to memory cells. The novel attention mechanism is not only faster than the standard one used in Deep Learning (Bahdanau et al., 2014), but it also facilities learning algorithms due to a built-in bias towards operating on intervals.
∗Equal contribution.
We show that an LSTM augmented with HAM is able to learn algorithms for tasks like merging, sorting or binary searching. In particular, it is the first neural network, which we are aware of, that is able to learn to sort from pure input-output examples and generalizes well to input sequences much longer than the ones seen during the training. Moreover, the learned sorting algorithm runs in time Θ(n log n). We also show that the HAM memory itself is capable of simulating different classic memory structures: a stack, a FIFO queue and a priority queue.
2 RELATED WORK
In this section we mention a number of recently proposed neural architectures with an external memory, which size is independent of the number of the model parameters.
Memory architectures based on attention Attention is a recent but already extremely successful technique in Deep Learning. This mechanism allows networks to attend to parts of the (potentially preprocessed) input sequence (Bahdanau et al., 2014) while generating the output sequence. It is implemented by giving the network as an auxiliary input a linear combination of input symbols, where the weights of this linear combination can be controlled by the network. Attention mechanism was used to access the memory in Neural Turing Machines (NTMs) proposed by Graves et al. (2014). It was the first paper, that explicitly attempted to train a computationally universal neural network and achieved encouraging results.
The Memory Network (Weston et al., 2014) is an early model that attempted to explicitly separate the memory from computation in a neural network model. The followup work of (Sukhbaatar et al., 2015) combined the memory network with the soft attention mechanism, which allowed it to be trained with less supervision. In contrast to NTMs, the memory in these models is non-writeable. Another model without writeable memory is the Pointer Network (Vinyals et al., 2015), which is very similar to the attention model of Bahdanau et al. (2014). Despite not having a memory, this model was able to solve a number of difficult algorithmic problems, like the Convex Hull and the approximate 2D TSP.
All of the architectures mentioned so far use standard attention mechanisms to access the memory and therefore memory access complexity scales linearly with the memory size.
Memory architectures based on data structures Stack-Augmented Recurrent Neural Network (Joulin & Mikolov, 2015) is a neural architecture combining an RNN and a differentiable stack. Grefenstette et al. (2015) consider extending an LSTM with a stack, a FIFO queue or a double-ended queue and show some promising results. The advantage of the latter model is that the presented data structures have a constant access time.
Memory architectures based on pointers In two recent papers (Zaremba & Sutskever, 2015; Zaremba et al., 2015) authors consider extending neural networks with nondifferentiable memories based on pointers and trained using Reinforcement Learning. The big advantage of these models is that they allow a constant time memory access. They were however only successful on relatively simple tasks.
Another model, which use a pointer-based memory and learns sub-procedures is the Neural Programmer-Interpreter (Reed & de Freitas, 2015). Unfortunately, it requires strong supervision in the form of execution traces. Different type of pointer-based memory was presented in Neural Random-Access Machine (Kurach et al., 2015), which is a neural architecture mimicking classic computers.
Parallel memory architectures There are two recent memory architectures, which are especially suited for parallel computation. Grid-LSTM (Kalchbrenner et al., 2015) is an extension of LSTM to multiple dimensions. Another recent model of this type is Neural GPU (Kaiser & Sutskever, 2015), which can learn to multiply long binary numbers.
3 HIERARCHICAL ATTENTIVE MEMORY
In this section we describe our novel memory module called Hierarchical Attentive Memory (HAM). The HAM module is generic and can be used as a building block of larger neural network architectures. For instance, it can be added to feedforward or LSTM networks to extend their capabilities. To make our description more concrete we will consider a model consisting of an LSTM “controller” extended with a HAM module.
The high-level idea behind the HAM module is as follows. The memory is structured as a full binary tree with the leaves containing the data stored in the memory. The inner nodes contain some auxiliary data, which allows us to efficiently perform some types of “queries” on the memory. In order to access the memory, one starts from the root of the tree and performs a top-down descent in the tree, which is similar to the hierarchical softmax procedure (Morin & Bengio, 2005). At every node of the tree, one decides to go left or right based on the auxiliary data stored in this node and a “query”. Details are provided in the rest of this section.
3.1 NOTATION
The model takes as input a sequence x1, x2, . . . and outputs a sequence y1, y2, . . .. We assume that each element of these sequences is a binary vector of size b ∈ N, i.e. xi, yi ∈ {0, 1}b. Suppose for a moment that we only want to process input sequences of length ≤ n, where n ∈ N is a power of two (we show later how to process sequences of an arbitrary length). The model is based on the full binary tree with n leaves. Let V denote the set of the nodes in that tree (notice that |V | = 2n− 1) and let L ⊂ V denote the set of its leaves. Let l(e) for e ∈ V \ L be the left child of the node e and let r(e) be its right child. We will now present the inference procedure for the model and then discuss how to train it.
3.2 INFERENCE y1
LSTM
HAM
x1 . . . xm
y2
LSTM
HAM
y3
LSTM
HAM
. . .
Figure 1: The LSTM+HAM model consists of an LSTM controller and a HAM module. The execution of the model starts with the initialization of HAM using the whole input sequence x1, x2, . . . , xm. At each timestep, the HAM module produces an input for the LSTM, which then produces an output symbol yt. Afterwards, the hidden states of the LSTM and HAM are updated.
The high-level view of the model execution is presented in Fig. 1. The hidden state of the model consists of two components: the hidden state of the LSTM controller (denoted hLSTM ∈ Rl for some l ∈ N) and the hidden values stored in the nodes of the HAM tree. More precisely, for every node e ∈ V there is a hidden value he ∈ Rd. These values change during the recurrent execution of the model, but we drop all timestep indices to simplify the notation.
The parameters of the model describe the inputoutput behaviour of the LSTM, as well as the following 4 transformations, which describe the HAM module: EMBED : Rb → Rd, JOIN : Rd × Rd → Rd, SEARCH : Rd × Rl → [0, 1] and WRITE : Rd × Rl → Rd. These transformations may be represented by arbitrary function approximators, e.g. Multilayer Perceptrons (MLPs). Their meaning will be described soon.
The details of the model are presented in 4 figures. Fig. 2a describes the initialization of the model. Each recurrent timestep of the model consists of three phases: the attention phase described in Fig. 2b, the output phase described in Fig. 2c and the update phase described in Fig. 2d. The whole timestep can be performed in time Θ(log n).
The HAM parameters describe only the 4 mentioned transformations and hence the number of the model parameters does not depend on the size of the binary tree used. Thus, we can use the model to process the inputs of an arbitrary length by using big enough binary trees. It is not clear that the same set of parameters will give good results across different tree sizes, but we showed experimentally that it is indeed the case (see Sec. 4 for more details).
h1
h2 h3
h4 h5 h6 h7
h8 h9 h10 h11 h12 h13 h14 h15
x1 x2 x3 x4 x5 x6
EMBED EMBED EMBED EMBED EMBED EMBED
JOIN
JOIN JOIN
JOIN JOIN JOIN JOIN
(a) Initialization of the model. The value in the i-th leaf of HAM is initialized with EMBED(xi), where EMBED is a trainable feed-forward network. If there are more leaves than input symbols, we initialize the values in the excessive leaves with zeros. Then, we initialize the values in the inner nodes bottom-up using the formula he = JOIN(hl(e), hr(e)). The hidden state of the LSTM — hLSTM is initialized with zeros.
h1
h2 h3
h4 h5 h6 h7
h8 h9 h10 h11 h12 ha h14 h15
SEARCH(h1, hLSTM) = 0.95
SEARCH(h3, hLSTM) = 0.1
SEARCH(h6, hLSTM) = 1
(b) Attention phase. In this phase the model performs a top-down “search” in the tree starting from the root. Suppose that we are currently at the node c ∈ V \ L. We compute the value p = SEARCH(hc, hLSTM). Then, with probability p the model goes right (i.e. c := r(c)) and with probability 1 − p it goes left (i.e. c := l(c)). This procedure is continued until we reach one of the leaves. This leaf is called the attended or accessed leaf and denoted a.
ha hLSTM yt
(c) Output phase. The value ha stored in the attended leaf is given to the LSTM as an input. Then, the LSTM produces an output symbol yt ∈ {0, 1}b. More precisely, the value u ∈ Rb is computed by a trainable linear transformation from hLSTM and the distribution of yt is defined by the formula p(yt,i = 1) = sigmoid(ui) for 1 ≤ i ≤ b. It may be beneficial to allow the model to access the memory a few times between producing each output symbols. Therefore, the model produces an output symbol only at timesteps with indices divisible by some constant η ∈ N, which is a hyperparameter.
h1
h2 h3
h4 h5 h6 h7
h8 h9 h10 h11 h12 ha h14 h15
hLSTM ha := WRITE(ha, hLSTM)
JOIN
JOIN
JOIN
(d) Update phase. In this phase the value in the attended leaf a is updated. More precisely, the value is modified using the formula ha := WRITE(ha, hLSTM). Then, we update the values of the inner nodes encountered during the attention phase (h6, h3 and h1 in the figure) bottom-up using the equation he = JOIN(hl(e), hr(e)).
Figure 2: The model. One timestep consists of three phases presented in Figures (b)–(d).
We decided to represent the transformations defining HAM with MLPs with ReLU (Nair & Hinton, 2010) activation function in all neurons except the output layer of SEARCH, which uses sigmoid activation function to ensure that the output may be interpreted as a probability. Moreover, the network for WRITE is enhanced in a similar way as Highway Networks (Srivastava et al., 2015), i.e. WRITE(ha, hLSTM) = T (ha, hLSTM) ·H(ha, hLSTM) + (1− T (ha, hLSTM)) · ha, where H and T are two MLPs with sigmoid activation function in the output layer. This allows the WRITE transformation to easily leave the value ha unchanged.
3.3 TRAINING
In this section we describe how to train our model from purely input-output examples using REINFORCE (Williams, 1992). In Appendix B we also present a different variant of HAM which is fully differentiable and can be trained using end-to-end backpropagation.
Let x, y be an input-output pair. Recall that both x and y are sequences. Moreover, let θ denote the parameters of the model and let A denote the sequence of all decisions whetherto go left or right made during the whole execution of the model. We would like to maximize the log-probability of producing the correct output, i.e.
L = log p(y|x, θ) = log (∑ A p(A|x, θ)p(y|A, x, θ) ) .
This sum is intractable, so instead of minimizing it directly, we minimize a variational lower bound on it: F = ∑ A p(A|x, θ) log p(y|A, x, θ) ≤ L.
This sum is also intractable, so we approximate its gradient using the REINFORCE, which we briefly explain below. Using the identity ∇p(A|x, θ) = p(A|x, θ)∇ log p(A|x, θ), the gradient of the lower bound with respect to the model parameters can be rewritten as:
∇F = ∑ A p(A|x, θ) [ ∇ log p(y|A, x, θ) + log p(y|A, x, θ)∇ log p(A|x, θ) ] (1)
We estimate this value using Monte Carlo approximation.For every x we sample Ã from p(A|x, θ) and approximate the gradient for the input x as∇ log p(y|Ã, x, θ) + log p(y|Ã, x, θ)∇ log p(Ã|x, θ). Notice that this gradient estimate can be computed using normal backpropagation if we substitute the gradients in the nodes1 which sample whether we should go left or right during the attention phase by
log p(y|Ã, x, θ)︸ ︷︷ ︸ return ∇ log p(Ã|x, θ).
This term is called REINFORCE gradient estimate and the left factor is called a return in Reinforcement Learning literature. This gradient estimator is unbiased, but it often has a high variance. Therefore, we employ two standard variance-reduction technique for REINFORCE: discounted returns and baselines (Williams, 1992). Discounted returns means that our return at the t-th timestep has the form ∑ t≤i γ
i−t log p(yi|Ã, x, θ) for some discount constant γ ∈ [0, 1], which is a hyperparameter. This biases the estimator if γ < 1, but it often decreases its variance.
For the lack of space we do not describe the baselines technique. We only mention that our baseline is case and timestep dependent: it is computed using a learnable linear transformation from hLSTM and trained using MSE loss function. The whole model is trained with the Adam (Kingma & Ba, 2014) algorithm. We also employ the following three training techniques:
Different reward function During our experiments we noticed that better results may be obtained by using a different reward function for REINFORCE. More precisely, instead of the log-probability of producing the correct output, we use the percentage of the output bits, which have the probability of being predicted correctly (given Ã) greater than 50%, i.e. our discounted return is equal∑ t≤i,1≤j≤b γ i−t [ p(yi,j |Ã, x, θ) > 0.5 ] . Notice that it corresponds to the Hamming distance be-
tween the most probable outcome accordingly to the model (given Â) and the correct output.
Entropy bonus term We add a special term to the cost function which encourages exploration. More precisely, for each sampling node we add to the cost function the term αH(p) , where H(p) is the entropy of the distribution of the decision, whether to go left or right in this node and α is an exponentially decaying coefficient. This term goes to infinity, whenever the entropy goes to zero, what ensures some level of exploration. We noticed that this term works better in our experiments than the standard term of the form −αH(p) (Williams, 1992).
Curriculum schedule We start with training on inputs with lengths sampled uniformly from [1, n] for some n = 2k and the binary tree with n leaves. Whenever the error drops below some threshold, we increment the value k and start using the bigger tree with 2n leaves and inputs with lengths sampled uniformly from [1, 2n].
1 For a general discussion of computing gradients in computation graphs, which contain stochastic nodes see (Schulman et al., 2015).
4 EXPERIMENTS
In this section, we evaluate two variants of using the HAM module. The first one is the model described in Sec. 3, which combines an LSTM controller with a HAM module (denoted by LSTM+HAM). Then, in Sec. 4.3 we investigate the “raw” HAM (without the LSTM controller) to check its capability of acting as classic data structures: a stack, a FIFO queue and a priority queue. It would be also interesting to get some insight into the algorithms learned by the model. In Appendix A we present an example execution on the Sort task.
4.1 TEST SETUP
For each test that we perform, we apply the following procedure. First, we train the model with memory of size up to n = 32 using the curriculum schedule described in Sec. 3.3. The model is trained using the minibatch Adam algorithm with exponentially decaying learning rate. We use random search to determine the best hyper-parameters for the model. We use gradient clipping (Pascanu et al., 2012) with constant 5. The depth of our MLPs is either 1 or 2, the LSTM controller has l = 20 memory cells and the hidden values in the tree have dimensionality d = 20. Constant η determining a number of memory accesses between producing each output symbols (Fig. 2c) is equal either 1 or 2. We always train for 100 epochs, each consisting of 1000 batches of size 50. After each epoch we evaluate the model on 200 validation batches without learning. When the training is finished, we select the model parameters that gave the lowest error rate on validation batches and report the error using these parameters on fresh 2, 500 random examples.
We report two types of errors: a test error and a generalization error. The test error shows how well the model is able to fit the data distribution and generalize to unknown cases, assuming that cases of similar lengths were shown during the training. It is computed using the HAM memory with n = 32 leaves, as the percentage of output sequences, which were predicted incorrectly. The lengths of test examples are sampled uniformly from the range [1, n]. Notice that we mark the whole output sequence as incorrect even if only one bit was predicted incorrectly, e.g. a hypothetical model predicting each bit incorrectly with probability 1% (and independently of the errors on the other bits) has an error rate of 96% on whole sequences if outputs consist of 320 bits.
The generalization error shows how well the model performs with enlarged memory on examples with lengths exceeding n. We test our model with memory 4 times bigger than the training one. The lengths of input sequences are now sampled uniformly from the range [2n+ 1, 4n].
During testing we make our model fully deterministic by using the most probable outcomes instead of stochastic sampling. More precisely, we assume that during the attention phase the model decides to go right iff p > 0.5 (Fig. 2b). Moreover, the output symbols (Fig. 2c) are computed by rounding to zero or one instead of sampling.
4.2 LSTM+HAM
We evaluate the model on a number of algorithmic tasks described below:
1. Reverse: Given a sequence of 10-bit vectors, output them in the reversed order., i.e. yi = xm+1−i for 1 ≤ i ≤ m, where m is the length of the input sequence.
2. Search: Given a sequence of pairs xi = keyi||valuei for 1 ≤ i ≤ m− 1 sorted by keys and a query xm = q, find the smallest i such that keyi = q and output y1 = valuei. Keys and values are 5-bit vectors and keys are compared lexicographically. The LSTM+HAM model is given only two timesteps (η = 2) to solve this problem, which forces it to use a form of binary search.
3. Merge: Given two sorted sequences of pairs — (p1, v1), . . . , (pm, vm) and (p′1, v ′ 1), . . . , (p ′ m′ , v ′ m′), where pi, p ′ i ∈ [0, 1] and vi, v′i ∈ {0, 1}5, merge them. Pairs
are compared accordingly to their priorities, i.e. values pi and p′i. Priorities are unique and sampled uniformly from the set { 1300 , . . . , 300 300}, because neural networks cannot easily distinguish two real numbers which are very close to each other. Input is encoded as
xi = pi||vi for 1 ≤ i ≤ m and xm+i = p′i||v′i for 1 ≤ i ≤ m′. The output consists of the vectors vi and v′i sorted accordingly to their priorities 2.
4. Sort: Given a sequence of pairs xi = keyi||valuei sort them in a stable way3 accordingly to the lexicographic order of the keys. Keys and values are 5-bit vectors.
5. Add: Given two numbers represented in binary, compute their sum. The input is represented as a1, . . . , am,+, b1, . . . , bm,= (i.e. x1 = a1, x2 = a2 and so on), where a1, . . . , am and b1, . . . , bm are bits of the input numbers and +,= are some special symbols. Input and output numbers are encoded starting from the least significant bits.
Every example output shown during the training is finished by a special “End Of Output” symbol, which the model learns to predict. It forces the model to learn not only the output symbols, but also the length of the correct output.
We compare our model with 2 strong baseline models: encoder-decoder LSTM (Sutskever et al., 2014) and encoder-decoder LSTM with attention (Bahdanau et al., 2014), denoted LSTM+A. The number of the LSTM cells in the baselines was chosen in such a way, that they have more parameters than the biggest of our models. We also use random search to select an optimal learning rate and some other parameters for the baselines and train them using the same curriculum scheme as LSTM+HAM.
Table 1: Experimental results. The upper table presents the error rates on inputs of the same lengths as the ones used during training. The lower table shows the error rates on input sequences 2 to 4 times longer than the ones encountered during training. LSTM+A denotes an LSTM with the standard attention mechanism. Each error rate is a percentage of output sequences, which contained at least one incorrectly predicted bit.
test error LSTM LSTM+A LSTM+ HAM Reverse 73% 0% 0% Search 62% 0.04% 0.12% Merge 88% 16% 0% Sort 99% 25% 0.04% Add 39% 0% 0%
2-4x longer inputs LSTM LSTM+ A LSTM+ HAM
Reverse 100% 100% 0% Search 89% 0.52% 1.68% Merge 100% 100% 2.48% Sort 100% 100% 0.24% Add 100% 100% 100%
Complexity Θ(1) Θ(n) Θ(log n)
The results are presented in Table 1. Not only, does LSTM+HAM solve all the problems almost perfectly, but it also generalizes very well to much longer inputs on all problems except Add. Recall that for the generalization tests we used a HAM memory of a different size than the ones used during the training, what shows that HAM generalizes very well to new sizes of the binary tree. We find this fact quite interesting, because it means that parameters learned from a small neural network (i.e. HAM based on a tree with 32 leaves) can be successfully used in a different, bigger network (i.e. HAM with 128 memory cells).
In comparison, the LSTM with attention does not learn to merge, nor sort. It also completely fails to generalize to longer examples, which shows that LSTM+A learns rather some statistical dependencies between inputs and outputs than the real algorithms.
The LSTM+HAM model makes a few errors when testing on longer outputs than the ones encountered during the training. Notice however, that we show in the table the percentage of output sequences, which contain at least one incorrect bit. For instance, LSTM+HAM on the problem Merge predicts incorrectly only 0.03% of output bits, which corresponds to 2.48% of incorrect output sequences. We believe that these rare mistakes could be avoided if one trained the model longer and chose carefully the learning rate schedule. One more way to boost generalization would be to simultaneously train the models with different memory sizes and shared parameters. We have not tried this as the generalization properties of the model were already very good.
2 Notice that we earlier assumed for the sake of simplicity that the input sequences consist of binary vectors and in this task the priorities are real values. It does not however require any change of our model. We decided to use real priorities in this task in order to diversify our set of problems.
3Stability means that pairs with equal keys should be ordered accordingly to their order in the input sequence.
4.3 RAW HAM
In this section, we evaluate “raw” HAM module (without the LSTM controller) to see if it can act as a drop-in replacement for 3 classic data structures: a stack, a FIFO queue and a priority queue. For each task, the network is given a sequence of PUSH and POP operations in an online manner: at timestep t the network sees only the t-th operation to perform xt. This is a more realistic scenario for data structures usage as it prevents the network from cheating by peeking into the future.
Raw HAM module differs from the LSTM+HAM model from Sec. 3 in the following way:
• The HAM memory is initialized with zeros. • The t-th output symbol yt is computed using an MLP from the value in the accessed leaf ha. • Notice that in the LSTM+HAM model, hLSTM acted as a kind of “query” or “command”
guiding the behaviour of HAM. We will now use the values xt instead. Therefore, at the t-th timestep we use xt instead of hLSTM whenever hLSTM was used in the original model, e.g. during the attention phase (Fig. 2b) we use p = SEARCH(hc, xt) instead of p = SEARCH(hc, hLSTM).
We evaluate raw HAM on the following tasks:
1. Stack: The “PUSH x” operation places the element x (a 5-bit vector) on top of the stack, and the “POP” returns the last added element and removes it from the stack.
2. Queue: The “PUSH x” operation places the element x (a 5-bit vector) at the end of the queue and the “POP” returns the oldest element and removes it from the queue.
3. PriorityQueue: The “PUSH x p” operations adds the element x with priority p to the queue. The “POP” operation returns the value with the highest priority and remove it from the queue. Both x and p are represented as 5-bit vectors and priorities are compared lexicographically. To avoid ties we assume that all elements have different priorities.
Table 2: Results of experiments with the raw version of HAM (without the LSTM controller). Error rates are measured as a percentage of operation sequences in which at least one POP query was not answered correctly.
Task Test Error GeneralizationError Stack 0% 0% Queue 0% 0%
Priority Queue
0.08% 0.2%
Model was trained with the memory of size up to n = 32 with operation sequences of length n. Sequences of PUSH/POP actions for training were selected randomly. The t-th operation out of n operations in the sequence was POP with probability tn and PUSH otherwise. To test generalization, we report the error rates with the memory of size 4n on sequences of operations of length 4n.
The results presented in Table 2 show that HAM simulates a stack and a queue perfectly with no errors whatsoever even for memory 4 times bigger. For the PriorityQueue task, the model generalizes almost perfectly to large memory, with errors only in 0.2% of output sequences.
5 COMPARISON TO OTHER MODELS
As far as we know, our model is the first one which is able to learn a sorting algorithm from pure input-output examples. Although this problem was considered in the original NTM paper, the error rate achieved by the NTM is in fact quite high – the log-likelihood of the correct output was equal around 20 bits on outputs consisting of 128 bits. In comparison our model learns to solve almost perfectly - only 0.04% of the outputs produced by our model contain at least one incorrect bit.
Reed & de Freitas (2015) shown that an LSTM is able to learn to sort short sequences, but it fails to generalize to inputs longer than the ones seen during the training. It is quite clear that an LSTM cannot learn a “real” sorting algorithm, because it uses a bounded memory independent of the length of the input. The Neural Programmer-Interpreter (Reed & de Freitas, 2015) is a neural network architecture, which is able to learn bubble sort, but it requires strong supervision in the form of
execution traces. In comparison, our model can be trained from pure input-output examples, which is crucial if we want to use it to solve problems for which we do not know any algorithms.
An important feature of neural memories is their efficiency. Our HAM module in comparison to many other recently proposed solutions is effective and allows to access the memory in Θ(log(n)) complexity. In the context of learning algorithms it may sound surprising that among all the architectures mentioned in Sec. 2 the only ones, which can copy a sequence of length n without Θ(n2) operations are: Reinforcement-Learning NTM (Zaremba & Sutskever, 2015), the model from (Zaremba et al., 2015), Neural Random-Access Machine (Kurach et al., 2015), and Queue-Augmented LSTM (Grefenstette et al., 2015). However, the first three models have been only successful on relatively simple tasks. The last model was successful on some synthetic tasks from the domain of Natural Language Processing, which are very different from the tasks we tested our model on, so we cannot directly compare the two models.
6 CONCLUSIONS
We presented a new memory architecture for neural networks called Hierarchical Attentive Memory. Its crucial property is that it scales well with the memory size — the memory access requires only Θ(log n) operations. This complexity is achieved using a new attention mechanism based on a binary tree. The model proved to be successful on a number of algorithmic problems. The future work is to apply this or similar architecture to very long real-world sequential data like books or DNA sequences.
A EXAMPLE: HAM SORTING
We present some insights into the algorithms learned by the LSTM+HAM model, by investigating the hidden representations he learned for a variant of the problem Sort in which we sort 4-bit vectors lexicographically4. For demonstration purposes, we use a small tree with n = 8 leaves and each node’s hidden state has size d = 6 values.
The trained network performs sorting perfectly. It attends to the leaves in the order corresponding to the order of the sorted input values, i.e. at every timestep HAM attends to the leaf corresponding to the smallest input value among the leaves, which have not been attended so far.
It would be interesting to exactly understand the algorithm used by the network to perform this operation. A natural solution to this problem would be to store in each hidden node e the smallest input value among the (unattended so far) leaves below e together with the information whether the smallest value is in the right or the left subtree under e.
In the Fig. 3 we present two timesteps of our model. The LSTM controller is not presented to simplify the exposition. The input sequence is presented on the left, below the tree: x1 = 0000, x2 = 1110, x3 = 1101 and so on. The 2x3 grids in the nodes of the tree represent the values he ∈ R6. White cells correspond to value 0 and non-white cells correspond to values > 0.
The lower-rightmost cells are presented in pink, because we managed to decipher the meaning of this coordinate for the inner nodes. This coordinate in the node e denotes whether the minimum in the subtree (among the values unattended so far) is in the right or left subtree of e. Value greater than 0 (pink in the picture) means that the minimum is in the right subtree and therefore we should go right while visiting this node in the attention phase.
In the first timestep the leftmost leaf (corresponding to the input 0000) is accessed. Notice that the last coordinates (shown in pink) are updated appropriately, e.g. the smallest unattended value at the beginning of the second timestep is 0101, which corresponds to the 6-th leaf. It is in the right subtree under the root and accordingly the last coordinate in the hidden value stored in the root is high (i.e. pink in the figure).
4 In the problem Sort considered in the experimental results, there are separate keys and values, which forces the model to learn stable sorting. Here, for the sake of simplicity, we consider the simplified version of the problem and do not use separate keys and values.
B USING SOFT ATTENTION
One of the open questions in the area of designing neural networks with attention mechanisms is whether to use a soft or hard attention. The model described in the paper belongs to the latter class of attention mechanisms as it makes hard, stochastic choices. The other solution would be to use a soft, differentiable mechanism, which attends to a linear combination of the potential attention targets and do not involve any sampling. The main advantage of such models is that their gradients can be computed exactly.
We now describe how to modify the model to make it fully differentiable ("DHAM"). Recall that in the original model the leaf which is attended at every timestep is sampled stochastically. Instead of that, we will now at every timestep compute for every leaf e the probability p(e) that this leaf would be attended if we used the stochastic procedure described in Fig. 2b. The value p(e) can be computed by multiplying the probabilities of going in the right direction from all the nodes on the path from the root to e.
As the input for the LSTM we then use the value ∑ e∈L p(e)·he. During the write phase, we update the values of all the leaves using the formula he := p(e)·WRITE(he, hROOT)+(1−p(e))·he. Then, in the update phase we update the values of all the inner nodes, so that the equation he = JOIN(hl(e), hr(e)) is satisfied for each inner node e. Notice that one timestep of the soft version of the model takes time Θ(n) as we have to update the values of all the nodes in the tree. Our model may be seen as a special case of Gated Graph Neural Network (Li et al., 2015).
This version of the model is fully differentiable and therefore it can be trained using end-to-end backpropagation on the log-probability of producing the correct output. We observed that training DHAM is slightly easier than the REINFORCE version. However, DHAM does not generalize as well as HAM to larger memory sizes.

This paper proposes to use a hierarchical softmax to speed up attention based memory addressing in memory augmented network (e.g. NTM, memNN…).
The model build a hierarchical softmax on top of the input sequence then at each time step SEARCH for the most relevant input to predict the next output (this search is discrete), and use its corresponding embedding to update the state of an LSTM that will then produce the output. Finally the embedding of the used input is update by a WRITE function (an LSTM working that takes hidden state of the other LSTM as an input). The model has a discrete component (the SEARCH) and is thus trained with REINFORCE. In the experimental section they test their approach on several algorithmic tasks such as search, sort...
The main advantage of replacing the full softmax by a hierarchical softmax is that during inference, the complexity goes from O(N) to O(log(N)). It would be great to see if the gain in complexity allows to tackle problem which are a few orders of magnitude bigger than the one addressed with full softmax. However the authors only test on toy sequences up to 32 tokens, which is quite small.
The model requires a relatively complex search mechanism that can only be trained with REINFORCE. While this seems to work on problems with relatively small and simple sequences, it would be great to see how performance changes with the size of the problem.
Overall, while the idea of replacing the softmax in the attention mechanism by a hierachical softmax is appealing, this work is not quite convincing yet. Their approach is not very natural, may be hard to train and may not be that simple to scale. The experiment section is very weak.


REVISITING DISTRIBUTED SYNCHRONOUS SGD
Distributed training of deep learning models on large-scale training data is typically conducted with asynchronous stochastic optimization to maximize the rate of updates, at the cost of additional noise introduced from asynchrony. In contrast, the synchronous approach is often thought to be impractical due to idle time wasted on waiting for straggling workers. We revisit these conventional beliefs in this paper, and examine the weaknesses of both approaches. We demonstrate that a third approach, synchronous optimization with backup workers, can avoid asynchronous noise while mitigating for the worst stragglers. Our approach is empirically validated and shown to converge faster and to better test accuracies.
1 INTRODUCTION
The recent success of deep learning approaches for domains like speech recognition (Hinton et al., 2012) and computer vision (Ioffe & Szegedy, 2015) stems from many algorithmic improvements but also from the fact that the size of available training data has grown significantly over the years, together with the computing power, in terms of both CPUs and GPUs. While a single GPU often provides algorithmic simplicity and speed up to a given scale of data and model, there exist an operating point where a distributed implementation of training algorithms for deep architectures becomes necessary.
Currently, popular distributed training algorithms include mini-batch versions of stochastic gradient descent (SGD) and other stochastic optimization algorithms such as AdaGrad (Duchi et al., 2011), RMSProp (Tieleman & Hinton, 2012), and ADAM (Kingma & Ba, 2014). Unfortunately, bulksynchronous implementations of stochastic optimization are often slow in practice due to the need to wait for the slowest machine in each synchronous batch. To circumvent this problem, practitioners have resorted to asynchronous approaches which emphasize speed by using potentially stale information for computation. While asynchronous training have proven to be faster than their synchronous counterparts, they often result in convergence to poorer results.
In this paper1, we revisit synchronous learning, and propose a method for mitigating stragglers in synchronous stochastic optimization. Specifically, we synchronously compute a mini-batch gradient with only a subset of worker machines, thus alleviating the straggler effect while avoiding any staleness in our gradients. The primary contributions of our paper are:
• Illustration of how gradient staleness in asynchronous training negatively impacts test accuracy and is exacerbated by deep models.
• Measurement of machine response times for synchronous stochastic optimization in a large deployment of 100 GPUs, showing how stragglers in the tail end affect convergence speed.
• Proposal of synchronous stochastic optimization with backup workers to mitigate straggler effects without gradient staleness.
• Establishing the need to measure both speed of convergence and test accuracy of optimum for empirical validation.
∗Joint first authors †UC Berkeley, Berkeley, CA, USA, xinghao@eecs.berkeley.edu 1This is an extension of our ICLR 2016 workshop extended abstract (Chen et al., 2016).
• Empirical demonstration that our proposed synchronous training method outperforms asynchronous training by converging faster and to better test accuracies.
The remainder of this paper is organized as follows. We briefly present preliminaries and notation in Section 1.1. Section 2 describes asynchronous stochastic optimization and presents experimental evidence of gradient staleness in deep neural network models. We present our approach in Section 3, and exhibit straggler effects that motivate the approach. We then empirically evaluate our approach in Sections 4. Related work is discussed in Section 5, and we conclude in Section 6.
1.1 PRELIMINARIES AND NOTATION
Given a dataset X = {xi : i = 1, . . . , |X |}, our goal is to learn the parameters θ of a model with respect to an empirical loss function f , defined as f(θ) ∆= 1|X | ∑|X | i=1 F (xi; θ), where F (xi; θ) is the loss with respect to a datapoint xi and the model θ.
A first-order stochastic optimization algorithm achieves this by iteratively updating θ using a stochastic gradient G ∆= ∇F (xi; θ) computed at a randomly sampled xi, producing a sequence of models θ(0), θ(1), . . . . Stochastic optimization algorithms differ in their update equations. For example, the update of SGD is θ(t+1) = θ(t) − γtG(t) = θ(t) − γt∇F (xi; θ(t)), where γt is the learning rate or step size at iteration t. A mini-batch version of the stochastic optimization algorithm computes the stochastic gradient over mini-batch of size B instead of a single datapoint, i.e., G
∆ = 1B ∑B i=1∇F (x̃i; θ(t)), where x̃i’s are randomly sampled from X . We will often evaluate performance on an exponential moving average θ̄(t) = αθ̄(t−1) + (1− α)θ(t) with decay rate α. Our interest is in distributed stochastic optimization using N worker machines in charge of computing stochastic gradients that are sent to M parameter servers. Each parameter server j is responsible for storing a subset θ[j] of the model, and performing updates on θ[j]. In the synchronous setting, we will also introduce additional b backup workers for straggler mitigation.
2 ASYNCHRONOUS STOCHASTIC OPTIMIZATION
An approach for a distributed stochastic gradient descent algorithm was presented in Dean et al. (2012), consisting of two main ingredients. First, the parameters of the model are distributed on multiple servers, depending on the architecture. This set of servers are called the parameter servers. Second, there can be multiple workers processing data in parallel and communicating with the parameter servers. Each worker processes a mini-batch of data independently of the others, as follows:
• The worker fetches from the parameter servers the most up-to-date parameters of the model needed to process the current mini-batch; • It then computes gradients of the loss with respect to these parameters; • Finally, these gradients are sent back to the parameter servers, which then updates the
model accordingly.
Since each worker communicates with the parameter servers independently of the others, this is called Asynchronous Stochastic Gradient Descent (Async-SGD), or more generally, Asynchronous Stochastic Optimization (Async-Opt). A similar approach was later proposed by Chilimbi et al. (2014). Async-Opt is presented in Algorithms 1 and 2.
In practice, the updates of Async-Opt are different than those of serially running the stochastic optimization algorithm for two reasons. Firstly, the read operation (Algo 1 Line 2) on a worker may be interleaved with updates by other workers to different parameter servers, so the resultant θ̂k may not be consistent with any parameter incarnation θ(t). Secondly, model updates may have occurred while a worker is computing its stochastic gradient; hence, the resultant gradients are typically computed with respect to outdated parameters. We refer to these as stale gradients, and its staleness as the number of updates that have occurred between its corresponding read and update operations.
Understanding the theoretical impact of staleness is difficult work and the topic of many recent papers, e.g. Recht et al. (2011); Duchi et al. (2013); Leblond et al. (2016); Reddi et al. (2015);
Algorithm 1: Async-SGD worker k Input: Dataset X Input: B mini-batch size
1 while True do 2 Read θ̂k = (θ[0], . . . , θ[M ]) from PSs. 3 G
(t) k := 0.
4 for i = 1, . . . , B do 5 Sample datapoint x̃i from X . 6 G
(t) k ← G (t) k + 1 B ∇F (x̃i; θ̂k).
7 end 8 Send G(t)k to parameter servers. 9 end
Algorithm 2: Async-SGD Parameter Server j Input: γ0, γ1, . . . learning rates. Input: α decay rate. Input: θ(0) model initialization.
1 for t = 0, 1, . . . do 2 Wait for gradient G from any worker. 3 θ(t+1)[j]← θ(t)[j]− γtG[j]. 4 θ̄(t)[j] = αθ̄(t−1)[j] + (1− α)θ(t)[j]. 5 end
De Sa et al. (2015); Mania et al. (2015), most of which focus on individual algorithms, under strong assumptions that may not hold up in practice. This is further complicated by deep models with multiple layers, since the times at which model parameters are read and which gradients are computed and sent are dependent on the depth of the layers (Figure 1). To better understand this dependence in real models, we collected staleness statistics on a Async-Opt run with 40 workers on a 18-layer Inception model (Szegedy et al., 2016) trained on the ImageNet Challenge dataset (Russakovsky et al., 2015), as shown in Table 1.
Despite the abovementioned problems, Async-Opt has been shown to be scale well up to a few dozens of workers for some models. However, at larger scales, increasing the number of machines (and thus staleness of gradients) can result in poorer trained models.
2.1 IMPACT OF STALENESS ON TEST ACCURACY
We explore how increased staleness contributes to training of poorer models. In order to mimic the setting on a smaller scale, we trained a state-of-the-art MNIST CNN model but simulated staleness by using old gradients for the parameter updates. Details of the model and training are provided in Appendix A.1.
The best final classification error on a test set was 0.36%, which increases to 0.47% with average gradient staleness of 20 steps, and up to 0.79% with 50 steps (see Figure 2).
Once the average simulated staleness was chosen to be more than 15 steps, the results started to significantly deteriorate and the training itself became much less stable. We had to employ following tricks to prevent the results from blowing up:
• Slowly increase the staleness over the first 3 epochs of training. This mimics increasing the number of asynchronous workers and is also very important in practice for some of the models we experimented with (e.g. large word-level language models). The trick was not relevant with a simulated staleness less than 15 but became crucial for larger values.
• Use lower initial learning rates when staleness is at least 20, which reduces a frequency of explosions (train error goes to 90%). This observation is similar to what we found in other experiments - we were able to use much larger learning rates with synchronous training and the results were also more stable.
• Even with above tricks the divergence occurs occasionally and we found that restarting training from random weights can lead to more successful runs. The best results were then chosen based on validation set performance.
3 REVISTING SYNCHRONOUS STOCHASTIC OPTIMIZATION
Both Dean et al. (2012) and Chilimbi et al. (2014) use versions of Async-SGD where the main potential problem is that each worker computes gradients over a potentially old version of the model. In order to remove this discrepancy, we propose here to reconsider a synchronous version of distributed stochastic gradient descent (Sync-SGD), or more generally, Synchronous Stochastic Optimization (Sync-Opt), where the parameter servers wait for all workers to send their gradients, aggregate them, and send the updated parameters to all workers afterward. This ensures that the actual algorithm is a true mini-batch stochastic gradient descent, with an effective batch size equal to the sum of all the mini-batch sizes of the workers.
While this approach solves the staleness problem, it also introduces the potential problem that the actual update time now depends on the slowest worker. Although workers have equivalent computation and network communication workload, slow stragglers may result from failing hardware, or contention on shared underlying hardware resources in data centers, or even due to preemption by other jobs.
To alleviate the straggler problem, we introduce backup workers (Dean & Barroso, 2013) as follows: instead of having only N workers, we add b extra workers, but as soon as the parameter servers receive gradients from any N workers, they stop waiting and update their parameters using the N gradients. The slowest b workers’ gradients will be dropped when they arrive. Our method is presented in Algorithms 3, 4.
Algorithm 3: Sync-SGD worker k, where k = 1, . . . , N + b
Input: Dataset X Input: B mini-batch size
1 for t = 0, 1, . . . do 2 Wait to read θ(t) = (θ(t)[0], . . . , θ(t)[M ]) from parameter servers. 3 G
(t) k := 0
4 for i = 1, . . . , B do 5 Sample datapoint x̃k,i from X . 6 G
(t) k ← G (t) k + 1 B ∇F (x̃k,i; θ(t)).
7 end 8 Send (G(t)k , t) to parameter servers. 9 end
Algorithm 4: Sync-SGD Parameter Server j Input: γ0, γ1, . . . learning rates. Input: α decay rate. Input: N number of mini-batches to aggregate. Input: θ(0) model initialization.
1 for t = 0, 1, . . . do 2 G = {} 3 while |G| < N do 4 Wait for (G, t′) from any worker. 5 if t′ == t then G ← G ∪ {G}. 6 else Drop gradient G. 7 end 8 θ(t+1)[j]← θ(t)[j]− γt
N ∑ G∈G G[j].
9 θ̄(t)[j] = αθ̄(t−1)[j] + (1− α)θ(t)[j]. 10 end
3.1 STRAGGLER EFFECTS
The use of backup workers is motivated by the need to mitigate slow stragglers while maximizing computation. We investigate the effect of stragglers on Sync-Opt model training here.
We ran Sync-Opt with N = 100 workers, b = 0 backups, and 19 parameter servers on the Inception model. Using one variable as a proxy, we collected for each iteration both the start time of the iteration and the time when the kth gradient of that variable arrived at the parameter server. These times are presented in Figure 3 for k = 1, 50, 90, 97, 98, 99, 100. Note that 80% of the 98th gradient arrives in under 2s, whereas only 30% of the final gradient do. Furthermore, the time to collect the final few gradients grows exponentially, resulting in wasted idle resources and time expended to wait for the slowest gradients. This exponential increase is also seen in Figure 4.
Thus, one might choose to drop slow stragglers to decrease the iteration time. However, using fewer machines implies a smaller effective mini-batch size and thus greater gradient variance, which in turn could require more iterations for convergence. We examine this relationship by running Sync-Opt2 with N = 50, 70, 80, 90, 100 and b = 6, and note the number of iterations required for convergence in Figure 5. Additional details of this training are provided in Appendix A.2. As N is doubled from 50 to 100, the number of iterations to converge nearly halves from 137.5e3 to 76.2e3.
2 Since we are interested in the gradient quality and convergence behavior but not running time in this experiment, the backups serve only to reduce our data collection time but do not affect our analysis.
Hence, there is a trade-off between dropping more stragglers to reduce iteration time, and waiting for more gradients to improve the gradient quality. Consider a hypothetical setting where we have N + b = 100 machines, and we wish to choose the best configuration of N and b to minimize running time to convergence3. For each configuration, we can estimate the iterations required from Figure 5 (linearly interpolating for values of N for which we did not collect data). We can multiply this with the mean iteration times (Figure 4) to obtain the running time required to converge for each setting of N and b. These results are shown in Figure 6, indicating that N = 96, b = 4 converges fastest. Therefore, this motivates our choice to use a few backup workers for mitigating stragglers.
4 EXPERIMENTS
In this section, we present our empirical comparisons of synchronous and asynchronous distributed stochastic optimization algorithms as applied to models such as Inception and PixelCNN. All experiments in this paper are using the TensorFlow system (Abadi et al., 2015).
4.1 METRICS OF COMPARISON: FASTER CONVERGENCE, BETTER OPTIMUM
We are interested in two metrics of comparison for our empirical validation: (1) test error or accuracy, and (2) speed of convergence3. We point out that for non-convex deep learning models, it is possible to converge faster to a poorer local optimum. Here we show a simple example with Inception using different learning rates.
We ran Sync-Opt on Inception with N = 100 and b = 6, but varied the initial learning rate γ0 between 1.125 and 9.0. (Learning rates are exponentially decreased with iterations.) Table 2 shows that smaller γ0 converge faster, but to poorer test precisions. Focusing on speed on an early phase of training could lead to misleading conclusions if we fail to account for eventual convergence. For example, Figure 3b shows that using γ0 = 1.125 reaches = 75% precision 1.5× faster than γ0 = 4.5, but is slower for = 77.75%, and fails to reach higher precisions.
4.2 INCEPTION
We conducted experiments on the Inception model (Szegedy et al., 2016) trained on ImageNet Challenge dataset (Russakovsky et al., 2015), where the task is to classify images out of 1000 categories. We used several configurations, varying N + b from 53 to 212 workers. Additional details of the training are provided in Appendix A.3. An epoch is a synchronous iteration for Sync-Opt, or a full pass of N updates for Async-Opt, which represent similar amounts of computation. Results of this experiment are presented in Figure 8.
Figure 8b shows that Sync-Opt outperforms Async-Opt in test precision: Sync-Opt attains ∼0.5% better test precision than Async-Opt for comparable N + b workers. Furthermore, Sync-Opt con-
3Convergence is defined as the point where maximum test accuracy or lowest test error is reached.
verges 6h and 18h faster than Async-Opt for 106 and 212 workers respectively, and is 3h slower when 53 workers are used, as seen in Figure 8d. This difference in speed is largely due to the fewer epochs (Figure 8c) needed by Sync-Opt, but comparable or better epoch time (Figure 8e).
4.3 PIXELCNN EXPERIMENTS
The second model we experimented on is PixelCNN (Oord et al., 2016), a conditional image generation deep neural network, which we train on the CIFAR-10 (Krizhevsky & Hinton, 2009) dataset. Configurations of N + b = 1, 8, 16 workers were used; for Sync-Opt, we always used b = 1 backup worker. Additional details are provided in Appendix A.4.
Convergence of the test negative log likelihood (NLL) on PixelCNN is shown in Figure 9a, where lower is better. Observe that Sync-Opt obtains lower NLL than Async-Opt; in fact, Async-Opt is even outperformed by serial RMSProp with N = 1 worker, with degrading performance as N increases from 8 to 16. Figure 9b further shows the time taken to reach test NLL. Sync-Opt reduces the time to reach = 2.145 from 247h to 58.3h; this NLL is not even achieved by Async-Opt.
5 RELATED WORK
Multicore and distributed optimization algorithms have received much attention in recent years. Asynchronous algorithms include Recht et al. (2011); Duchi et al. (2013); Zhang et al. (2015a); Reddi et al. (2015); Leblond et al. (2016). Implementations of asynchronous optimization include Xing et al. (2015); Li et al. (2014); Chilimbi et al. (2014). Attempts have also been made in Zinkevich et al. (2010) and Zhang & Jordan (2015) to algorithmically improve synchronous SGD.
An alternative solution, “softsync”, was presented in Zhang et al. (2015b), which proposed batching gradients from multiple machines before performing an asynchronous SGD update, thereby reducing the effective staleness of gradients. Similar to our proposal, softsync avoids stragglers by not forcing updates to wait for the slowest worker. However, softsync allows the use of stale gradients but we do not. The two solutions provide different explorations of the trade-off between high accuracy (by minimizing staleness) and fast throughput (by avoiding stragglers).
Watcharapichat et al. (2016) introduces a distributed deep learning system without parameter servers, by having workers interleave gradient computation and communication in a round-robin pattern. Like Async-Opt, this approach suffers from staleness. We also note that in principle, workers in Sync-Opt can double as parameter servers and execute the update operations and avoid the need to partition hardware resources between workers and servers.
Das et al. (2016) analyzes distributed stochastic optimization and optimizes the system by solving detailed system balance equations. We believe this approach is complimentary to our work, and could potentially be applied to guide the choice of systems configurations for Sync-Opt.
Keskar et al. (2016) suggests that large batch sizes for synchronous stochastic optimization leads to poorer generalization. Our effective batch size increases linearly with the number of workers N . However, we did not observe this effect in our experiments; we believe we are not yet in the large batch size regime examined by Keskar et al. (2016).
6 CONCLUSION AND FUTURE WORK
Distributed training strategies for deep learning architectures will become ever more important as the size of datasets increases. In this work, we have shown how both synchronous and asynchronous distributed stochastic optimization suffer from their respective weaknesses of stragglers and staleness. This has motivated our development of synchronous stochastic optimization with backup workers, which we show to be a viable and scalable strategy.
We are currently experimenting with different kinds of datasets, including word-level language models where parts of the model (the embedding layers) are often very sparse, which involves very different communication constraints. We are also working on further improving the performance of synchronous training like combining gradients from multiple workers sharing the same machine before sending them to the parameter servers to reduce the communication overhead. An alternative of using time-outs instead of backup workers is also being explored.
A DETAILS OF MODELS AND TRAINING
A.1 MNIST CNN, SECTION 2.1
The model used in our experiments is a 4-layer CNN that have 3x3 filters with max-pooling and weight normalization in every layer. We trained the model with SGD for 25 epochs and evaluated performance on the exponential moving average θ̄ using a decay rate of α = 0.9999. Initial learning rate was set to be 0.1 and linearly annealed to 0 in the last 10 epochs. We also used small image rotations and zooms as a data augmentation scheme.
A.2 INCEPTION, SECTION 3.1
For our straggler experiments, we trained the Inception (Szegedy et al., 2016) model on the ImageNet Challenge dataset (Russakovsky et al., 2015). 10 parameter servers were used, and each worker was equipped with a k40 GPU.
The underlying optimizer was RMSProp with momentum, with decay of 0.9 and momentum of 0.9. Mini-batch size B = 32 was used. Initial learning rates γ0 were set at 0.045N , which we found to provide good test precisions for Inception. Learning rates were also exponentially decreased with decay rate β = 0.94 as γ0βtN/(2T ), where T = |X |/B is the number of mini-batches in the dataset. Test precisions were evaluated on the exponential moving average θ̄ using α = 0.9999.
A.3 INCEPTION, SECTION 4.2
For experiments comparing Async-Opt and Sync-Opt on the Inception model in Section 4.2, each worker is equipped with a k40 GPU. For N + b = 53 workers, 17 parameter servers were used; for N + b = 106 workers, we used 27 parameter servers; and 37 parameter servers were used for N + b = 212.
In the asynchronous training mode, gradient clipping is also needed for stabilization, which requires each worker to collect the gradient across all layers of the deep model, compute the global norm ||G|| and then clip all gradient accordingly. However, synchronization turns out to be very stable so gradient clipping is no longer needed, which means that we can pipeline the update of parameters in different layers: the gradient of top layers’ parameters can be sent to parameter servers while concurrently computing gradients for the lower layers.
The underlying optimizer is RMSProp with momentum, with decay of 0.9 and momentum of 0.9. Mini-batch size B = 32 was used. Initial learning rates γ0 for Async-Opt were set to 0.045; for Sync-Opt, we found as a rule-of-thumb that a learning rate of 0.045N worked well for this model. Learning rates were then exponentially decayed with decay rate β = 0.94 as γ0βt/(2T ) for AsyncOpt, where T = |X |/B is the number of mini-batches in the dataset. For Sync-Opt, we learning rates were also exponentially decreased at rate of γ0βtN/(2T ), so that the learning rates after computing the same number of datapoints are comparable for Async-Opt and Sync-Opt.
Test precisions were evaluated on the exponential moving average θ̄ using α = 0.9999.
A.4 PIXELCNN, SECTION 4.3
The PixelCNN (Oord et al., 2016) model was trained on the CIFAR-10 (Krizhevsky & Hinton, 2009) dataset. Configurations of N + b = 1, 8, 16 workers each with a k80 GPU, and 10 parameter servers were used. For Sync-Opt, we always used b = 1 backup worker. The underlying optimizer is RMSProp with momentum, using decay of 0.95 and momentum of 0.9. Initial learning rates γ0 were set to 1e − 4 and slowly decreased to 3e − 6 after 200,000 iterations. Mini-batch size B = 4 was used.

This paper was easy to read, the main idea was presented very clearly.
The main points of the paper (and my concerns are below) can be summarized as follows:
1. synchronous algoriths suffer from some struggeling nodes, for which the algorithm has to wait. From my own experience, this has never happend for me on e.g. Amazon EC2 cloud, however, it happens on our own cluster at my university, if the cluster is shared and some users make some nodes very busy. So maybe if the nodes would be dedicated to just user's job, it wouldn't be such a big concer (I am not sure what kind of cluster was used to produce Figure 3 and 4). Also how many experiments have you run? In my own experience, most of the time I get the gradient on time from all nodes equality fast, but maybe just in less than 0.1% of iterations I observe that it took maybe twice as long for some node. Also the increasing shape of the curve is somehow implying some weird implementation of communication. Isn't it only because you are somehow serialize the communication? And it would be maybe much faster if a "MPI_Reduce" would be used (even if we wait for the slowest guy)?
2. asynchronous algorithms are cutting the waiting time, however, the convergence speed may be slower. Moreover, those algorithms can be divergence it special care is not given to stale gradients. Also they have a nice guarantees for convex functions, but the non-convex DNN may cause pain.
3.they propose to take gradient from the first "N" workers out of "N+b"
workers available. My concern here is that they focused only on the
workers, but what if the "parameter server" will became to slow? What
if the parameter server would be the bottleneck? How would you address
this situation? But still if the number of nodes (N) is not large, and
the deep DNN is used, I can imagine that the communciation will not
take more than 30% of the run-time.
My largest concern is with the experiments. Different batch size
implies that different learning rate should be chosen, right? How did
you tune the learning rates and other parameters for e.g. Figure 5 you
provide some formulas in (A2) but clearly this can bias your Figures,
right? meaning, that if you tune "\gamma, \beta" for each N, it could
be somehow more representative? also it would be nicer if you run the
experiment many times and then report average, best and worst case
behaviour. because now it can be just coinsidence, right?


AN ACTOR-CRITIC ALGORITHM FOR LEARNING RATE LEARNING
Stochastic gradient descent (SGD), which updates the model parameters by adding a local gradient times a learning rate at each step, is widely used in model training of machine learning algorithms such as neural networks. It is observed that the models trained by SGD are sensitive to learning rates and good learning rates are problem specific. To avoid manually searching of learning rates, which is tedious and inefficient, we propose an algorithm to automatically learn learning rates using actor-critic methods from reinforcement learning (RL). In particular, we train a policy network called actor to decide the learning rate at each step during training, and a value network called critic to give feedback about quality of the decision (e.g., the goodness of the learning rate outputted by the actor) that the actor made. Experiments show that our method leads to good convergence of SGD and can prevent overfitting to a certain extent, resulting in better performance than human-designed competitors.
1 INTRODUCTION
While facing large scale of training data, stochastic learning such as stochastic gradient descent (SGD) is usually much faster than batch learning and often results in better models. An observation for SGD methods is that their performances are highly sensitive to the choice of learning rate LeCun et al. (2012). Clearly, setting a static learning rate for the whole training process is insufficient, since intuitively the learning rate should decrease when the model becomes more and more close to a (local) optimum as the training goes on over time Maclaurin et al. (2015). Although there are some empirical suggestions to guide how to adjust the learning rate over time in training, it is still a difficult task to find a good policy to adjust the learning rate, given that good policies are problem specific and depend on implementation details of a machine learning algorithm. One usually needs to try many times and adjust the learning rate manually to accumulate knowledge about the problem. However, human involvement often needs domain knowledge about the target problems, which is inefficient and difficult to scale up to different problems. Thus, a natural question arises: can we automatically adjust the learning rate? This is exactly the focus of this work and we aim to automatically learn the learning rates for SGD based machine learning (ML) algorithms without human-designed rules or hand-crafted features.
By examining the current practice of learning rate control/adjustment, we have two observations. First, learning rate control is a sequential decision process. At the beginning, we set an initial learning rate. Then at each step, we decide whether to change the learning rate and how to change it, based on the current model and loss, training data at hand, and maybe history of the training process. As suggested in Orr & Müller (2003), one well-principled method for estimating the ideal learning rate that is to decrease the learning rate when the weight vector oscillates, and increase it when the weight vector follows a relatively steady direction. Second, although at each step some immediate reward (e.g., the loss decrement) can be obtained by taking actions, we care more about the performance of the final model found by the ML algorithm. Consider two different learning rate
control policies: the first one leads to fast loss decrease at the beginning but gets saturated and stuck in a local minimum quickly, while the second one starts with slower loss decrease but results in much smaller final loss. Obviously, the second policy is better. That is, we prefer long-term rewards over short-term rewards.
Combining the two observations, it is easy to see that the problem of finding a good policy to control/adjust learning rate falls into the scope of reinforcement learning (RL) Sutton & Barto (1998), if one is familiar with RL. Inspired by the recent success of RL for sequential decision problems, in this work, we leverage RL techniques and try to learn the learning rate for SGD based methods.
We propose an algorithm to learn the learning rate within the actor-critic framework Sutton (1984); Sutton et al. (1999); Barto et al. (1983); Silver et al. (2014) from RL. In particular, an actor network is trained to take an action that decides the learning rate for current step, and a critic network is trained to give feedbacks to the actor network about long-term performance and help the actor network to adjust itself so as to perform better in the future steps. The main contributions of this paper include:
• We propose an actor-critic algorithm to automatically learn the learning rate for ML algorithms.
• Long-term rewards are exploited by the critic network in our algorithm to choose a better learning rate at each step.
• We propose to feed different training examples to the actor network and the critic network, which improve the generalization performance of the learnt ML model.
• A series of experiments validate the effectiveness of our proposed algorithm for learning rate control.
2 RELATED WORK
2.1 IMPROVED GRADIENT METHODS
Our focus is to improve gradient based ML algorithm through automatic learning of learning rate. Different approaches have been proposed to improve gradient methods, especially for deep neural networks.
Since SGD solely rely on a given example (or a mini-batch of examples) to compare gradient, its model update at each step tends to be unstable and it takes many steps to converge. To solve this problem, momentum SGD Jacobs (1988) is proposed to accelerate SGD by using recent gradients. RMSprop Tieleman & Hinton (2012) utilizes the magnitude of recent gradients to normalize the gradients. It always keeps a moving average over the root mean squared gradients, by which it divides the current gradient. Adagrad Duchi et al. (2011) adapts component-wise learning rates, and performs larger updates for infrequent and smaller updates for frequent parameters. Adadelta Zeiler (2012) extends Adagrad by reducing its aggressive, monotonically decreasing learning rate. Instead of accumulating all past squared gradients, Adadelta restricts the window of accumulated past gradients to some fixed size. Adam Kingma & Ba (2014) computes component-wise learning rates using the estimates of first and second moments of the gradients, which combines the advantages of AdaGrad and RMSProp.
Senior et al. (2013); Sutton (1992); Darken & Moody (1990) focus on predefining update rules to adjust learning rates during training. A limitation of these methods is that they have additional free parameters which need to be set manually. Another recent work Daniel et al. (2016) studies how to automatically select step sizes, but it still requires hand-tuned features. Schaul et al. (2013) proposes a method to choose good learning rate for SGD, which relies on the square norm of the expectation of the gradient, and the expectation of the square norm of the gradient. The method is much more constrained than ours and several assumption should be met.
2.2 REINFORCEMENT LEARNING
Since our proposed algorithm is based on RL techniques, here we give a very brief introduction to RL, which will ease the description of our algorithm in next section.
Reinforcement learning Sutton (1988) is concerned with how an agent acts in a stochastic environment by sequentially choosing actions over a sequence of time steps, in order to maximize a cumulative reward. In RL, a state st encodes the agents observation about the environment at a time step t, and a policy function π(st) determines how the agent behaves (e.g., which action to take) at state st. An action-value function (or, Q function) Qπ(st, at) is usually used to denote the cumulative reward of taking action at at state st and then following policy π afterwards.
Many RL algorithms have been proposed Sutton & Barto (1998); Watkins & Dayan (1992), and many RL algorithms Sutton (1984); Sutton et al. (1999); Barto et al. (1983); Silver et al. (2014) can be described under the actor-critic framework. An actor-critic algorithm learns the policy function and the value function simultaneously and interactively. The policy structure is known as the actor, and is used to select actions; the estimated value function is known as the critic, and it criticizes the actions made by the actor.
Recently, deep reinforcement learning, which uses deep neural networks to approximate/represent the policy function and/or the value function, have shown promise in various domains, including Atari games Mnih et al. (2015), Go Silver et al. (2016), machine translation Bahdanau et al. (2016), image recognition Xu et al. (2015), etc.
3 METHOD
In this section, we present an actor-critic algorithm that can automate the learning rate control for SGD based machine learning algorithms.
Many machine learning tasks need to train a model with parameters ω by minimizing a loss function f defined over a set X of training examples:
ω∗ = arg min ω fω(X). (1)
A standard approach for the loss function minimization is gradient descent, which sequentially updates the parameters using gradients step by step:
ωt+1 = ωt − at∇f t, (2) where at is the learning rate at step t, and ∇f t is the local gradient of f at ωt. Here one step can be the whole batch of all the training data, a mini batch of tens/hundreds of examples, or a random sample.
It is observed that the performance of SGD based methods is quite sensitive to the choice of at for non-convex loss function f . Unfortunately, f is usually non-convex with respect to the parameters
w in many ML algorithms, especially for deep neural networks. We aim to learn a learning rate controller using RL techniques that can automatically control at.
Figure 1 illustrates our automatic learning rate controller, which adopts the actor-critic framework in RL. The basic idea is that at each step, given the current model ωt and training sample x, an actor network is used to take an action (the learning rate at, and it will be used to update the model ωt), and a critic network is used to estimate the goodness of the action. The actor network will be updated using the estimated goodness of at, and the critic network will be updated by minimizing temporal difference (TD) Sutton & Barto (1990) error. We describe the details of our algorithm in the following subsections.
3.1 ACTOR NETWORK
The actor network, which is called policy network in RL, plays the key role in our algorithm: it determines the learning rate control policy for the primary ML algorithm1 based on the current model, training data, and maybe historical information during the training process.
Note that ωt could be of huge dimensions, e.g., one widely used image recognition model VGGNet Simonyan & Zisserman (2014) has more than 140 million parameters. If the actor network takes all of those parameters as the inputs, its computational complexity would dominate the complexity of the primary algorithm, which is unfordable. Therefore, we propose to use a function χ(·) to process and yield a compact vector st as the input of the actor network. Following the practice in RL, we call χ(·) the state function, which takes ωt and the training data x as inputs:
st = χ(ωt, X). (3)
Then the actor network πθ(·) parameterized by θ yields an action at:
πθ(s t) = at, (4)
where the action at ∈ R is a continuous value. When at is determined, we update the model of the primary algorithm by Equation 2.
Note that the actor network has its own parameters and we need to learn them to output a good action. To learn the actor network, we need to know how to evaluate the goodness of an actor network. The critic network exactly plays this role.
3.2 CRITIC NETWORK
Recall that our goal is to find a good policy for learning rate control to ensure that a good model can be learnt eventually by the primary ML algorithm. For this purpose, the actor network needs to output a good action at at state st so that finally a low training loss f(·) can be achieved. In RL, the Q function Qπ(s, a) is often used to denote the long term reward of the state-action pair s, a while following the policy π to take future actions. In our problem, Qπ(st, at) indicates the accumulative decrement of training loss starting from step t. We define the immediate reward at step t as the one step loss decrement:
rt = f t − f t+1. (5)
The accumulative value Rtπ of policy π at step t is the total discounted reward from step t:
Rtπ = Σ T k=tγ k−tr(sk, ak),
where γ ∈ (0, 1] is the discount factor. Considering that both the states and actions are uncountable in our problem, the critic network uses a parametric function Qϕ(s, a) with parameters ϕ to approximate the Q value function Qπ(s, a).
1Here we have two learning algorithms. We call the one with learning rate to adjust as the primary ML algorithm, and the other one which optimizes the learning rate of the primary one as the secondary ML algorithm.
3.3 TRAINING OF ACTOR AND CRITIC NETWORKS
The critic network has its own parameters ϕ, which is updated at each step using TD learning. More precisely, the critic is trained by minimizing the square error between the estimation Qϕ(st, at) and the target yt:
yt = rt + γQϕ(s t+1, at+1). (6)
The TD error is defined as:
δt = yt −Qϕ(st, at) = rt + γQϕ(s t+1, πθ(s t+1))−Qϕ(st, at)
(7)
The weight update rule follows the on-policy deterministic actor-critic algorithm. The gradients of critic network are:
∇ϕ = δt∇ϕQϕ(st, at), (8)
The policy parameters θ of the actor network is updated by ensuring that it can output the action with the largest Q value at state st, i.e., a∗ = arg maxaQϕ(st, a). Mathematically,
∇θ = ∇θπθ(st+1)∇aQϕ(st+1, at+1)|a=πθ(s). (9)
Algorithm 1 Actor-Critic Algorithm for Learning Rate Learning Require: Training steps T ; training set X; loss function f ; state function χ; discount factor: γ ; Ensure: Model parameters w, policy parameters θ of the actor network, and value parameters ϕ of
the critic network; 1: Initial parameters ω0, θ0, ϕ0; 2: for t = 0, ..., T do 3: Sample xi ∈ X, i ∈ 1, ..., N . 4: Extract state vector: sti = χ(ω
t, xi). 5: //Actor network selects an action. 6: Computes learning rate ati = πθ(s t i). 7: //Update model parameters ω. 8: Compute∇f t(xi). 9: Update ω: ωt+1 = ωt − ati∇f t(xi).
10: //Update critic network by minimizing square error between estimation and label. 11: rt = f t(xi)− f t+1(xi) 12: Extract state vector: st+1i = χ(ω t+1, xi) 13: Compute Qϕ(st+1i , πθ(s t+1 i )), Qϕ(s t i, a t i) 14: Compute δt according to Equation 7: δt = rt + γQϕ(s t+1 i , πθ(s t+1 i ))−Qϕ(sti, ati) 15: Update ϕ using the following gradients according to Equation 8 : ∇ϕ = δt∇ϕQϕ(sti, ati) 16: // Update actor network 17: Sample xj ∈ X, j ∈ 1, ..., N, j 6= i. 18: Extract state vector: st+1j = χ(ω
t+1, xj). 19: Compute at+1j = πθ(s t+1 j ). 20: Update θ from Equation 9: ∇θ = ∇θπθ(st+1j )∇aQϕ(s t+1 j , a t+1 j )|a=πθ(s) 21: end for 22: return ω, θ, ϕ;
3.4 THE ALGORITHM
The overall algorithm is shown in Algorithm 1. In each step, we sample an example (Line 3), extract the current state vector (Line 4), compute the learning rate using the actor network (Line 6), update the model (Lines 8-9), compute TD error (Lines 11-14), update the critic network (Line 15), and sample another example (Line 17) to update the actor network (Line 18-20). We would like to make some discussions about the algorithm.
First, in the current algorithm, for simplicity, we consider using only one example for model update. It is easy to generalize to a mini batch of random examples.
Second, one may notice that we use one example (e.g., xi) for model and the critic network update, but a different example (e.g., xj) for the actor network update. Doing so we can avoid that the algorithm will overfit on some (too) hard examples and can improve the generalization performance of the algorithm on the test set. Consider a hard example2 in a classification task. Since such an example is difficult to be classified correctly, intuitively its gradient will be large and the learning rate given by the actor network at this step will also be large. In other words, this hard example will greatly change the model, while itself is not a good representative of its category and the learning algorithm should not pay much attention to it. If we feed the same example to both the actor network and the critic network, both of them will encourage the model to change a lot to fit the example, consequently resulting in oscillation of the training, as shown in our experiments. By feeding different examples to the actor and critic networks, it is very likely the critic network will find that the gradient direction of the example fed into the actor network is inconsistent with its own training example and thus criticize the large learning rate suggested by the actor network. More precisely, the update of ω is based on xi and the learning rate suggested by the actor network, while the training target of the actor network is to maximize the output of the critic network on xj . If there is big gradient disagreement between xi and xj , the update of ω, which is affected by actor’s decision, would cause the critic’s output on xj to be small. To compensate this effect, the actor network is forced to predict a small learning rate for a too hard xi in this situation.
4 EXPERIMENTS
We conducted a set of experiments to test the performance of our learning rate learning algorithm and compared with several baseline methods. We report the experimental results in this section.
4.1 EXPERIMENTAL SETUP
We tested our method on two widely used image classification datasets: MNIST LeCun et al. (1998) and CIFAR-10 Krizhevsky & Hinton (2009). Convolutional neural networks (CNNs) are the standard model for image classification tasks in recent years, and thus the primary ML algorithm adopted the CNN model in all our experiments.
We specified our actor-critic algorithm in experiments as follows. Given that stochastic mini-batch training is a common practice in deep learning, the actor-critic algorithm also operated on minibatches, i.e., each step is a mini batch in our experiments. We defined the state st = χ(ωt, Xi) as the average loss of learning model ωt on the input min-batch Xi. We specified the actor network as a two-layer long short-term memory (LSTM) network with 20 units in each layer, considering that a good learning rate for step t depends on and correlates with the learning rates at previous steps while LSTM is well suited to model sequences with long-distance dependence. We used the absolute value activation function for the output layer of the LSTM to ensure a positive learning rate. The LSTM was unrolled for 20 steps during training. We specified the critic network as a simple neural network with one hidden layer and 10 hidden units. We use Adam with the default setting in TensorFlow optimizer toolbox Abadi et al. (2015) to train the actor and critic networks in all the experiments.
We compared our method with several mainstream SGD algorithms, including SGD, Adam Kingma & Ba (2014), Adagrad Duchi et al. (2011) and RMSprop Tieleman & Hinton (2012). For each of these algorithms and each dataset, we tried the following learning rates 10−4, 10−3, ..., 100. We report the best performance of these algorithms over those learning rates. If an algorithm needs some other parameters to set, such as decay coefficients for Adam, we used the default setting in TensorFlow optimizer toolbox. For each benchmark and our proposed method, five independent runs are averaged and reported in all of the following experiments.
4.2 RESULTS ON MNIST
MNIST is a dataset for handwritten digit classification task. Each example in the dataset is a 28×28 black and white image containing a digit in {0, 1, · · · , 9}. The CNN model used in the primary
2For example, an example may has an incorrect label because of the limited quality of labelers.
ML algorithm is consist of two convolutional layers, each followed by a pooling layer, and finally a fully connected layer. The first convolutional layer filters each input image using 32 kernels of size 5 × 5. The max-pooling layer following the first convolutional layer is performed over 2 × 2 pixel windows, with stride 2. The second convolutional layer takes the outputs of the first max-pooling layer as inputs and filters them with 64 kernels of size 5 × 5. The max-pooling layer following the second convolutional layer is performed over 2 × 2 pixel windows, with stride 2. The outputs of second max pooling layer are fed to a fully connected layer with 512 neurons. Dropout was conducted on the fully connect layer with a dropout rate of 0.5. ReLU activation functions are used in the CNN model. There are 60,000 training images and 10,000 test images in this dataset. We scaled the pixel values to the [0,1] range before inputting to all the algorithms. Each mini batch contains 50 randomly sampled images.
Figure 2 shows the results of our actor-critic algorithm for learning rate learning and the baseline methods, including the curves of training loss, test loss, and test accuracy. The final accuracies of these methods are summarized in Table 1. We have the following observations.
• In terms of training loss, our algorithm has similar convergence speed to the baseline methods. One may expect that our algorithm should have significantly faster convergence speed considering that our algorithm learns both the learning rate and the CNN model while the baselines only learn the CNN model and choose the learning rates per some predefined rules. However, this is not correct. As discussed in Section 3.4, we carefully design the algorithm and feed different samples to the actor network and critic network. Doing so we can focus more on generalization performance than training loss: as shown in Figure 4, our algorithm achieves the best test accuracy.
10-2 10-1 100 101
Epoch
0.010
0.015
0.020
0.025
0.030
0.035
0.040
0.045
0.050
0.055
Le a rn
in g r
a te
CIFAR-10
Figure 5: The learning rate learned by actor network for CIFAR-10.
• Our algorithm achieves the lowest error rate on MNIST. Although the improvement looks small, we would like to point out that given that the accuracy of CNN is already close to 100%, it is a very difficult task to further improve accuracy, not to mention that we only changed learning rate policy without changing the CNN model.
4.3 RESULTS ON CIFAR-10
CIFAR-10 is a dataset consisting of 60000 natural 32 × 32 RGB images in 10 classes: 50,000 imagesfor training and 10,000 for test. We used a CNN with 2 convolutional layers (each followed by max-pooling layer) and 2 fully connected layers for this task. There is a max pooling layer which performed over 2× 2 pixel windows, with stride 2 after each convolutional layer. All convolutional layers filter the input with 64 kernels of size 5× 5. The outputs of the second pooling layer are fed to a fully connected layer with 384 neurons. The last fully connected layer has 192 neurons. Before inputting an image to the CNN, we subtracted the per-pixel mean computed over the training set from each image.
Figure 3 shows the results of all the algorithms on CIFAR-10, including the curves of training loss, the test loss and test accuracy. Table 2 shows the final test accuracy. We get similar observations as MNIST: our algorithm achieves similar convergence speed in terms of training loss and slightly better test accuracy than baselines. Figure 5 shows the learning rate learned by our method on CIFAR-10. To further understand the generalization performance of our algorithm, we ran all the
algorithms on two subsets of training data on CIFAR-10: one with only 20% training data The curves of training loss and test loss are shown in Figure 4. As can be seen from the figure, those baseline methods are easy to overfit and their test loss increases after 5000 steps (mini batches). In contrast, our algorithm is relatively robust and can prevent overfitting to some extent.
As we explained in Section 3.4, feeding different examples to the actor and critic networks is important to guarantee generalization ability. Here we conducted another experiment to verify our intuitive explanation. Figure 6 shows the results of two different implementations of our actor-critic algorithm on CIFAR-10. In the first implementation, we fed the sample examples to the two net-
works, i.e., xi = xj in the algorithm, and in the second implementation, the input xj of the critic network is different from the input xi of the actor network. It is easy to see from the figure that setting xi = xj tends to oscillate during training and leads to poor test performance. Thus, we need to feed different training data to the actor network and the critic network to ensure the performance of the algorithm.
4.4 COMPARISON WITH OTHER ADAPTIVE LEARNING RATE METHOD
We also compare our method with “vSGD” from previous by work Schaul et al. (2013), which can automatically adjust learning rates to minimize the expected error. This method tries to compute learning rate at each update by optimizing the expected loss after the next update according to the square norm of the expectation of the gradient, and the expectation of the square norm of the gradient. Note that our method learns to predict a learning rate at each time step by utilizing the long term reward predicted by a critic network.
For a fair comparison, we followed the experiments settings of Schaul et al. (2013), which designed three different network architectures for MNIST task to measure the performance. The first one is denoted by ‘M0’ which is simple softmax regression (i.e. a network with no hidden layer). The second one (‘M1’) is a fully connected multi-layer perceptron, with a single hidden layer. The third one (denoted ‘M2’) is a deep, fully connected multi-layer perceptron with two hidden layers. The vSGD has three variants in their paper. We referred to the results reported in their paper and compared our method with all of three variants of their algorithm (vSGD-l, vSGD-b, vSGD-g). The learning rates of SGD are decreased according to a human designed schedule, and the hyperparameters of SGD, ADAM, Adagrad, RMSprop are carefully determined by their lowest test error among a set of hyper-parameters. All hyper-parameters can be found in Schaul et al. (2013).
The experimental results are reported in Table 3. It shows that our proposed method performs better than vSGD and other baseline methods, and is stable across different network architectures.
5 CONCLUSIONS AND FUTURE WORK
In this work, we have studied how to automatically learn learning rates for gradient based machine learning methods and proposed an actor-critic algorithm, inspired by the recent success of reinforcement learning. The experiments on two image classification datasets have shown that our method (1) has comparable convergence speed with expert-designed optimizer while achieving better test accuracy, and (2) can successfully adjust learning rate for different datasets and CNN model structures.
For the future work, we will explore the following directions. In this work, we have applied our algorithm to control the learning rates of SGD. We will apply to other variants of SGD methods. We have focused on learning a learning rate for all the model parameters. We will study how to learn an individual learning rate for each parameter. We have considered learning learning rates using RL techniques. We will consider learning other hyperparameters such as step-dependent dropout rates for deep neural networks.
A APPENDIX
A method of automatically controlling learning rate is proposed in the main body of the paper. The learning rate controller adjusts itself during training to control the learning rate. Here, we propose an improved version that can leverage experiences from several repeated training runs to learn a fixed learning rate controller. Empirically, this algorithm can achieve better performance than the previous one. Given that it requires more time for training the learning rate controller, this method is more suitable for training offline models.
In this algorithm, during every training run, we fix the actor network and compute the weighted sum of the gradients of its parameter θ. The parameter is updated after each run (modified from Equation 9):
∇θ = ΣTt=1h(t)∇θπθ(st+1)∇aQϕ(st+1, at+1)|a=πθ(s). (10)
h(t) is weighted function which is used to amplify the feedback signal from the initial training stage. It is defined as h(t) = 1/t in our experiments. An error rate of 0.48% was achieved with 5 repeated training runs in MNIST experiment (the same setting as Table 1), and in CIFAR-10 experiment (the same setting as Table 2), 80.23% accuracy was achieved with 10 training runs. This method showed better performance in both experiments.

The authors present a method for adaptively setting the step size for SGD by treating the learning rate as an action in an MDP whose reward is the change in loss function. The method is presented against popular adaptive first-order methods for training deep networks (Adagrad, Adam, RMSProp, etc). The results are interesting but difficult to assess in a true apples-to-apples manner. Some specific comments:
-What is the computational overhead of the actor-critic algorithm relative to other algorithms? No plots with the wall-time of optimization are presented, even though the success of methods like Adagrad was due to their wall-time performance, not the number of iterations.
-Why was only a single learning rate learned? To accurately compare against other popular first order methods, why not train a separate RL model for each parameter, similar to how popular first-order methods adaptively change the learning rate for each parameter.
-Since learning is a non-stationary process, while RL algorithms assume a stationary environment, why should we expect an RL algorithm to work for learning a learning rate?
-In figure 6, how does the proposed method compare to something like early stopping? It may be that the actor-critic method is overfitting less simply because it is worse at optimization.