Large Product Key Memory for Pretrained Language Models (2020)

Abstract: Product key memory (PKM) proposed by Lample et al. (2019) enables to improve prediction accuracy by increasing model capacity efficiently with insignificant computational overhead. However, their empirical application is only limited to causal language modeling. Motivated by the recent success of pretrained language models (PLMs), we investigate how to incorporate large PKM into PLMs that can be finetuned for a wide variety of downstream NLP tasks. We define a new memory usage metric, and careful observation using 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 effective 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 memory utilization and downstream performance. Code and pretrained weights are available at https://github.com/clovaai/pkm-transformers .

## introduction larger model capacity has brought improvement in accuracy by enabling better modeling of data. however, increasing model capacity causes a significant increase in computational cost at both training and inference time despite better accuracy. to address this issue, @xcite propose product key memory (pkm) that enables very efficient and exact nearest neighbor search in a large number of learnable memory slots. they substitute a feed-forward network (ffn) in a transformer block @xcite with a pkm layer. augmenting large pkm layers to networks allows we measure various memory utilization metrics to analyze how many memory slots contribute to the model prediction. careful examination about memory utilization during and after the training demonstrates that only a few memory slots are be-ing used importantly ( § 3.1). we attribute this phenomenon, called a catastrophic drift, to the sparsely updated memory parameters. the lower memory utilization implies that model capacity from memory is not fully exploited. it promotes us to develop methods that can overcome this issue. we found that initialization from weights pretrained without memory is essential for pretraining pkm-augmented plms. moreover, rather than replacing an ffn to a pkm as lample et al. (2019) do, we show that adding pkm to a transformer layer @xcite with a residual connection @xcite without removing ffn is advantageous. both the initialization ( § 4.1) and our proposed residual memory (resm, § 4.2) prevent a sudden change of transformer parameters, thus allow to train memory parameters better by less suffering from the catastrophic drift. consequently, we obtain pkm-augmented-bert base having comparable accuracy and faster than bert large . as demonstrated in the main contributions of this work are summarized as follows. first, we explore how to incorporate pkm to plms to be finetuned for downstream tasks and find that simple application does not work well. secondly, we attribute this to a catastrophic drift during the training by careful monitoring of memory utilization. lastly, we propose simple yet effective solutions to tackle the observed catastrophic drift problem: (1) weight initialization without pkm and (2) the residual memory layer. we empirically verify that both of them are crucial to achieve improved accuracy. in our knowledge, this is the first work that successfully applies pkm to plms. ## memory utilization analysis as shown in our experiment ( ## finetuning results ## conclusion and future work this work starts from unexpected results that directly applying pkm to plms does not work well in downstream tasks, contrary to ( lample @xcite . in this paper, we successfully augment pkm to plms with two ingredients, weight initialization and residual connection, based on the observation of memory utilization and catastrophic drift during the training. consequently, we encourage to utilize memory architecture such as pkm for plms in practical use. although our approach mitigates the catastrophic drift problem somehow, we leave further study on it during both pretraining and finetuning as future work. one possible solution is to regularize a pkm memory by a structured dropout on the memory keys like drophead @xcite . it would also help to prune unnecessary memory slots on-demand during the inference time.