The Devil in Linear Transformer

TL;DR

This paper introduces transNormer, a novel linear transformer that stabilizes gradients through normalization and enhances local attention, leading to improved performance and efficiency on various NLP tasks and benchmarks.

Contribution

It proposes a new linear attention mechanism that replaces scaling with normalization and employs diagonal attention for local context, addressing key issues in existing linear transformers.

Findings

Outperforms vanilla and existing linear transformers on NLP benchmarks.

Demonstrates superior performance on Long-Range Arena benchmark.

Achieves better space-time efficiency while maintaining high accuracy.

Abstract

Linear transformers aim to reduce the quadratic space-time complexity of vanilla transformers. However, they usually suffer from degraded performances on various tasks and corpus. In this paper, we examine existing kernel-based linear transformers and identify two key issues that lead to such performance gaps: 1) unbounded gradients in the attention computation adversely impact the convergence of linear transformer models; 2) attention dilution which trivially distributes attention scores over long sequences while neglecting neighbouring structures. To address these issues, we first identify that the scaling of attention matrices is the devil in unbounded gradients, which turns out unnecessary in linear attention as we show theoretically and empirically. To this end, we propose a new linear attention that replaces the scaling operation with a normalization to stabilize gradients. For the issue of attention dilution, we leverage a diagonal attention to confine attention to only neighbouring tokens in early layers. Benefiting from the stable gradients and improved attention, our new linear transformer model, transNormer, demonstrates superior performance on text classification and language modeling tasks, as well as on the challenging Long-Range Arena benchmark, surpassing vanilla transformer and existing linear variants by a clear margin while being significantly more space-time efficient. The code is available at https://github.com/OpenNLPLab/Transnormer .