Improved Algorithms for Kernel Matrix-Vector Multiplication Under Sparsity Assumptions

Thanks a lot for the review.

->It may be worth emphasizing the new assumption in the title, since the title suggests an unconditionally faster algorithm for kernel matrix-vector multiplication, which is somewhat misleading.

Done. The new title is “Improved Algorithms for Kernel Matrix-Vector Multiplication Under Sparsity Assumptions”.

->How does the running time of the algorithm scale with the constant $c$ made in assumption A?

The running time scales linearly with $c$, i.e., $\widetilde O((1+c)\cdot n^{1.89}/\varepsilon^2)$. It comes into play in the following parts of the analysis where the corresponding runtime are proportional to $\|K\|_1$ which is bounded by $O(n+cn)= O((1+c)n)$:

1. In FindHeavy (Lemma 4.4), it is used in bounding the number of hash collisions per hash table across all queries (line 708-711), which governs the “false positive” overhead of identifying the heavy keys $\{S_i:i\in[n]\}$. The expected number of collisions is $\widetilde O((1+c)n^{1+\alpha+o(1)})$ per table, and since we have $T=\widetilde O(n^{\alpha})$ tables, we enumerate in total over $\widetilde O((1+c)n^{1+2\alpha+o(1)}$ collisions. Lines 708-711 show the total number of collisions/runtime described in this point is proportional to $\|K\|_1\leq O((1+c)n)$
2. In ApproxLight (Lemma 4.5), it is used in bounding the number of averaging repetitions of the estimator (line 744-747), by $O((1+c)n^{1-\alpha})$. Hence the total number of samples for the estimator is $\widetilde O((1+c)n^{2+\gamma-\alpha}/\varepsilon^2)$ (line 662, eq. (2)). Line 744-747 show the number of averaging repetitions of the estimator described in this point is proportional to $\|K\|_1\leq O((1+c)n)$.
3. In ApproxLight, it is also used for bounding the number samples needed for estimating the number of repetitions (line 801-803), by $\widetilde O((1+c)n^{2+\gamma-\beta}/\varepsilon^2)$. Hence, the total number of samples in ApproxLight is like eq. (3) (line 806) times $(1+c)$. Line 801-803 show that the number of repetitions described in this point is proportional to $\|K\|_1\leq O((1+c)n)$. These are the terms that govern the running time.

->Theorem 1.1: you state the theorem as a constant probability statement, but is there a way to boost the probability of this procedure with logarithmic dependence on the failure probability?

This is indeed possible, by increasing the number of repetitions already built into the algorithm: given $\delta\in(0,1)$, set $T=10n^{\alpha}\log(n/\delta)$ in FindHeavy line 3, and take median over $10\log(n/\delta)$ averages in ApproxLight line 9 (the current settings of these parameters are the same without $\delta$). This boosts the probability of achieving the error guarantee to $1-\delta$, and the expected running time is $\log(1/\delta)$ times the currently stated running time.

->Line 151: should be $\lVert …\rVert_2^2$ rather than $\lVert …\rVert_2$ Done.

->Do your results imply improved results for the special case of KDE? power method?

Our results do not improve over KDE, as they rely on a KDE data structure as a black-box (the one from Charikar et al. 2020). They do, however, improve over the Kernel Noisy Power Method (KNPM) from Backurs et al. 2021, for the task of approximating the top eigenpair of the kernel matrix, in some cases.

The running time Backurs et al. 2021 provided for Gaussian KNPM is $O(n^{1.173+o(1)}/\varepsilon^{7.692+o(1)})$ (this is their Corollary 13, with $p=0.173+o(1)$ for the Gaussian kernel due to Charikar et al. 2020), although we believe their analysis in fact gives $O(n^{1.173+o(1)}/\varepsilon^{6.019+o(1)})$.

If we plug our KMV algorithm into KNPM, the running time we get is $O(n^{1.89}/(\varepsilon^5\lambda_1^2)))$, where $\lambda_1$ is the top eigenvalue of the kernel matrix, which is always between $1$ and $n$. The algorithm achieves this running time without needing to know $\lambda_1$ (or any bound on it) in advance. This improves over Backurs et al. 2021 whenever $\lambda_1$ is sufficiently large (and the matrix satisfies our assumption A). Furthermore, the two algorithms can be combined into a single one whose running time is the minimum (up to constants) of the two algorithms.

->Maybe too much to ask for of the authors since this is largely a theoretical contribution, but it would have been more convincing if there were experiments which showed that this new algorithm could be used to speed up the KDEformer application of Zandieh et al. Again, that's ok though since this is a theoretical paper.

We are working on an implementation of the algorithm, building on the code base of KDEFormer. We hope to be able to present results of an empirical evaluation in the final version of the paper.

Thanks a lot for the review.

->In places, contributions seem overstated.

Thank you for the comments. We rephrased the paper to account for them.

->How practical is the algorithm?

Indeed, we are currently working on an implementation of the algorithm, building on the code base of KDEFormer. The details are non-trivial, but we hope to be able to include the results of an empirical evaluation in the final version of the paper.

->Do you really test assumption A? We have added additional experiments for the scaling behavior of the constant with $n$. Please refer to Section 4 and the Appendix.

->Another nice idea (which might be harder) would be proving whether your assumption holds under different data distributions (queries and keys that are uniform on a hypersphere, Gaussian etc.).

Our assumption does not hold for random keys and queries on the hypersphere of large enough dimension as a function of n (i.e. at least polylogarithmic in n), as these are uncorrelated with high probability and the entries of the matrix will be fairly uniform. It is indeed a good idea to add such a claim. We will work out the details and add a formal proof in the final version of the paper.

->I would consider bringing Alg. 3 or a related schematic to Sec. 1.3.

We tried to do that, but found it difficult to present the algorithm clearly without the accompanying text of Section 3. Thus, we decided to keep the presentation as is. However, we can revisit this issue if the reviewer thinks it is important.

->I would also swap the order of Secs 3 and 4

Done.

->What is the space complexity of your algorithm?

The space complexity is $\widetilde O(dn^{1.782}/\varepsilon^{0.346})$. It is dominated by the KDE data structures created in ApproxLight, whose total space complexity is proportional to their construction time, $\widetilde O(dn^{1.346+2\beta}/\varepsilon^{0.346})$, and $\beta=0.218$ (see lines 716-717). Another source of space usage is the hash tables created in FindHeavy, whose space complexity is $dnT=\widetilde O(dn^{1+\alpha})$, and $\alpha=1/3$, so this is dominated by the KDE data structures.

->Lemma 3.1 is applied to the attention matrix to obtain a related Gaussian matrix. Why not simply use $K_{S.M.} = \text{diag}(\exp(q_i^2/2)) K_{G} \text{diag}(\exp(k_j^2/2))$? Diagonal matrix multiplication is trivially linear. Your guarantees might need to be tweaked slightly. Would this be simpler, or are there good reasons to prefer the current formulation?

Thanks for the great suggestion. This suggested approach would also work. However, let us explain why Lemma 3.1 is preferable. In short, it leads to a Gaussian kernel matrix with smaller $\ell_1$ norm, which improves the runtime of our algorithm (the error bounds are the same for both approaches). We provide more details below.

First, our reduction also uses diagonal matrix multiplication after Gaussian kernel mat vec: let $A$ be the unnormalized attention matrix and $x$ be the vector for which we want to compute $Ax$. Lemma 3.1 produces a diagonal matrix $S$ with its $i^{th}$ diagonal entry $e^{|q_i|^2/2\sqrt{d}} \cdot e^{\max_{j\in [n]}|k_j|^2/2\sqrt{d}}$ and Gaussian kernel matrix $K$ such that $A=SK$. In the reduction you propose, let us denote by $K_G$ the Gaussian kernel matrix $K_G$ and by $D_1,D_2$ the diagonal matrices whose $i^{th}$ entries contain $e^{|q_i|^2/(2\sqrt{d})}$ and $e^{|k_i|^2/(2\sqrt{d})}$ respectively, such that $A = D_1 K_G D_2$.

Our algorithm’s running time improves (linearly) with the $\ell_1$ of the matrix. This can be seen in response to reviewer uqhw about the runtime dependence on $c$: in that response, $(1+c)$ serves as an upper bound on the $\tfrac1n\|K\|_1$. We will observe that $\|K\|_1\leq \|K_G\|_1$, and therefore, the reduction in Lemma 3.1 promotes faster runtime for our algorithm.

By the fact that $A = SK = D_1 K_G D_2$, $A_{i,j} = e^{|q_i|^2/2\sqrt{d}} \cdot e^{\max_{j\in [n]}|k_j|^2/2\sqrt{d}} \cdot K_{i,j} = e^{|q_i|^2/2\sqrt{d}} \cdot e^{|k_j|^2/2\sqrt{d}} (K_G)_{i,j}$ for all $i,j\in [n]$. Now since $e^{\max |k_j|_2^2/(2\sqrt(d))}$ $\geq e^{|k_j|_2^2/(2\sqrt{d})}$, we have that the $i,j$ entry of $K$ is smaller than the $i,j$ entry of $K_G$ for all $i,j$, thus $\|K\|_1 \leq \|K_G\|_1$.

At the same time, both approaches lead to the same error for matrix vector multiplication with $A$. We have added this argument in another comment below.

->Presumably one can trade off the time complexity and probability that the estimate is accurate, changing the exponent on $n$ and the value of $p$. Can the authors comment more on this?

The probability of failure can be reduced to $p$ at the expense of a log(1/p) increase in runtime, which does not affect the exponent of n as long as p is not exponentially small in n. We will make this dependence explicit in the final version of the paper.

You don't need to implement with KDEFormer to provide some wall-clock times, or even just number of FLOPs, for your algorithm on real or synthetic queries, keys and values. This is a much simpler experiment and should be straightforward if you've already implemented the algorithm, no?

Thanks a lot for the review.

->However, the experiment in this paper only includes a pre-trained BERT model, whereas the introduction claims that, “we show empirically that our modeling assumption holds for kernel matrices that arise in modern transformer-based language models.” Beyond the BERT model, I am uncertain whether this assumption would apply to other transformer-based large language models (LLMs) as well.

Thanks a lot for the comment. We have added additional experiments for RoBERTA and GPT models. Please refer to Section 4 and the Appendix.

->While the experimental results support, to some extent, that the assumption made in this paper holds in practice, there are no experimental results verifying the improvement in running time. Since the main theoretical contribution is this improvement in running time, it would be preferable if the paper included experimental evidence to support it.

Indeed, we are currently working on an implementation of the algorithm, building on the code base of KDEFormer. The details are non-trivial, but we hope to be able to include the results of an empirical evaluation in the final version of the paper.

->In Zandieh et al. (2023), their results (Theorem 3.5) hold for any arbitrary $Q, K, V \in \mathbb{R}^{n \times d}$ to approximate the attention computation, and it does not seem that their Theorem 3.5 rely on the non-negative constraint. How does the result of this paper help with the attention approximation? My concern here is that if existing work already provides a nearly linear time algorithm to approximate attention computation, why is there a need to develop kernel density estimation for attention computation in a less efficient running time?

The algorithm of Alman and Song (2023) crucially uses the assumption $d=O(\log n)$. This is very restrictive as there is no dimensionality reduction primitive that works in this setting – higher dimension really is important. Our result, on the other hand, applies in a much more general setting than Alman & Song (2023), namely as long as $d=n^{o(1)}$. We added a remark to that effect in the paper.

The work of Zandieh at al. (2023) only gives a fast algorithm under a stable rank assumption on the attention matrix. This assumption is incomparable to ours. In particular, near permutation matrices, which are often observed in attention computation, have close-to-full rank. Such matrices are covered by our assumption.

->Alman & Song (2023) was published in NeurIPS 2023, but in the paper, the authors cite it as Alman & Song (2024). We used Google Scholar Latex citations feature, which apparently set the year to 2024. Fixed now.

Thanks a lot for the review.

->The core techniques such as finding heavy keys using LSH and random sampling, are slight adaptations of the ones developed in [Charikar et al. FOCS20] (or earlier).

Thank you for the comments. We would like to point out that our paper introduces several new ideas to the space of kernel density estimation based primitives for kernel matrix vector multiplication. The central ones are perhaps the idea of adaptively allocating sampling budgets to the n rows of the matrix, and the observation that the actual budget calculation is a kernel density like problem itself. This allows our algorithm to go beyond what is achievable by n independent invocations of kernel density estimation on the rows.

->Several fine-grained lower bounds for either the KMV problem or the Attention subroutine are known [Backurs et al. NeurIPS2017, Alman-Song NeurIPS2024, Alman-Guan CCC2024]. Showing how the new algorithm manages or fails to get around the known lower bounds might shed more light on the power and usefulness of the new subclass of kernel matrices.

Our results are for the low-precision regime, i.e. we have $poly(1/\epsilon)$ dependence on $\epsilon$ in our running time. The lower bounds you mentioned are against high precision algorithms, i.e. algorithms achieving a $poly(\log(1/\epsilon))$ dependence. To the best of our knowledge, there are no lower bounds known for KMV or the attention subroutine in the low-precision regime.

->A few other subclasses of kernel matrices were proposed in recent years that correspond to easier KMV problems. One example is kernel matrices with stable rank constraints. How are those matrices compared to ones with weight distribution studied in this paper?

Our assumption on the kernel matrix is incomparable with stable rank constraints. Near permutation matrices, which are often observed in attention computation, have close-to-full rank, but are covered by our assumption. On the other hand, the “all ones” matrix has low stable rank (rank 1), but does not fit our model

->Several variants of KDE problems with different levels of hardness are mentioned in this paper – KMV with vector entries being (1) all 1, (2) all positive or (3) arbitrary. For the relative-error setting, (2) is strictly easier than (3). Is this true in the additive-error setting as well?

In the additive error setting, the settings (a) with the vector being non-negative vs (b) with the vector being arbitrary have equal difficulty. This is because if one has an algorithm that can handle the case when the vector $x$ is non-negative, then one can write $x$ as $x= x_{pos} - x_{neg}$ where both $x_{pos}$ and $x_{neg}$ are non-negative vectors, and contain the absolute values of the positive numbers and negative numbers in $x$ respectively. Then applying the algorithm to $x_{pos}$ and $x_{neg}$ separately and subtracting the result gives an approximation to $Kx$ in $\ell_2$ norm with an error $\epsilon$ (|x_{pos}|2 + |x{neg}|_2) $\leq 2 \epsilon |x|_2$. Thus adjusting $\epsilon$ by a factor 2 gives us an output satisfying the guarantees we want for arbitrary vector $x$.