Fine-Tuning Pre-trained Transformers into Decaying Fast Weights (2022)

Abstract: Autoregressive Transformers are strong language models but incur O(T) complexity during per-token generation due to the self-attention mechanism. Recent work proposes kernel-based methods to approximate causal self-attention by replacing it with recurrent formulations with various update rules and feature maps to achieve O(1) time and memory complexity. We explore these approaches and find that they are unnecessarily complex, and propose a simple alternative - decaying fast weights - that runs fast on GPU, outperforms prior methods, and retains 99% of attention’s performance for GPT-2. We also show competitive performance on WikiText-103 against more complex attention substitutes.

## introduction autoregressive transformers @xcite have demonstrated strong performance on text generation @xcite . the success of selfattention in transformers over recurrent models @xcite ) can be attributed to its parallelizability @xcite and its effective gradient propagation over many time steps @xcite . however, self-attention has a high computation and memory cost. during inference sampling, it consumes o(t ) time and memory and grows linearly per token generated. these drawbacks motivated recent work to convert or fine-tune attention into recurrent formulations with o(1) memory and time complexity for auto-regressive generation. kernel-based methods for self-attention @xcite learn approximations of the exponential similarity function using m-dimensional feature maps to reformulate attention as a recurrent computation. they replace attention with "unlimited capacity" with fixedcapacity fast weights @xcite , where the memory-accuracy tradeoff @xcite is controlled by m. several works explored different feature maps and recurrent formulations (i.e., update rules). @xcite propose feature maps to maintain positive outputs, while @xcite ; @xcite carefully ensure their random feature maps are unbiased estimates of the softmax attention kernel. @xcite ; @xcite propose more sophisticated update rules to forget information in the recurrent state to improve performance. recently, @xcite showed that pre-trained transformers can be fine-tuned into a recurrent formulation using learned relu feature maps with minor degradations. while promising, it is unclear which update rules or feature maps are critical for successful fine-tuning. in this work, we investigate various update rule configurations to fine-tune pre-trained transformers into rnns for fast inference. we find that prior proposals contain unnecessary operations, leading us to propose a simple element-wise decay update rule with no feature map. we fine-tune gpt-2 @xcite into our recurrent formula-tion to demonstrate that our rule outperforms prior methods and recovers 99% of self-attention's performance. we also show competitive performance on wikitext-103 @xcite compared to more complex attention alternatives. our results support the idea @xcite that it is unnecessary for attention alternatives to maintain a close analogy to self-attention, and it is more important to focus on designing an expressive update rule. ## decay update rule we propose the decay update rule, which replaces the self-attention mechanism in transformers with decaying fast weights @xcite that evolves with linear dynamics. we modify the additive update rule (eq. 3) by adding a low-rank decay matrix g @xmath0, 1) d×m to forget information. our rule is based on the gated rule @xcite and gated fast weights @xcite . unlike the gated rule, our gate g t is a learned matrix and not a scalar, which enables more control when decaying s t as feature size m increases. compared to @xcite ; @xcite , we do not gate the v t ϕ(k t ) ⊺ term as it did not bring performance gains given our initialization scheme. unlike the delta rule @xcite , our rule is a pure element-wise operation, where state dimensions m do not mix @xcite . this can enable more efficient parallelization on gpu @xcite (fig. 1 ). finally, we choose a linear projection feature map and remove attention normalization @xcite and its associated vector z t from eq. 4. in practice, this feature map can be subsumed into w k , w q @xcite and is equivalent to the identity function (i.e., no feature map). ## gpt-2 fine-tuning experiments we perform fine-tuning experiments to speed up existing pre-trained transformers in a similar setting to transformer-to-rnn (t2r) @xcite . we choose gpt-2 small @xcite as our candidate model to fine-tune, as it has a direct scale-up to large models such as gpt-3 @xcite . to facilitate reproducibility on publicly available datasets, we first fine-tune gpt-2 on the pile @xcite ) -a large general corpus of text. this fine-tuned gpt-2 serves as our baseline target model (14.5 ppl, ## wikitext-103 fine-tuning experiments we perform a similar experiment to fine-tune a pretrained transformer into our decay rule with @xmath2, 64} on the wikitext-103 @xcite language modeling dataset. we fine-tune from the publicly available checkpoint from @xcite (base) by swapping attention with our decay rule and tuning the entire model, with similar hyperparameters as @xcite . we compare against recently proposed self-attention approximations including a causal variant of lin-former @xcite and attention with bounded memory control (abc mlp ) @xcite evaluated under the same setting of 0 and 480 token context sizes on the wikitext-103 test set. our method is simpler than these approaches and performs competitively foot_3 under similar state memory sizes ( ## limitations to compare update rules, our experiments use the simple feature map from @xcite for all rules. however, we acknowledge that some of these rules are jointly proposed with their feature maps, which may be required for good performance. for example, when training peng et al. ( 2021 )'s gated rule (baseline), we attempted various initialization schemes and tried adding small constants to the normalization divisor to help the model converge. however, none of our efforts worked. we note that @xcite 's original proposal of the gated rule does not use the more recently proposed learned relu feature maps @xcite , which may be incompatible with their original proposal. similarly, @xcite 's delta rule was also proposed along with their feature map. our results suggest that their proposed update rules may not be robust to alternative feature maps or the fine-tuning setting. the majority of our experiments are performed on the pile with gpt-2 and assumes a large dataset pre-training setup. while this setup is typical of many nlg applications, it may not generalize to settings without a large dataset, such as lowresource nlp settings. in these cases, more complex approaches may provide better inductive biases than our simple approach. our work focuses on the specific case of finetuning pre-trained auto-regressive transformer language models into fast recurrent variants for inference applications. however, we do not explore training these models from scratch or on different architectures. it may be the case that our approach only works in the fine-tuning setting and that self-attention pre-training is required @xcite . our work also exclusively focuses on language modeling and does not consider evaluation on downstream tasks. while language modeling performance typically correlates to downstream task performance @xcite , future work should further validate our proposal on these practical areas of interest. finally, as mentioned in sec. 3.1, our proposed update rule requires more memory (i.e., o(t dm)) than the add and delta rule (i.e., o(t d)) during training. this may lead to memory issues when training state sizes with larger m, in which the delta rule may be preferred over the decay rule. future work should explore extensions of our approach with better space complexity. ## ethical considerations our work uses datasets crawled from the public web, including wikitext-103 @xcite and the pile @xcite ) and may contain sensitive information or contain undesirable biases. we refer readers to the dataset descriptions from their respective papers for details. our work focuses on improving the inference speed of generating from language models by fine-tuning pre-trained models, which may benefit downstream applications that require the deployment of language models. our work does not explore the issues of bias and disinformation in language models nor specifically aims to mitigate these issues. our models will likely exhibit the same biases and issues that large language models exhibit @xcite . practitioners who deploy models with our proposal should not assume our method mitigates these issues. text sampled from these models may contain offensive content and biases and we advise practical uses of these models to involve some form of human supervision. pile consists of 825 gb of diverse text (englishdominant) including stories, websites, code and mathematical questions. it is intended for largescale language model pre-training. due to the scale of the pile, our fine-tuning only sees a small subset of the training set. we evaluate on the provided validation set from the pile. gpt-2 small @xcite has approximately 124m parameters. all models are finetuned with a batch size of 32 for 100,000 iterations with a learning rate of 6 × 10 -4 and gradient clipping of 1.0 on a single nvidia a100 gpu for 13 hours. training was performed on pytorch 1.11. we use mixed precision training @xcite except for computing normalizer z t (eq. 3), which is prone to numerical overflows. we evaluate on the validation set every 4000 training iterations and report the best validation perplexity achieved for a single run.