Navigating with Graph Representations for Fast and Scalable Decoding of Neural Language Models
Neural language models (NLMs) have recently gained a renewed interest by achieving state-of-the-art performance across many natural language processing (NLP) tasks. However, NLMs are very computationally demanding largely due to the computational cost of the decoding process, which consists of a softmax layer over a large vocabulary. We observe that in the decoding of many NLP tasks, only the probabilities of the top-K hypotheses need to be calculated preciously and K is often much smaller than the vocabulary size. This paper proposes a novel softmax layer approximation algorithm, called Fast Graph Decoder (FGD), which quickly identifies, for a given context, a set of K words that are most likely to occur according to a NLM. We demonstrate that FGD reduces the decoding time by an order of magnitude while attaining close to the full softmax baseline accuracy on neural machine translation and language modeling tasks. We also prove the theoretical guarantee on the softmax approximation quality.
Update after author response:  Thanks for the detailed response! It's a strong submission and I vote for an accept.  ===========  This paper aims to speed up the computation of the softmax over a large vocabulary, which is quite common in some NLP tasks like e.g., language modeling. Specifically, the proposed method formulates the problem into a nearest neighbor search in a small world graph, and applies a log time algorithm to find the approximate top K predictions. The resulting time complexity reduces to logarithmic in the vocabulary size in expectation, in contrast to the linear one in a standard softmax. The proposed method is empirically compared to both full softmax and a state-of-the-art baseline on language modeling and neural machine translation. Experimental results show that the method achieves significant speed-up, without sacrificing too much accuracy, and that it scales to larger-size models.  The paper is overall well-motivated and clearly-written, and I find it interesting to read. My major concern is perhaps how the method generalizes to training,  detailed below.  Concerns:  - Small world graph construction: It would be nice to see some discussions on the complexity of the graph construction step, and whether it could become a potential problem if the algorithm is applied to training.  - Probabilities: Are there any good reasons to not report perplexities in the language modeling task? It appears to me that in a FGD, it is less straightforward to assign probabilities to the corpus. One thing on top of my head is to attribute most of the mass to the retrieved top K, with some additional tricks to avoid numerical issues. Have the authors tried this?  - The title is a bit confusing, since ``language models'' usually mean the language modeling task, but here the algorithm applies to many other tasks involving a large softmax.  - There are several statements that look like ``finding top K in sublinear time complexity,'' especially around lines 30--40. I find such statements a bit misleading, since essentially one would need approximations to achieve this.