Momentum Transformer: Closing the Performance Gap Between Self-attention and Its Linearization
This work addresses the efficiency-accuracy trade-off in transformers for sequence modeling tasks, offering a method to enhance linear transformers without quadratic costs, though it is incremental as it builds on existing linear attention techniques.
The paper tackles the performance gap between self-attention and its linear approximations in transformers by introducing momentum into linear attention and residual connections, resulting in improved accuracy while maintaining linear complexity, as demonstrated in experiments on image generation and machine translation tasks.
Transformers have achieved remarkable success in sequence modeling and beyond but suffer from quadratic computational and memory complexities with respect to the length of the input sequence. Leveraging techniques include sparse and linear attention and hashing tricks; efficient transformers have been proposed to reduce the quadratic complexity of transformers but significantly degrade the accuracy. In response, we first interpret the linear attention and residual connections in computing the attention map as gradient descent steps. We then introduce momentum into these components and propose the \emph{momentum transformer}, which utilizes momentum to improve the accuracy of linear transformers while maintaining linear memory and computational complexities. Furthermore, we develop an adaptive strategy to compute the momentum value for our model based on the optimal momentum for quadratic optimization. This adaptive momentum eliminates the need to search for the optimal momentum value and further enhances the performance of the momentum transformer. A range of experiments on both autoregressive and non-autoregressive tasks, including image generation and machine translation, demonstrate that the momentum transformer outperforms popular linear transformers in training efficiency and accuracy.