Matrix Product Sketching via Coordinated Sampling
Authors: Majid Daliri
Juliana Freire
Danrong Li
Christopher Musco
Venue: ICLR 2025 Poster (2025)
Keywords: Sketching Algorithm
Matrix Multiplication
Model Compression
Data Discovery
Efficient resource
Primary area: optimization

Abstract: We revisit the well-studied problem of approximating a matrix product, $\bv{A}^T\bv{B}$, based on small space sketches $\mathcal{S}(\bv{A})$ and  $\mathcal{S}(\bv{B})$ of $\bv{A} \in \R^{n \times d}$ and $\bv{B}\in \R^{n \times m}$. We are interested in the setting where the sketches must be computed independently of each other, except for the use of a shared random seed. We prove that, when $\bv{A}$ and $\bv{B}$ are sparse, methods based on \emph{coordinated random sampling} can outperform classical linear sketching approaches, like Johnson-Lindenstrauss Projection or CountSketch. For example, to obtain Frobenius norm error $\epsilon\|\bv{A}\|_F\|\bv{B}\|_F$, coordinated sampling requires sketches of size $O(s/\epsilon^2)$ when $\bv{A}$ and $\bv{B}$ have at most $s \leq d,m$ non-zeros per row. In contrast, linear sketching leads to sketches of size $O(d/\epsilon^2)$ and $O(m/\epsilon^2)$ for $\bv{A}$ and $\bv{B}$. We empirically evaluate our approach on two applications: 1) distributed linear regression in databases, a problem motivated by tasks like dataset discovery and augmentation, and 2) approximating attention matrices in transformer-based language models. In both cases, our sampling algorithms yield an order of magnitude improvement over linear sketching.

Published as a conference paper at ICLR 2025
MATRIX PRODUCT SKETCHING VIA COORDINATED SAMPLING
Majid Daliri', Juliana Freire!, Danrong Li", Christopher Museo
NNew York University
"Penssylvania State University
ABSTRACT
We revisit the well-studied problem of approximating matrix product, A ABB, based on small space sketches S(A) and (BB) of A E Rand and B E RR We are interested in the setting where the sketches must be computed independently of each other, except for the use of a shared random seed. We prove that, when A and B are sparse, methods based on coordinated random sampling can outperform classical linear sketching approaches, like
Projection or CountSketch. For example, to obtain Frobenius norm error < || [||[[[|r coor- dinated sampling requires sketches of size Ossfc²) when A and B have at most 8 S d, m non-zeros per row. In contrast, linear sketching leads to sketches of size (d/e²) and (mm/e² for A and B. We empirically evaluate two applications: 1) distributed linear regression in databases, a problem motivated by tasks like dataset discovery and augmentation, and 2) approximating atten- tion matrices in transformer- based language models. both cases, our sampling algorithms yield order of magnitude improvement over linear sketching.
approach on
INTRODUCTION
Over the past 20 years, sketching and sampling methods have emerged as powerful tools for solv- ing massive linear algebraic problems that arise in machine learning, data science, and scientific
computing (Woodruf) 2014 Drineas & Mahoney 2016 Martinsson & Tropp 2020). Matrix and
vector sketching has also been widely applied in federated learning (Rothchild et al. [2020 Konečný et al.) 2020), distributed learning (Jiang et al. 2018), and beyond (Cohen et al. 2015ap. One of the most fundamental problems where randomization has found success is matrix-matrix multiplication: are given A RRxd and B E Raxm and hope to compute an approximation to the product A E ddxm Naively, it takes O(drm) time to compute A exactly, or bit less if fast (rectan- gular) matrix multiplication methods are used (Le Gall, 2012). We are also interested in minimizing communication costs: when A and B are stored in separate machines, computing A A requires communicating at least O(d min m)) numbers (either all of A or all of B). We explore the question of whether this bound can be improved.
In particular, we consider solving matrix matrix multiplication in the widely studied sketching setting (Nelson & Nguyen] 2013).. The goal is to compute small space sketches S( A and S(B), from which the matrix product can be approximated using some routine F(S(A) (A) (B)) and S S(B) should be much smaller than the space required to store A and B - i.c. smaller than ((d)) and O(nm) space for dense matrices, or smaller than ((nz((A)))) A)) and O (mnz(B)) space when the matrices sparse with nnz denoting the number of non-zeros. This reduction in size allows the sketches to be processed, stored, and communicated with less cost than the original matrices.
) R A BB. Ideally, S(A)
We emphasize that, in the sketching setting, while S(A) and S(B) can be computed using shared source of random bits (e.g., for constructing hash functions), no additional communication is allowed between the processes computing S(A) and S(B). This restriction is crucial in settings where communication between processes is expensive, slow, or impossible, as is often the case in distributed systems, or in applications where A and B are processed at different times. Moreover, the restriction is critical in applications that require repeated matrix multiplications. For example, if we want to approximate A A B B,..., A B for set of matrices A A!.... ,.... A using the same sketch S(B), that

---

Published as a conference paper at ICLR 2025
sketch cannot be tailored to any one particular A'. We further discuss applications of matrix- productt sketching to central problems like dataset discovery and multi-vector retrieval Section
1.1
PRIOR WORK
The best existing sketching methods for matrix product approximation are based on linear sketches that compress A and B via multiplication by a random matrix generated with a shared random seed. In particular, below we state seminal result of Sarlós that was later refined and generalized (see Cohen et al. (2016) or Woodrulf (2014) for more detailed discussions).
Fact (Srróós 2006 Kane & Nelson 2014; Cohen et al. 2016).. Ler II E R be a scaled random Gaussian matrix, random sign matrix, CountSketch matrix (Churikar et al. 2002), or any of a variety of other randomized linear embeddings. If k- O (2+)) then with probability at least 1 1-8, 8,
- Sc|AA|
IIA R and IIB E Rkxm are sketches of size Old/28) and (mm/²5) respectively, For dense matrices, these sketches are smaller than and B. respectively, whenever 2 n.
The above result can be strengthened if additional assumptions are made on A and B. For example, tighter bound is possible if the matrices are low-rank or nearly low rank (Cohen et al.) 2016). However, if no additional assumptions are made on the matrices' spectra, then Fact the best result known.
1.2 SAMPLING BASED MATRIX PRODUCT APPROXIMATION
Our goal is to devise an alternative approach to matrix product sketching that improves on Fact I the important setting where A and B are sparse, providing better sketches even for dense matrices. To do so, we build on another classical approach for using randomization to speed up matrix-matrix multiplication: subsampling. A seminal paper by Drineas, Kannan, and Mahoney (Drineas & Kannan 2001 Drineas et al., (2066)) proves that it is possible to sample and reweight O(1/="0) rows of A and B so that the product of the subsampled matrices À and B satisfies A ATB S
||||[||||] F. with probability at least 8.
Notably, this approximation guarantee exactly matches Fact However, since À and B consist of subsample of rows from A and B, they can be much more compact to store than IIA and IIB. For example, if A and B have most non-zeros per row, Á and B take (ss/²) space to store, instead of Odd/e) and Omm/=)) respectively.
Importantly, however, the row subsamples Á and B guaranteed by Drineas et al. (2006a) cannot be computed in the sketching setting. The challenge is that, obtain theoretical accuracy bound, rows must be sampled with non-uniform probabilities. Intuitively. A 's 2thh row A, E Rd and B's 7th row B, E IRm have large magnitude, they contribute more the matrix product A AB, which can be written as sum rank-one outerproducts: ATB Σi= A, So, high- nigh-magnitude rows must be sampled with higher probability. The original analysis from Drineas et al. (2006a) suggests sampling with probabilities proportional to ~ ||A:||2]| B:Il2l where ||x||2 denotes the Euclidean norm of a vector x. While these probabilities were shown satisfy an optimal variance property, they cannot be computed without access to both A and B, which is impossible in the sketching setting, since the sketches S(A) and S(B) must be computed independently of each other. It can be shown that it also suffices to use probabilities proportional to either || or ~ ||B.l|2 (so, only taking one matrix into account),
the choice must be consistent, so the issue remains that either the machine sketching A or the machine sketching B does not know the right probabilities.
We emphasize that this issue cannot simply be resolved by communicating probabilities between the processes computing the matrix sketches (which would be relatively inexpensive). In partic- ular, the applications we consider in Section I.4 require computing all pairwise matrix products B Using single collection of sketches
between two sets of matrices A ...... A and B Even if probabilites could be communicated, it is not clear
S(A') ",....,.(AA.,,, S(AA),S(B)) S(B'), ..... S(B)). what choice of probabilities should be used to make all pairs of sketches compatible.
2

---

Published as a conference paper at ICLR 2025
1.3 OUR RESULTS
Our main contribution is to show that, despite the limitations above, sampling can in fact be applied effectively in the sketching setting by drawing on coorlinated random sampling methods. Such methods include MinHash (Broder et al. 1998; Manasse et al. 2010; loffe [offe]2010 (2100) the k-minimum values (KMV) sketch (Beyer et al. 2007), conditional random sampling (Li et al.) 2006), and coordinated variants of PPSWOR sampling (Cohen & Kaplan 2007, 2013). We refer the reader to the recent survey of Cohen (2023) for more complete review of prior work.
The idea behind coordinated sampling is to use shared random seed to draw samples on two different machines that are likely to contain significant number of shared indices, even if the exact same sampling probabilities are not used. In our setting, A rows will be sampled using probabilities proportional to ||A and B's rows with probabilities proportional to |B, Our samples ((A) and S(B) only likely to contain A, and B, for shared index i (which can then be used estimate A A via the sum A Σi=1 A, if both ||A, and |[B, 12 are large. Fortunately, it turns out that this suffices to prove a bound equivalent to Fact
Formally, we use a method for coordinated sampling without replacement called Priority Sampling OOhlsson 1998 Duffield et al. 2004 2007) / to prove the following main theoretical result:
Theorem 2 (Main Result). Consider A E gaxd, B Rax= and any e, å E (0, There is a
sketching procedure (Algorithm ( that constructs sketches S(A) and S(B) consisting of at most
k o + rows from A and B, and there is corresponding estimation procedure (Algorithm) 2
that, using the information in these sketches, returns an estimate W such that, with probability 1
A S < ||A||||| B||p.
Theorem 2 matches the guarantee of the sampling method from Drineas et al. (2006a) up a constant, albeit method computes S(A) and S(B) completely independently from each other. As such, we match the oftte-of-the-arr Fact guarantee for linear sketching in the worst case, and improve on it whenever A and B are sparse. For example, in our experiments in Section 4 we consider some applications involving matrices with only 2% of entries in each row non-zero. For these applications, linear sketch IIA with height k (
storing subsample of size & = O(15c²) translates to 50x savings in space.
Priority Sampling and related methods have recently been leveraged to give sketching procedures for estimating inner products, which improve on linear sketching methods like
projection for that problem (Bessa et al. 2023 Daliri et al. 2024a b). Inner products represent special case of matrix products when d = m = I. Our work significantly extends these results by addressing general matrix multiplication. In doing so, we encounter several technical challenges, in- cluding the fact that Priority Sampling beetinng-beinngg without-replacemmnt sampling procedure- generates non-i.i.d. row samples from and B. We address these challenges in Section 2
1.4 EXAMPLE APPLICATIONS
As mentioned, sketching methods are most useful in distributed computing environments, or in settings where we wish to compute many pairs of matrix-matrix products from fixed collection of sketches
Multi-vector Retrieval. One such setting arise the "vector set search" or "multi- retrieval" problem, which has received recent attention (Engels et al. 2023 Dhulipala et al., 2024). This problem generalizes standard vector similarity search to matrices: we have a database of matrices A !!..... A (each representing a document or media item) and another matrix B that represents a query. These matrices can be viewed as collections of column vectors, hence the name "vector set scarch". The goal is to find min, d(A' B), where is some distance function that depends on the matrix matrix-matrix product A B. Approximate matrix maatix-multiplicationn can used to speed up the computation of A " B and thus, the distance computation. For example, Johnson-Lindeenstraus sketch is used in (Dhulipala et al.) 2024). Here, the sketching setting is key: A A!...., A are preprocessed into sketches that are computed before the query B is issued and S(B) cannot be chosen to depend on any particular A' or S(A'), )), as it will be used to estimate B B's matrix product with all g matrices A ......A A².
3

---

Published as a conference paper at ICLR 2025
Regression-based Dataset Search. Another motivation of our work is to develop efficient methods for dataset search and discovery, a problem that has received significant interest in recent years Chepurko et al. x Castelo et al. [202T) Liu al. (2022) Lonescu et al. (2022).
In particular, suppose we have a data lake consisting of many datasets A ,...,A. A and we want to support queries where a user provides data vector b and the system returns all candidate datasets A' that are predictive of b. Lee., for which minx || bll2 is small. The sketching setting can be used to support such queries efficiently: b is sketched by the user and is sent to the dataset search system. It is then compared to precomputed sketches for each of A , A to approximate minx | bll2-
It turns out that this problem can be solved with modified version of our matrix- -rroucct sketching method. Concretely, restricting our attention
single A R and vector b E R", our goal is to find X which is a near minimizer of the standard least squares problem: minx || || bll The optimal x has the form x" = (A A) A) A b,, where A A A is relatively small, d matrix. So, the challenge is approximating the matrix-vector product A b using compact sketches
Approximating A b directly using Theorem 2 does not suffice, as to ensure an accurate X, we need small error with respect to different norm than the standard Frobenius norm. Instead, we introduce a variant of Algorithm| that collects row samples from A based on the matrix's statistical leverage scores. Entries from b are sampled based on their squared magnitude. Our main result is follows:
Theorem 3 (Sketched Regression). There is procedure that constructs sketches S(A) and S(b) consisting of ((d/c) row samples from A E Tuaxd and b E RR such that, using only the information in those sketches, we can compute x E TRd satisfying, with probability at least 99/ 19010
| bl2 < Ax bll2 b
Sketching algorithms for regression have been studied in prior work. For the problem above, the best existing result is based on linear sketching (e.g.. Boghnson-Lindenstenn
projection). Linear sketching methods achieve the same guarantee as Theorem 3 with a sketch of size O(d /E) for A and sketch
of size (dd/e) for (specifically, the d x (d/e ( matrix IIA and the vector IIb) Sarlós (2006); Woodrufi (2014) Theorem 3 improves these bounds when A is sparse. In particular, if A has
8 5 d non-zeros per row, we require a sketch of size (sd/cc) for A and of size ((d/c) for b
Theorem is proven in Section We remark that leverage score sampling has already been widely applied to regression problems outside of the distributed sketching setting (Drineas et al.) 20066] Cohen et al. 20155] Chen & Price 2019). It is well known that, if the rows of A and b are sampled with probability proportional to the leverage scores of A, then a guarantee matching Theorem holds as long as ((d/e) samples are taken Sarlós (2006). However, standard leverage score sampling cannot be applied in our sketching setting since A's leverage scores cannot be used when subsampling b. Again, this not an issue with simply needing to communicate the scores. We want a sketch of b that is compatible with each matrix in a collection A !!...., A", which might have very different leverage score distributions. This challenge necessitates both new algorithm and new analysis.
Our work builds on recent line of work that uses sketching methods for efficient dataset search in general. For example, sketching methods for estimating inner products have been applied to finding individual columns in a datalake that are highly correlated with given query vector b (Santos et al. 2021 2022; Daliri et al. 2024b). Our sketching methods for regression allow for more advanced search queries that go beyond pairwise correlation.
1.5 NOTATION AND PRELIMINARIES
Before proceeding, we briefly review notation used throughout the paper.
Linear Algebra Notation. The ith row of matrix A is denoted by A4, The entry in the 7thh row and 5th column of A is denoted by Aijj For a vector x; denotes the th entry. We use ||x||2 to denote
'Linear sketching methods actually ensure a stronger guarantee: the additive error el|b|la can be replaced with the residual <|| Ax" bll2. which is always smaller. In some applications, the difference is not significant. For example. in dataset search, most matrices A will be unrelated to b, so we expect ||Ax* - bbl2 ~ |b]l Additive error e| e||b]] should suffice to at least rule out bad candidates. However nice question for future work understand if more compact sketches can be obtained when targeting the stronger residual error guarantee.
4

---

Published as a conference paper at ICLR 2025
the standard Euclidean norm of vector x and ||]|| to denote the Frobenius norm of of matrix A. The transpose of a matrix A is denoted by A and the inverse of is denoted by A -1 provided exists. The zero vector is denoted by O, with dimension clear from context.
Other Notation. The expected value of a random variable X is denoted by B[ and its variance is denoted by Var(X). We use the notation [ni] to represent the set {1, (1....,n.. n.).
2 MATRIX PRODUCT SKETCHING WITH PRIORITY SAMPLING
Our main approach to matrix product sketching is based on subsampling. We can write any matrix product A ABB sum of outer-products A B Σi A,B." We will estimate this sum as Liet W; A,B, where T is small subset of (1,.... n and w; is appropriately chosen weight. Typically 'I is selected via importance sampling: indices that correspond to larger norm rows in A or B are sampled with higher probability (Drineas et al., 2006a). The challenge in the sketching setting is that A and B must be sampled independently from each other, without knowledge of the other matrices row norms.
We address this issue by using coordinated sampling technique known as Priority Sampling, which has been widely used for subsampling data streams Duffield et al. (2007), and more recently for subsampling vectors for inner product estimation (Dalini et al. [2024b). Pseudocode for the method included in Algorithm|] [ To give better intuition for the method, we informally describe another closely related algorithm called Threshold Sampling, which gives the same guarantees as Priority Sampling for our problem, but has the disadvantage of producing a sketch whose size can only be bounded in expectation.
Threshold Sampling works as follows: 1) using shared random bits, we select a random hash function h: {1.... n} [0, 1] that assigns a uniformly random number between [0, 1] to any index iF 2) we collect in the sketch ((A) any row A, for which hii) I k- || l and in the sketch S(B) any row b; for which h(i) S k. |B BB]. Equivalently, A, sampled if the reweighted hash value h(t)/|| || falls below a fixed threshold A/| I/ (and likewise for B₂).
It is easy that each sketch contains i rows in expectation. Moreover, since hash function, it can be checked that, for any index i, we have that both A; E (AA) and B, E S(B) with probability:
shared
= min (1 / // k- /B,
Let T denote the set of indices that appear in both sketches. We return the unbiased estimate W -EeT p; A,B.. To show that this estimate is accurate, can follow an analysis similar to the original paper subsampled randomized matrix multiplication, which bounds the expected squared error TE| W - A BI before applying Markov's 's inequality (Drineas et al., 2006a). The only difference is that must show that it suffices to sample indices with probability proportional to the minimum of | | and ||B / | instead of the product of these numbers. Perhaps the fact that this suffices is intuitive: for the outerproduct A, to make a significant contribution to A A, neither A, nor B, can have small magnitude. A full analysis of Threshold Sampling is given in Appendix B
The method we propose, Priority Sampling, is almost identical to Threshold Sampling. However, instead of fixing the threshold k/ ||?, which leads to a random number of indices being sampled, we dynamically set the threshold to collect exactly samples. Doing so does not change the method in spirit, but complicates the analysis since samples no longer independent.
Nevertheless, drawing inspiration from a new, simple analysis of Priority Sampling for sampling numbers from stream (Daliri al. 2024a), we able prove the following bound:
Theorem 4. Let A Raxd B E Raxmm and let S(A) = {IA, VA,TA) and ((B (IB, VB, TB} be sketches produced by Algorithm 7 with input k and shared seed s. Suppose W is the approximate
'In practice, / can be substituted with pseudorandom function mapping to large discrete subset of (0, 1 For simplicity, we assume access to real-valued, perfect hash function, as is standard in the literature (Cormode al. 2011)
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 Priority Sampling
Input: Matrix A size n x d, random seed s, number of row samples, k.
Output: Sketch S(A) (IA, VA, TA), where ZA is subset of row indices from (1 (1...., n) and selected
VA contains A, for all E IA, 7A the threshold used to determine whether a row in ZA is
1: Use random seed select uniformly random hash function h {1...., [0, 1].
2: Initialize ZA and VA to be empty lists.
3: Compute rank R h(s) for all such that A, + O.
4: Set TA equal to the (k + 1)ah smallest value R2 or set TA = 0c if A has < k non-zero rows.
5: for such that A; * 0 do
6: if R₂ < TA then
7: Append ZA. append A, VA.
8: return S(A) (IA, TA
Algorithm 2 Approximate Matrix Multiplication
Input: Sketches S(A) {ZA. VA. TA). S(B) (IB, VB, TB} constructed by Algorithm Output: Estimate W for A
1: Compute T IA In. Note that for all i E T, VA and VB contain A; and B;
2: return
A,
W
Σ
min(1,
A 2 TA, B TB)
SET
matrix of A calculated using Algorithm 2 on these sketches. Then, [W] A and
2 E [ - A -- A
Additionally, [TA] S. and [T]] S k. Lee, each sketch contains no more than h rows from A and B, respectively. If each matrix has at least k non-zero we have that [IAl [nn k.
We prove Theorem via Lemma 5 which is proven in Appendix A due to space limitations.
Lemma 5. Let A, B, and W be as in Theorem
For any xr E [d] x [m]]
we have:
n2
A B
|
B [
E W z.y B) z. and E [W 2.y [A 2
k
|
Lemma 5 givee an entrywise guarantee the of the approximation W, which we can then use to give an overall bound on the Frobenius norm error. That analysis is given below.
Proof of Theorem 4 The entrywise expectation guarantee Lemma immediately gives E W ATB. We left bound the expected Frobenius error, which can be written sum entries:
A
B
||A
B
E W - A A
Σ E (W 2,9
[ D.Y.
+
k
A
B
2031
30,91
1=1
B
d
m
B
|
B
1,2
1.
y=1
|
B
|A
|
B.
Σ
|
B
2
BB
2 4
14
BIB
With Theorem in place, our main result, Theorem 2 follows immediate corollary:

---

Published as a conference paper at ICLR 2025
Proof of Theorem 2 Our main result, Theorem follows an immediate corollary of Theorem A -
4
In particular, if we set 2/4 + 1, then we have that E [[[ Applying Markov's inequality proves the theorem.
| B
SKETCHED REGRESSION
In this section, we focus on proving Theorem 3 To do so, we first prove a simpler version of the result that, instead of just storing subsample of rows from A in S(A), also explicitly stores the x d covariance matrix A T ". The pseudocode for this method included as Algorithm 3 It leads to sketch of size (d + ds/c) when A has non-zeros per row. We later show that the sketch can be modified to have size (ds//c) by replacing A T with another subsample of rows from A.
Proof of Theorem Our goal is to compute a vector x that approximates x" = (ATA) A In particular, we wish to obtain an upper bound || bl1.. Observing that Ax* - b is orthogonal to any vector in the column span of A, we can apply Pythagorean theorem to write:
| - bll ||Ax'
bl12
||A - Ax"1
or equivalently:
| bl2 - ||Ax - b12 ||A A(A
(1)
We claim that the sketching and regression procedure in Algorithm returns x such that A is exactly
equal F(S(A((( A ((b)), where S() denotes the sketching procedure of Algorithm 1 and FF-) denotes the estimation procedure of Algorithm 2 However, it does so in an implicit way. without every explicitly forming the large n and possibly dense matrix A(A
T A) A." If
immediately follows. In particular, by the guarantee of Theorem 2 as long as we choose sketch size (d/c), we would have:
this claim holds, then the main guarantee of Theorem 3
É | - A(A"A)-AA"DI S á [IA(A A) A 1
E d
d b = </bb|/
Above we have used that d. Plugging into equation
= tr(Ia)
! would then prove the theorem.
So. it is left to establish that Algorithm 3 returns x such that Ax is exactly equal to F(S(A(ATA) ")"",,",, S(b)). For this to be the case, it can be checked that it suffices to sim- ply sample from A with probabilities proportional to the squared row norms in A (A" A Then, multiplying the sampled rows by (AT AT) produce x, and again by A to produce Ax exactly reproduces F(S( F(S(A(AT A) AT) S(b)).
The Th squared row norm of A(A can be written ||A(A-A) A l| where e; denotes the 1th standard basis vector. We then have:
AA(A"A)
= A?
The quantity above is exactly equal to the hh statistical leverage score A, which is the quantity that Theorem 3 uses when Priority Sampling, so the claim holds. We note that, besides the naive approach, efficient algorithms are known for more quickly computing the statistical leverage scores of matrix A, although our focus here on sketch size opposed to construction time (Mahoney et al. 2012 Clarkson & Woodruff" 2013)
The approach above immediatly gives an O(d + ds/e) space sketch for least squares regression when A has s-sparse rows. This already improves on the space complexity of (dd² /c) achieved by linear sketching methods. However, in settings where d is large, it would be better to avoid quadratic dependence on d entirely. To do so, instead of explicitly storing the d x matrix A T A in our sketch, we can store SA, where S E Rexd is matrix that selects and reweights
Optimized Method.
z rows from A. Instead returning x = (A" A) I W as in Line 6 of Algorithm 3
we would return:
x (A S -'W.
7

---

Published as a conference paper at ICLR 2025
Algorithm 3 Sketching for Regression (not optimized)
Input:
Matrix Anxd. matrix b randomness seed s, and target error E.
1: Compute & as the leverage score of A: 4 (A"A)-
2: Construct sketches ((A) with target sampling size of O (rows) using Algorithm 1 with
shared seed s. However, compute the rank (line 2 of Algorithm I as R. A()
2,
3: Construet sketches S(b) with a target sampling size of O (1) (rows) using Algorithm with shared seed 8. Output: (S(A), A A). S(b)
Procedure REGRESSION ( (ATA,S(A) S(b)).
4: Compute Compute & as the leverage score of A: & A,(ATA)-1A? for any A, S(A).
5: Compute T = IAA nIp.
A.b. 6: Compute W JET min(L.ArT..bbm.
Output: & = (A"A)-'W
is well known that there exist choices of S with m = O(d/c) rows such that A is Ve-relative error spectral approximation A A'A Batson et al. 2012). Lee.
(1 -e) 3 A SA 3 (1 Vē) A A and
(1- - (A 3 (A"S'SA)-'
(
where denotes the Loewner order.
check that A lla I vc, where || 2l2 denotes the operator norm.
Given the second guarantee,
JA(A A) A
Now, observe that Ax' A(A"STSA) - 'A Aă, and thus:
|AX' A (A(AT AT A(A"A) AT S Ve|lAR|aa S 2
In the last step, we used that || blI2 < < b|bll to loosely upper bound || Finally, put everything together. As proven earlier, A Ax 12 S Vellb/l2- Applying triangle inequality, we
thus that that ||AX' Ax" 1|2 < bl| Applying Pythagorean theorem as before, we conclude
that AX bl3 & Axx b - be][b]| Adjusting e by a constant gives the desired result. O
We remark that way of producing a matrix S satisfying the spectral approximation guarantee above is to subsample and reweight m = O log d/e) rows from A by leverage scores. While worse by log ol factor than the deterministic methods given e.g. in Batson et al. (2012), the advantage of such an approach is that we can reuse the samples used to approximate A A when approximating A b, for which we also sample via leverage scores. This "single sketch" procedure is what implement in our experimental section.
0.10
Prionly Sampligg
0.10
Prionty Sampling
0.10
Prionly Sampling
Jl Sketch
Sketch
J. Sketch
Enor |iic
Threshold
Error L1re ira--foos dfil
0.06
Throshald
Eror ar
0.06
Threshold Samping
0.00s
1.,00
00.00
0.06
0,01
0.04
0.002
0.02
(a) 80% Sparsity.
(b) 40% Sparsity.
(c) 10% Sparsity.
Figure 1: Performance of matrix product sketching synthetic data with varying sparsity levels (10%, 40%, and 80%). Priority sampling and threshold sampling are depicted on top of each other and both methods outperform the JL sketch the level of sparsity increases.
4 EXPERIMENTS
We experimentally evaluate our sampling based matrix product sketches variety of settings. First, synthetic data to directly compare the method to the standard linear sketching approach:
basedd
8

---

Published as a conference paper at ICLR 2025
random project (which we denote JL Sketch in our plots). As predicted by our theory, the methods perform similarly for dense matrices, but our Priority Sampling method obtains more compact sketches for sparse matrices. We also apply real-world problems, including to regression tasks, as outlined in Section 3 We show that for two popular datasets, our method can outperform over the best existing linear sketch.
method
number of
Additionally, we apply our Priority Sampling method to approximating matrix multiplications that arise in transformer-basee large language models. Deploying auto regressive language models involves performing attention decoding in an online setting, where key and value embeddings from each transformer layer are cached to eliminate redundant computations. More precisely, during each token generation phase, the stream of tokens is encapsulated by three matrices known as the query (Q), key (K), and value (V) embeddings. At the heart of this process, each iteration involves calculating the attention matrix as Att dimension. Recent studies (Hooper et al. 2024; Zirui Liu al. 2023) have introduced methods that apply vector quantization to the key and value embeddings, replacing the full matrices with a quantized matrix. In this study, we employ Priority Sampling to sketch Q and K. We assess the performance of our approach in approximating QK" compared to linear sketching techniques (see appendix for detailed results). Moreover, recent work Zandich et al. (2024) introduced quantized JL-bbased method for KV cache compression, outperforming previous baselines. However, our results show that Priority Sampling achieves even greater accuracy and efficiency, highlighting its potential for matrix product sketching in KV cache sketching.
Softmax(())
/vd)V where d is the embedding
1.0
1.0
Prirtyy Sampiing
Prorty Sampinng
0.0
Sketch
0.0
Optiall
Optimal
Error 112--
20.5
0.58
00.5
Prionty Sammpinn JI Sketch
Eror
0.7
0.7
Ertor
Optmal
0.81
0,6
0.6
0.5
0.03
0.04
0.16
3.66
0,37
0,08
0.09
0.05
0.06
0.30
0,05
0.5
0,000
0,075
0,060 Sketch Size
00.55
0,00
Sketch Size
Sketch Size
dimension 128.
(b) dimension 256.
(c) dimension 512.
Figure 2: Comparison of Regression Sketching Methods on the IMDB Dataset: The plots illustrate the approximation error of different sketching methods across various sketch sizes. The matrix A is generated using TF-IDF on 10,000 random reviews, keeping the top 256, 512, and 1024 features. As the dimensionality increases, the matrices become more sparse. The matrix b represents the sentiment scores (positivity or negativity) of the reviews.
Prirty Samping
1.0
Pronty Sampinn
Prinniy Semping
JL Sketch Optinal
Optial
0.55
Optimal
Error (:2-81
Eroo
0.6
Enor
0.4
0.4
0.4
0.2
0.025
0.000
0.2
0.005
01000
0.1100
0.846
0.2
0.064
11.06
0L070
Sketch Size
Skatch
Sketch
dimension 128.
(b) dimension 256.
dimension $12.
Figure 3: Sketched Regression Methods on the Android Review Dataset: The plots illustrate the approximation error of different sketching methods across various sketch sizes. The matrix A is generated using sparse transformer SPLADE (Formal et al.) 2022) over 10,000 random reviews, retaining the top 128, 256, and 512 important features. The matrix represents the review scores.
Datasets For the task of sketched regression, we use two primary datasets. The first dataset includes Android application reviews with user ratings, which reflect the quality of the applications (Grano et al. 2017). The second dataset contains IMDB movie reviews, labeled as positive or negative (Maas
al.
201

---

Published as a conference paper at ICLR 2025
We transform the reviews into sparse vector embeddings using TF-IDF and SPLADE (Formal et al. 2022). Regression analysis is then performed to explore the relationship between the vectorized reviews and their ratings (Figure
3 or sentiment labels (Figure 2).
Additionally, we evaluate our model on synthetic datasets (Figure T for the task of approximate matrix multiplication BB The entries of the matrices A and B are generated from Gaussian distribution (0, 1). However, 10% of the dataset includes outliers, with values that are 10 times higher. The choice of 10% outliers reflects a moderate level of noise typically observed in real-world datasets, where outliers, while not dominating the data, can still have a substantial impact on the outcome. Testing with outliers allows us to assess the robustness of our model under such conditions. To examine how varying the sparsity of the matrices A and B affects approximation, we modify the number of non-zero entries. Both matrices are flattened, and we keep only designated percentage of entries non-zero.
Alongside these datasets, we use the (Bai et al. 2023) dataset to produce long text prompt from its MultiFieldi dataset for the task of KV cache in transformers. We compare our sketch as quantization method for the Key to reduce cache usage (Figure 4
Sketching Size For sampling-basee sketches, it is necessary store the indices of the sampled rows. We account for both the size of each index and the need to store selected items in full precision. In contrast, for linear sketches, only the output of the projected matrix IIA needs be stored. We report the total number of bits required for storage across all approaches and present the relative size of the sketches compared to the original matrices in bits, referred to the compression ratio. Additionally, for threshold sampling, the sketch size is not fixed due to the nature of the algorithm. Therefore, instead of reporting the expected value, we provide the average compression ratio.
Estimation Error
For the plots of matrix multiplication A and B, report the absolute difference between the estimated product and the true product, divided by || [B]].. As stated in theorem 2 this term appears on the right-hand side of the accuracy guarantee for approximate matrix multipli- cation. For all plots regarding the regression between A and b, we report the absolute difference between the estimated value and the true value, divided by ||b||2-
Interpretation of Experimental Results As we can for both the tasks of matrix multiplication (Figure D Figure 4 and sketched regression (Figure 2 Figure 3), we have observed that as the matrices become sparser, our method improves over the best-known linear sketching methods. Even though our method needs to allocate some of its budget to store indices, it still outperforms JL methods.
0.355
Priorlly Sempling
Prioriy Semplnn
Priority Samping
JL SSktth
JL Sketch
0,20
JL Skkakthhh Threshoke Sampling
0.00
Threshold Sampling
Erao 106" -pot
Threchold Sampling
Lo Ixte
0.15
0.15
0.15
0,10
0.10
0.05
0.05
0.05
0.6
Size
(a) Layer
(b) Layer 15
(c) Layer
Figure 4: Comparison of KV Cache Sketching Methods the LongBench for Mu MultiFieldgy The plots show the accuracy of different sketching methods approximating QK" across various sketch sizes. The matrices Q and K are generated from prompt tokens, and the approximation errors are displayed. Layers refer to the individual layers (hidden layers) of the Transformer architecture in the I.LaMA 2 model (meta-1l amma / ama - -2-7b-ct h-)
ACKNOWLEDGMENTS
This work supported by NSF Award 11S-2106888.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding, arXiv preprint arXiv:2308. 14508, 2023.
Joshua Batson, Daniel A. Spielman, and Nikhil Srivastava. Swice-Ramanujat sparsifiers. SIAM Journal on Computing. 41(6): 704-1111, 2012. doi: 10.1 1371/33094911 17. Preliminary version in the s1st Annual ACM Symposium on Theory of Computing (STOC).
Aline Bessa, Majid Daliri, Juliana Freire, Cameron Musco, Christopher Musco, Aécio Santos, and Haoxiang Zhang. Weighted minwise hashing beats linear sketching for inner product estimation. In Proceedings of the 42nd Symposium on Principles of Database Systems (PODS), pp. 169-181, 2023.
Kevin Beyer, Peter J. Haas, Berthold Reinwald, Yannis Sismanis, and Rainer Gemulla. On synopses for distinct-value estimation under multiset operations. In Proceedings of the 2007 ACM SIGMOD International Conference on Management of Data, 2007.
Andrei Z. Broder, Moses Charikar, Alan M. Frieze, and Michael Mitzenmacher. Min-wise indepen- dent permutations (extended abstract). In Proceedings of the 30th Annual ACM Symposium on Theory of Computing (STOC), 1998.
Sonia Castelo, Rémi Rampin, Aécio Santos, Aline Bessa, Fernando Chirigati, and Juliana Freire. Auctus: a dataset search engine for data discovery and augmentation. Proc. VLDB Endow., 14(12), 2021.
Moses Charikar, Kevin Chen, and Martin Farach-Colton. Finding frequent items in data streams. In Proceedings of the 29th International Colloquium on Automata, Languages and Programming 2002.
(ICALP), pp.
Xue Chen and Eric Price. Active regression via linear-sample sparsification active regression via linear-sample sparsification. In Proceedings of the 32nd Annual Conference on Computational Learning Theory (COLT), 2019.
Nadiaa Chepurko, Ryan Mareus, Emanuel Zgraggen, Raul Castro Fernandez, Tim Kraska, and David Karger. Arda: automatic relational data augmentation for machine learning. Proc. VLDB Endow., 13(9), 2020.
Kenneth L. Clarkson and David P. Woodruff. Low rank approximation and regression input sparsity time. In Proceedings of the 45th Annual ACM Symposium on Theory of Computing (STOC), pp. 81-90, 2013.
Edith Cohen. Sampling big ideas in query optimization. In Proceedings of the 42nd Symposium on Principles of Database Systems (PODS), 2023.
Edith Cohen and Haim Kaplan. Summarizing data using bottom-k sketches. In Proceedings of the
2007 ACM Symposium on Principles of Distributed Computing (PODC), 2007.
Edith Cohen and Haim Kaplan. What you can do with coordinated samples. In Proceedings of the 16th International Workshop on Approximation Algorithms for Combinatorial Optimization Problems (APPROX), 2013.
Michael Cohen, Sam Elder, Cameron Musco, Christopher Musco, and Madalina Persu. Dimension- ality reduction for k-meaas clustering and low rank approximation. In Proceedings of the 47th Annual ACM Symposium on Theory of Computing (STOC), pp. 163-172, 2015a.
Michael B. Cohen, Yin Tat Lee, Cameron Musco, Christopher Musco, Richard Peng, and Aaron Sidford. Uniform sampling for matrix approximation. In Proceedings of the 6th Conference on Innovations in Theoretical Computer Science (ITCS), pp. 181-190, 2015b.
Michael B. Cohen, Jelani Nelson, and David P. Woodruff. Optimal approximate matrix product in terms of stable rank. Proceedings of the 43rd International Colloquium on Automata, Languages and Programming (ICALP), 2016.
11

---

Published as a conference paper at ICLR 2025
Graham Cormode, Minos Garofalakis, Peter Haas, and Chris Jermaine. Synopses for Massive Data: Samples, Histograms, Wavelets, Sketches. Foundations and Trends in Databases. NOW publishers, 2011.
Majid Daliri, Juliana Freire, Christopher Musco, Aécio Santos, and Haoxiang Zhang. Simple analysis of priority sampling. In 2024 Symposium on Simplicity in Algorithms (SOSA), pp. 224-229. SIAM, 2024a.
Majid Daliri, Juliana Freire, Christopher Musco, Aécio Santos, and Haoxiang Zhang. Sampling methods for inner product sketching, 2024b.
Laxman Dhulipala, Majid Hadian, Rajesh Jayaram, Jason Lee, and Vahab Mirrokni. MUVERA: multi-vector retrieval via fixed dimensional encodings, (rrXiv2405.79500 2024.
Petros Drineas and Ravi Kannan. Fast Monte-Carlo algorithms for approximate matrix multiplication. In Proceedings of the 42nd Annual IEEE Symposium on Foundations of Computer Science (FOCS),
pp. 452 459, 2001.
Petros Drineas and Michael W. Mahoney. RandNLA: Randomized numerical linear algebra. Commun. ACM, 59(6), 2016.
Petros Drineas, Ravi Kannan, and Michael W. Mahoney. Fast Monte Carlo algorithms for matrices I: Approximating matrix multiplication. SIAM Journal on Computing, 36(1):132-155) 2006a. Preliminary version in the 42nd Annual IEEE Symposium on Foundations of Computer Science (FOCS).
Petros Drineas, Michael W. Mahoney, and S. Muthukrishnan. Sampling algorithms for &2 regression and applications. In Proceedings of the 17th Annual ACM-SIAM Symposium on Discrete Algorithms (SODA), pp. 1127-1136, 2006b.
Nick Duffield, Carsten Lund, and Mikkel Thorup. Flow sampling under hard resource constraints. In Proceedings of the Joint International Conference on Measurement and Modeling of Computer Systems (SIGMETRICS 2004), 2004.
Nick G. Duffield, Carsten Lund, and Mikkel Thorup. Priority sampling for estimation of arbitrary subset sums. J. ACM, 54(6), 2007.
Joshua Engels, Benjamin Coleman, Vihan Lakshman, and Anshumali Shrivastava. DESSERT: an efficient algorithm for vector search with vector set queries. In Advances in Neural Information Processing Systems 36 (NeurIPS), 2023.
Thibault Formal, Carlos Lassance, Benjamin Piwowarski, and Stéphane Clinchant. From distillation to hard negative sampling: Making sparse neural models more effective. In Proceedings of the 45th international ACM SIGIR conference research and development in information retrieval, pp. 2353-2359, 2022.
Giovanni Grano, Andrea Di Sorbo, Francesco Mercaldo, Corrado A Visaggio, Gerardo Canfora, and Sehastiano Panichella. Software applications user reviews. 2017.
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization, 2024.
Sergey loffe. Improved consistent sampling, weighted minhash and 11 sketching. In Proceedings of the 2010 IEEE International Conference on Data Mining (ICDM), 2010.
Andra Ionescu, Rihan Hai, Marios Fragkoulis, and Asterios Katsifodimos. Join path-based data augmentation for decision trees. In 2022 IEEE 38th International Conference on Data Engineering Workshops (ICDEW). IEEE, 2022.
Jiawei Jiang, Fangcheng Fu, Tong Yang, and Bin Cui. SketchML: accelerating distributed machine learning with data sketches. In Proceedings of the 2018 ACM SIGMOD International Conference on Management of Data, pp. 1269-1222 1284, 2018.
12

---

Published as a conference paper at ICLR 2025
Daniel M. Kane and Jelani Nelson. Sparser Johnson--indeestrrass transforms. Journal of the ACM, 61(1):4, 2014. Preliminary version in the 23rd Annual ACM-SIAM Symposium on Discrete Algorithms (SODA).
Jakub Konečný, H. Brendan McMahan, Felix X. Yu, Peter Richtárik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency.
arXiv:2005. 10009 2020.
François Le Gall. Faster algorithms for rectangular matrix multiplication. In Proceedings of the 53rd Annual IEEE Symposium on Foundations of Computer Science (FOCS), pp. 514-523, 2012.
Ping Li, Kenneth Church, and Trevor Hastie, Conditional random sampling: A sketch-based sampling technique for sparse data. In Advances in Neural Information Processing Systems 19 (NeurIPS), volume 19, 2006.
Jiabin Liu, Chengliang Chai, Yuyu Luo, Yin Lou, Jianhua Feng, and Nan Tang. Feature augmentation with reinforcement learning. In Proceedings of the 38th IEEE International Conference on Data Engineering (ICDE), 2022.
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, 2011.
Michael W. Mahoney, Petros Drineas, Malik Magdon-Ismaie and David Woodruff. Fast approxi- mation of matrix coherence and statistical leverage. Journal of Machine Learning Research, 13: 3475-3506, 2012. Preliminary version in the 29th International Conference on Machine Learning (ICML).
Mark Manasse, Frank McSherry, and Kunal Talwar. Consistent weighted sampling. Technical Report MSR-TR-2010-11 Microsoft Research, June 2010. URL URL wttp:/wwwwwwwwwwwpwpptptt com/ http:///essarrollllllaaaaaooo/oa
Per-Gunnar Martinsson and Joel A. Tropp. Randomized numerical linear algebra: Foundations and algorithms. Acta Numerica, 29:403-572, 2020.
Jelani Nelson and Huy L. Nguyen. OSNAP: Faster numerical linear algebra algorithms via sparser subspace embeddings. In Proceedings of the 54th Annual IEEE Symposium on Foundations of Computer Science (FOCS), pp. 117-126, 2013.
Esbjörn Ohlsson. Sequential poisson sampling. Journal of Official Statistics, 14(2), 1998.
Daniel Rothchild, Ashwinee Panda, Enayat Ullah, Nikita Ivkin, Ion Stoica. Vladimir Braverman, Joseph Gonzalez, and Raman Arora. FetchSGD: communication--oe -fffcciinn federated learning with sketching. In Proceedings of the 37th International Conference Machine Learning (ICML), 2020.
Aécio Santos, Aline Bessa, Fernando Chirigati, Christopher Musco, and Juliana Freire. Correlation sketches for approximate join-correlation queries. Proceedings of the 2021 ACM SIGMOD International Conference on Management of Data, 2021.
Aécio Santos, Aline Bessa, Christopher Musco, and Juliana Freire. A sketch-based index for correlated dataset search. In Proceedings of the 38th IEEE International Conference on Data Engineering (ICDE), 2022.
Tamás Sarlós. Improved approximation algorithms for large matrices via random projections. In Proceedings of the 47th Annual IEEE Symposium Foundations of Computer Science (FOCS), pp. 152, 2006.
David P. Woodruff. Sketching as a tool for numerical linear algebra. Foundations and Trends in Theoretical Computer Science, 10(1-2):1-157, 2014.
Amir Zandieh, Majid Daliri, and Insu Han. Qjl: 1-bit quantized jl transform for kv cache quantization with zero overhead, 2024. httpst haattps://axxiv.orrt
2406. 03482
13

---

Published as a conference paper at ICLR 2025
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi Plug-and-plas 2bit kv cache quantization with streaming asymmetric quantization. 2023. doi:
URL
13140/RG 21..4/R.....2220 2772222277
A SUPPORTING PROOFS
We begin by proving Lemma | which was the key result used proving our main matrix product sketching result.
Proof of Lemma 5
For any 2, y E [d] x [m]] we can write the æ, entry of our estimate W as:
A
W a.y
min |A
TA
B
1 .B)
Recall that we aim to prove the following:
A B 4V
|
BB
E (Wa.a) z.y [ 2.
and E[(W z.y [A]
}3]
2.
1:
|A
B
For any i. let be a 0-1 indicator random variable for the event that index selected for both S(A) and S(B). I..,, for the event that i E T.
For any i E [n]] let T denote the hthh smallest value of h(j)/llA, over all [n]] {i). If [n|] {i}
has fewer than & values, define =å oc. Define 7 analogously. The probability that i T
IA n IB conditioned on (A) and r' (B) is equal to the probability that both h(1)/||, 11 T
and B(i)/||BBBB < TB. Le., the conditional probability is equal to min ||A T BB 12 TB)
Additionally, observe that, for all sampled i E T = IA n IB. T -A and T So, we have:
Ai. Biy
E
1
min(1,
A,
TA.
B
T
A B k.
E E min (1, A,
Pr E I r
T
|
T
-i K.y
By linearity of expectation, follows that E [Wayl
E E
A Big min( - BBB. 3 1
Σ. A Bi.y = A z,y" as desired.
The next step is to bound the expected squared error. As mentioned earlier, this made difficult by the fact that 11, and 1, are not independent random variables. Fortunately, however, that appropriate (random) scalings of these random variables are uncorrelated, a technique that is standard in prior analyses of priority sampling for other problems.
show
In particular, define P: min ||A TA |B, TB) and define 7 to be the (k: 1)a ..j
A(k) smallest value among TA for all & E [n] {i,j}. Define TB analogously. We can see that Pr|i.j E T Moreover, conditioned on j E T. we have that Tir(A) = (AA) and T (B) TBB). So, in
J.j B
min(1, 1, ||A 12 T | I T min(1, ||
A |By T
A
particular, conditioned on i.j j E T, min 1, A 1
T B -
1 1
E
= E E
P: Pj
min 1, A.
1 TH B.B.B
TB min( 1, | 18 T |B, +
Pr ijET
E
E
P3
14

---

Published as a conference paper at ICLR 2025
Since the random variables 1i and 1 pairwise uncorrelated for all i.j. we have that A.,. and j 3.9"
Pi Pj
1s are pairwise uncorrelated well. So, apply the linearity of variance to Ps
conclude:
A B Ly
E
W
z. ) [
Var W s.yl
Var
1i
a.y
min (1, |A, TAi |
2
TB
A Biyy
Var
1
min (1, A 11 TA |
BB
1=1
So, suffices to establish individual bounds on Var
A Biy
1/
In order to do
min(1. 13
-B)
so, first observe that, conditioned on A BB
A L.a Big
E
11
TA TB
min(1 (1, A 11
TA
B TB)
A Biy
2
-min(, (1, || T | T
minn1.
A 11 T
B.1 TB)
A B
A2 B kM
max 1, 1A:1..2 -
|A, A B 7
B
B
We can thus write:
A B 1.
Var
11
min
1,, ||A
13
TA
B.112 TB
1
A. B 1.9
E
max
A. , B
A
12
TA
122
TB
)
A B
1,2 1.9
E
max
A
12
1,
B
A
B
S 2 B2
1,2
E
max
|
B
TB
A 2 B 1.9
E
2
B
E
||A 1
T
TB
A B
A
B
E
|
B
T
|
BB
We can apply Claim 5 from Daliri et al.
(024a) bound
and
E
So
we have:
A
B TV
2
B 4.9
E (W z.y
[
E
E
2.91
|
7
B
TB
A2 1,
B
||A
113
A? aB
B
B
1
|| T.
1.9
1
B
1e
1
A B
(IA |AA
B
k 1
[
as desired.
B MATRIX PRODUCT SKETCHING WITH THRESHOLD SAMPLING
We motivated our Priority Sampling method from Section 2 via a simpler matrix sketching method based on Threshold Sampling. We include full analysis of this method here for pedagogical purposes,
15

---

Published as a conference paper at ICLR 2025
since the method is much easier analysis. However, in general, we believe that Priority Sampling preferable since it offers a fixed size sketch. In terms of accuracy, recent work on inner product sketching finds that both methods perform nearly identically to each other (Daliri al. 2024b). Experiments suggest the same is true for general matrix-matio product sketching.
Sketching. As discussed, Threshold Sampling uses shared hash function h [n] - [0, 11, which is assumed to be uniformly random. As shown in the pseudocode in Algorithm + the method selects all from A for which hu(i)|l || A, ll falls below fixed "global threshold". TA | l| Here, is parameter that determines the size of the sketch S(A) produced by Algorithm There will be most rows selected in expectation, but the exact number depends on the random choice of h.
Algorithm 4 Threshold Sampling
Input: Matrix A of size d, random seed 8, target number of row samples, k: Output: Sketch ((A) [IA. VA, where ZA is a subset of row indices from {1..... ......... n2 and VA contains A; for all E ZA-
1: Use random seed to select uniformly random hash function h (I, nz - [0, 1].
2: Initialize ZA and VA to be empty lists.
3: for E I...., do
I|A
4: Set threshold T3 k
|A
then 6: Append to IA. append A, to VA.
5: if' h(i) S
7: return S(A) = (IA, VA. FA} where TA = k/
Estimation. Similar Priority Sampling, after constructing our sketches S(A) and S(B), we approximate the matrix product of A and B by computing weighted sum of outerproducts of rows included in both S(A) and S(B). fact, we can use the exact same procedure defined in Algorithm 2 from Section 2
Unlike Priority Sampling, Threshold Sampling ensures that the probability of sampling any given row A; is an independent random event. There is no dependence on the event that another row gets sampled, In particular, we can easily compute the probability of index being included both
||A. sketches S(A) and S(B). It is exactly equal to pt min (1, TAI
B.II.
A
BB
Guarantees. Our primary theoretical guarantee for Threshold Sampling can be stated as follows: Theorem 6. Let A Raxd, B E Raxm, and let S(A) {IA, TA) and S(B) (I,, Vn, TB) be sketches produced by Algorithm 4 with imput and a shared seed 8. Suppose W is the approximate matrix of A calculated using Algorithm! 2
on these sketches. Then, [W]- ATB and
2
[|W
ATB
Additionally, E [[AA]] & k and E [B]]] k. Lee, each sketch contains no more than k row indices in expectation.
Theorem@essenteee matches Theorem 4 although is actually a bit tighter, the 2/(K: 1) prefactor replaced with 2/k. The only disadvantage of the theorem is that we do not have a fixed upper bound on [A and [Za] which are equal the number of rows sampled from A and B, respectively. We also remark that, as for Theorem Theorem Ga be combined with Markov's 's inequality to give high probability bound: if we set k with probability
2/6 then we achieve error JW AT - |||A||||[B]]]
2
Proof of Theorem 6 Let 1, denote the indicator random variable for the event that is included in both ZA and ZB. 1, = 1 if this event occurs and 0 if it does not. Note that, for i / j, 1; is independent from L,. since the hash values h(i)) and h(j) drawn uniformly and independently from [0, 1]. Moreover, we claim that 1, is equal to with probability:
pi min (1. k- |
k- B
min(1,
1A B 12)
(2)
A
16

---

Published as a conference paper at ICLR 2025
This is because for index to be included in both IA and ZB, h(i) must be less than both
|| and ||B.]// [B] simultaneously (see line 3 of Algorithm 4). Given the probability
of sampling each item, can find the expectation of the approximation W.
A.B. Pi
E W
Σ
Pr
A. A
(3)
This proves the desired claim the expection of W. is left to bound the E [IW A.
We do so by bounding the squared error of each entry in W separately. In particular, for the I, entry, we write:
A A Σ
n
Ai B.y
W s.y (E
(.
Var [ (aay Var
1,
Var 1;
Ps
P
Above we use the fact that 1..... In are independent to apply linearity of variance. Let H. denote the the set of all for which p; o and pi + 1. Using that Var[1.] p;(1 pi). have:
B
A 2 L,
B
E (
[A )
Pe(1
p.) -
Pi
iE
So we can bound E
||
A |
by:
d m
B I
A
ATB
E
K.y
[A
[ATB].
B) (x.p))
y=1
m
A2 Lan
B
Σ
1
A
Σ
B²
Pi
Pi ==1
IEH
y=1
| 12
|
Σ
iCH Ps
NE
min
K- & &
TAI
B
Recall that we defined p. = min 11, 1. 22
B
so in the last step, we have used the fact that
A
B
Ps * for i H. Continuing, we can bound:
B
IAA.II
E -W S ΣIAR
B
A
A-
B.B.E
iEH
TB
max
11///1
I
/BBB///BB))))
|
k
|
B
|
B B
2
2 K
+
BAE
A
NEH
In the second to last step, we upper bounded the maximum but the sum.
C FURTHER EXPERIMENTS ON ATTENTION MODELS
One key reason for not sketching the Query matrix in this application is that is not quantized. unlike the Key and Value matrices. This means there is no need to apply sketching techniques to the Query matrix, as is recalculated for each input token and does not benefit from compression methods used on more static, larger matrices. An important advantage of using the sampling method is that it allows selective sketching of matrices that are both large and have a static or slow-changing nature, such as the Key matrix. In contrast, linear sketching requires projecting both the Query and Key matrices, regardless of their size or dynamism. In our experiments, we had access to the entire Queries matrix while applied Priority Sampling to sketch the considerably larger Key matrix. This approach effectively demonstrates the efficiency of sampling in handling large-scale data while preserving the dynamic properties of the Query matrix in real-time applications.
17

---

Published as a conference paper at ICLR 2025
Piictty
Prioriy Samplng
Priority Semplngg
0.21
Sketch
0.30
Skatch
0.20
Threshold Sampling
Enar 108"-0xT1
Threshold Sampinng
Thresheld Sampling
0.15
0.15
0.15
0.16
0.10
0.05
0.05
0.2
0.4
0.2
0.4
0.6
0.8
0.2
0.1
0.5
0.5
Sketch Size
Sketch Size
Skatch Size
(a) Layer 0, Key Sketched
(b) Layer 15, Key Sketched
(c) Layer 31, Key Sketched
Figure 5: Comparison of KV Cache Sketching Methods on the LongBench for MultiFieldga. The plots illustrate the accuracy of various sketching methods in approximating QK" across different sketch sizes. The Query matrix remains untouched, and only the Key matrices are sketched using Priority Sampling and Threshold Sampling, whereas the JL sketch requires the projection of both matrices Q K. Layers refer to the individual layers (hidden layers) the Transformer architecture in the L.LMMA 2 model.
18


Reviews:
[{"review_id": "UX41dLNNto", "reviewer": "Reviewer_St9U", "summary": "This paper addresses the problem of approximating the matrix product $A^T B$ using sketching-based techniques in a new setting with two key requirements: (1) preserving the potential sparsity of $A$ and $B$ and (2) ensuring that sketches of $A$ and $B$ can be computed independently—motivated by scenarios like distributed computing and multi-vector retrieval. Traditional linear sketches, often relying on dense random matrices, fail to preserve sparsity. However, sparsity preservation (requirement 1) can be achieved using importance sampling techniques, where rows of $A$ and $B$ are subsampled and reweighted. Unfortunately, existing importance sampling methods require knowledge of both $A$ and $B$ simultaneously, which contradicts requirement (2).\n\nTo address this, the paper proposes a new approach called coordinated random sampling, which only requires shared randomness to independently compute sketches of $A$ and $B$ by row subsampling, without further interaction between the matrices. This method is based on priority sampling: the same rows of $A$ and $B$ are selected only when they both contribute significantly to $A^TB$, even though rows are sampled independently within each matrix.\n\nThe paper shows that subsampling $O(1/\\epsilon^2 \\delta)$ rows of $A$ and $B$ suffices to achieve $\\epsilon$ error in Frobenius norm with probability $1-\\delta$, matching the best known result from dense linear sketches. This paper demonstrates one application of their proposed algorithm in sketched linear regression. Lastly, the paper empirically compares the effectiveness of coordinated random subsampling against linear sketches in terms of approximation accuracy with varying degrees of sparsity in the input matrices, and in two applications: 1) sketched linear regression, 2) approximation to the attention layer in autoregressive language models.", "questions": "Why is Theorem 3 presented as a constant probability guarantee rather than a high-probability guarantee? What specific challenges arise in deriving a high-probability guarantee in this context? Additionally, what makes it difficult to achieve a multiplicative $(1+\\epsilon)$ guarantee instead of the additive error guarantee stated in the footnote on page 4?\n\nIn the context of sketched linear regression problems, the guarantee provided by Theorem 3 is relatively weak. A discussion on the challenges of obtaining a stronger guarantee would be valuable and could provide useful insights for future work.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. I found the presentation of this paper to be very clear and engaging. The problem setting, requirements, and notation are all well-defined, and the core idea is explained thoroughly, making it easy for readers to follow. The proof is presented in a clean and structured manner, enhancing readability.\n\n2. The paper provides strong motivation for studying the problem of computing independent sketches and discusses several potential applications, demonstrating their proposed algorithm in a one important application: sketched linear regression. \n\n3. The setting of preserving sparsity of the inputs and computing independent sketches seems to be novel.", "weaknesses": "1. One concern I have regarding the experiments is that while vector quantization—a nonlinear compression technique—has been widely studied and applied in practice for approximating the computation of the key matrix in the attention layer, it remains unclear whether using linear compression techniques, such as approximate matrix products, to approximate $QK^T$ or just the key matrix $K$ could degrade model performance significantly in downstream applications. I suggest that the authors cite works that empirically investigate this issue (if there is any) or include explicit discussion and caveats when discussing the use of approximate matrix products for approximating the attention layer. The application discussed in Section 1.4 on multi-vector information retrieval seems more reasonable and well-suited for demonstrating the applicability of the proposed algorithm in experiments.\n\n2. Additionally, some aspects of the experiment section need clarification. For instance, does the x-axis labeled “Sketch Size” in Figures 1-5 represent the percentage of rows subsampled from the input matrices? What does “Layer x” in Figures 4 and 5 refer to? Furthermore, what architecture is used for the autoregressive language model in the experiment? Providing these details would enhance the clarity of the experimental setup.\n\n3. It would also be helpful to include empirical comparisons of memory usage between linear sketches (e.g., JL sketches) and the proposed method.\n\n4. Minor issues: Bad citation in line 109. “is include as” => “is included as” in line 324.\n\n5. To improve clarity for readers, it would be helpful to include a description of $\\tau_A$ in the output of Algorithm 1, explaining that it represents the threshold used to determine whether a row in the input matrix $A$ should be retained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RixPrrpDpA", "reviewer": "Reviewer_NkFE", "summary": "Overview: The paper presents a method for sketching two matrices A and B, each held by a different party, such that the product of the matrices can be approximated up to an additive term of eps ||A||||B|| from the sketches alone. Each sketch consists of k rows/columns of the sketched matrix, for k=O(1/eps^2). For sparse matrices, such sketches can be more compact than the state-of-the-art linear sketches obtained via random projections. \n\nTechniques:  the paper presents two sketching methods: Threshold Sampling and Priority Sampling. Both methods are variants of coordinated sampling, where a shared random seed is used to sample columns/rows in a correlated manner. Threshold Sampling is simpler to analyze, while Priority Sampling offers better control over the sketch size.  In addition, the paper presents two applications of the proposed sketches, to multi-vector retrieval and regression-based dataset search.\n\nThe authors complement the theoretical development with empirical evaluation of their algorithms. In particular, both sketching methods are shown to improve (for sparse matrices) over the standard random projection approach, as predicted theoretically.", "questions": "Q: In the paper you write \"In doing so, we encounter several technical challenges, including the fact that Priority Sampling—being a without-replacement sampling procedure—generates non-i.i.d. row samples from A and B.\". Can you list some of the other technical challenges as well ?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "S1: Interesting problem.\nS2: Elegant solutions.\nS3: Solid experiments.", "weaknesses": "W1: The result and the approach are not very surprising, given the prior work of Bessa et al and Daliri et al\nW2: The analysis of one of the algorithms (Threshold Sampling) seems fairly straightforward.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cEFA9TCQp9", "reviewer": "Reviewer_ev8D", "summary": "The paper revisits the approximation matrix product problem. That is, given two matrix $A, B$, we want to compute a matrix $W$ in a fast time where $\\|W - A^T B\\|_F^2 \\le \\|A\\|_F^2 \\|B\\|_F^2$. The paper studies the case where sketches must be computed independently of each other, except for the use of a shared random seed. The paper provides a new method that is based on coordinated random sampling and has sketch size $O(s/\\epsilon^2)$ assuming $A$ and $B$ have $s$ non-zero entries per row. Based on this, the paper next proposes a new algorithm for the $\\ell_2$ regression problem. Finally, the paper gives an empirical evaluation that demonstrates the advantage of the proposed algorithm.", "questions": "Approximation Matrix Product: I am not sure I fully understand the motivation of the model considered, as mentioned by the authors, sampling to either $\\|A_i\\|_2$ and $\\|B_i\\|_2 $ can also get a good bound. The communication about the probability can be somewhat expensive but we can also sample from $A$ and then tell the other part the indices of the sampled rows, in which case we will have much better communication cost? \n\nRegression: there are some recent works that study the active regression problem where the goal is to access as fewer number of the entries of $b$ as possible, see, e.g, [1, 2]. In particular, it has been shown that sample $O(d/\\epsilon)$ rows of $A$(where the sample probability is only computed from the matrix) can give a $(1 + \\epsilon)$ relative error of the solutions. Can the authors give a comparison of both models?\n\nMinor:\n\n1.Line 238: that is $k \\cdot |B_i|^2 / |B|_F^2$ ?\n\n[1] Cameron Musco et al. Active Linear Regression for $\\ell_p$ Norms and Beyond.\n\n[2] Xue Chen et al. Active regression via linear-sample sparsification active regression via linear-sample sparsification.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The theoretical analysis of this paper is solid. The paper gives a new sketching algorithm with size $O(s^2 /\\epsilon^2)$. This bound will be better for sparse matrix compared to the previous methods, which is interesting to me.\n\n- The paper gives a detailed experiment that demonstrates the advantage of the proposed algorithms.\n\n- The presentation of the paper is good. The paper has a nice introduction section.", "weaknesses": "- I still do not understand the motivation of the new model the paper discusses well (see the questions below). Maybe the authors can give more explanation about this？ \n\n- It will be better if the experiments can also give a comparison to the previous sampling-based method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0rg6aBsfob", "reviewer": "Reviewer_Ah4Z", "summary": "Given two matrices $\\mathbf{A} \\in \\mathbb{R}^{n \\times d}$ and $\\mathbf{B} \\in \\mathbb{R}^{n \\times m}$, the problem of approximately computing the product $\\mathbf{A}^\\top \\mathbf{B}$ using sketching techniques has been extensively studied, with a wide range of applications in machine learning and numerical linear algebra. The existing work includes two types of sketching methods: (1) Generating a random matrix $\\boldsymbol{\\Pi} \\in \\mathbb{R}^{k \\times n}$ such that $(\\boldsymbol \\Pi \\mathbf{A})^\\top (\\boldsymbol \\Pi \\mathbf{B}) \\approx \\mathbf{A}^\\top \\mathbf{B}$; (2) Sampling and reweighting the rows of $\\mathbf{A}$ and $\\mathbf{B}$ to obtain $\\widetilde{\\mathbf{A}}$ and $\\widetilde{\\mathbf{B}}$ such that $\\widetilde{\\mathbf{A}}^\\top \\widetilde{\\mathbf{B}} \\approx \\mathbf{A}^\\top \\mathbf{B}$. This paper studies matrix multiplication in the $\\textit{sketching setting}$, which requires the sketches $\\mathcal{S}(\\mathbf{A})$ and $\\mathcal{S}(\\mathbf{B})$ of $\\mathbf{A}$ and $\\mathbf{B}$ are computed independently. Nevertheless, the classical sampling method does not work as the rows are sampled with non-uniform probabilities. This paper applies the sampling method effectively in the sketching setting by virtue of the $\\textit{coordinated sampling method}$ (aka $\\textit{priority sampling}$), which ensures that the number of commonly sampled row indices from $\\mathbf{A}$ and $\\mathbf{B}$ is guaranteed. Furthermore, the proposed algorithm achieves the same guarantee as the existing sampling method and is more efficient when the two input matrices are sparse. In addition, the developed method can be extended to $\\ell_2$ regression. The experimental results prove the effectiveness of the proposed algorithm on synthetic and real datasets.", "questions": "In Algorithm 3, it requires computing the leverage score of each row of $\\mathbf{A}$ (i.e., $\\mathbf{A}_i (\\mathbf{A}^\\top \\mathbf{A})^{-1} \\mathbf{A}_i^\\top$) and $(\\mathbf{A}^\\top \\mathbf{A})^{-1}$.   \n(1) What's their running time?  \n(2) Can you compare this algorithm with the linear regression in the work of Clarkson and Woodruff [CW'13]?  \n(3) For sketched regression, do we only care about the size of sketches? If so, in Section 1.4, it needs to compute $\\min {\\mathbf{x}} ||\\mathbf{A}^i \\mathbf{x} - \\mathbf{b}||_2$ for each $i$. Then the running time could be more crucial.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "(1) This paper studies the approximate matrix multiplication in the sketching setting by the coordinated sampling technique. The proposed algorithm matches the guarantee of the existing sampling-based method and is more efficient when the input matrices are sparse.    \n(2) The developed algorithm also works for $\\ell_2$ regression problem and improves the existing bound when the input matrix is sparse.        \n(3) This paper is well-written and easy to understand.", "weaknesses": "(1) In essence, the coordinated sampling technique is nearly the same as the threshold sampling method and has also been applied in other sketching work [DLT'07, DFMSZ'24]. From the technique perspective, this paper's contribution is kind of limited.    \n(2) The restriction in sketching setting often occurs when communications matter, like in distributed systems. Although this paper proposes an algorithm that works in the sketching setting, there are no specific applications in distributed sketching.     \n(3) This paper claims that priority sampling is more complicated than threshold sampling whose sketch size can only be bounded in expectation. However, in Figure 1 and Figure 4, we can see that priority sampling does not outperform threshold sampling.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)