Fine-Tuning Pre-trained Transformers into Decaying Fast Weights

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 10236 - 10242 December 7-11, 2022 ©2022 Association for Computational Linguistics Fine-Tuning Pre-trained Transformers into Decaying Fast Weights Huanru Henry Mao Jenni henry@jenni.ai Abstract Autoregressive Transformers are strong lan- guage 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 for- mulations with various update rules and fea- ture 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 com- petitive performance on WikiText-103 against more complex attention substitutes. 1 Introduction Autoregressive Transformers (Vaswani et al., 2017) have demonstrated strong performance on text gen- eration (Brown et al., 2020). The success of self- attention in Transformers over recurrent models (Hochreiter and Schmidhuber, 1997) can be at- tributed to its parallelizability (Hooker, 2021) and its effective gradient propagation over many time steps (Ke et al., 2018). 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 con- vert or fine-tune attention into recurrent formula- tions with O(1) memory and time complexity for auto-regressive generation. Kernel-based meth- ods for self-attention (Tay et al., 2021) learn ap- proximations of the exponential similarity function using m-dimensional feature maps to reformulate attention as a recurrent computation. They replace attention with “unlimited capacity” with fixed- capacity fast weights (Schmidhuber, 1992; Peng et al., 2022), where the memory-accuracy trade- off (Kerg et al., 2020) is controlled by m. Several Figure 1: Plot of memory usage and execution time of our decay rule, delta rule and attention when gen- erating the next token at various sequence lengths on Quadro RTX 4000. Decay and delta rule use approxi- mately the same peak memory (overlapped in plot). works explored different feature maps and recur- rent formulations (i.e., update rules). Katharopou- los et al. (2020) propose feature maps to maintain positive outputs, while Choromanski et al. (2021); Peng et al. (2021) carefully ensure their random feature maps are unbiased estimates of the softmax attention kernel. Schlag et al. (2021a); Peng et al. (2021) propose more sophisticated update rules to forget information in the recurrent state to improve performance. Recently, Kasai et al. (2021) 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 Transform- ers into RNNs for fast inference. We find that prior proposals contain unnecessary operations, lead- ing us to propose a simple element-wise decay up- date rule with no feature map. We …

Content truncated to 500 words