Fast Attention Requires Bounded Entries

In modern machine learning, inner product attention computation is a fundamental task for training large language models such as Transformer, GPT-1, BERT, GPT-2, GPT-3 and ChatGPT. Formally, in this problem, one is given as input three matrices Q, K, V ∈ [-B,B]n × d, and the goal is to construct the matrix Att(Q,K,V) := diag(A bf 1n)-1 A V ∈ ℝn × d, where A = exp(QK^top/d) is the `attention matrix', and exp is applied entry-wise. Straightforward methods for this problem explicitly compute the n × n attention matrix A, and hence require time Omega(n2) even when d = no(1) is small. In this paper, we investigate whether faster algorithms are possible by implicitly making use of the matrix A. We present two results, showing that there is a sharp transition at B = Theta(√log n). bullet If d = O(log n) and B = o(√log n), there is an n1+o(1) time algorithm to approximate Att(Q,K,V) up to 1/poly(n) additive error. bullet If d = O(log n) and B = Theta (√log n), assuming the Strong Exponential Time Hypothesis from fine-grained complexity theory, it is impossible to approximate Att(Q,K,V) up to 1/poly(n) additive error in truly subquadratic time n2 - Omega(1). This gives a theoretical explanation for the phenomenon observed in practice that attention computation is much more efficient when the input matrices have smaller entries.