(Wu et al., 2016), summarization (Rush, Chopra, and Weston, 2015), image captioning (Xu et al., 2015) and poetry generation (Ghazvininejad et al., 2016). Decoding can be time consuming: 1) Most tasks generate target sentences in an on-line fashion, one token at a time, unlike batch mode used during training; 2) Decoding time is proportional to the beam size , which is usually around 10 in machine translation and 50 in poetry generation; 3) Vocabulary size can be very large (tens of thousands), and the computational complexity of the two major components, Softmax and Beam expansion, are proportional to . As Table 1 shows, these two parts occupy 46% and 30% of the decoding time respectively.
The major bottleneck for Softmax is the matrix multiplication between hidden states and the word embedding matrix . As for Beam expansion, tremendous time is spent on transferring the output of Softmax from device memory to host memory and the following heap-sort on CPU, with complexity .
This work aims to speed up the beam search by reducing the runtime vocabulary size, using Locality Sensitive Hashing (Gionis et al., 1999) . We first hash each high-dimensional word embedding into different buckets so that the embeddings which are closer under a certain distance measure will be in the same bucket with high probability. We choose the winner-take-all (WTA) hash
. We first hash each high-dimensional word embedding into different buckets so that the embeddings which are closer under a certain distance measure will be in the same bucket with high probability. We choose the winner-take-all (WTA) hash(Yagnik et al., 2011) due to its robustness against the perturbations of numerical value and correlation with dot-product distance. Then during decoding, we construct a candidate word list by retrieving the words that share the same bucket with the hidden state under the same hash function. Finally an actual dot-product is calculated between and the shrunken word embedding matrix .
|Total||GPU+CPU||100 %||1178.5 s||574.3 s||2.05|
|Source side||GPU||7 %||88 s||88.1 s||1.00|
|Target side||GPU||63 %||735.5 s||387.2 s||1.90|
|-- Softmax||GPU||43 %||505.3 s||157 s||3.22|
|-- 2nd layer||GPU||10 %||113.7 s||113.7 s||1.00|
|-- 1st layer||GPU||10 %||115.2 s||115.2 s||1.00|
|Beam Expansion||GPU+CPU||30 %||352.4 s||96.4 s||3.66|
|-- Device2Host data transfer||GPU+CPU||12 %||138.1 s||25.8 s||5.35|
|-- Heapsort||CPU||15 %||176 s||31.9 s||5.52|
|-- Hidden states reorder||GPU||3 %||38.3 s||38.7 s||0.99|
|Runtime vocab size||-||-||40,000||5,792||6.91|
Vijayanarasimhan et al. successfully applies this idea to speed up both training and inference of deep neural networks with large output space on multi-core CPU. However, beam search on GPU poses several unique and hard challenges:
The LSH schema used by Vijayanarasimhan et al.
is not GPU-friendly: a) It uses a hash table on CPU to store the bucket key and word list, and the underlying data structure is usually a balanced binary search tree or linked list, both of which are hard to transport to GPU. b) It requires sorting to get the candidate lists, which can not easily parallelize on GPU. c) It processes each hidden vector in the batch one by one, whereas the matrix dot-product on GPU is calculated across the whole batch to fully take advantage of the GPU parallelism.
Beam search generally requires high recall of the top words according to the actual dot-product, because the error will accumulate fast as the sentence is built up. Whereas in practice, we find LSH alone does not delivery an adequate recall/speedup trade-off.
Our main contribution is to re-design the LSH algorithm on GPU for beam search to address the above-mentioned challenges. After fully considering the computational capabilities of CUDA-enabled GPU, we propose the following algorithm for LSH:
We implement a parallel Cuckoo hash table (Pagh and Rodler, 2004) for LSH code lookup on GPU. The Cuckoo hash table can achieve worst-case constant-time lookup, offering an excellent load balance that is important for GPU parallelism.
The candidate word list is shared across different beams so that the GPU can calculate the actual matrix multiplication between and in a batch mode. We also use a threshold to select the candidate list to avoid the expensive sorting operation.
To solve the low-recall problem, we always merge the top frequent words into the candidate list.
We conduct experiments on the task of neural machine translation (NMT). We train NMT models on 4 language pairs, and our LSH algorithm can achieve a consistent 2x overall speedup over the full vocabulary decoding with less than 0.4 BLEU drop.
2 Related Work
Several approaches have been proposed to speed up beam search for RNN-based generation tasks. The first line of research is to use specialized hardware, like Tensor Processing Unit (TPU) and low precision (Low-p) calculation
Several approaches have been proposed to speed up beam search for RNN-based generation tasks. The first line of research is to use specialized hardware, like Tensor Processing Unit (TPU) and low precision (Low-p) calculation(Wu et al., 2016). This method will usually speedup all parts of the neural models.
The second line tries to compress the original large model to a small model by weight pruning (WP) (See, Luong, and Manning, 2016) or sequence-level knowledge distillation (KD) (Kim and Rush, 2016). These methods require additional fine-tuning.
The third line is to modify the Softmax layer to speed up the decoding. Noise-contrastive estimation (NCE) (Gutmann and Hyvärinen, 2010) discriminates between the gold target word and () other sampled words. It has been successfully applied on several NLP tasks (Mnih and Teh, 2012; Vaswani et al., 2013; Williams et al., 2015; Zoph et al., 2016). Morin and Bengio introduces hierarchical softmax (H-softmax) where binary classifications are performed rather than a single -way classification. However, these two methods can only speedup training and still suffer at the decoding phase. Chen, Grangier, and Auli propose differentiated softmax (D-softmax) based on the idea that more parameters should be assigned for embeddings of frequent words and fewer for rare words. It can achieve speedups on both training and decoding.
The fourth line of research uses word alignments (WA) trained on the parallel corpus to construct a small runtime vocabulary for each sentence (Jean et al., 2015; Mi, Wang, and Ittycheriah, 2016; L’Hostis, Grangier, and Auli, 2016; Shi and Knight, 2017). However, this approach is only suitable for tasks where sensible alignments can be extracted, such as machine translation and summarization, and do not benefit tasks like image caption or poem generation.
Table 2 compares different speed up methods. Compared to these existing methods, LSH has the following advantages:
It is orthogonal to the first two lines of research. The first two lines of approaches do not decrease the ratio of the number of word embedding parameters to the number of the rest parameters. Thus, LSH can be applied for further speedup.
It is a machine learning free (ML-free) method, which means it can be used in plug-and-play style, without requiring additional tuning processes or alignment information, once the original model is done training.
|Speedup Train||Speedup Decode||ML Free|
3 GPU Computing Concept
In this section, we describe the basic concepts of CUDA-enabled GPU computing to better motivate our decisions in re-designing the LSH algorithm.
Kernels are functions executed on GPU. Each kernel will be executed by many GPU threads in parallel. These threads are further grouped into warps, where each warp consists of 32 threads. All 32 threads in a wrap will execute the same instruction concurrently. However, due to the branches inside the code, some threads in a warp will be diverged and the rest of the threads in the same wrap will have to idle in that cycle. We should make every effort to avoid warp divergence to maximize GPU usage. Figure 1 shows an example of warp divergence.
3.2 Memory Hierarchy
The bandwidth of GPU global memory is 208 GB/s for Tesla K20 GPU, and it takes 400-800 cycles for global memory access. Another faster but limited memory is shared memory, whose bandwidth is more than 2 TB/s and only takes 3-4 cycles for each access.
The way to access the global memory also strongly affects the speed. Coalesced access, where threads in the same wrap will access consecutive address, can take the full bandwidth, i.e. around 200 GB/s. Whereas the bandwidth of random access can be as low as 20 GB/s.
In practice, we will load data from global memory in a coalesced way to shared memory, then manipulate the data on shared memory, and finally write back to global memory in a coalesced way again.
3.3 Latency Hiding
The global memory access can lead to a large latency (400-800 cycles). However, the GPU scheduler has a smart strategy to hide the latency: when a warp needs to access global memory, GPU will put it into wait, and switch to another warp to execute. In order to hide the latency completely, each kernel should launch enough threads (more than 1000) in parallel.
4 Locality Sensitive Hashing
At each step during beam search with beam size , given the hidden state from the top RNN layer, the probability distribution
from the top RNN layer, the probability distributionover will be calculated by softmax:
where is the ith row of and is the word embedding matrix. The computational intensive part is the matrix product in 2, whose complexity is .
For each beam entry, we are only interested in the top words according to the probability/logit. We can reduce the complexity down to
words according to the probability/logit. We can reduce the complexity down toif we can inexpensively construct a much smaller vocabulary set that also contains the top words.
is proportional to the dot-product between and . Thus, finding the top words with highest probability is equivalent to finding the nearest neighbors of from all embedding vectors under the dot-product distance measure. LSH (Gionis et al., 1999) is an efficient tool for the nearest neighbor problem. LSH will construct a small candidate vocabulary set which will contains the top words with a high expectation.
4.1 LSH on CPU
Vijayanarasimhan et al. successfully utilize winner-take-all (WTA) LSH to speed up the softmax calculation on CPU:
Hash every word embedding into hash code where is the dimension of the hash space. Organize these hash codes in hash tables.
For each hidden vector , apply the same hash function to get .
Given , select the top collisions in the hash table, to construct the candidate vocabulary set .
Calculate the dot-product .
Repeat step 2-4 for each entry in the batch.
The tasks Vijayanarasimhan et al. that sped up are Image Classification, Skipgram Word2Vec, and Video Identification, which all involve only one step of the softmax calculation. When conducting inference in batch mode, the test entries inside a batch can be very different, thus the candidate list will differ a lot. Therefore, step 2-4 must be executed independently for each entry in the batch.
This will lead to less speedup when batch size is large: Vijayanarasimhan et al. reports that speedup decreases from 6.9x to 1.6x when batch size increase from 8 to 64. This is problematic as beam size could be more than 50 for certain task.
4.2 LSH on GPU
A similar reduction in speedup will also happen on GPU, especially because GPUs prefer large-scale calculation to hide latency, as described in Section 3.3. Table 3 shows a comparison of matrix multiplications at different scales. Even though calculation is reduced to 1/12th, the time spent is only shrunk to one third.
On GPU, another drawback of processing each beam one by one, is that we must first construct embedding matrix that occupies contiguous space in global memory to do matrix multiplication. This expensive data movement can further downgrade the speedups.
To solve this problem, we propose to share the candidate vocabulary set across different entries in the batch. Only one embedding matrix will be constructed and the following matrix multiplication will be calculated in batch mode in a single kernel launch. This idea is motivated by the intuition that during beam search, at each step, different will have a big portion of words in common, thus . Although the amount computation in matrix multiplication will increase (), it makes every step of our LSH algorithm on GPU executed in batch mode, saving a lot of time in practice.
Another issue that beam search poses is that the error will accumulate fast as the sentence is built. Therefore, missing a correct word at a certain step will lead to a catastrophe. In practice, we find that even the combined candidate list can miss out some important words. We solve this problem by further merging the top frequent words into the candidate list:
Figure 2 illustrates the detailed pipeline of our LSH algorithm on GPU. Table 4 shows the time breakdown of each step in LSH softmax. Although every step is running nicely in batch mode, we still need to carefully design each step as a naive implementation leads to large overhead on GPU. The naive implementation of cuckoo lookup and construct candidate list experience 3.4x and 1.7x slowdown respectively, compared to our optimized version.
|Full Vocab (ms)||Percent||LSH (ms)||Percent||Naive slowdown|
|Softmax||120.97||100.0 %||44.09||100.0 %|
|-- LSH overhead||-||16.53||37.5 %|
|-- -- WTA-hash||-||3.72||8.4 %|
|-- -- Cuckoo lookup||-||7.29||16.5 %||3.4x|
|-- -- Construct candidate list||-||2.51||5.7 %||1.7x|
|-- -- Construct||-||3.01||6.8 %|
|-- Matrix multiply||108.16||89.4 %||22.43||50.9 %|
|-- Normalization||12.33||10.2 %||2.74||6.2 %|
|Runtime vocab size||40,000||6,177|
4.2.1 Winner-Take-All hashing
WTA will convert a -dimension real value hidden vector into a -dimension int value hash code. Each is the index of the maximum value of the first elements of the permuted . Thus, the hash code is actually an ordinal embedding of the original hidden vector, and we use the ordinal similarity as a proxy for the dot-product similarity. We can further group these hash codes into bands, and convert into a -dimension band code:
where each band code is the concatenation of hash codes. Table 5 shows an example of WTA hash. We can represent in bits, and we make sure so that we can store each band code using an int32 on GPU. The hyper parameters for WTA hash are .
4.2.2 Cuckoo lookup
Given , this step will calculate the hit matrix , where
counts how many band codes are the same between and , which estimates the dot-product similarity between and .
First, we hash all word embeddings . For each band, we will build a hash table:
where the key is the band code and the value is a list containing all the words whose band code is equal to the key. We re-organize into an flat array on GPU: word ids with same band code are stored in continuous span, and we build a cuckoo hash table for each to store the starting position and corresponding length of each span:
Second, we launch a total of GPU threads to calculate , and each thread follows Algorithm 1. To look up certain key in cuckoo hash table, it hashes and compares the key at most twice. Thus the warp(a group of 32 threads) won’t diverge at line 2.
However, at line 3, because different threads will have different values, the execution time of the warp will depend on the largest length. There will be a serious warp divergence at line 3-6. To solve this problem, we re-arrange the execution order so that the 32 threads will first process the for-loop of thread together, then the for-loop of thread together, until that of thread. Such re-arrangement will speed up this step by 3.4x. Figure 4 illustrates the two different thread arrangements.
4.2.3 Construct candidate list
Given the hit matrix and a threshold , this step selects the final candidate vocabulary set , where:
We use threshold to avoid the inefficient sorting on GPU. is a sparse matrix after filtering with , whereas should be a dense array. This is the canonical Stream Compaction problem, and one can simply use copy_if function provided in thrust library. To further improve the efficiency, we re-design the algorithm by taking advantage of shared memory and coalesced access. The new algorithm is illustrated in Figure 5.
Hyper-parameters that define a WTA LSH beam search are , where is beam size, is the number of top frequent words to merge and is the threshold to select .
Training configurations of different language pairs. The attention model is based onLuong, Pham, and Manning. Data sources: ASPEC Japanese-English Corpus (Nakazawa et al., 2016), French-English Corpus from WMT2014 (Bojar et al., 2014), and Uzbek-English Corpus (Linguistic Data Consortium, 2016).
We conduct our experiment on 4 machine translation models: Japanese to English (J2E), English to Japanese (E2J), French to English (F2E) and Uzbek to English (U2E). The statistics and training parameters are shown in Table 6.
Overall speedup As Table 7 shows, our LSH decoding achieves up to 4x speedup on the softmax, and a consistent 2x overall speedup for J2E, E2J and F2E with tiny BLEU score loss. For U2E, the speedup without BLEU loss is 1.78x, due to the small original target vocabulary size (25,000).
We compare our algorithm with two other decoding acceleration methods: Decoding using only the top frequent words (TOP) and decoding with word alignments (WA) (Shi and Knight, 2017). We conduct a grid search of the LSH hyper parameters and find that and generally deliver good performance/speedup trade-off. We vary other two hyper parameters to get different speedup and BLEU score. Figure 6 shows the BLEU/speedup curve of the three decoding methods on 4 translation directions.
Our LSH decoding always obtain a higher BLEU with a large margin than TOP decoding at the same speedup level. Table 4 shows that the optimized LSH overhead can take up to 37.5% of the total time to calculate softmax. Thus, to achieve the same speedup with TOP decoding, the runtime vocabulary size of LSH can not exceed half of that of TOP decoding, which demonstrates that our LSH algorithm indeed selects a smaller yet more accurate vocabulary set.
The WA decoding achieves higher speedup with the same BLEU. However, this approach can only work in the context where a sensible alignment information can be provided.
Effects of beam size Another easy way to speed up decoding is to just reduce the beam size, which is also orthogonal to our LSH decoding. Figure 7 demonstrates that LSH decoding with reduced beam size can achieve even better speed/performance trade-off.
|Beam size||Speedup||BLEU loss|
On the other hand, the speed and performance at larger beam size is also important because certain application, like poem generation, requires the beam size larger than 50 to work in practice. Table 8 shows the speedup and BLEU loss of our LSH decoding over the full vocabulary decoding at larger beam size (batch size). Unlike the algorithm proposed by Vijayanarasimhan et al., our algorithm will maintain or even obtain a higher speedup with the same level of BLEU when beam size (batch size) increases, which can be further explained by that large batch size will saturate the GPU and fully exploit its parallel power.
Effects of T The merge of top frequent words into the candidate word list is necessary to obtain good performance. Figure 8 shows the BLEU/speedup curve for different on the French-to-English translation task. Having too small will result in low BLEU whereas having too large will limit the highest speedup the method can achieve.
We re-design the LSH algorithm for beam search on GPU. The candidate vocabulary set is shared across the beams to execute every step in batch mode. Several key functions are optimized by using a cuckoo hash table, taking advantage of shared memory, and avoiding warp divergence. Top frequent words are merged into to further improve the performance. Our LSH algorithm is a machine-learning-free acceleration method that achieves 2x speedup on 4 machine translation tasks, and delivers better BLEU/speedup trade-off than TOP decoding.
- Bojar et al. (2014) Bojar, O.; Buck, C.; Federmann, C.; Haddow, B.; Koehn, P.; Machacek, M.; Monz, C.; Pecina, P.; Post, M.; Saint-Amand, H.; Soricut, R.; and Specia, L., eds. 2014. Proc. Ninth Workshop on Statistical Machine Translation.
- Chen, Grangier, and Auli (2016) Chen, W.; Grangier, D.; and Auli, M. 2016. Strategies for training large vocabulary neural language models. Proc. ACL.
Ghazvininejad et al. (2016)
Ghazvininejad, M.; Shi, X.; Choi, Y.; and Knight, K.
Generating topical poetry.
Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, 1183--1191.
- Gionis et al. (1999) Gionis, A.; Indyk, P.; Motwani, R.; et al. 1999. Similarity search in high dimensions via hashing. In VLDB, volume 99, 518--529.
- Gutmann and Hyvärinen (2010) Gutmann, M., and Hyvärinen, A. 2010. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proc. AISTATS.
- Jean et al. (2015) Jean, S.; Cho, K.; Memisevic, R.; and Bengio, Y. 2015. On using very large target vocabulary for neural machine translation. In Proc. ACL.
- Kim and Rush (2016) Kim, Y., and Rush, A. M. 2016. Sequence-level knowledge distillation. In Proc. EMNLP.
- L’Hostis, Grangier, and Auli (2016) L’Hostis, G.; Grangier, D.; and Auli, M. 2016. Vocabulary selection strategies for neural machine translation. arXiv preprint arXiv:1610.00072.
- Linguistic Data Consortium (2016) Linguistic Data Consortium. 2016. Bolt lrl uzbek representative language pack v1.0. ldc2016e29.
- Luong, Pham, and Manning (2015) Luong, M.-T.; Pham, H.; and Manning, C. D. 2015. Effective approaches to attention-based neural machine translation. In Proc. EMNLP.
- Mi, Wang, and Ittycheriah (2016) Mi, H.; Wang, Z.; and Ittycheriah, A. 2016. Vocabulary Manipulation for Neural Machine Translation. In Proc. ACL.
- Mnih and Teh (2012) Mnih, A., and Teh, Y. W. 2012. A fast and simple algorithm for training neural probabilistic language models. Proc. ICML.
- Morin and Bengio (2005) Morin, F., and Bengio, Y. 2005. Hierarchical probabilistic neural network language model. In Proc. AISTATS.
- Nakazawa et al. (2016) Nakazawa, T.; Yaguchi, M.; Uchimoto, K.; Utiyama, M.; Sumita, E.; Kurohashi, S.; and Isahara, H. 2016. Aspec: Asian scientific paper excerpt corpus. In Proc. LREC.
- Pagh and Rodler (2004) Pagh, R., and Rodler, F. F. 2004. Cuckoo hashing. Journal of Algorithms 51(2):122--144.
- Rush, Chopra, and Weston (2015) Rush, A. M.; Chopra, S.; and Weston, J. 2015. A neural attention model for abstractive sentence summarization. In Proc. EMNLP.
- See, Luong, and Manning (2016) See, A.; Luong, M.-T.; and Manning, C. D. 2016. Compression of neural machine translation models via pruning.
- Shi and Knight (2017) Shi, X., and Knight, K. 2017. Speeding up neural machine translation decoding by shrinking run-time vocabulary. In Proc. ACL.
- Vaswani et al. (2013) Vaswani, A.; Zhao, Y.; Fossum, V.; and Chiang, D. 2013. Decoding with large-scale neural language models improves translation. In Proc. EMNLP.
- Vijayanarasimhan et al. (2015) Vijayanarasimhan, S.; Shlens, J.; Monga, R.; and Yagnik, J. 2015. Deep networks with large output spaces. In Proc. ICLR.
- Williams et al. (2015) Williams, W.; Prasad, N.; Mrva, D.; Ash, T.; and Robinson, T. 2015. Scaling recurrent neural network language models. In Proc. ICASSP.
- Wu et al. (2016) Wu, Y.; Schuster, M.; Chen, Z.; Le, Q. V.; Norouzi, M.; Macherey, W.; Krikun, M.; Cao, Y.; Gao, Q.; Macherey, K.; et al. 2016. Google’s neural machine translation system: Bridging the gap between human and machine translation. arXiv preprint arXiv:1609.08144.
- Xu et al. (2015) Xu, K.; Ba, J.; Kiros, R.; Cho, K.; Courville, A. C.; Salakhutdinov, R.; Zemel, R. S.; and Bengio, Y. 2015. Show, attend and tell: Neural image caption generation with visual attention. In ICML, volume 14, 77--81.
- Yagnik et al. (2011) Yagnik, J.; Strelow, D.; Ross, D. A.; and Lin, R.-s. 2011. The power of comparative reasoning. In Proc. ICCV.
- Zoph et al. (2016) Zoph, B.; Vaswani, A.; May, J.; and Knight, K. 2016. Simple, fast noise-contrastive estimation for large rnn vocabularies. In Proc. NAACL-HLT.