Meta-Learning Fast Weight Language Models (2022)

Abstract: Dynamic evaluation of language models (LMs) adapts model parameters at test time using gradient information from previous tokens and substantially improves LM performance. However, it requires over 3x more compute than standard inference. We present Fast Weight Layers (FWLs), a neural component that provides the benefits of dynamic evaluation much more efficiently by expressing 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 gradient updates. FWLs can easily be added on top of existing transformer models, require relatively little extra compute or memory to run, and significantly improve language modeling perplexity.

## introduction a key challenge in language modeling is representing the contextual information from previous tokens. transformer language models use attention to pass along this information, but constantly referring back to the previous text is a cognitively implausible model of working memory. an appealing alternative is using fast weight neural networks @xcite . 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 @xcite , reinforcement @xcite , and few-shot @xcite learning. dynamic evaluation @xcite uses a variant of fast weights to improve language models at inference time. after scoring (or generating) a chunk of text, dynamic evaluation applies a gradient update to the model coming from the lm loss over that chunk before continuing. intuitively, this process improves performance because an update that makes the model better at predicting previous tokens will likely also make it better at predicting future ones. dynamic evaluation substantially improves lm perplexity, but has numerous drawbacks. 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 evaluation 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 neural component that provides the benefits of dynamic 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 parameters using gradient information from previous tokens, but fwls employ three key ideas to improve efficiency. first, fwls are added on top of the transformer after the last attention layer, which avoids having to backpropagate through the whole transformer and circumvents the complexities of backpropagation through time. secondly, fwls compute gradients in parallel rather than recurrently. lastly, fwls leverage the property that gradient matrices are rank one to compute their outputs efficiently. crucially, the efficient design means fwls can be used at training time so the model learns to obtain beneficial updates from previous gradients. another benefit is that while dynamic evaluation needs hyperparameter search to find a good step size for the gradient update, fwls can instead learn the step size. training fwls can be viewed as applying gradient-based meta-learning (finn et al., 2017) to language modeling, where the support set contains tokens seen so far and the query set contains future tokens, a perspective which helps explain some of the behaviors of fwls. fwls scale well to long sequences and are complementary to existing long-text generation methods such as sparse attention @xcite or recurrent processing @xcite . we evaluate fwls at language modeling on the wikitext-103 dataset @xcite . they substantially improve perplexities over strong baselines; for example lowering transformer-xl's perplexity from 18.1 to 16.6. this gain is comparable to the one from dynamic evaluation, but fwls add less than 20% compute overhead to the model compared to the over 200% of dynamic evaluation. ablations show fwls achieve superior compute vs perplexity trade-offs compared to alternative fast weight methods. we also analyze how fwls improve results and find they are especially effective at modeling rare tokens, repeated tokens, and long documents. ## method our models first run a transformer decoder (or other left-to-right neural network) over an input sequence [x 1 , ..., x t ], yielding vector representations [h 1 , ..., h t ]. fwls use a learned function f θ with parameters θ to further process the text. these parameters (and the transformer's parameters) are "slow" weights learned during training. we chose to use a hidden layer, projection layer, and layernorm @xcite for f θ , which worked well in initial experiments: the fwl's output can be passed into a prediction layer such as an output softmax: in addition to the slow weights, an fwl also employs "fast" weights that let it adapt to contextual information. first, the fwl does a forward and backward pass over the input sequence with its slow weights θ to obtain a gradient ∇ θ l t from each position t. importantly, these t gradients can be computed in a single backwards pass because l t depends only on h t . if we added attention or other temporal dependence to the fwl, backpropagation through time would instead only compute the summed gradient t @xmath0. this limitation is why dynamic evaluation has to process the input in chunks, which means the model does not get updates from recent tokens within the same chunk; fwls do not have this drawback. the fwl then does a second forward pass using evolving fast weights θ . the fast weights are initialized with the slow weights θ. conceptually, the second pass processes the sequence in a leftto-right order, although we will show this can be parallelized in practice. 1 for each position t, it re-runs the fwl and output softmax on h t , but now using θ instead of θ to produce a new loss l t . then the fwl updates its fast weights as θ ← θα • ∇ θ l t where α consists of learned step sizes. we learn a separate step size for each weight matrix/vector in the fwl, which performs slightly better than having one global step size. although we tried more sophisticated update rules such as one based on adam, they did not outperform this simpler one in initial experiments. an overview of fwls is shown in figure 1 . efficient fast weight computation. we will now show how the fwl outputs can be computed without needing to store the t gradients of l 1:t in memory (which would have prohibitive memory requirements) or recurrently computing θ (which would be slow on modern accelerators). we will illustrate this for the second matrix-multiply in f θ ; other parts of the fwl can be computed analogously. we use v t to denote the input, o t to denote the output, and the symbol to differentiate activations in the second (fast weight) forward pass from those in the first (slow weight) pass. the slow weight output at position t is simply o @xmath1. the fast weight output at position t is the first term is just a regular matrix-multiply and can be computed easily. interestingly, the sum in the second term can be interpreted as linear attention (i.e., without a softmax) using query v t , keys v i , and values ∇ o i l i . the partial sum up to t corresponds to causal masking in the attention. computing the fwl output using (2) means we only need to store gradient vectors ∇ o i l i in memory (which is no worse than storing model activations) rather than gradient matrices ∇ w l i . it also can be computed fully in parallel using causal attention. however, naive attention scales quadratically with sequence length, which is an obstacle for applying (2) to long documents. luckily, because the attention is linear, there exist more efficient methods for computing it. in particular, we use the mixed chunk attention method from @xcite , which improves efficiency through breaking the document into chunks while still computing the attention scores exactly. viewing the fast weight output as linear attention makes it clearer how fwls work. for each previous position i, the error signal ∇ o i l i is the direction for adjusting the output o i to lower the loss, and the attention score v t v t i measures the similarity between position i and the current position. this means fwls adjust the current output in a way that mimics the loss-reducing directions of similar past positions. computing fast weights for vectors (e.g. biases or layernorm weights) rather than matrices is simpler because the gradients require less memory. for example, the fast weight output after adding the bias term b to o t is o t + bα b @xmath2. this can be computed in parallel for all t using the cumsum operation over the matrix of gradients [∇ b l 1 , ..., ∇ b l t ], which can be obtained using standard backpropagation and stored in memory. the parallelized fwl updates are not usable during generation, which inherently has to produce one token at a time. instead the model at each step (1) samples xt from its output distribution using the current weights θ , (2) computes the gradient ∇ θ lt of the slow weights predicting xt , and (3) updates the fast weights as θ ← θ -α•∇ θ lt . this process is still relatively efficient because the additional backward/forward passes are only applied to the fwl, not the whole transformer. fwl gradient updates may first seem strange in that the model is "training" on test sequences during perplexity evaluation or on its own outputs during generation. we would like to emphasize that this way of conditioning on previous tokens is just as valid as standard transformer lms attending over gold-standard tokens (from teacher forcing) or their own generated tokens (during sampling). indeed, (2) shows that fwls can be viewed as adding another attention layer to the model, but with a pre-specified function computing the values (the gradient) rather than a learned matrix-multiply. empirically, we did not find fwls to produce degenerate repetitive outputs @xcite ) more than baselines. training. training jointly optimizes the transformer, softmax, and fwl parameters, as well as the step sizes α, to minimize the combined loss t @xmath3. we compute second-order gradients for the fwl parameters, backpropagating through the gradient updates ∇ θ l t . intuitively, this means that in addition to learning to expect the gradient update and adapt quickly, the model also learns to produce effective gradient updates. while second-order gradients are expensive to compute for some models, they are not for fwls because there is no backpropagation through time for the fwl parameters. runtime. fwls add relatively modest compute and memory overhead to the model. in our experiments, fwls @xmath4. the majority of this extra compute comes from the two extra passes through the output layer (the backward pass and fast weight pass), although using an adaptive softmax somewhat alleviates this cost. in contrast, dynamic evaluation adds over 200% extra compute due to needing an extra backward and forward pass through the whole transformer. connection to meta-learning. fwls can be viewed as applying gradient-based meta-learning to lms. specifically, language modeling is treated as a few-shot learning task where the support set contains the tokens seen so far and the query set contains the next token. the fwl training is essentially using maml @xcite , where there is a single inner loop optimization step over the support set that adapts the model parameters using a gradient update. this connection helps explain a surprising property of fwls: although it applies the fast weight gradient update to only a small subset of the network, it performs comparably to dynamic evaluation, which updates all transformer parameters. @xcite show that maml works just as well when the inner-loop update is only applied to the last layer of the network, which is similar to how fwls only update a few weights on top of the transformer. ## experiments we experiment on wikitext-103 @xcite , a standard benchmark for language modeling consisting of approximately 100m tokens from english wikipedia. we report perplexity using the standard tokenization and splits. we consider two baseline models: the sparse local attention transformer from @xcite and transformer-xl @xcite . the sparse transformer model has 12 layers with 768 hidden units (121m parameters). for transformer-xl, we use the large model (257m parameters). transformer-xl processes the text recurrently: at each step it trains on one chunk while attending over but not backpropagating into a previous chunk. 2 we use an analogous trick with fwls to efficiently use them with the recurrent model: we keep a running fast weight update ∆ θ from the tokens in previous chunks and similarly don't backpropagate into this update. more specifically, when processing a chunk of text [ ## conclusion fast weight layers provide the benefits of dynamic evaluation at a fraction of the compute cost and memory usage. they can easily be added to existing language models and yield strong results on language modeling benchmarks. applying fwls to few-shot learning tasks is one interesting future direction: doing one (or perhaps a small number) of gradient updates on few-shot examples might offer a nice middle ground incontext learning where the model parameters are fixed and full fine-tuning. indeed, @xcite show that hidden state optimization, a method closely related to dynamic evaluation, can improve few-shot lm performance.