Paper: Algorithm and Hardness for Dynamic Attention Maintenance in Large Language Models (ICLR-2024)

Reviewer comment (Theory):
The online setting dealt in this paper has less reality; in the dynamic setting of this problem, only one entry of input matrices is changed by one update. However, in real situation, although it can be assumed that the values of entries change slowly, a substantial portion of entries are subject to change. Therefore, when applied to real situations, polynomially many update operations may occur, which eventually result in quadratic computation time.

Author response:
Your careful reading and the insightful suggestions are greatly appreciated. In response to the concerns raised regarding the practical implications of the dynamic attention maintenance problem, we emphasise that entry-wise update of the weight matrix is prevalent in real-world scenarios. Deep neural network architectures frequently exhibit significant redundancy, and empirical evidence supports the capacity of deep neural networks to tolerate substantial levels of sparsity [1,2]. In downstream fine-tuning tasks, the dimensions of the model often make the fine tuning infeasible.
Sparse neural networks have demonstrated enhanced parameter and computational efficiency relative to dense networks, and, in many cases, can significantly decrease wall clock inference times [3,4]. Over the past few years, numerous techniques for inducing sparsity have been proposed to sparsify the neural network such as magnitude pruning [1,5], RegL [6] and dynamic sparse reparameterization [7]. For attention based models, leveraging the sparsity of the attention can help make the model more efficient, as in [8,9,10]. 
For example, to save the computation time and memory of the transformer, [11] only computes the top-k entries for each row in the attention matrix. Their work implies that we do not need to update the entire query and key matrix. A promising avenue for future exploration involves the adaptation of this concept to implement rank-1 updates, potentially surpassing entry-wise updates in effectiveness. We believe that it is an interesting direction for enhancing practical training efficiency. Our approach acknowledges the value of simplicity as an initial foundation, drawing from a range of theoretical papers [12,13,14,15]. Starting with a simpler version allows us to build a solid starting point for understanding the broader context. 

[1] Han, Song, et al. "Learning both weights and connections for efficient neural network." Advances in Neural Information Processing Systems 28 (2015).

[2] Gale, Trevor, Erich Elsen, and Sara Hooker. "The state of sparsity in deep neural networks." arXiv preprint arXiv:1902.09574 (2019).

[3] Theis, Lucas, et al. "Faster gaze prediction with dense networks and Fisher pruning." arXiv preprint arXiv:1801.05787 (2018).

[4] Kalchbrenner, Nal, et al. "Efficient neural audio synthesis." International Conference on Machine Learning. PMLR, 2018.

[5] Zhu, Michael, and Suyog Gupta. "To prune, or not to prune: exploring the efficacy of pruning for model compression." arXiv preprint arXiv:1710.01878 (2017).

[6] Evci, Utku, et al. "Rigging the lottery: Making all tickets winners." International Conference on Machine Learning. PMLR, 2020.

[7] Mostafa, Hesham, and Xin Wang. "Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization." International Conference on Machine Learning. PMLR, 2019.

[8] Roy, Aurko, et al. "Efficient content-based sparse attention with routing transformers." Transactions of the Association for Computational Linguistics 9 (2021): 53-68.

[9] Choromanski, Krzysztof, et al. "Rethinking attention with performers." arXiv preprint arXiv:2009.14794 (2020).

[10] Kaiser, Łukasz, and Samy Bengio. "Discrete autoencoders for sequence models." arXiv preprint arXiv:1801.09797 (2018).

[11] Gupta, Ankit, et al. "Memory-efficient Transformers via Top-$ k $ Attention." arXiv preprint arXiv:2106.06899 (2021).

[12] Alman, Josh, and Zhao Song. "Fast attention requires bounded entries." arXiv preprint arXiv:2302.13214 (2023).

[13] Deng, Yichuan, Sridhar Mahadevan, and Zhao Song. "Randomized and deterministic attention sparsification algorithms for over-parameterized feature dimension." arXiv preprint arXiv:2304.04397 (2023).

[14] Kitaev, Nikita, Łukasz Kaiser, and Anselm Levskaya. "Reformer: The efficient transformer." arXiv preprint arXiv:2001.04451 (2020).

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

Reviewer comment (Theory):
As described in "weaknesses" section, I think this work has of theoretical interest, and so the worst-case update time instead of amortized update time is preferable if possible. As described in Remark B.2, since the proposed algorithm relies on the periodic rebuilding of data structures, we can make the update time fit for worst-case analysis theoretically by dividing the rebuilding of data structures into updates. So, the question is, why the main result (Theorem 1.3 or Theorem B.1) is posed on amortized update time? For clarity?

Author response:
Thanks for pointing that out. In our paper, we don’t go into detail for worst-case as the worst case time can be derived from our results by simply using standard techniques as stated in Remard B.2. We believe showing the proof of the worst case would confuse the readers.

Reviewer comment (Theory):
I'm interested in the situation that we are in a static setting but we only have to compute one entry of attention matrix. In such a case, there is possibility to lower time complexity, but it is not straightforward because it involves $n\times n$ matrix's inverse. Do you have any idea for this? If so, it complements this work since in the dynamic setting we also have to compute only one entry of attention for a query.

Author response:
One entry can be computed statically in O(nd) time: Compute one row of QK, then do entry-wise exponentiation and compute inner product with column of V. Normalization works because we have the necessary row of exp(QK).

Reviewer comment (Presentation):
The paper structure is a little bit unbalanced, with the first introduction section taking up more than half of the pages. The related work section (Sec 1.2) can be compressed.

Author response:
**A2**: Thanks for pointing it out. To make the structure more balance, we will make our introduction and related work more compressed and also move some of the related work to our appendix. We will be sure to update that in our camera ready version.

Reviewer comment (Novelty):
The biggest question I have is why we need dynamic attention multiplication. There are many works for sparse attention and low-rank attention [1, 2], which provide potential solutions for the long-range attention problem. And when we do fine-tuning on LLMs, we can also freeze part of its parameters and only tune the rest. It seems not very common to do targeted update for specific entries of K, V, or query specific positions of the attention matrix. So I'm not sure what is the scenario that the proposed method is mostly suitable for.

Author response:
**A3** : The question you've raised about the necessity of dynamic attention multiplication in light of existing sparse and low-rank attention models is indeed pertinent. 
The significance of our approach lies in its ability to leverage the redundancy often found in deep neural networks. As you mentioned, many methods have been proposed to leverage the sparsity to speed up the computation [1,2,3,4,5,6]. For example, in the top-k method [6], that only updates the top-k entries that have the highest attention score. In a lazy fashion, we can accumulate the update of them and recompute the attention after enough iterations. Our proposed method opens the door to future explorations, such as implementing rank-1 updates, which could potentially be more effective than entry-wise updates. It acknowledges the importance of starting with a simpler model as a foundation. This approach not only allows for a clearer understanding of the complex dynamics within LLMs but also provides a solid base for further enhancements in training efficiency, drawing from various theoretical perspectives [7,8,9,10].

[1] Han, Song, et al. "Learning both weights and connections for efficient neural network." Advances in Neural Information Processing Systems 28 (2015).

[2] Gale, Trevor, Erich Elsen, and Sara Hooker. "The state of sparsity in deep neural networks." arXiv preprint arXiv:1902.09574 (2019).

[3] Roy, Aurko, et al. "Efficient content-based sparse attention with routing transformers." Transactions of the Association for Computational Linguistics 9 (2021): 53-68.

[4] Choromanski, Krzysztof, et al. "Rethinking attention with performers." arXiv preprint arXiv:2009.14794 (2020).

[5] Kaiser, Łukasz, and Samy Bengio. "Discrete autoencoders for sequence models." arXiv preprint arXiv:1801.09797 (2018).

[6] Gupta, Ankit, et al. "Memory-efficient Transformers via Top-$ k $ Attention." arXiv preprint arXiv:2106.06899 (2021).

[7] Alman, Josh, and Zhao Song. "Fast attention requires bounded entries." arXiv preprint arXiv:2302.13214 (2023).

[8] Deng, Yichuan, Sridhar Mahadevan, and Zhao Song. "Randomized and deterministic attention sparsification algorithms for over-parameterized feature dimension." arXiv preprint arXiv:2304.04397 (2023).

[9] Kitaev, Nikita, Łukasz Kaiser, and Anselm Levskaya. "Reformer: The efficient transformer." arXiv preprint arXiv:2001.04451 (2020).

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

Reviewer comment (Experiments):
It might be more interesting to consider a row or a column, or even a submatrix of fixed size, being updated in each iteration.

Author response:
Your careful reading and the insightful suggestions are greatly appreciated. In response to the concerns raised regarding the practical implications of the dynamic attention maintenance problem, we emphasise that entry-wise update of the weight matrix is prevalent in real-world scenarios. Deep neural network architectures frequently exhibit significant redundancy, and empirical evidence supports the capacity of deep neural networks to tolerate substantial levels of sparsity [1,2]. In downstream fine-tuning tasks, the dimensions of the model often make the fine tuning infeasible.

Sparse neural networks have demonstrated enhanced parameter and computational efficiency relative to dense networks, and, in many cases, can significantly decrease wall clock inference times [3,4]. Over the past few years, numerous techniques for inducing sparsity have been proposed to sparsify the neural network such as magnitude pruning [1,5], RegL [6] and dynamic sparse reparameterization [7]. For attention based models, leveraging the sparsity of the attention can help make the model more efficient, as in [8,9,10]. 
For example, to save the computation time and memory of the transformer, [11] only computes the top-k entries for each row in the attention matrix. Their work implies that we do not need to update the entire query and key matrix. A promising avenue for future exploration involves the adaptation of this concept to implement rank-1 updates, potentially surpassing entry-wise updates in effectiveness. We believe that it is an interesting direction for enhancing practical training efficiency. Our approach acknowledges the value of simplicity as an initial foundation, drawing from a range of theoretical papers [12,13,14,15]. Starting with a simpler version allows us to build a solid starting point for understanding the broader context. 

[1] Han, Song, et al. "Learning both weights and connections for efficient neural network." Advances in Neural Information Processing Systems 28 (2015).

[2] Gale, Trevor, Erich Elsen, and Sara Hooker. "The state of sparsity in deep neural networks." arXiv preprint arXiv:1902.09574 (2019).

[3] Theis, Lucas, et al. "Faster gaze prediction with dense networks and Fisher pruning." arXiv preprint arXiv:1801.05787 (2018).

[4] Kalchbrenner, Nal, et al. "Efficient neural audio synthesis." International Conference on Machine Learning. PMLR, 2018.

[5] Zhu, Michael, and Suyog Gupta. "To prune, or not to prune: exploring the efficacy of pruning for model compression." arXiv preprint arXiv:1710.01878 (2017).

[6] Evci, Utku, et al. "Rigging the lottery: Making all tickets winners." International Conference on Machine Learning. PMLR, 2020.

[7] Mostafa, Hesham, and Xin Wang. "Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization." International Conference on Machine Learning. PMLR, 2019.

[8] Roy, Aurko, et al. "Efficient content-based sparse attention with routing transformers." Transactions of the Association for Computational Linguistics 9 (2021): 53-68.

[9] Choromanski, Krzysztof, et al. "Rethinking attention with performers." arXiv preprint arXiv:2009.14794 (2020).

[10] Kaiser, Łukasz, and Samy Bengio. "Discrete autoencoders for sequence models." arXiv preprint arXiv:1801.09797 (2018).

[11] Gupta, Ankit, et al. "Memory-efficient Transformers via Top-$ k $ Attention." arXiv preprint arXiv:2106.06899 (2021).

[12] Alman, Josh, and Zhao Song. "Fast attention requires bounded entries." arXiv preprint arXiv:2302.13214 (2023).

[13] Deng, Yichuan, Sridhar Mahadevan, and Zhao Song. "Randomized and deterministic attention sparsification algorithms for over-parameterized feature dimension." arXiv preprint arXiv:2304.04397 (2023).

[14] Kitaev, Nikita, Łukasz Kaiser, and Anselm Levskaya. "Reformer: The efficient transformer." arXiv preprint arXiv:2001.04451 (2020).

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

Reviewer comment (Writing):
In the abstract, $\omega(1,1,\tau)$ and $\tau$ are undefined notations that make it hard for the readers to understand the core results of the paper at first glance. While one might be able to infer that $\omega$ means the matrix multiplication exponent from the context, it is hard to understand $\omega(1,1,\tau)$. Also, without properly introducing $\tau$, it is unclear whether the run time results hold for a specifically chosen $\tau$ (might be too complicated to write out in the abstract) or for any $\tau \in (0, 1]$.

Author response:
Thanks for your careful reading and pointing them out. To fix the presentation issues, we have defined $\omega(1,1, \tau)$ to be a matrix multiplication exponent and $\tau$ to be a constant in $(0,1]$ in the abstract.

Reviewer comment (Writing):
In Theorem 1.3, $\delta$ is undefined. The initial impression might be some notion of failure probability, but it turns out in later sections that $\delta$ is indeed the update to the attention matrix.

Author response:
Thanks for your careful reading and pointing them out. We have defined $\delta$ to be the update of the matrix in Theorem 1.3.

Reviewer comment (Writing):
Section 3.1 does not give a precise and clear overview of the lazy update based algorithm that solves ODAMV. For example, it is unclear what “List_C, List_D, List_V” in “Lazy Update” store. In “Fast Query”, it is confusing why there are two stacked matrices, $\Delta_{V,1}$ and $\Delta_{V,2}$, from List_V.

Author response:
Thanks for your careful reading. We kindly argue that, as you can see in Line 16,17 of Algorithm 2 and Line 5 of Algorithm 3, the update that store in  "List_C, List_D, List_V" is a bit lengthy to write them down in technique overview. In Section 3.1, we just want to give reader a brief idea of how our algorithm works.

Reviewer comment (Theory):
In Section 3.2 and 4 on the Lower Bound (LB), it is unclear why this work presents the LB result for a simpler problem called OAMV, which is not the true problem ODAMV considered, instead of directly presenting the LB results of ODAMV. Section 4 is essentially a more detailed / formalized version of Section 3.2, and these two sections might be combined.

Author response:
Due to space limitations, we provide our result and proof sketch for the simplified version of ODAMV, which is OAMV, in the main content and defer the proof of the true problem ODAMV to appendix. Section 3.2 is the technique overview that gives a proof sketch of our main results and Section 4 is the formal results for the OAMV problem.

Reviewer comment (Writing):
Minor issue: In Theorem 1.3, “This operation has the same … as K update” should be “This operation has the same … as UpdateK”.

Author response:
Thanks for pointing this out. We have fixed them as you suggested.

Reviewer comment (Writing):
Minor issue: Below Definition 1.2, “When then complement our result …” should be “We then complement our result”.

Author response:
Thanks for pointing this out. We have fixed them as you suggested.

Reviewer comment (Writing):
Minor issue: In the paragraph “Fast Query” in Section 3.1, “denote the lates $D^{-1}$” should be “denote the latest $D^{-1}$”.

Author response:
Thanks for pointing this out. We have fixed them as you suggested.