Kshiteej Sheth

ICLR Conference 2025 Conference Paper

Improved Algorithms for Kernel Matrix-Vector Multiplication Under Sparsity Assumptions

• Piotr Indyk
• Michael Kapralov
• Kshiteej Sheth
• Tal Wagner

Motivated by the problem of fast processing of attention matrices, we study fast algorithms for computing matrix-vector products for asymmetric Gaussian Kernel matrices $K\in \mathbb{R}^{n\times n}$. $K$'s columns are indexed by a set of $n$ keys $k_1,k_2\ldots, k_n\in \mathbb{R}^d$, rows by a set of $n$ queries $q_1,q_2,\ldots,q_n\in \mathbb{R}^d $, and its $i,j$ entry is $K_{ij} = e^{-\|q_i-k_j\|_2^2/2\sigma^2}$ for some bandwidth parameter $\sigma>0$. Given a vector $x\in \mathbb{R}^n$ and error parameter $\epsilon>0$, our task is to output a $y\in \mathbb{R}^n$ such that $\|Kx-y\|_2\leq \epsilon \|x\|_2$ in time subquadratic in $n$ and linear in $d$. Our algorithms rely on the following modelling assumption about the matrices $K$: the sum of the entries of $K$ scales linearly in $n$, as opposed to worst case quadratic growth. We validate this assumption experimentally, for Gaussian kernel matrices encountered in various settings such as fast attention computation in LLMs. Under this assumption, we obtain the first subquadratic time algorithm for kernel matrix-vector multiplication for unrestricted vectors.

SODA Conference 2024 Conference Paper

Sublinear Time Low-Rank Approximation of Toeplitz Matrices

• Cameron Musco
• Kshiteej Sheth

SODA Conference 2023 Conference Paper

Toeplitz Low-Rank Approximation with Sublinear Query Complexity

• Michael Kapralov
• Hannah Lawrence
• Mikhail Makarov
• Cameron Musco
• Kshiteej Sheth

We present a sublinear query algorithm for outputting a near-optimal low-rank approximation to any positive semidefinite Toeplitz matrix T ∈ ℝ d×d. In particular, for any integer rank k ≤ d and ε, δ > 0, our algorithm makes Õ ( k 2 · log(1/δ) · poly(1/ε)) queries to the entries of T and outputs a rank Õ ( k · log(1/δ)/ε) matrix d×d such that || T – || F ≤ (1 + ε) · || T - T k || F + δ|| Τ || F. Here, || · || F is the Frobenius norm and T k is the optimal rank- k approximation to T, given by projection onto its top k eigenvectors. Õ (·) hides polylog( d ) factors. Our algorithm is structure-preserving, in that the approximation is also Toeplitz. A key technical contribution is a proof that any positive semidefinite Toeplitz matrix in fact has a near-optimal low-rank approximation which is itself Toeplitz. Surprisingly, this basic existence result was not previously known. Building on this result, along with the well-established off-grid Fourier structure of Toeplitz matrices [Cybenko'82], we show that Toeplitz with near optimal error can be recovered with a small number of random queries via a leverage-score-based off-grid sparse Fourier sampling scheme.