Stack Long Short-Term Memory (StackL-STM) is useful for various applications such as parsing and string-to-tree neural machine translation, but it is also known to be notoriously difficult to parallelize for GPU training due to the fact that the computations are dependent on discrete operations. In this paper, we tackle this problem by utilizing state access patterns of StackLSTM to homogenize computations with regard to different discrete operations. Our parsing experiments show that the method scales up almost linearly with increasing batch size, and our parallelized PyTorch implementation trains significantly faster compared to the Dynet C++ implementation.
Introduction Tree-structured representation of language has been successfully applied to various applications including dependency parsing (Dyer et al., 2015) , sentiment analysis (Socher et al., 2011) and neural machine translation (Eriguchi et al., 2017) . However, most of the neural network architectures used to build tree-structured representations are not able to exploit full parallelism of GPUs by minibatched training, as the computation that happens for each instance is conditioned on the input/output structures, and hence cannot be naïvely grouped together as a batch. This lack of parallelism is one of the major hurdles that prevent these representations from wider adoption practically (e.g., neural machine translation), as many natural language processing tasks currently require the ability to scale up to very large training corpora in order to reach state-of-theart performance. We seek to advance the state-of-the-art of this problem by proposing a parallelization scheme for one such network architecture, the Stack Long Short-Term Memory (StackLSTM) proposed in Dyer et al. (2015) . This architecture has been successfully applied to dependency parsing (Dyer et al., 2015 (Dyer et al., , 2016 Ballesteros et al., 2017) and syntax-aware neural machine translation (Eriguchi et al., 2017) in the previous research literature, but none of these research results were produced with minibatched training. We show that our parallelization scheme is feasible in practice by showing that it scales up near-linearly with increasing batch size, while reproducing a set of results reported in (Ballesteros et al., 2017) . StackLSTM StackLSTM (Dyer et al., 2015) is an LSTM architecture (Hochreiter and Schmidhuber, 1997) augmented with a stack H that stores some of the hidden states built in the past. Unlike traditional LSTMs that always build state h t from h t−1 , the states of StackLSTM are built from the head of the state stack H, maintained by a stack top pointer p(H). At each time step, StackLSTM takes a realvalued input vector together with an additional discrete operation on the stack, which determines what computation needs to be conducted and how the stack top pointer should be updated. Throughout this section, we index the input vector (e.g. word embeddings) x t using the time step t it is fed into the network, and hidden states in the stack h j using their position j in the stack H, j being defined as the 0-base index starting from the stack bottom. The set of input discrete actions typically contains at least Push and Pop operations. When these operations are taken as input, the corresponding computations on the StackLSTM are listed below: 1 (Dyer et al., 2015) . S0 and B0 refers to the token-level representation corresponding to the top element of the stack and buffer, while S1 and B1 refers to those that are second to the top. We use a different notation here to avoid confusion with the states in StackLSTM, which represent non-local information beyond token-level. Reflecting on the aforementioned discussion on parallelism, one should notice that StackLSTM falls into the category of neural network architectures that is difficult to perform minibatched training. This is caused by the fact that the computation performed by StackLSTM at each time step is dependent on the discrete input actions. The following section proposes a solution to this problem. Parallelizable StackLSTM Continuing the formulation in the previous section, we will start by discussing our proposed solution under the case where the set of discrete actions contains only Push and Pop operations; we then move on to discussion of the applicability of our proposed solution to the transition systems that are used for building representations for dependency trees. The first modification we perform to the Push and Pop operations above is to unify the pointer update of these operations as p(H) ← p(H) + op, where op is the input discrete operation that could either take the value +1 or -1 for Push and Pop operation. After this modification, we came to the following observations: Now, what remains to homogenize Push and Pop operations is conducting the extra computations needed for Push operation when Pop is fed in as well, while guaranteeing the correctness of the resulting hidden state both in the current time step and in the future. The next observation points out a way for this guarantee: Observation 2 In a StackLSTM, given the current stack top pointer position p(H), any hidden state h i where i > p(H) will not be read until it is overwritten by a Push operation. What follows from this observation is the guarantee that we can always safely overwrite hidden states h i that are indexed higher than the current stack top pointer, because it is known that any read operation on these states will happen after another overwrite. This allows us to do the extra computation anyway when Pop operation is fed, because the extra computation, especially updating h p(H)+1 , will not harm the validity of the hidden states at any time step. Algorithm 1 gives the final forward computation for the Parallelizable StackLSTM. Note that this algorithm does not contain any if-statements that depends on stack operations and hence is homogeneous when grouped into batches that are consisted of multiple operations trajectories. In transition systems (Nivre, 2008; Kuhlmann et al., 2011) used in real tasks (e.g., transition-based parsing) as shown in Table 1 , it should be noted that more than push and pop operations are needed for the StackLSTM. Fortunately, for Arc-Eager and h_prev ← h p(H) ; h ← LSTM(x t , h_prev); h p(H)+1 ← h; p(H) ← p(H) + op; return h p(H) ; Arc-Hybrid transition systems, we can simply add a hold operation, which is denoted by value 0 for the discrete operation input. For that reason, we will focus on parallelization of these two transition systems for this paper. It should be noted that both observations discussed above are still valid after adding the hold operation. Experiments Setup We implemented 2 the architecture described above in PyTorch (Paszke et al., 2017) . We implemented the batched stack as a float tensor wrapped in a non-leaf variable, thus enabling in-place operations on that variable. At each time step, the batched stack is queried/updated with a batch of stack head positions represented by an integer vector, an operation made possible by gather operation and advanced indexing. Due to this implementation choice, the stack size has to be determined at initialization time and cannot be dynamically grown. Nonetheless, a fixed stack size of 150 works for all the experiments we conducted. We use the dependency parsing task to evaluate the correctness and the scalability of our method. a test set accuracy of 97.47%. We use exactly the same pre-trained English word embedding as Dyer et al. (2015) . We use Adam (Kingma and Ba, 2014) as the optimization algorithm. Following Goyal et al. (2017) , we apply linear warmup to the learning rate with an initial value of τ = 5 × 10 −4 and total epoch number of 5. The target learning rate is set by τ multiplied by batch size, but capped at 0.02 because we find Adam to be unstable beyond that learning rate. After warmup, we reduce the learning rate by half every time there is no improvement for loss value on the development set (ReduceLROnPlateau). We clip all the gradient norms to 5.0 and apply a L 2 -regularization with weight 1 × 10 −6 . We started with the hyper-parameter choices in Dyer et al. (2015) but made some modifications based on the performance on the development set: we use hidden dimension 200 for all the LSTM units, 200 for the parser state representation before the final softmax layer, and embedding dimension 48 for the action embedding. We use Tesla K80 for all the experiments, in order to compare with Neubig et al. (2017b); Dyer et al. (2015) . We also use the same hyper-parameter setting as Dyer et al. (2015) for speed comparison experiments. All the speeds are measured by running through one training epoch and averaging. close to linear, which means there is very little overhead associated with our batching scheme. Quantitatively, according to Amdahl's Law (Amdahl, 1967) , the proportion of parallelized computations is 99.92% at batch size 64. We also compared our implementation with the implementation that comes with Dyer et al. (2015) , which is implemented in C++ with DyNet (Neubig et al., 2017a) . DyNet is known to be very optimized for CPU computations and hence their implementation is reasonably fast even without batching and GPU acceleration, as shown in Figure 1 . 4 But we would like to point out that we focus on the speed-up we are able to obtain rather than the absolute speed, and that our batching scheme is framework-universal and superior speed might be obtained by combining our scheme with alternative frameworks or languages (for example, the torch C++ interface). The dependency parsing results are shown in Table 2. Our implementation is able to yield better test set performance than that reported in Ballesteros et al. (2017) for all batch size configurations except 256, where we observe a modest performance loss. Like Goyal et al. (2017) ; Keskar et al. (2016) ; Masters and Luschi (2018) , we initially observed more significant test-time performance deterioration (around 1% absolute difference) for models trained without learning rate warmup, and concurring with the findings in Goyal et al. (2017) , we find warmup very helpful for stabilizing largebatch training. We did not run experiments with batch size below 8 as they are too slow due to 4 Measured on one core of an Intel Xeon E7-4830 CPU. Results Python's inherent performance issue. Related Work DyNet has support for automatic minibatching (Neubig et al., 2017b) , which figures out what computation is able to be batched by traversing the computation graph to find homogeneous computations. While we cannot directly compare with that framework's automatic batching solution for StackLSTM 5 , we can draw a loose comparison to the results reported in that paper for BiLSTM transition-based parsing (Kiperwasser and Goldberg, 2016) . Comparing batch size of 64 to batch size of 1, they obtained a 3.64x speed-up on CPU and 2.73x speed-up on Tesla K80 GPU, while our architecture-specific manual batching scheme obtained 60.8x speed-up. The main reason for this difference is that their graph-traversing automatic batching scheme carries a much larger overhead compared to our manual batching approach. Another toolkit that supports automatic minibatching is Matchbox 6 , which operates by analyzing the single-instance model definition and deterministically convert the operations into their minibatched counterparts. While such mechanism eliminated the need to traverse the whole computation graph, it cannot homogenize the operations in each branch of if. Instead, it needs to perform each operation separately and apply masking on the result, while our method does not require any masking. Unfortunately we are also not able to compare with the toolkit at the time of this work as it lacks support for several operations we need. Similar to the spirit of our work, Bowman et al. (2016) attempted to parallelize StackLSTM by using Thin-stack, a data structure that reduces the space complexity by storing all the intermediate stack top elements in a tensor and use a queue to control element access. However, thanks to Py-Torch, our implementation is not directly dependent on the notion of Thin-stack. Instead, when an element is popped from the stack, we simply shift the stack top pointer and potentially re-write the corresponding sub-tensor later. In other words, there is no need for us to directly maintain all the intermediate stack top elements, because in PyTorch, when the element in the stack is re-written, its underlying sub-tensor will not be destructed as there are still nodes in the computation graph that point to it. Hence, when performing back-propagation, the gradient is still able to flow back to the elements that are previously popped from the stack and their respective precedents. Hence, we are also effectively storing all the intermediate stack top elements only once. Besides, Bowman et al. (2016) didn't attempt to eliminate the conditional branches in the StackLSTM algorithm, which is the main algorithmic contribution of this work. Conclusion We propose a parallelizable version of StackLSTM that is able to fully exploit the GPU parallelism by performing minibatched training. Empirical results show that our parallelization scheme yields comparable performance to previous work, and our method scales up very linearly with the increasing batch size. Because our parallelization scheme is based on the observation made in section 1, we cannot incorporate batching for neither Arc-Standard transition system nor the token-level composition function proposed in Dyer et al. (2015) efficiently yet. We leave the parallelization of these architectures to future work. Our parallelization scheme makes it feasible to run large-data experiments for various tasks that requires large training data to perform well, such as RNNG-based syntax-aware neural machine translation (Eriguchi et al., 2017) .
1 2 3 4.1 4.2 5 6
Introduction StackLSTM Parallelizable StackLSTM Setup Results Related Work Conclusion
Benchmark
Transition-based dependency parsing on Stanford Dependency Treebank
PyTorch, Single K80 GPU
Transition-based dependency parsing on Stanford Dependency Treebank
PyTorch, Single K80 GPU