Summary
The attention mechanism, whose runtime complexity is approximately quadratic in the sequence length $n$, is the main bottleneck to LLM inference. It is thus quite reasonable to study whether this can be *approximated* with sufficient accuracy in subquadratic time. In particular, for this paper the authors define and study the problem of approximate attention computation (AAttC), wherein one seeks a result that is accurate within a pre-specified entrywise error tolerance $\epsilon_a$.
Under what I think are reasonable assumptions, the authors find that whether AAttC can run in subquadratic time hinges on the bound $B$ on the magnitudes of the entries of $Q$/$K$/$V$. This is slightly surprising and, in my opinion, a neat story for the paper. This conclusion consists in two results:
1. Theorem 3.8: an algorithm that computes AAttC in near-linear time provided that $B = o(\sqrt{\log n})$. The high-level idea is
- Approximate the $\exp$ in $\exp(QK^T/d)$ with a polynomial $P$ (Lemma 2.1 from [AA22]).
- Compute a low-rank decomposition $U_1 U_2^T = P(QK^T/d)$ directly in linear time based on knowledge of $P$ and the entries of $Q$ and $K$ (Lemma 3.2).
These lemmas are the main workhorses of Theorem 3.8, but effort is also made stitching everything together. The crucial thread is that $B^2$ determines the degree of $P$, which in turn determines the rank of the decomposition $U_1 U_2^T$. $B$ thus needs to grow slowly in $n$ to make the algorithm subquadratic.
2. Theorem 4.6, which shows that if $B = \Theta(\log n)$, there is no subquadratic-time algorithm for AAttC (assuming the strong exponential time hypothesis, SETH). This result is based on a reduction of approximate nearest neighbor search to AAttC.
Strengths
- Although past work has produced some closely related results (ZHDK23, KWH23), the authors' focus on $B$, the entrywise bound, is original and insightful. Furthermore, the proof techniques in this paper bring up interesting connections to the KDE literature which may be interesting for the community.
- They provide a near-linear time attention algorithm based on $B = o(\sqrt{\log n})$, and show that the rate at which $B$ scales cannot be improved upon. In some sense their work is a complete description of the problem, conditional on their (reasonable) assumptions on how other quantities scale with $n$ ($d = O(\log n), \epsilon_a = 1/\text{poly}(n)$). This is quite satisfying.
- For a theoretical paper, this is well-written and easy to understand. The authors do a great job of motivating, roadmapping, and explaining their work.
- This work offers a key insight about an important problem.
Weaknesses
This is a pair of complexity results. It is not clear to me how the insights provided here can be operationalized to improve AAttC in a practical setting. (The authors do highlight practical examples in ZBIW19 and KVPF20, but the connections between those works and the theory developed here do not seem particularly deep).
Typos
- L117: Hedge -> Hegde
- L133: algoriths -> algorithms
- L134: I believe should read "To design our algorithm for *Theorem 1.4*" not Theorem 1.3.
- L242: entires -> entries
- L232, L257, Algorithm 1 L4: I believe $B$ in $\frac{\log(1/\epsilon)}{\log(\log(1/\epsilon)/B)}$ was meant to be $B^2$
- Appendix L434: roamdap -> roadmap
- Appendix L453: $v_i$ -> $u_i$ under the product sign
Questions
I am curious to hear how you think your results might lead to practical speedups in attention computation for future work.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Limitations
The authors do highlight their main limitation: Theorem 4.6's reliance on the strong exponential time hypothesis (SETH). As far as I can tell, no important limitations were omitted.