Randomized and Deterministic Attention Sparsification Algorithms for Over-parameterized Feature Dimension

Large language models (LLMs) have shown their power in different areas. Attention computation, as an important subroutine of LLMs, has also attracted interests in theory. Recently the static computation and dynamic maintenance of attention matrix has been studied by [Alman and Song 2023] and [Brand, Song and Zhou 2023] from both algorithmic perspective and hardness perspective. In this work, we consider the sparsification of the attention problem. We make one simplification which is the logit matrix is symmetric. Let n denote the length of sentence, let d denote the embedding dimension. Given a matrix X ∈ ℝn × d, suppose d gg n and | X X^top |∞ < r with r ∈ (0,0.1), then we aim for finding Y ∈ ℝn × m (where mll d) such that beginalign* | D(Y)-1 exp( Y Y^top ) - D(X)-1 exp( X X^top) |∞ ≤ O(r) endalign* We provide two results for this problem. bullet Our first result is a randomized algorithm. It runs in widetildeO(nnz(X) + nω ) time, has 1-δ succeed probability, and chooses m = O(n log(n/δ)). Here nnz(X) denotes the number of non-zero entries in X. We use ω to denote the exponent of matrix multiplication. Currently ω ≈ 2.373. bullet Our second result is a deterministic algorithm. It runs in widetildeO(minsumi∈[d]nnz(Xi)2, dnω-1 + nω+1) time and chooses m = O(n). Here Xi denote the i-th column of matrix X. Our main findings have the following implication for applied LLMs task: for any super large feature dimension, we can reduce it down to the size nearly linear in length of sentence.