PolySketchFormer: Fast Transformers via Sketches for Polynomial Kernels

• “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.

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

• 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$.

• 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.

• 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.

• "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.

• 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.

• "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.

• "Fast Conditional Generation": Since the KV cache for linear transformers essentially reduces to storing a single matrix with size independent of context length as opposed to in the vanilla softmax models which need to store key, value pairs for all the tokens in the context to perform inference, the linear transformers are significantly less memory hungry during inference. The focus of our paper was to mainly improve training latency, so we did not include the improved memory requirements/inference latency for our models.

• 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.

• 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.

• 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.

• "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.

• "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.

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

• “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.

• 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.

• 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.

• Rematerialization: Rematerialization is used in Our model to recompute the output of the polysketch attention mechanism. Recomputation is also used in the Performer implementation and FlashAttention implementation.

• 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.

• Speedup is Not Large: We note that we do obtain significant speedups and the speedups improve as we increase the context lengths. For example, we report that on a 4 layer model with a context length 32k, we obtain a 4.5x reduction in training latency compared to FlashAttention and on the main model we study (similar parameters to GPT2-small), we get a 2x reduction. Thus, we can train our models with twice as many iterations in the same amount of time as we can train a Flash Attention model.

• 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”.

• Novelty: We do agree that the use of tensorsketch and other polynomial sketches is not new. But we believe our new finding that the lack of “non-negativity” for the polynomial sketches is a hindrance to model convergence is an important point. And our simple fix to this using tensoring to ensure non-negativity is also we believe an important contribution and a useful technique in general.

While the block-based algorithm to compute LT($AB^T$)$C$ is a simple modification of the existing algorithm, we want to make an important point that the slowness of kernel based linear transformers is not intrinsic and just changing the algorithm to ours can make it significantly faster in practice. The slowness of the cumulative sum algorithm was cited as the most important reason in the “FLASH : Transformer Quality in Linear Time” paper to explore other types of models but as we show we can significantly improve the running time just using a block based algorithm.

• Intuition for why polynomial works: Yes, you are correct that the functions behave very differently as x goes to -infty. Our intuition is that the polynomial attention learns to make the key vectors, that it doesn’t want to put a probability mass on, close to orthogonal to the query vector and hence will assign a low probability under the polynomial distribution.

• Memory footprint of attention: 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.

• “RNN style Sequential Dependence”: 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.

• Missing Citations and Typos: 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.

• Double Sketch: 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.

• Theorem 3.1: 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.

• Explicitly Computing Q^{tensor 2} and K^{tensor 2}: 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.

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

Thanks for going through the rebuttal!

• Base Sketch: In our implementations for the paper, we used the base sketch. After submitting the paper, we realized that applying the base sketch was not necessary and can be removed. We want to note that including/removing the base sketch did not affect the quality of the models in our experiments.

• Empirical Evaluation: Yes, perplexity as a single metric is not very useful in determining the quality of a model since models with small context lengths can also achieve good perplexities. In general, specific benchmarks such as LongBench (Bai et al., 2023), which includes tasks such as long document summarizing, document question-answering, token recall capability, are used to measure how good large state-of-the-art models are at handling long contexts. In our case, since we train the models from scratch on a limited amount of data, we do not expect these small models to be good at handling benchmarks in LongBench. Our main aim was to show, using perplexities as a proxy for model quality as has been done in previous efficient attention mechanism papers such as "FLASH : Transformer Quality in Linear Time", that polysketch attention is a promising mechanism to scale transformers to long contexts, in the sense that the perplexity trends are similar to the softmax attention transformers while being significantly faster to train.