Fast Locality Sensitive Hashing for Beam Search on GPU
Xing Shi, Shizhen Xu, Kevin Knight

TL;DR
This paper introduces a GPU-optimized LSH algorithm using WTA hashing and CUDA architecture design to accelerate beam search in sequence models, achieving significant speedups without loss of accuracy.
Contribution
The paper presents a novel GPU-based LSH algorithm with architecture-aware optimizations for faster beam search in neural sequence models.
Findings
Up to 4x speedup on softmax module
2x overall speedup on GPU
No degradation in BLEU scores
Abstract
We present a GPU-based Locality Sensitive Hashing (LSH) algorithm to speed up beam search for sequence models. We utilize the winner-take-all (WTA) hash, which is based on relative ranking order of hidden dimensions and thus resilient to perturbations in numerical values. Our algorithm is designed by fully considering the underling architecture of CUDA-enabled GPUs (Algorithm/Architecture Co-design): 1) A parallel Cuckoo hash table is applied for LSH code lookup (guaranteed O(1) lookup time); 2) Candidate lists are shared across beams to maximize the parallelism; 3) Top frequent words are merged into candidate lists to improve performance. Experiments on 4 large-scale neural machine translation models demonstrate that our algorithm can achieve up to 4x speedup on softmax module, and 2x overall speedup without hurting BLEU on GPU.
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsAdvanced Image and Video Retrieval Techniques · Algorithms and Data Compression · Genomics and Phylogenetic Studies
MethodsSPEED: Separable Pyramidal Pooling EncodEr-Decoder for Real-Time Monocular Depth Estimation on Low-Resource Settings · Softmax
