Abstract
Product key memory (PKM) proposed by Lample et al. (2019) enables to improveprediction accuracy by increasing model capacity efficiently with insignificantcomputational overhead. However, their empirical application is only limited tocausal language modeling. Motivated by the recent success of pretrainedlanguage models (PLMs), we investigate how to incorporate large PKM into PLMsthat can be finetuned for a wide variety of downstream NLP tasks. We define anew memory usage metric, and careful observation using this metric reveals thatmost memory slots remain outdated during the training of PKM-augmented models.To train better PLMs by tackling this issue, we propose simple but effectivesolutions: (1) initialization from the model weights pretrained without memoryand (2) augmenting PKM by addition rather than replacing a feed-forwardnetwork. We verify that both of them are crucial for the pretraining ofPKM-augmented PLMs, enhancing memory utilization and downstream performance.Code and pretrained weights are available athttps://github.com/clovaai/pkm-transformers.