Meta-Learning Fast Weight Language Models

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 9751 - 9757 December 7-11, 2022 ©2022 Association for Computational Linguistics Meta-Learning Fast Weight Language Models Kevin Clark Kelvin Guu Ming-Wei Chang Panupong Pasupat Geoffrey Hinton Mohammad Norouzi Google Research {kevclark,kguu,mingweichang,ppasupat,geoffhinton,mnorouzi}@google.com Abstract Dynamic evaluation of language models (LMs) adapts model parameters at test time using gradient information from previous to- kens and substantially improves LM perfor- mance. However, it requires over 3x more compute than standard inference. We present Fast Weight Layers (FWLs), a neural com- ponent that provides the benefits of dynamic evaluation much more efficiently by express- ing gradient updates as linear attention. A key improvement over dynamic evaluation is that FWLs can also be applied at training time so the model learns to make good use of gradi- ent updates. FWLs can easily be added on top of existing transformer models, require rela- tively little extra compute or memory to run, and significantly improve language modeling perplexity. 1 Introduction A key challenge in language modeling is repre- senting the contextual information from previous tokens. Transformer language models use atten- tion to pass along this information, but constantly referring back to the previous text is a cogni- tively implausible model of working memory. An appealing alternative is using fast weight neural networks (Hinton and Plaut, 1987; Schmidhuber, 1992). Inspired by short-term plasticity in the brain, these models have parameters that change on-the-fly based on input data (previous tokens for LMs) in addition to standard “slow” weights learned during training. Fast weights have proven successful for supervised (Ba et al., 2016a), rein- forcement (Munkhdalai et al., 2019), and few-shot (Munkhdalai and Yu, 2017) learning. Dynamic evaluation (Mikolov et al., 2010; Krause et al., 2018) uses a variant of fast weights to improve language models at inference time. Af- ter scoring (or generating) a chunk of text, dy- namic evaluation applies a gradient update to the model coming from the LM loss over that chunk before continuing. Intuitively, this process im- proves performance because an update that makes the model better at predicting previous tokens will likely also make it better at predicting fu- ture ones. Dynamic evaluation substantially im- proves LM perplexity, but has numerous draw- backs. It requires an extra forward and backward pass through the model to compute the gradients, and the sequential gradient updates over chunks cannot be parallelized. Furthermore, it is very memory intensive because a separate copy of the evolving model weights has to be stored for each example in a minibatch. Lastly, dynamic evalua- tion is only used at test-time, so the model does not learn to make good use of gradient updates. We present Fast Weight Layers (FWLs), a neu- ral component that provides the benefits of dy- namic evaluation with none of these downsides. They can be added to any LM without requiring changes to the training or evaluation loop. Like dynamic evaluation, FWLs also update their pa- rameters using gradient information from previous tokens, but FWLs employ …

Content truncated to 500 words