Elamur Research⌕
☰
home / papers / 2021 / Efficient Nearest Neighbor Language Models

Efficient Nearest Neighbor Language Models

Paper · 2021 · Junxian He, Graham Neubig and Taylor Berg-Kirkpatrick

Summary

Non-parametric language models such as the kNN-LM (Khandelwal et al., 2019) improve perplexity by consulting an external datastore of context–token pairs. However, searching that datastore at every prediction step makes inference far slower than with a standard parametric NLM. He, Neubig and Berg-Kirkpatrick study how to cut this cost through three complementary techniques:

  • Adaptive retrieval: a small learned adaptor decides when retrieval can be skipped.
  • Datastore pruning: entries are compressed, for example by greedy merging.
  • Dimension reduction: PCA, using the FAISS implementation, shrinks the stored context vectors.

They evaluate on the WikiText-103 benchmark and on the Law-MT domain-adaptation setting. Each technique is tuned on validation data to give the largest speed-up while losing at most about 0.1 perplexity against the vanilla kNN-LM. Combined, the techniques make the kNN-LM roughly 5–6.6x faster with essentially unchanged or even better perplexity. The analysis offers practical guidance for building and deploying retrieval-based language models at lower cost.

Key points

  • On the WikiText-103 test set, combining all three methods gives 16.67 ppl at 1835 tokens/s, compared with 16.65 ppl at 277 tokens/s for the vanilla kNN-LM, a 6.6x speed-up. The parametric NLM baseline reaches 18.66 ppl.
  • PCA from 1024 to 512 dimensions slightly improves WikiText-103 perplexity (16.40) and runs 3.6x faster. Reducing below 256 dimensions hurts badly, and below 128 dimensions the kNN-LM loses its edge over the NLM.
  • Adaptive retrieval skips 50% of lookups on WikiText-103 for nearly 2x speed. On Law-MT it removes only 10%, because the domain datastore matters much more there.
  • Greedy merging prunes 40% of the datastore, costing 0.2 ppl on WikiText-103 and 0.7 ppl on Law-MT. The smaller Law-MT datastore (19M entries vs. 103M) appears less redundant.
  • On Law-MT, the combined system runs 5.4x faster than the kNN-LM and improves perplexity (12.29 vs. 12.64). Code is available at github.com/jxhe/efficient-knnlm.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

Cite

Cite as Junxian He, Graham Neubig and Taylor Berg-Kirkpatrick. 2021. Efficient Nearest Neighbor Language Models. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2021.emnlp-main.461

@inproceedings{he2021efficient,
  title = {{Efficient Nearest Neighbor Language Models}},
  author = {Junxian He, Graham Neubig and Taylor Berg-Kirkpatrick},
  booktitle = {Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing},
  year = {2021},
  doi = {10.18653/v1/2021.emnlp-main.461},
  url = {https://doi.org/10.18653/v1/2021.emnlp-main.461}
}

References

36 references, 10 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Roee Aharoni and Yoav Goldberg. 2020. Unsupervised domain clusters in pretrained language models. In Proceedings of ACL. our card
  2. Rami Al-Rfou, Dokook Choe, Noah Constant, Mandy Guo, and Llion Jones. 2019. Character-level language modeling with deeper self-attention. In Proceedings of AAAI. find
  3. Fabien André, Anne-Marie Kermarrec, and Nicolas Le Scouarnec. 2019. Quicker adc: Unlocking the hidden potential of product quantization with simd. IEEE transactions on pattern analysis and machine intelligence. find
  4. Alexei Baevski and Michael Auli. 2019. Adaptive input representations for neural language modeling. In Proceedings of ICLR. find
  5. Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2015. Neural machine translation by jointly learning to align and translate. In Proceedings of ICLR. our card
  6. Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. arXiv preprint arXiv:2005.14165. our card
  7. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL. our card
  8. Martin Ester, Hans-Peter Kriegel, Jörg Sander, and Xiaowei Xu. 1996. A density-based algorithm for discovering clusters in large spatial databases with noise. In Proceedings of KDD. find
  9. Aristides Gionis, Piotr Indyk, Rajeev Motwani, et al. 1999. Similarity search in high dimensions via hashing. In Proceedings of VLDB. find
  10. Yunchao Gong, Svetlana Lazebnik, Albert Gordo, and Florent Perronnin. 2012. Iterative quantization: A procrustean approach to learning binary codes for large-scale image retrieval. IEEE transactions on pattern analysis and machine intelligence, 35(12):2916–2929. find
  11. Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar. 2020. Accelerating large-scale inference with anisotropic vector quantization. In Proceedings of ICML. find
  12. Kelvin Guu, Tatsunori B Hashimoto, Yonatan Oren, and Percy Liang. 2018. Generating sentences by editing prototypes. Transactions of the Association for Computational Linguistics, 6:437–450. our card
  13. Junxian He, Taylor Berg-Kirkpatrick, and Graham Neubig. 2020. Learning sparse prototypes for text generation. In Proceeings of NeurIPS. find
  14. Sepp Hochreiter and Jürgen Schmidhuber. 1997. Long short-term memory. Neural computation, 9(8):1735–1780. find
  15. Herve Jegou, Matthijs Douze, and Cordelia Schmid. 2010. Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence, 33(1):117–128. find
  16. Alexander Johansen and Richard Socher. 2017. Learning when to skim and when to read. In Proceedings of the 2nd Workshop on Representation Learning for NLP. find
  17. Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2019. Billion-scale similarity search with gpus. IEEE Transactions on Big Data. our card
  18. Urvashi Khandelwal, Angela Fan, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2020. Nearest neighbor machine translation. arXiv preprint arXiv:2010.00710. — [arXiv:2010.00710](https://arxiv.org/abs/2010.00710) find
  19. Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2019. Generalization through memorization: Nearest neighbor language models. In Proceedings of ICLR. find
  20. Diederik P Kingma and Jimmy Ba. 2015. Adam: A method for stochastic optimization. In Proceedings of ICLR. find
  21. Philipp Koehn and Rebecca Knowles. 2017. Six challenges for neural machine translation. In Proceedings of the First Workshop on Neural Machine Translation. find
  22. Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2017. Pointer sentinel mixture models. In Proceedings of ICLR. 2019. XLNet: Generalized autoregressive pretraining for language understanding. In Proceedings of NeurIPS. our card
  23. Tomáš Mikolov, Martin Karafiát, Lukáš Burget, Jan Černockỳ, and Sanjeev Khudanpur. 2010. Recurrent neural network based language model. In Eleventh annual conference of the international speech communication association. find
  24. Dani Yogatama, Cyprien de Masson d’Autume, and Lingpeng Kong. 2021. Adaptive semiparametric language models. Transactions of the Association for Computational Linguistics, 9:362–373. find
  25. Sewon Min, Jordan Boyd-Graber, Chris Alberti, Danqi Chen, Eunsol Choi, Michael Collins, Kelvin Guu, Hannaneh Hajishirzi, Kenton Lee, Jennimaria Palomaki, et al. 2020. find
  26. Neurips 2020 efficientqa competition: Systems, analyses and lessons learned. arXiv preprint arXiv:2101.00133. — [arXiv:2101.00133](https://arxiv.org/abs/2101.00133) find
  27. Marius Muja and David G Lowe. 2009. Fast approximate nearest neighbors with automatic algorithm configuration. In Proceedings of the International Conference on Computer Vision Theory and Applications. find
  28. Nathan Ng, Kyra Yee, Alexei Baevski, Myle Ott, Michael Auli, and Sergey Edunov. 2019. Facebook fair’s wmt19 news translation task submission. In Proceedings of the Fourth Conference on Machine Translation (Volume 2: Shared Task Papers, Day 1). our card
  29. Matthew Peters, Mark Neumann, Mohit Iyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. 2018. Deep contextualized word representations. In Proceedings of NAACL. our card
  30. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9. find
  31. Alexander M Rush, Sumit Chopra, and Jason Weston. 2015. A neural attention model for abstractive sentence summarization. In Proceedings of EMNLP. find
  32. Josef Sivic and Andrew Zisserman. 2003. Video google: A text retrieval approach to object matching in videos. In Proceedings of ICCV. find
  33. Martin Sundermeyer, Ralf Schlüter, and Hermann Ney. 2012. Lstm neural networks for language modeling. In Thirteenth annual conference of the international speech communication association. find
  34. Nenad Tomasev, Milos Radovanovic, Dunja Mladenic, and Mirjana Ivanovic. 2013. The role of hubness in clustering high-dimensional data. IEEE transactions on knowledge and data engineering, 26(3):739–751. find
  35. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Proceedings of NeurIPS. our card
  36. Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/he-2021-efficient-nearest-neighbor-language-models/