DenseAttention: No-Compromise Exact All $N \times N$  Interactions Algorithm with $O(N)$ Space and Time Complexity

Authors: Andrew Argatkiny
Primary area: foundation or frontier models, including LLMs
Keywords: self-attention, deep learning, transformer architecture, nlp, efficient transformers, DenseAttention, long context, Long Range Arena

Abstract:
The ubiquitous Transformer architecture suffers from two main bottlenecks: 1) low computational and memory efficiency, leading to suboptimal hardware utilization, and 2) quadratic time complexity with respect to sequence length $N$, making it slow and costly for large data contexts. We propose a novel DenseAttention Network architecture, a straightforward simplification of the standard Transformer block that addresses these issues and serves as a drop-in replacement for language modeling tasks. We eliminate memory-bound components in DenseAttention, including Softmax, masking, one skip connection, and both LayerNorms, as well as key, value, and output projection matrices, as they become redundant. Despite these removals, it maintains exact $N \times N$ pairwise interactions between tokens. By exploiting the associativity of matrix multiplications, DenseAttention can be computed with $O(N^2d)$ or $O(Nd^2)$ time and space complexity, depending on the context. To handle the absence of Softmax and prevent numerical instability, we introduce MaxNormActivation at both ends of the Transformer block. We also devise Cosine Relative Positional Embeddings as a computationally efficient replacement for RoPE, and simple LocalAttention variations of the block to help the model focus on details in extremely long contexts. 

DenseAttention competes with FlashAttention in speed on small sequences and outperforms it by orders of magnitude on large contexts. We pre-train encoder language models on sequences up to 16K in length, which perform similarly or better than baseline BERT-large, while significantly improving speed and efficiency.  Finally, we achieve state-of-the-art on the LRA benchmark among the Transformer-based architectures.

Decision: reject
Reviewer scores: 5, 3, 5, 5

Reviews and discussion:
<thread>
creation date:
1729383657457

thread_type:
official_review

<thread_object>
creator: reviewer_rvwa

<content>
summary:
The authors propose a new architecture which they call DenseAttention Network, which is a variation of the standard transformer architecture which is specifically tuned to perform well on long sequences. The changes include:
1) Take the softmax away in the attention block, writing $QK^{\top}V$ instead of $\mathrm{softmax}(QK^{\top})V$, and using associativity to compute the matrix product in linear time. (They name this as DenseAttention mechanism/block.)
2) Use a MaxNormActivation block instead of LayerNorm, which scales each token feature by its maximum absolute value.
3) Use a novel positional embedding called Cosine RelPE, which is claimed to perform similarly but more efficiently computable than RoPE.
4) For very long contexts, use a hand-rolled local attention implementation suited for DenseAttention.
They show that this architecture has general improvements over a basic transformer + RoPE implementation in Long Range Arena, both in performance (across a few tasks) and efficiency (more broadly, as the usual attention mechanism does not have linear-in-$N$ time complexity). It also shows improvements against S4-V1 in Long Range Arena, and BERT in terms of Masked Language Modeling.

strengths:
- The MaxNormActivation block is potentially useful as a method to stabilize LLM training.
- The empirical results in tables 1-3 show that DenseAttention has some promise empirically on long-context sequence modeling tasks as compared to the standard transformer, S4-V1, and BERT.
- The efficiency results in table 4 show that it's possible to get out-of-the-box performance increases at long context compared to the usual BERT model. Specifically, all changes seem to be architectural, and no specialized kernels are needed to get better performance, as a result of using `torch.compile` and potentially fusing linear operations together.

weaknesses:
- The main mechanism behind DenseAttention, i.e., removing the softmax and using associativity to compute the product in linear-in-$N$ time, has been studied before; see for example [1]. It is acknowledged that the paper cites [1], but the paper suggests that the mechanism in [1] has poor efficiency; however DenseAttention is strictly a special case of the mechanism in [1] (using the notation of [1], set $\phi$ as the identity mapping). So this claim does not make sense, and the novelty of DenseAttention seems limited.

- The reasoning behind using MaxNormActivation seems lacking. In particular, since all norms are equivalent (i.e., bounded by each other up to multiplicative constants, possibly dependent on dimension) in finite-dimensional vector space, the boundedness of the maximum norm is equivalent to the boundedness of the $\ell^{2}$ norm. So if the argument in the paper goes through, it should mean that $\ell^{2}$ normalization should also work (then why not LayerNorm? But LayerNorm doesn't work, as reported in the paper, so something else is going on.). Although the MaxNormActivation is an interesting and potentially useful contribution, it may not work for the reason explained in the paper. Also there's a potential typo in the equation defining MaxNormActivation: it should be $\frac{X_{i}}{\max_{j}|X_{ij}| + \epsilon}$ on the RHS (note the absolute value).

- Not much motivation is given for the two other modifications, e.g., CosineRelPE and the local attention proposal - they seem to have a flavor of "we tried it and it works," potentially with some ablation, and without context of why such an approach may or may not make sense or generalize to other architectures.

- The results in Long Range Arena are promising insofar as they match up against a standard transformer and an SSM, but this may not be a fair comparison. Given that the authors start with a regular transformer and apply modifications to show improvement on long-context, they could also compare against more recent models specialized for long context. For example, the authors omit comparison with S5, whose numbers are publicly available on [PapersWithCode](https://paperswithcode.com/dataset/lra), as well as S4 V2 and a long list of other models benchmarked on Long Range Arena but not necessarily added there.

- The result on efficiency compared to BERT also may seem to not be a fair comparison. BERT is trained with an encoder-only architecture, while DenseAttention Network is trained with a decoder-only architecture. A fairer comparison would pit DenseAttention Network against a regular decoder-only transformer (as well as BERT if desired, along with, say, an SSM), under the same experimental setting, and allow readers to observe trends in the different approaches as different scaling parameters vary.


[1] Katharopoulos, Angelos, et al. "Transformers are RNNs: Fast autoregressive transformers with linear attention." International conference on machine learning. PMLR, 2020.

questions:
- What is the specific motivation of designing CosineRelPE? 

- Is there anything that suggests that any new block (DenseAttention, CosineRelPE, MaxNormActivation) can generalize to other architectures and improve either performance or efficiency (while not degrading the other)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730190707932

thread_type:
official_review

<thread_object>
creator: reviewer_ogkz

<content>
summary:
This paper proposes a novel neural network architecture called DenseAttention Network as alternative to Transformer networks with Self-Attention. 
The core innovation of DANet is the novel DenseAttention mechanism, which removes Softmax and projection layers from original Self-Attention.
Additionally the authors modify the surrounding network block: They replace Layernorm or RMS with their novel “MaxNormActivation”, they remove some skip connections and modify the Rotary Positional Embeddings. 
The paper performs experiments on the Long Range Arena Benchmark and masked language modeling with BERT-large sized models. 

The paper claims to outperform a BERT baseline on masked language modeling. 
It claims to set a new SOTA on “Transformer-based” models on LRA and to outperform 4 of 6 State Space model baselines.

strengths:
- The usage of MaxNormActivation seems to be well motivated by a theoretical variance analysis. However, this could also be supported with experiments empirically too.
- Code provided.

weaknesses:
In general I believe this paper is not ready for publication as there are several weaknesses in terms of the new architecture, the experiments and the presentation in the paper. My main concerns are summarized below: 

- A related work section is missing in which the authors put DANet in relation to other Linear Attention variants (e.g. GLA https://arxiv.org/abs/2312.06635 ), State space models (e.g. Mamba https://arxiv.org/abs/2405.21060) or other RNNs variants (e.g. xLSTM (https://arxiv.org/abs/2405.04517 ) or RWKV (https://arxiv.org/abs/2305.13048 )). Also a relation to embedding models other than BERT is missing, e.g. Monarch Mixer (https://arxiv.org/abs/2310.12109). 
- Since DANet seems to be a hybrid architecture (Section 3.3), also a relation to hybrid architectures (e.g. https://arxiv.org/abs/2402.19427, https://arxiv.org/abs/2406.07522) is interesting.
- There are so many architecture changes (e.g. Layernorm, Positional Encoding, Block structure, Attention mechanism, Block order / hybrid variants) that leave the reader unclear of what brings performance gains. A careful ablation study could help here.
- While the paper demonstrates large throughput benefits in the long context regime compared to Transformers, it has not been shown in the paper that DANet performs well in the long context regime.
- Regarding Cosine RelPE: It is not clear why the authors made the modification to the original Rotary Positional Embedding. It seems to be motivated by efficiency gains, but this claim is not supported sufficiently. An experiment on this could help.
- A conclusion is missing.

questions:
- Why do you still consider DANet as Transformer-based? The only part of transformers that is left, is the Feeforward layers which is now inside the block.
- You train your model with 4 stages, but the original BERT was trained on 2 stages. Could you also train the baseline in the same way?
- On page 10, line 494 (key highlights) you hint at the fact that DANet outperforms the baseline due to a soft-capping of output logits that you use. Why did you not try this for the baseline as well? 
- L.400: The authors find that local attention is effective. Do you use the Transformer Self-Attention here? An ablation on this would be interesting.
- Why do you use float16 in the experiments?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700192667

thread_type:
official_review

<thread_object>
creator: reviewer_hatu

<content>
summary:
The paper proposes a DenseAttention Network (DANet), which addresses inefficiencies in the Transformer architecture, especially its high memory and computational cost - O(N^2), with respect to sequence length - N. DANet uses a new MaxNormActivation and Cosine Relative Positional Embeddings, capturing N x N interactions at O(N) space and time complexity. Experimental results demonstrate that DANet outperforms FlashAttention on long sequences on the Long Range Arena benchmark.

strengths:
1) The paper proposes an interesting approach to eliminate projection matrices in attention, considering that the multiplication $W_QW_K^{\top}$ can be replaced with a single parameter, which I don't think exists in previous literature.
2) The paper also proposes using local attention in conjunction with the proposed attention function.

weaknesses:
1) The paper lacks comparison against mamba in experiments. Mamba-I and Mamba-II are fast approaches for long range sequence modeling. 
2) This is not the first paper which captures NXN correlations with O(N) complexity. Linear attention [1] uses linear approximations of attention. A fair comparison with this paper would be great.
3) The mathematical writing in this paper is inconsistent. Here are some instances:


Better Notation:
1. Standard operators max, var should be mentioned in times new roman using \DeclareMathOperator
2. Defined operators such as MaxNormActivation can be put in \text{MaxNormActivation}, as done in 200-204. 
3. Line 240: has a typo open bracket.
4. Line 284: << should be \ll.
5. Line 246: why is fp16 and bf16 bolded?

Major readability issues:
1. Inconsistent definition of $X_i$ in line 247, 300 and 311.

If the above issues are resolved I am willing to increase my score.

[1] Katharopoulos, Angelos, et al. "Transformers are rnns: Fast autoregressive transformers with linear attention." International conference on machine learning. PMLR, 2020.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730748251435

thread_type:
official_review

<thread_object>
creator: reviewer_rubj

<content>
summary:
In this paper, the authors propose a new architecture, DenseAttention Network, which could potentially replace Transformer. The motivation of this new design is to alleviate the quadratic time complexity in sequence length as well as the memory-bound operations in the vanilla Transformer (e.g. softmax and layer normalization). Specifically, they propose to linearize the original multi-head attention layer with naive matrix multiplications. To stabilize the forward pass, the inputs of each layer are scaled to have the same $\ell_\infty$ norm. The authors further propose a replacement for the rotary embedding and sliding window that is compatible with their approach.

**Strengths**
1. Given the popularity of Transformer models, the topic of their efficiency becomes more and more important. The proposed solution is also well-motivated.
2. The paper is well-written and easy to follow.

**Weaknesses**
1. The major flaw of this paper is the thin experiments.
2. The paper lacks several important previous papers.

In summary, this paper proposes a potential solution to accelerate Transformer models. However, the experiments are not convincing enough. Therefore, I would recommend a clear rejection unless there is further evidence.

strengths:
1. Given the popularity of Transformer models, the topic of their efficiency becomes more and more important. The proposed solution is also well-motivated.
2. The paper is well-written and easy to follow.

weaknesses:
1. The major flaw of this paper is the thin experiments. The Transformer model is known to perform well on a wide range of tasks. In addition, it also demonstrates a promising scaling effect. On the other hand, this paper only contains limited experiments: (1) the testbeds are limited. Currently, the only benchmark is Long Range Arena (for causal LMs); (2) the baselines are limited. There is only one Transformer model that serves as the baseline without specifying how the model is trained; (3) the scaling effect is not studied. The authors do not analyze how the parameter number affects the results. It is unclear if the method could be scaled to larger-scale applications.
2. The paper lacks several important previous papers. In fact, linearizing attention has been heavily studied before [1, 2, 3]. This paper has no comparisons or discussions. 

[1] Random Feature Attention, ICLR 2021 \
[2] Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention, ICML 2020 \
[3] Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel, EMNLP 2019

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732693619464

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response. Part 1. Discussion of Linear Transformers in relation to DenseAttention

comment:
Dear Reviewers,

We sincerely thank you for the precious time, insightful feedback and constructive suggestions which helped to clarify and enhance the exposition of our work.

Here we would like to address several key commonly shared comments.

---

**Discussion of Linear Transformers in relation to DenseAttention**

The majority of reviewers suggested to put DenseAttention into perspective with Linear Transformer class of algorithms. Here’s the extended discussion comparing Linear Transformer and its derivatives with DenseAttention Network from the architectural standpoint. We present it below in full and will include it in the revised version of the paper.

Given entries $Q_i, K_j, V_j \in \mathbb{R}^{1 \times d}$ of  matrices $\mathbf{Q}, \mathbf{K}$ and $\mathbf{V}$, standard softmax attention for input i can be reformulated as


$ 
\begin{equation}
 A_i = \frac{\sum_{j=1}^N \text{Sim}(Q_i, K_j)V_j}{\sum_{j=1}^N \text{Sim}(Q_i, K_j)} \in \mathbb{R}^{1 \times d}
\end{equation},
$

where $\text{Sim}(Q_i, K_j)=\text{exp}(Q_i K_j^\top)$. Conceptually, linear attention class of algorithms, described in [1] and built upon in numerous subsequent works, approximates or replaces this similarity function with separable kernel $\text{Sim}(Q_i, K_j)=\mathcal{K}(Q_i, K_j)=\phi(Q_i) \phi(K_j^\top)$, where $\phi: \mathbb{R}^d \to \mathbb{R}_{+}^{r}$ maps query and key vectors to non-negative vectors with possibly different dimension r.

Hence, the attention mechanism becomes:

 $A_i = \frac{\sum_{j=1}^N \phi(Q_i) \phi(K_j^\top)V_j}{\sum_{j=1}^N \phi(Q_i) \phi(K_j^\top)} = \frac{\phi(Q_i) \sum_{j=1}^N  \phi(K_j^\top)V_j}{\phi(Q_i) \sum_{j=1}^N  \phi(K_j^\top)}, \quad (1)$ 

which can be computed in linear time.

The function $\phi(\cdot)$ can take various forms, such as 1 + ELU [1], ReLU [2], squared ReLU [3], Taylor [4-6] or Random Feature [7-8] expansions, and even MLPs trained to mimic softmax attention [6]. They aim to approximate softmax without its explicit calculation when being applied jointly to queries and keys, or to retain its properties, most importantly, non-negativity of resulting dot products $\phi(Q_i) \phi(K_j^\top)$. 


The latter property, together with reweighting attention scores (denominator in the formula 1) are defining for Linear Transformer algorithms. Absence of scaling by $\frac{1}{\phi(Q_i) \sum_{j=1}^n  \phi(K_j^\top)}$ leads to numerical instabilities, and the scaling factor itself is not guaranteed to be bounded without non-negative $\phi(\cdot)$.  However, both mappings $\phi(\cdot)$  and memory intensive non-MatMul operations for reweighting contribute to subpar speed and computational efficiency in comparison with ordinary and fast self-attention algorithms on all but large context sizes.

We forgo both transforming $\mathbf{Q}, \mathbf{K}$ by $\phi(\cdot)$ and reweighting in DenseAttention as we believe the main factor of success of Transformer is the ability of all $N \times N$ interactions between tokens. It results in an improved computational efficiency and simpler design which can be expressed entirely by matrix multiplications:

$\mathbf{A} = \mathbf{Q} \mathbf{K}^\top \mathbf{V} $

---

**References**

[1] Katharopoulos et al., "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention." ICML 2020

[2] Qin et al., "cosFormer: Rethinking Softmax in Attention." ICLR 2022

[3] Hua et al., "Transformer Quality in Linear Time." ICML 2022 

[4] Keles et al., "On The Computational Complexity of Self-Attention." International Conference on Algorithmic Learning Theory 2023

[5] Arora et al., "Simple Linear Attention Language Models Balance the Recall-Throughput Tradeoff." ICLR 2024

[6] Zhang et al., "The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry." ICLR 2024

[7] Choromanski et al., "Rethinking Attention with Performers." ICLR 2021

[8] Peng et al., "Random Feature Attention." ICLR 2021
</content>
</thread_object>

</thread>
<thread>
creation date:
1733299597745

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General summary

comment:
We sincerely thank reviewers for their time, feedback, and constructive suggestions which led to additional experiments and helped to clarify and strengthen our work. We would like to emphasize the main points of the paper, highlighted strengths, and key updates to the revision.  

---

**Recap**

In this paper we propose DenseAttention Network, a simplification of Self-Attention and the vanilla Transformer, which runs both in $O(N)$ and $O(N^2)$ time depending on what’s more computationally efficient.  Among other modifications to the architecture, we completely eliminate softmax from attention with no substitutions to it. We show it’s able to work if the standard LayerNorm is replaced with a proposed MaxNormActivation. DenseAttention in both regimes runs faster than highly-optimized low-level Transformer implementations, despite being coded in plain PyTorch. It retains or outperforms the modeling quality of the Transformer and its multiple variations across the wide range of context sizes, as shown in experiments on language modeling.

Additionally, we suggest two extensions, Cosine RelPE as a computationally efficient alternative to RoPE, and local-global attention scheme, to aid the model performance on extremely long sequences while maintaining computational efficiency. It allows us to achieve state-of-the-art results on the Long Range Arena suite of benchmarks among all Transformer-based models and the overall best result on the Pathfinder-256 task (65K context size). 

We are grateful to reviewers for acknowledging the following strengths in our work:

* The paper being **clearly written and easy to follow** as well as proposing a solution to **important** problem of **efficiency** of Transformer architectures (ruBJ). 

* The solution being  **well-motivated** (ruBJ, oGkz), in particular, MaxNormActivation, which is also acknowledged to be **useful** primitive **to stabilize LLM training** (rVwA).

* Originality, novelty and efficiency of the approach to unite and **fuse** projection matrices of Transformer into **single parameter** (HAtu, rVwA) and of the proposition to **use local attention** together with the new architecture (HAtu).

* Strong empirical results for speed/ **efficiency** gains with just **architectural** changes and **no specialized kernels** as well as for promising modeling performance in **long-context sequence modeling tasks** (rVwA). 

* **Provision of code** (oGkz) which ensures easy reproducibility. 

**Summary of updates**

Here’s the summary of the key updates we made based on reviewers’ comments and suggestions. 

* Added a discussion of differences of Linear Transformers in relation to DenseAttention in **Appendix D** of the revised manuscript (also in **General Response. Part 1** for convenience) (ruBJ, HAtu, oGkz, rVwA). Additionally provided an exposition of other sub-quadratic algorithms for sequence processing in **General Response. Part 3** (please note: version of this exposition in the manuscript is a draft which will be updated). (oGkz). 

* Added an “Extended comparison with Transformer-based models” on the LRA for 25+ models in **Appendix E.2** (also in **General Response. Part 2**), which highlights strong performance of DANet (ruBJ, HAtu, rVwA). Provided an extended “Discussion of the LRA tasks” clarifying the scope and difficulty for each task in **Appendix E.1** (ruBJ, oGkz).

* Conducted additional experimental studies on extremely long sequences performance on Pathfinder-256 benchmark in **Appendix H.1** (ruBJ, oGkz), and scaling effects in  **Appendix H.3** (ruBJ, rVwA), both with positive outcomes.

* Reproduced all DANet-BERT pretraining experiments with architecture equivalent to original BERT in **General Response. Part 4** and arrived at results and conclusions similar to former ones (oGkz). 

* Performed additional ablation studies on local-global attention in DANet-BERT for extremely long context sizes in **General Response. Part 4** (oGkz), speed gains of Cosine PE in **Appendix H.2** (oGkz, rVwA), and use of MaxNormActivation in vanilla Transformer in **General Response. Part 4** (rVwA), all to expected/ favorable outcomes.

* Added a “Conclusion and Future Work” section in **Appendix A** (draft) and **General Response. Part 3** (final version) to clarify the scope of the work and future directions (ruBJ, oGkz).
</content>
</thread_object>

</thread>


Metareview:
metareview:
The primary goal of this paper is to develop an attention mechanism which, given an input of N tokens, allows for the computation of the full matrix of NxN pairwise attention weights in less than N^2 time (i.e., linear in N).  This is done in a manner similar to the Linear Attention mechanism, but as discussed by the reviews and authors is still somewhat distinct from Linear Attention in a few technical details regarding issues like normalization, which the authors suggest is important for numerical stability.  Beyond this, the authors also make several changes to the overall architecture in terms of other normalization operators, positional encoding, etc and achieve good performance on several benchmarks.

Overall, the reviewers are fairly unanimous in leaning to rejection, and I agree that the paper is perhaps not ready for publication just yet.  As has been noted in the reviews, the proposed mechanism is very similar to Linear Attention, and while the current work may not be strictly captured by Linear Attention, this proximity combined with the numerous other changes to the architecture make it difficult to establish which elements are critical to performance and for readers to understand the key aspects of the contribution and why the proposed mechanism is needed relative to Linear Attention.  I would encourage the authors to take these aspects and additional comments from the reviewers into account when preparing a new manuscript for a future conference which better highlights the advantages of the proposed approach relative to existing techniques (such as Linear Attention) as well as justifying the additional modification to the overall architecture.