Source-linked AI summary
cosFormer: Rethinking Softmax in Attention
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, Yiran Zhong
TL;DR
Softmax attention captures long-range dependencies but has quadratic time and space complexity, limiting efficient scaling to long sequences. COSFORMER replaces it with a linear operator plus cosine-based re-weighting that preserves non-negativity and concentrated attention weighting. Across language modeling, text understanding, and Long-Range Arena evaluations, it achieves competitive performance with better efficiency and ranks first on Long-Range Arena.
Problem
Softmax attention captures long-range dependencies but its quadratic time and space complexity limits scaling to long inputs, while efficient approximations can lose performance.
Method
COSFORMER uses a linear projection with ReLU-enforced non-negativity and a cosine-based re-weighting mechanism that preserves key softmax attention properties.
Results
COSFORMER achieves better inference speed and smaller memory footprint with performance on par with the vanilla transformer, and ranks first on Long-Range Arena.
Takeaways & Limitations
The results support COSFORMER as an efficient linear-time and linear-space transformer for modeling long-range dependencies in causal and cross attention.
Abstract
from arXiv · showhide
Transformer has shown great successes in natural language processing, computer vision, and audio processing. As one of its core components, the softmax attention helps to capture long-range dependencies yet prohibits its scale-up due to the quadratic space and time complexity to the sequence length. Kernel methods are often adopted to reduce the complexity by approximating the softmax operator. Nevertheless, due to the approximation errors, their performances vary in different tasks/corpus and suffer crucial performance drops when compared with the vanilla softmax attention. In this paper, we propose a linear transformer called cosFormer that can achieve comparable or better accuracy to the vanilla transformer in both casual and cross attentions. cosFormer is based on two key properties of softmax attention: i). non-negativeness of the attention matrix; ii). a non-linear re-weighting scheme that can concentrate the distribution of the attention matrix. As its linear substitute, cosFormer fulfills these properties with a linear operator and a cosine-based distance re-weighting mechanism. Extensive experiments on language modeling and text understanding tasks demonstrate the effectiveness of our method. We further examine our method on long sequences and achieve state-of-the-art performance on the Long-Range Arena benchmark. The source code is available at https://github.com/OpenNLPLab/cosFormer.
1 INTRODUCTION
Transformers capture long-range dependencies effectively, but softmax attention makes long-sequence scaling costly because its time and space complexity is quadratic. COSFORMER replaces softmax with a linear design preserving non-negativity and nonlinear re-weighting, achieving improved efficiency and competitive performance.
- Softmax attention captures long-range dependencies but incurs quadratic time and space complexity as sequence length grows.This computational overhead becomes especially prohibitive for long inputs.
- Existing efficient attention methods reduce complexity through sparse matrices, low-rank representations, or kernel approximations.These approaches report comparable performance on selected tasks or corpora, but often rely on additional assumptions or constrained approximations.
- Approximation errors and violated assumptions can produce performance deficiencies, particularly in linear-complexity transformer variants.Performer, RFA, and Reformer are cited as examples with less satisfactory performance.
- COSFORMER preserves softmax attention’s non-negativity and nonlinear re-weighting properties using a linear operator and cosine-based re-weighting.The method applies ReLU to features and uses cosine re-weighting to stabilize attention weights and amplify local correlations.
- COSFORMER delivers faster inference and lower memory footprint while achieving performance on par with the vanilla transformer across autoregressive and bidirectional settings.The paper evaluates language modeling and text understanding benchmarks, including Long-Range Arena.
2 OUR METHOD
COSFORMER replaces softmax attention with a linear attention design that preserves non-negativity and nonlinear re-weighting. Its decomposable computation reduces sequence-length complexity while cosine re-weighting encourages locality.
- Method: COSFORMER replaces non-decomposable softmax with a decomposable linear operation and cosine-based nonlinear re-weighting.The method is designed to retain two empirically important softmax properties: non-negative attention values and concentrated attention weights.
- Scope: The method is presented as applicable to causal and cross attention while retaining linear time and space complexity for modeling long-range dependencies.The supplied method description states this scope but provides no quantitative result in this section.
- Motivation: Softmax attention captures long-range dependencies but has quadratic space and time complexity in sequence length, creating a bottleneck for long inputs.The paper motivates linearization because computing the full attention matrix scales as O(N^2).
- Linearization: Kernelized attention rewrites similarity as φ(Q_i)φ(K_j)^T, enabling matrix reordering so φ(K)^T V is computed before multiplication by φ(Q).This avoids explicitly constructing the N×N attention matrix.
- Linearization: When d ≪ N, the reordered computation costs O(Nd^2), which the paper treats as O(N) time and space complexity with respect to sequence length.Figure 2 illustrates the corresponding linearized computation for input length N and feature dimension d.
- COSFORMER Components: ReLU maps queries and keys to non-negative representations, while cosine re-weighting concentrates weights on neighboring tokens and reinforces locality.The cosine mechanism is decomposable and is intended to stabilize attention weights without softmax normalization.
3 EXPERIMENTS
Experiments evaluate COSFORMER in causal and bidirectional language modeling, downstream text understanding, long-sequence benchmarks, efficiency, and re-weighting ablations. Across these settings, it generally matches or exceeds competing methods while retaining linear computation.
- Experimental settings: COSFORMER is evaluated in autoregressive and bidirectional WikiText-103 language modeling, downstream text classification, Long-Range Arena, and efficiency experiments.The study also tests the effect of cos-based re-weighting through composite downstream and benchmark metrics.
- Language modeling: COSFORMER outperforms the quadratic-computation baseline in autoregressive language modeling while using linear computation.It achieves comparable validation perplexity to other methods and significantly outperforms competing methods on the test set.
- Language modeling: COSFORMER converges faster than vanilla transformer in bidirectional language modeling while reaching comparable or lower loss values.The comparison uses a sequence length of 512 and keeps the RoBERTa structure otherwise unchanged.
- Long-range evaluation: COSFORMER achieves competitive Long-Range Arena results across five tasks, with best performance on ListOps and Document Retrieval.The benchmark includes ListOps, byte-level classification, document retrieval, pixel-sequence image classification, and Pathfinder.
- Ablation: The cos-based re-weighting improves convergence speed and final autoregressive results, while raising average scores on bidirectional fine-tuning and Long-Range Arena.The ablation compares COSFORMER with and without the re-weighting mechanism using composite metrics over five datasets and five tasks.
4 RELATED WORK
Prior efficient-transformer methods mainly reduce attention cost through patterned sparsity or kernel-based linearization. These approaches trade computation against assumptions about attention structure or approximation quality.
- Pattern based methods: Pattern-based methods sparsify attention using handcrafted or learnable patterns.Examples include inducing points, strided patterns, local patterns, and fixed diagonal sliding windows with global attention.
- Kernel based methods: Kernel-based methods reduce self-attention complexity from quadratic to linear by using decomposable approximations.Representative approaches assume similar attention distributions for neighboring queries or approximate softmax with Gaussian kernels.
5 CONCLUSION
COSFORMER is a linear-time, linear-space transformer built around non-negative attention and locality-focused re-weighting. Across language modeling, text understanding, and long-sequence evaluation, it matches or surpasses vanilla transformer performance.
- Conclusion: COSFORMER has linear time and space complexity and applies to both causal and cross attentions.Its attention output is naturally decomposable, avoiding accumulated approximation errors associated with previous linear transformers.
- Conclusion: COSFORMER uses ReLU to preserve non-negative attention contributions and cos-based re-weighting to impose locality bias.The design aims to exclude negatively correlated contextual information and concentrate attention on nearby sequence elements.
- Conclusion: COSFORMER achieves comparable or better performance than vanilla transformer across causal pre-training, bidirectional pre-training, and downstream text-understanding tasks.It also achieves state-of-the-art performance on the Long-Range Arena benchmark.
A.1 MATHEMATICAL DERIVATION OF cos-BASED RE-WEIGHTING
The derivation explains how COSFORMER attention produces each output position in linear form. It presents the resulting attention computation as linear rather than explicitly quadratic.
- Mathematical derivation: The derivation obtains the output at each sequence position by following Equation 11.The accompanying statement identifies the proposed COSFORMER attention output as obtainable in a linear manner.
A.2 PSEUDO CODE OF COSFORMER
The COSFORMER attention algorithm maps query, key, and value matrices to an output through cosine-based accumulations and linear-time processing.
- Inputs and outputs: COSFORMER accepts Q, K, and V matrices and outputs O with dimensions determined by the query count and value width.The inputs are Q ∈ R^{N×d1}, K ∈ R^{M×d1}, V ∈ R^{M×d2}, and O ∈ R^{N×d2}.
- Initialization: The procedure initializes attention weights, output entries, cosine and sine accumulators, and corresponding normalization terms before processing keys.These accumulators support the cosine-based attention computation described by the algorithm.
- Computation: The pseudocode iterates over the key dimension and then the query dimension to construct the attention output.The supplied algorithm excerpt explicitly shows loops over M and N.
A.3 ALGORITHM TO VISUALIZE ATTENTION MATRIX
The visualization algorithm converts multiple attention matrices into a thresholded aggregate heatmap by retaining the highest-scoring entries in each row.
- Inputs and output: The algorithm takes n matrices M_k and a threshold, then returns an aggregated matrix M for heatmap visualization.Each M_k has shape d × d, and the output M has the same shape.
- Row-wise selection: For every row, entries are sorted in descending order and accumulated until the cumulative score exceeds the threshold.The algorithm uses argsort and a running sum p to select the retained indices.
- Aggregation: The selected entries are counted across matrices, normalized by n, and displayed as a heatmap.The final matrix is divided by n before visualization.
A.4 INTRODUCTION OF DATASET
The experiments use WikiText-103 for language modeling and several text-classification datasets, with Long-Range Arena covering diverse long-sequence abilities.
- Language modeling: WikiText-103 is used for autoregressive and bidirectional language modeling before downstream fine-tuning.The dataset is split by tokens, and the resulting models are fine-tuned on text-classification tasks.
- Text classification: The downstream classification datasets include QQP, SST-2, IMDB, AMAZON, and MNLI, covering semantic equivalence, sentiment, scoring, and natural-language inference.QQP tests question-pair equivalence; SST-2 and IMDB test sentiment; AMAZON tests product-review scoring; MNLI uses entailment, contradiction, and neutral labels.
- Long-sequence evaluation: Long-Range Arena contains five datasets testing hierarchical parsing, text classification, spatial reasoning, and compressed-representation retrieval.Its tasks include ListOps, IMDB, CIFAR-10, Pathfinder, and AAN.
- Dataset statistics: The dataset statistics are reported in Table 7, including a small electronics-category subset of AMAZON.The table caption specifies that a subset of the small AMAZON dataset is used.
A.5 QUALITATIVE RESULTS OF LRA
Qualitative Long-Range Arena examples compare COSFORMER with the vanilla transformer on ListOps and document retrieval, showing differing predictions against ground truth.
- ListOps: ListOps requires ten-way prediction over hierarchically structured sequences with operators such as MAX, MEAN, MEDIAN, and SUM MOD.The task requires access to all tokens and modeling of the inputs’ logical structure.
- Document Retrieval: Document Retrieval assigns a binary similarity label to pairs of long documents and tests compressed representations for matching and retrieval.The examples are shortened for display because the original LRA samples are too long.
- Qualitative comparisons: In the displayed ListOps examples, COSFORMER predictions match the ground truth while the transformer predictions differ.The listed examples report COSFORMER outputs of 0 and 9 matching ground truth, versus transformer outputs of 9 and 3.
- Qualitative comparisons: In the displayed document-retrieval example, COSFORMER outputs False matching ground truth, while the transformer outputs True.The example is labeled as a byte-level document-retrieval case.