Fine-Tuning Pre-trained Transformers into Decaying Fast Weights

TL;DR

This paper introduces a simple decaying fast weights method for autoregressive Transformers that achieves near-attention performance with significantly reduced computational complexity, outperforming prior kernel-based approaches.

Contribution

The paper proposes a straightforward decaying fast weights approach as an alternative to complex kernel-based methods for efficient Transformer inference.

Findings

Decaying fast weights run efficiently on GPU.

They retain 99% of GPT-2 attention performance.

Achieve competitive results on WikiText-103.

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.