Large Product Key Memory for Pretrained Language Models

Findings of the Association for Computational Linguistics: EMNLP 2020, pages 4060–4069 November 16 - 20, 2020. c⃝2020 Association for Computational Linguistics 4060 Large Product Key Memory for Pretrained Language Models Gyuwan Kim∗ Clova AI, NAVER Corp. gyuwan.kim@navercorp.com Tae-Hwan Jung∗† Kyung Hee University nlkey2022@gmail.com Abstract Product key memory (PKM) proposed by Lample et al. (2019) enables to improve pre- diction accuracy by increasing model capac- ity efficiently with insignificant computational overhead. However, their empirical applica- tion is only limited to causal language model- ing. Motivated by the recent success of pre- trained language models (PLMs), we investi- gate how to incorporate large PKM into PLMs that can be finetuned for a wide variety of downstream NLP tasks. We define a new mem- ory usage metric, and careful observation us- ing this metric reveals that most memory slots remain outdated during the training of PKM- augmented models. To train better PLMs by tackling this issue, we propose simple but ef- fective solutions: (1) initialization from the model weights pretrained without memory and (2) augmenting PKM by addition rather than replacing a feed-forward network. We verify that both of them are crucial for the pretraining of PKM-augmented PLMs, enhancing mem- ory utilization and downstream performance. Code and pretrained weights are available at https://github.com/clovaai/pkm-transformers. 1 Introduction Larger model capacity has brought improvement in accuracy by enabling better modeling of data. However, increasing model capacity causes a sig- nificant increase in computational cost at both train- ing and inference time despite better accuracy. To address this issue, Lample et al. (2019) propose product key memory (PKM) that enables very effi- cient and exact nearest neighbor search in a large number of learnable memory slots. They substi- tute a feed-forward network (FFN) in a transformer block (Vaswani et al., 2017) with a PKM layer. Augmenting large PKM layers to networks allows ∗Equal contribution. †TJ was an intern at Clova AI while doing this work. Model # Layers # Params Inference Speed (batch/sec) BERTBASE 12 110M 79.8 BERTBASE +PKM 12 506M 61.4 BERTBASE +ResM 12 515M 59.3 BERTLARGE 24 340M 43.1 BERTLARGE +PKM 24 860M 37.2 BERTLARGE +ResM 24 876M 36.1 Table 1: Comparison of inference speed between dif- ferent model sizes and the memory layers. We run each model for the classification task with batch size 1, and measure inference speed on a single V100 GPU. We follow the model size settings of BERT (Devlin et al., 2018). We use two memory layers with the recom- mended setting of PKM hyper-parameters following Lample et al. (2019) as described in §5. As marked bold, BERTBASE with our proposed residual memory (ResM) is much faster than BERTLARGE, while having more parameters. increasing model capacity, with only a slight in- crease in inference time. Lample et al. (2019) prove the efficiency of PKM on causal language models (CLMs) in terms of the superior trade-off between perplexity and inference speed. For instance, they achieve a PKM-augmented CLM with only 12 lay- ers that is more accurate and twice faster than …

Content truncated to 500 words