Fast attention mechanisms: a tale of parallelism

Thank you for your review. We appreciate all the feedback provided.

Novelty of the paper?: the ingredients – LSH-based attention and the correspondence transformers MPC – are known. The paper’s contribution is to tighten this correspondence from $O(N^2)$ to $O(N^{1+o(1)})$ machines and, through this, to prove that ANNA strictly subsumes low-rank attention. It would help if the authors highlighted more explicitly which lemmas are new, e.g. the almost-linear-machine simulation (Theorem 4.2) and the reduction from low-rank to ANNA (Theorem 4.4).

Thanks for your suggestion! We will highlight these connections more in the main results part of the paper. The ANNA simulation of MPC (Theorem 4.1) is new, because ANNA is a more restrictive and different model from transformer, and thus the correspondence between transformer and MPC does not apply to ANNA. Put differently, we propose ANNA not as an approximator to standard transformers (implemented via, say, LSH), but rather as an alternative, more efficient model in itself (and hence the model to be trained to start with).

The authors should add a limitation section: the experiments distill from a soft-max model; the paper does not yet demonstrate an end-to-end differentiable training scheme for ANNA, although such schemes exist for related LSH attentions (Reformer, Routing Transformer).

Thank you for the suggestion; we will add a limitation section discussing these issues. Note that Reformer and Routing Transformer are also not differentiable, so they also do not have end-to-end differentiable training schemes. (Cf. "Differentiable Approximations for Distance Queries" by Abdelkader and Mount, which gives a scheme for approximate differentiable distance query computations that is efficient in low-dimensional spaces.) But it is a beautiful open question!

Memory & constant factors: in line 230, the authors say that the runtime is sub-quadratic in $O(N^{1+3\rho})$. With c=10 (which is already pretty aggressive), we get $\rho=0.01$, so the exponent is 1.03 and $\ell=N^{0.03}\log N$. For N=16 000 that is 64 hash tables per layer – still a noticeable overhead, though far from the “900 tables” suggested by a literal reading of the bound. A discussion of how these constants affect wall-clock speed and memory would be valuable.

Our theoretical analysis is not fine-grained enough to pin down the constants suppressed by the big-O notation, so wall-clock time and the precise number of hash tables is beyond the scope of our paper. Please also see our reply to Reviewer XSKf regarding (lack of) memory overhead.

We'd love to discuss more if you have any further questions.

Thank you for your review, and we appreciate the suggested change.

On the weak side, as in prior work [51,53], this paper assumes that the element-wise operations (Q, K, V, ψ) can be arbitrary functions. This deviates from the typical transformer architectures that are used in practice.

Such assumptions are necessary to establish an equivalence to MPC, since the MPC model also allows arbitrary computation by each machine on their local memory content. This is a reasonable abstraction because our goal is to understand the capabilities/limitations due to the attention (or attention-like) mechanism. (That said, many concrete MPC algorithms do have a simple local algorithm.)

Line 355-356, I think there is an error. In this example, sigma(w, 9) = 7, so sigma(w, sigma(w,9)) = sigma(w, 7) = 6, and hence w_{sigma(w, sigma(w,9))} = a, not b.

Thank you for catching this error! We will fix the example in the final version of the paper.

We'd love to clarify more you have any further questions.

Thank you for the review. We appreciate all the feedback and questions provided.

W2 The paper does not discuss in practical terms the implications of important factors for the architecture such as memory overheads of locality sensitive hashing tables, failure probabilities, and trade-offs involving the approximation factor c for ANNA.

Q1 Could you discuss the practical implications of key architectural parameters in ANNA, such as the memory overhead of locality sensitive hashing tables, the failure probability, and the trade-offs involved in selecting the approximation factor c?

Memory overhead: In fact, Algorithm 1 can be implemented using linear memory ($O(N)$) with the same time complexity. Instead of maintaining $\ell$ hash tables, one can just store 1 hash table of size $O(N)$ with each entry responsible for tracking the values for each query. For each round of hashing ($\ell$ rounds in total), hash all queries using the hash functions and creates empty buckets for them. Then, hash each key, and if the key hashes to an existing query bucket, its value is added (along with a count). After processing keys, each query accumulates the values and counts from its corresponding bucket. We will add a remark about the memory complexity and this linear memory algorithm to the final version of the paper.

Failure probability: The failure probability decays at a polynomial rate with respect to the context length $N$. Context length $N$ is our main scaling parameter for our asymptotic analysis, so the failure probability tends to zero rapidly with $N$. (Note that $N=128000$ for GPT-4 and $N=1000000$ for Gemini.)

Trade-offs for selecting approximation factor: whether there are any trade-offs for selecting $c$ is highly problem-dependent. For example, for the Match2 and k-hop tasks in this paper, even a very large $c$ (e.g., $c=1000$) works based on the theoretical results in this paper. However, for problems that are conjectured to require quadratic time, such as Orthogonal Vectors, even $c=1.1$ is probably bad.

W1 The empirical evaluation is limited, and it mostly serves as proof-of-concept. Performance evaluation on real-world benchmarks or large-scale datasets is not provided.

Q2 Do you have any empirical results on real-world benchmarks or large-scale datasets to support the practical effectiveness of ANNA beyond the current proof-of-concept experiments?

The primary goal of our paper is to make theoretical contributions to the formal understanding of transformers and related architectures, and, as a result, we did not perform empirical studies on real-world benchmarks or large-scale datasets.

We'd love to discuss more if you have any further questions.