PaperlayerView exact revision on arXiv

Source-linked AI summary

The Devil in Linear Transformer

Zhen Qin, XiaoDong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, Yiran Zhong

arXiv:2210.10340v1cs.CLcs.LG

TL;DR

Linear transformers offer linear complexity but often underperform because of unbounded gradients and attention dilution. TRANSNORMER replaces attention scaling with normalization and adds early-layer diagonal attention, outperforming competing linear transformers across language modeling, classification, and Long-Range Arena evaluations. The paper supports improved efficiency and performance, while leaving applicability to linear vision transformers for future work.

  • Problem

    Kernel-based linear transformers reduce quadratic complexity but often show degraded performance because of unbounded gradients and attention dilution.

  • Method

    TRANSNORMER uses normalized attention to stabilize gradients and diagonal attention in early layers to preserve neighbouring information.

  • Results

    TRANSNORMER outperforms competing linear transformers on autoregressive and bidirectional language modeling, text classification, and Long-Range Arena benchmarks.

  • Takeaways & Limitations

    The model combines linear-transformer efficiency with the local and global attention behavior associated with vanilla transformers.

  • Takeaways & Limitations

    Whether the identified issues and method extend to linear vision transformers remains unresolved.

Abstract

from arXiv · show

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 .

1 Introduction

Kernel-based linear transformers reduce the quadratic cost of vanilla transformers but can lose performance because of unstable gradients and diluted attention. TRANSNORMER addresses these issues with normalized attention and early-layer diagonal attention, achieving favorable results across language, classification, and long-range tasks.

  • Kernel-based linear transformers retain linear space-time complexity but often underperform vanilla transformers on practical tasks.
  • Unbounded gradients from attention scaling destabilize optimization and impair convergence in existing linear transformers.
  • Attention dilution spreads scores across long sequences instead of preserving neighbouring information, contributing to inferior linear-transformer performance.
  • TRANSNORMER achieves better performance than vanilla transformer across broad tasks while remaining significantly faster at runtime.
  • NORMATTENTION removes attention-matrix scaling and applies post-attention normalization, producing more stable gradients and more consistent convergence.
  • TRANSNORMER combines early-layer diagonal attention with later NORMATTENTION to capture local and global language context while preserving linear complexity.
  • TRANSNORMER obtains lower WikiText-103 perplexities, higher GLUE classification accuracy, and favorable Long-Range Arena results with faster, more scalable processing.

2 Background and related work

Vanilla self-attention provides strong modeling but scales quadratically with sequence length. Efficient-transformer research therefore uses sparse attention patterns or kernel decompositions to reduce complexity, commonly trading performance for efficiency.

  • Vanilla self-attention becomes prohibitive for long inputs because its softmax computation has quadratic space-time complexity.
  • Efficient transformers broadly divide into pattern-based methods that sparsify attention and kernel-based methods that decompose softmax attention.
  • Kernel-based methods achieve theoretical linear complexity by mapping queries and keys through kernel functions and reordering key-value computation.
  • Existing kernel methods vary in kernel design, including ELU-based activations, low-rank assumptions, softmax approximations, and empirically motivated alternatives.
  • These efficient methods often trade performance for efficiency relative to vanilla transformers, motivating analysis of the sources of the gap.

3 The devil in linear attention

The paper identifies unbounded gradients and attention dilution as two causes of degraded linear-transformer performance. The analysis links unstable gradients to attention scaling and shows that linear attention distributes less mass locally than vanilla attention.

  • 3.1 Unbounded gradients: The gradient analysis represents attention through token similarities and shows that vanilla and linear attention differ in how those similarities are computed.
  • 3.1 Unbounded gradients: Kernel-based linear attention can have gradients with neither upper nor lower bounds, which leads to less stable optimization and worse convergence.
  • 3.2 Attention dilution: Vanilla attention is known to emphasize neighbouring tokens, whereas this locality property does not directly carry over to linear-transformer variants.
  • 3.2 Attention dilution: Locally accumulated attention measures the score assigned within a neighbourhood of ratio r around a token; higher values indicate stronger local concentration.
  • 3.2 Attention dilution: Vanilla attention has more concentrated local scores than linear attention, whose scores are distributed more evenly across local and distant tokens.

4 Method

TRANSNORMER combines diagonal attention in early layers with normalized linear attention later, targeting attention dilution, unstable gradients, and preserving linear complexity. Its ablations support normalization for stable training and local-to-global attention placement.

  • 4.1 The overall architecture: TRANSNORMER combines DIAGATTENTION in early layers with NORMATTENTION in later layers to capture local and global context while preserving linear space-time complexity.The architecture uses block-local attention early and normalized linear attention later.
  • 4.2 NORMATTENTION: Replacing scaling with normalization addresses unbounded gradients while avoiding the critical performance drop caused by simply removing scaling.The proposed attention removes scaling and applies normalization after the attention layer; LayerNorm and RMSNorm are both suitable choices.
  • 4.2 NORMATTENTION: NORMATTENTION produces more stable training gradients, with deviations closer to vanilla attention than existing linear methods over 50k iterations.The comparison uses RoBERTa on WikiText103 and evaluates relative gradient standard deviation.
  • 4.1 The overall architecture: The architecture overview places block-local DIAGATTENTION early and NORMATTENTION late, assigning neighbouring focus to early computation and stable linear attention to later stages.This arrangement is designed to combine local structure with later global processing.
  • 4.3 DIAGATTENTION: Replacing selected linear-attention layers with vanilla attention improves language-modeling performance, especially when local information is captured in early layers.The ablation supports using local attention earlier rather than later.
  • 4.3 DIAGATTENTION: DIAGATTENTION restricts attention to non-overlapping blocks and has complexity O(nwd), which scales linearly with sequence length when d ≪ n.Here n is sequence length, w is block size, and d is feature dimension.
  • 4.3 DIAGATTENTION: Applying DIAGATTENTION in later stages hurts performance, indicating that later layers require a global field of view and motivating NORMATTENTION there.The placement choice is supported by the later-stage ablation.

5 Experiments

Experiments evaluate TRANSNORMER on language modeling, GLUE, Long-Range Arena, speed, and ablations, showing strong accuracy and efficiency across settings.

  • 5.1 Settings: Experiments use two variants, T1 and T2, differing in DIAGATTENTION and NORMATTENTION choices, and evaluate autoregressive, bidirectional, and long-range tasks under matched configurations.The study uses WikiText-103, GLUE, and Long-Range Arena, with shared competitor settings where specified.
  • 5.2 Results: TRANSNORMER variants achieve comparable or better perplexity than vanilla attention and outperform existing linear models; TRANSNORMER T2 lowers validation and test perplexity by 2.31 and 1.58, respectively.
  • 5.2 Results: TRANSNORMER achieves superior average GLUE performance, with comprehensive gains on SST-2, MRPC, and CoLA, including 4.62 higher on CoLA.TRANSNORMER T1 also outperforms vanilla attention by a notable margin.
  • 5.2 Results: Both TRANSNORMER variants rank first or second across all five Long-Range Arena tasks and outperform competing methods, including vanilla transformer, overall.These results evaluate the model’s ability to encode long sequences.
  • 5.3 Speed comparison: TRANSNORMER is over 300% faster during training and 150% faster during inference than FLASH-quad, while maintaining comparable GPU memory footprints.The vanilla transformer exhausts GPU memory beyond sequence lengths of 3k, whereas TRANSNORMER remains faster than other efficient transformers.
  • 5.4 Ablation study: Using DIAGATTENTION in early layers and NORMATTENTION later performs better than using either attention throughout or reversing their order.The ablation supports local attention early and long-range attention later; the balanced attention structure is also reported as best.
  • 5.4 Ablation study: GLU performs better than the traditional FFN, while applying the two attentions in parallel doubles computation without improving performance.The ablations compare FFN choices and sequential or parallel attention combinations.
  • 5.4 Ablation study: Larger DIAGATTENTION blocks improve performance, but increase O(nwd) computation, motivating a block size of 64 as a performance–cost trade-off.The ablation reports clear gains with increased block size.

6 Conclusion

The paper identifies unbounded gradients and attention dilution as key issues in linear transformers, then addresses them with NORMATTENTION and DIAGATTENTION. TRANSNORMER outperforms competing linear transformers across language modeling, text classification, and Long-Range Arena tasks.

  • 6 Conclusion: The paper identifies unbounded gradients and attention dilution as two issues behind inferior linear-transformer performance.NORMATTENTION targets gradient stability, while DIAGATTENTION concentrates attention on neighboring tokens.
  • 6 Conclusion: TRANSNORMER combines the proposed attentions and outperforms competing linear transformers on autoregressive and bidirectional language modeling, text classification, and Long-Range Arena.The conclusion frames the model as combining strengths of vanilla and linear transformers.

Limitations

The analysis is limited to natural language processing tasks, and whether linear vision transformers share the identified issues remains untested.

  • Limitations: The paper studies linear transformers in natural language processing, leaving whether linear vision transformers share the same issues for future validation.The authors state that they will validate the method on linear vision transformers in future work.

Ethics Statement

The paper motivates resource-efficient language models while noting that public-web training data may contain problematic content requiring assessment and regulation before deployment.

  • Ethics Statement: The proposed technique is intended to reduce computing-resource demand for large-scale, environmentally friendlier language models.
  • Ethics Statement: Public-web training corpora may contain biased, explicit, or improper content, so deployment requires further assessment and regulation.
  • A Mathematical Notations: The notation section defines matrices with bold uppercase letters, vectors with bold lowercase letters, and scalars with lowercase letters.
  • A Mathematical Notations: Vectors are represented as columns, matrix rows are denoted explicitly, and ∥.∥2 and ∥.∥F denote l2 and Frobenius norms.
  • A Mathematical Notations: The main inputs are X ∈ Rn×d and the query, key, and value matrices, with projection parameters WQ, WK, WV ∈ Rd×d.

B Proof of gradients’ upper bound

This proof establishes the bounds needed for the paper’s gradient analysis using the geometric–arithmetic mean inequality and selected substitutions.

  • The proof reduces the target bounds to inequalities that can be established using the geometric mean–arithmetic mean relationship.
  • Substituting a = pik and b = 1 − pik completes the inequality argument for the required bound.
  • The first bound is obtained directly after invoking the preceding inequality, while the second begins from an additional stated fact.

C Proof of Proposition 3.1

The analysis contrasts gradient behavior in vanilla, kernel-based linear, and normalized attention, concluding that linear attention can have unbounded gradients while NORMATTENTION remains bounded.

  • D Analyze the gradient of each method: The analysis treats one-layer Transformers first and states that the multi-layer conclusion can be obtained by induction.
  • D.1 Notations: The proof framework assumes that the image of φ contains vectors arbitrarily close to zero, described as common for kernel functions.
  • D.2.3 NORMATTENTION: NORMATTENTION incorporates RMSNorm into the attention formulation, with RMSNorm defined as a normalization of x by its root-mean-square scale plus ϵ.
  • D.2.1 Preliminary: The preliminary propositions bound relevant quantities before the analysis derives method-specific attention and gradient expressions.
  • D.2.2 Vanilla/Linear attention: Vanilla and linear attention are unified through a common formulation whose scaling parameter differs between the two methods.
  • D.2.3 NORMATTENTION: ∂L/∂sij is bounded in vanilla attention and NORMATTENTION but unbounded in linear attention, making linear-transformer training unstable.

D.3 Proof of the proposition

The appendix proves bounded gradients for the relevant attention module, then records experiment configurations and pseudocode used to analyze local attention concentration.

  • D.3 Proof of the proposition: The proof analyzes both forward and backward passes and propagates gradient bounds through the model’s activation and parameter mappings.
  • D.3 Proof of the proposition: XAttention denotes vanilla or normalized attention in the proof notation.
  • D.3 Proof of the proposition: The proof concludes that the gradient passed to the XAttention module is bounded, with c1 = a3 < ∞.
  • E Experiment configs: The LRA setup uses a two-layer Transformer with 64 hidden dimensions, two attention heads, 85 GLU dimensions, and Swish activation, with task-specific batch sizes and learning rates.
  • F Pseudocode for visualization.: The visualization pseudocode expands a token’s neighborhood around its center and records the smallest local span whose accumulated attention exceeds threshold levels.
  • F Pseudocode for visualization.: Table 13 documents experiment configurations, including total batch size and the fact that attention dropout is used only for vanilla attention.
Loading 2210.10340v1…