Paper: PolySketchFormer: Fast Transformers via Sketches for Polynomial Kernels (ICLR-2024)

Reviewer comment (Reproducibility):
Typically in matrix sketching, we take the product of $A^{\otimes p} \in \mathbb{R}^{n \times h^p}$ and $S \in \mathbb{R}^{h^p \times r}$ explicitly. However, the proposed algorithm does not perform this explicit multiplication and avoids the exponential runtime dependency on $p$. The algorithm described in *Ahle et al. (2020)* uses the Fast Fourier Transform to compute $A^{\otimes p} S$ efficiently with the tree-based algorithm. This paper does not indicate how this product is computed. It is important to outline this process for completeness, as it is a significant component that enables the removal of exponential dependency on degree $p$.

Author response:
**“does not have any indication of how this product is computed”**: We use the algorithm as exactly described in the paper of Ahle et al., which avoids explicit computation of the large matrix to apply the sketch. We try to explain this algorithm in Figure 3 (Definitions of SRHT and TensorSRHT are presented in Definition 2.3 and 2.4). We will include a more detailed description of how exactly it is to be implemented to avoid the exponential dependence in p.

Reviewer comment (Writing):
In theorem 2.2, what is $m$?

Author response:
**Theorem 2.2**: Our apologies for the typo. It should be r instead of m.

Reviewer comment (Theory):
On page 12, in the proof of theorem 2.5, does the JL-moment property of $S$ imply that $\langle S^Tc_i, S^Td_j \rangle.\langle c_i, d_j \rangle \le \langle c_i, d_j \rangle^2$? It is unclear how $(\langle S^Tc_i, S^Td_j \rangle - \langle c_i, d_j \rangle + 2 \langle c_i, d_j \rangle )^2 (\langle S^Tc_i, S^Td_j \rangle - \langle c_i, d_j \rangle)^2 \le (1+C) (\langle S^Tc_i, S^Td_j \rangle - \langle c_i, d_j \rangle)^4 + 4(1+1/C) \langle c_i, d_j \rangle^2 (\langle S^Tc_i, S^Td_j \rangle - \langle c_i, d_j \rangle)^2$ holds for any $C \ge 1$ otherwise.

Author response:
**Proof of Theorem 2.5**:  We use AM-GM inequality to obtain the following inequality for any C > 0: for any $a, b \ge 0$, $(a + b)^2 = a^2 + b^2 + 2ab = a^2 + b^2 + 2(a/\sqrt{C})(b\cdot\sqrt{C}) \le a^2 + b^2 + (a^2/ C) + (b^2 C) \le (1 + 1/C)a^2 + (1 + C)b^2$  
  
We use this to obtain that $(\langle S^T c_i, S^T d_j \rangle - \langle c_i, d_j\rangle + 2\langle c_i, d_j\rangle)^2 \le (1 + C)(\langle S^T c_i, S^T d_j\rangle - \langle c_i, d_j\rangle)^2 + (1 + 1/C)(2\langle c_i, d_j\rangle)^2$ from which the inequality follows. So we do not need the assumption that $\langle S^Tc_i, S^Td_j\rangle \cdot \langle c_i, d_j\rangle \le \langle c_i, d_j\rangle^2$.

Reviewer comment (Experiments):
Even though the block-based method is targeted at decoder-only models, the approximate algorithm for polynomial kernel self-attention can be applied in general to any transformer-based architecture. Have the authors explored any other architectures with this approximate self-attention algorithm?

Author response:
**Other Self-Attention Mechanisms**: Yes, our sketch-based attention mechanism can definitely be used outside the causal language modeling that we focus on in this paper. We plan to explore this in future.

Reviewer comment (Theory):
In "On The Computational Complexity of Self-Attention [Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, Chinmay Hegde]" paper, a polynomial approximation to softmax self-attention is derived using the finite Taylor series which includes lower order terms as well. Technically, this should be able to get arbitrarily close to softmax self-attention as the order $p$ is increased. If we use Polysketch self-attention to approximate each term of the Taylor series approximation, and thereby approximate softmax self-attention, what would be the closest approximation one can get while still maintaining subquadratic complexity in context length and subexponential complexity in degree $p$?

Author response:
**Closest Approximation to Softmax using Polynomial**: It would definitely be interesting to study how close we can get to softmax using these techniques. But the lower bounds in the work of Alman and Song suggest that we shouldn’t be able to get too close to Softmax attention output in sub-quadratic time unless the entries of the query and key vectors are bounded.

Reviewer comment (Evaluation):
The empirical result is weak. It seems that the proposed method with 8k, 16k context is slower than the other method with 2k, 4k context and has higher perplexity. Doesn’t this mean the proposed method is not better?

Author response:
- **Empirical Results**: The table does suggest that our longer context models have worse perplexity. But we note that we have not tried to pick the best possible parameters for each context length and we trained at all context lengths using the same hyperparameters (learning rate schedule, peak learning rate, number of steps, etc.) which can be the reason for the trends.

In some situations such as when working with long documents, one might want to work with models trained on long contexts even if they have worse perplexities than models trained with shorter contexts since we cannot fit the whole document otherwise. We do also note that for longer contexts, our model trains significantly faster than vanilla transformer/Flash Attention so our models can be trained for significantly more steps using the same amount of compute time. Thus, training for a larger number of steps, we may be able to beat softmax transformers though we have not yet fully explored this direction yet.

Reviewer comment (Reproducibility):
Need to elaborate more on the data leak of Performer, so the reviewers can verify the statement.

Author response:
- **Data Leak in Performer**: This is an issue discovered by us since there was a mismatch between the eval perplexities we computed vs the quality of text generated by the model trained using performer attention. We will include a full explanation of the bug in the next version of the paper. Here we include a brief description of the issue. 

Consider a two token context length. Let $q_1$ and $q_2$ be the queries and $k_1$ and $k_2$ be the keys and $v_1$ and $v_2$ be the values. Performer first maps the queries and keys to $q’_1$ and $q’_2$ and $k’_1$ and $k’_2$ respectively using random projections. For the keys, they finally need to compute $\exp(-\|k_1\|_2^2) * \exp(k_1’)$ and $\exp(-\|k_2\|_2^2) * \exp(k_2’)$ where $\exp(k’)$ means applying the exponential function entrywise to coordinates of $k’$. So as to not create large entries and make the training stable they compute the maximum value of an entry in $k_1’$ and $k_2’$ and subtract the same value from all the entries in $k_1’$ and $k_2’$ essentially making the largest entry to be 0 and so when we apply exponential function, we will not get large values. Mathematically, subtracting the same quantity should not matter because after applying exponential, it is essentially just multiplying all the entries with the same quantity and hence the attention matrix they compute should not change. But observe that the entries in $k_1’$ may get subtracted from the largest entry in $k_2’$ and since we are working with limited precision, this creates a channel using which information from future keys leak to present attention calculation.

We observe that the model tries to exploit this channel and therefore gets better eval perplexities since it is using information from future tokens at test time to predict the same tokens. We reported this bug to the authors of Performer paper and they confirmed the bug. We do note that the bug is not present in the non-causal version of the Performer. But since our paper focuses on the causal version of the performer, we are unable to report the numbers for the correct version of the performer implementation. We did implement Performer attention removing the bug and using an alternate “normalization” to stabilize training but the training latency was significantly worse and we couldn’t finish training in a reasonable amount of time to be able to compare with our models. We note that the running times that we reported are for the buggy version of Performer which is publicly available.

Reviewer comment (Experiments):
The authors need to elaborate on why some newer methods that perform better on the long range arena does not need to be compared. For example, ChordMixer or LongFormer.

Author response:
- **Long Range Arena**: The benchmarks in Long Range Arena are for the non causal attention setting. Since we focused on improving causal language modeling (autoregressive decoding) in this work, we do not evaluate our mechanism in tasks in the Long Range Arena benchmark.

Reviewer comment (Reproducibility):
Can you briefly explain why Performer is slower than the softmax baseline in the experiments? This does not seem to be the case in the original paper. Is this because of bad implementation?

Author response:
- **Performer being Slow**: The runtimes we reported are directly using the Publicly available version of Performer (which as we note has a data leak.). Depending on the model parameters, the cumulative sum algorithm used in Performer can be significantly slow because of the large number of reads/writes it has to perform to compute the output of the attention mechanism. This was also observed by the authors of “FLASH : Transformer Quality in Linear Time”.

Reviewer comment (Theory):
According to Definition 2.4, PolySketchFormer can only express/approximate the polynomial $f(s)=s^p$ ($s:= \langle q,k \rangle$), where $p$ is **a power of 2**. This limits the flexibility to "approximate Softmax".

Author response:
**"Limited flexibility to approximate Softmax"**: We do not think of polynomial attention as trying to approximate softmax but instead think of it as an entirely different attention mechanism. Softmax attention puts more emphasis on large dot products and similarly a high degree polynomial attention will also put more emphasis on large dot products though the main difference is that polynomial attention also amplifies dot products with large negative values. Our hypothesis is that the query and key vectors obtained when training with polynomial attention don’t create large negative dot products and hence may not present large issues.

Reviewer comment (Writing):
A typo: On Page 3, the second to last line, it should be $q\in \mathbb{R}^h$, right?

Author response:
**Typo**: Yes, it should be $q \in \mathbb{R}^h$

Reviewer comment (Reproducibility):
In Table 1, what is exactly the Softmax model? Is it FlashAttention?

Author response:
**“In Table 1, what is exactly the Softmax model? Is it FlashAttention?”**: The softmax models are trained with a default implementation of dot product attention in Flax. This implementation does not have all the optimizations performed by FlashAttention. A properly implemented FlashAttention model should also achieve the same perplexity values.

Reviewer comment (Theory):
Even though the paper proposes a way to avoid the nonnegativity, another property of the softmax is that it upper-bounds a value to at most 1, while it’s less clear how the proposed model deals with this or whether this is even an issue.

Author response:
-  **"Upperbounds to 1"**: We normalize the attention matrix outputs to have a value 1 in our models as well. Concretely, given Q and K, we use sketches + tensoring to get Q’ and K’ and then to compute the final output we run our algorithm to get LT(Q’ K’^T)*[V 1] i.e., we append a column of 1s to the V matrix and the last column of the result gives the necessary values we need to use to normalize the rows of the output. So, the approximate attention matrix we implicitly apply to V to get the output also has the sum of each row as 1.

Reviewer comment (Experiments):
The paper mentions in the introduction that in practice most sota papers use the vanilla transformer in practice (i.e. large-scale experiments), but then there are no sota papers and large-scale experiments conducted in the experimental section. In addition, there are few to no ablations on the choices made in this work.

Author response:
-  **Comparison with SOTA models on NLP tasks**: As the focus of our work was on improving the training latency for language models, we could not compare on benchmarks with SOTA models that are quite large. We note that we train the models from scratch on two datasets for only 125k iterations and hence cannot be competitive with SOTA models (even with similar model configuration as us as they are trained for much longer on more diverse datasets) on usual long context benchmarks which study for e.g., how good the model is at summarizing long documents, how good the model is at retrieving tokens, or translation etc.

Reviewer comment (Experiments):
There are some related works in the direction of polynomial and efficient attention that might be worth comparing with: “Hydra Attention: Efficient Attention with Many Heads” (10/2022) and “Linear Complexity Self-Attention with 3rd Order Polynomials” (3/2023).

Author response:
- **"HydraAttention and Related Mechanisms"**: Thanks for the references! We looked at HydraAttention and it appears that the paper focuses on improving Vision Transformers instead of language modeling and at the first glance it seems to us that such models may be a bit too weak for causal language modeling which is our main focus. Our intuition is that the attention heads for language modeling seem to require a somewhat large (like 64/128) number of features to get something useful from the attention mechanism. But we agree it is something we should explore more.

Reviewer comment (Evaluation):
What is the perplexity of the model and why is this a good metric?

Author response:
- **Perplexity**: Perplexity is generally used as a measure of the quality of a language model. It is defined as the exponential of the average cross entropy loss of the model on the training dataset. One can roughly think of perplexity as the “inverse probability assigned by the model to the correct token as the next token” given the context. If a model has perplexity 100, then the model essentially only assigns on average (in the geometric mean sense) a probability of 1/100 to the correct token whereas a model that has perplexity 2, assigns on average a probability of ½ to the correct token.  

We agree a single measure, such as perplexity, is not good at predicting how good the model is but it is known to be a good proxy for the performance and is used as the main way to compare language models in earlier works on efficient mechanisms for training language models such as “FLASH : Transformer Quality in Linear Time” and references therein.

Reviewer comment (Theory):
Why is a degree 4 polynomial used?

Author response:
- **Degree 4**: We performed experiments with degrees 2, 4 and 8 and found that degree 2 polynomial attention is not close in performance to softmax attention whereas degree 4 is very close as we report in the paper. We also observed that going to higher degrees does not help much since degree 4 is already very close to softmax. So we settle to use degree 4 polynomial in our models.

Reviewer comment (Experiments):
Why are all experiments up until the 16k context window only? What’s the core limitation for extending it further?

Author response:
- **Experiments only up to 16k**: Since we trained all our models using data parallelism, the context length is limited only by the 40GB VRAM available in the A100 GPUs we use. There is nothing stopping us from extending it to longer contexts using multiple GPUs and using pipeline parallelism and the same improvements in the single device setting should carry over. We note that we do report training latency results for training a smaller 4-layer model with a context length of 32k in the paper.

Reviewer comment (Theory):
Why is a single-degree k polynomial (even more that it is restricted to be even) approximating the softmax? Are there any guarantees or insights on that?

Author response:
- **"Why is a single-degree k polynomial approximating the softmax?"**: We note that we are not claiming that polynomial attention is approximating softmax attention. We think of polynomial attention as an entirely new attention mechanism which also works. Intuitively, the softmax is used to convert a vector of numbers into a probability distribution. Nothing in principle is stopping us from using other functions, such as polynomials in our case, from converting a vector of numbers into a probability distribution. Our other intuition is that softmax forces the probability distribution to focus on the vector entry with the largest value which is also done by a high degree polynomial as well although even-degree polynomials also amplify large negative entries. So the key, query dot products function somewhat differently under softmax and polynomial attentions.

Reviewer comment (Evaluation):
In Fig. 1 I don’t understand how the vanilla softmax attention runs out of memory at context length of 4k. Shouldn’t the memory footprint just be roughly (4k)^2 doubles, or roughly 32 MBs? This is extremely small. I can compute a 4k x 4k attention matrix in a few seconds on my laptop.

Author response:
In order to train the models, we compute the gradient of the loss function with respect to the weights using standard backpropagation algorithm. The backpropagation algorithm needs to store all the results of the intermediate computation or recompute the activations from the checkpoints. In our model, with a context length of 4k, we have 12 layers each with 12 attention heads and hence the backpropagation algorithm needs to store 144 4k x 4k attention matrices and results of various other computations such as the key, query, value projections for each token in the context in each layer and the intermediate computations of the feedforward layers, etc. Our model with vanilla softmax attention runs out of memory to store all the required information for backpropagation at 4k context length.

Reviewer comment (Writing):
“because of an RNN-style sequential dependence of linear transformers which essentially renders the linear attention implementation to be memory-bandwidth bound” — I don’t understand what this means. Sequential dependence of what on what? Does this not depend on the algorithm used?

Author response:
The cumulative sum algorithm used by earlier works to compute LT($AB^T$)C computes rows of the output matrix one-by-one and inorder to perform these operations, we have to read/write $n$ matrices to the high-bandwidth memory of the GPU where $n$ is the number of rows of $A$. This is what the authors of “FLASH : Transformer Quality In Linear Time'' refer to as “RNN style Sequential Dependence of Cumulative Sum algorithm on the context length” since RNNs also need to read and write $n$ matrices to the high-bandwidth memory where $n$ is the context length.

	As you say, the sequential dependence (as in the number of read/writes performed) does depend on the algorithm used to compute LT($AB^T$)$C$ and is not an inherent property of Linear Transformers.

Reviewer comment (Novelty):
The work of Han, Avron, Shin on approximating entrywise transformed lower rank matrices with tensor sketch seems very related but is not cited. See: https://proceedings.mlr.press/v119/han20a.html

Author response:
Thanks for suggesting the paper to us. We will add a citation in the next version. In theorem 2.2, yes the parameter should be r and not m.

Reviewer comment (Writing):
The parameter m in Theorem 2.2 (in the runtime) seems to be undefined. Is this meant to be r?

Author response:
Thanks for suggesting the paper to us. We will add a citation in the next version. In theorem 2.2, yes the parameter should be r and not m.

Reviewer comment (Theory):
I didn’t follow the explanation at the bottom of page 5. If we break out exactly what this is doing it is computing:  
[AD_1 H P_1 D_2 H P_2 * AD_3 H P_3 D_4 H P_4]  
where D_1,P_1 and D_3,P_3 are the random matrices chosen for SRHT_1(A) and SRHT_2(A) respectively. And D_2,P_2 and D_4,P_4 are chosen by TensorSRHT_1. This doesn’t look right to me. What is the purpose of applying the ‘double sketch’ D_1 H P_1 D_2 H P_2. If one looks at the Ahl et al. paper (https://arxiv.org/pdf/1909.01410.pdf), the base sketch is applied already to the first level of tensor products, not to the original input. I.e. the base sketch sketches from h^2 to r dimensions. Not h to r dimensions. In particular, Def 15 of that paper defines TensorSHRT which matches Def 2.4 of this paper, and should be directly used for sketching A^{tensor 2}.

Author response:
In the work of Ahle et al., they apply base sketch to obtain input sparsity time algorithms which is not very relevant in this work. As you observe, we can skip applying the base sketch and directly proceed to applying the TensorSRHT sketch.

Reviewer comment (Theory):
I think as stated Theorem 3.1 is meaningless, as it doesn’t mention the memory tradeoff or even the algorithm used. So the theorem follows vacuously from noting that O(nr(b+d) >= O(nrd) and by Choromanski et al. we can compute mask(AB^T) C in O(nrb) time.

Author response:
We wanted to convey that there is a block-based algorithm exclusively operating on blocks with b rows and computes LT($AB^T$)$C$ using those numbers of floating point operations and which needs $O(n/b)$ intermediate sequential reads/writes. We will update the theorem statement.

Reviewer comment (Experiments):
In the end the tensor sketch is only actually applied for a degree 2 polynomial. Is it really the best option there? Why not explicitly compute Q^{tensor 2}, K^{tensor 2}. SVD them to get lower dimensional approximations tilde Q and tilde K and then compute tilde Q^{tensor 2} and tilde K^{tensor 2}. You are already explicitly tensoring up once in your last step anyways.

Author response:
In the model we use in our experiments, which is based on GPT-2 small, the $Q$ and $K$ matrices have 64 columns and in the current state of the art models (such as Palm-2), generally the $Q$, $K$ matrices have 128 columns. Explicitly tensoring these matrices would create matrices $Q’$ / $K’$ with 4096 (in 64 case) and 16384 (in 128 case) columns which is very slow to compute on GPUs and requires a large amount of memory as well. Our sketching technique will let us approximate attention even for larger head sizes such as 128 or more without such quadratic increase in the time/memory requirements.

Reviewer comment (Theory):
Related to the above, what is the final runtime complexity of your algorithm? The naive tensoring method is O(n h^5). Yours is what in terms of n, h, and the sketch size r? Roughly O(n(h+r^2))?

Author response:
Yes, the final complexity of sketch + tensoring would be $O(n(h \log h) + nr^2)$ using a fast Hadamard product implementation.