Matrix Product Sketching via Coordinated Sampling

We appreciate your thoughtful suggestions and detailed comments to help enhance the clarity of the paper.

Weakness:

1- We address this question in our response to Reviewer NkFE as well. That response is pasted below for convenience.

Regarding Priority Sampling, the main difference between the prior analysis for inner products is as follows: The matrix-matrix product $A^TB$ can be viewed as a set of $d\times m$ inner products between the $d$ columns of $A$ and $m$ columns of $B$. Consider two fixed columns, $a$ from $A$ and $b$ from $B$. In contrast to the prior work, we do not sample the $i^\text{th}$ entries $a_i$ and $b_i$ with probability proportion to their squared magnitude – the sampling probability depends on the norm of the entire $i^\text{th}$ row in $A$ and $B$. As a result, we cannot prove our result by proving a “per inner product” guarantee – for some column pairs, the method will undersample heavy entries, and thus have high error. Instead, we must bound the total variance across all $d\times m$ inner products. Ultimately, doing so is not too difficult: the analysis of priority sampling is still < 2 pages, and all “from first principles”. However, we hope the simplicity of the methods and analysis can be an advantage of the result.

2- We have also addressed this question in our response to Reviewer ev8D. For your convenience, we have included that response below.

The issue with communication arises when we wish to compress multiple different matrices and to compute approximate matrix-matrix products between multiple pairs of sketches. For example, suppose we have a database of matrices $A_1,\ldots, A_q$ that need to be “pre-compressed” (as is the case in dataset search or multi-vector retrieval). At some later time, we receive a “query matrix” $B$, and wish to compute $A_1^TB, \ldots, A_q^TB$ based on a sketch of $B$. Because $A_1, \ldots, A_q$ could have very different row norms, a norm based sampling method will not sample the same set of rows from the matrices. So it is not clear what set of indices to communicate that will allow us to create a single sketch of $B$ that can be used when approximating each product $A_1^TB, \ldots, A_q^TB$. In fact, given how different $A_1, \ldots, A_q$ could be, our intuition is that, to achieve worst case guarantees, the sketch of $B$ should not depend on any individual $A_i$, necessitating communication-free sketching. We will think about if this claim can be made formal, but the main point is that naive “index sharing” approaches do not work in such settings. There might be alternative methods that use communication in a more careful way, but this would be surprising to us.

We will work on clarifying this further in the introduction, as these settings are partially addressed in Section 1.4.

3- The reviewer is correct that, in terms of (expected) sketch size vs. accuracy, Priority Sampling is not better than Threshold Sampling. However, it has the advantage of a fixed sketch size, which is very important in applications, where a fixed amount of memory would typically be required for each sketch. For example, in vector-search applications, sketch sizes are often carefully chosen due to memory considerations (e.g., a certain number of sketched vectors should fit in L1 or L2 cache). We expect the same sort of requirements to arise when sketching matrices, which is why we analyze priority sampling.

Questions:

1- Naively, computing $(A^T A)^{-1}$ takes $O(n \cdot d^2 + d^3)$, and calculating the leverage scores for all rows takes $O(n \cdot d^2)$ , resulting in a total running time of $O(d^3 + n \cdot d^2)$. This runtime can be accelerated with more advanced techniques (e.g., the Clarkson-Woodruff 2013 paper mentioned by the reviewer), ultimately resulting in the $nd^2$ term being replaced with with $nnz(A)\log(n)$, where $nnz(A)$ is the number of non-zero entries in $A$. We will add some discussion about this to the paper.

2- This is a good question. The regression algorithm from Clarkson-Woodruff 2013 is a “linear sketching” method. As discussed below Theorem 3, it would lead to a sketch size of $O(d^2/\epsilon)$ for the regression problem, whereas we achieve size $O(sd/\epsilon)$. [CW13] does have an advantage over other linear sketching methods in that the sketching matrix, $S$, can be applied extremely quickly (roughly in $O(nnz(A))$ time). However, the sketch size is no better (and in fact, slightly worse) than if a slower random Gaussian or random sign sketch was used instead. We are primarily focused on sketched size for this paper, as in the settings we consider such as multi-vector retrieval, sketch time is less important (it is a preprocessing step).

3- As mentioned, for this paper, we are primarily focused on sketch size: in practical applications, there is often a strong connection between reducing memory requirements and improving wall-clock runtime. By working with smaller sampled matrices $S(A_i)$ and $S(b)$, instead of the full matrices $A_i$ and $b$, the computational cost of solving $\min_x || S(A_i)^T x - S(b) ||_2$ is significantly reduced. This reduction in matrix size directly translates to faster computations, making sketching not just memory-efficient but also runtime-efficient. That said, in future work, it might be interesting to consider if there are even faster estimation procedures that avoid computing $A_i^+$ for each $A_i$.

We appreciate the reviewer's constructive feedback and respond to specific questions below:

Weakness:

Regarding comparing to “the previous sampling-based method”, we are not quite sure what method the reviewer refers to. The previous methods based on weighted sampling cannot be implemented in the “sketching model” where communication is not allowed. Uniform row sampling can be implemented in the sketching model – we will consider adding an experiment comparing our approach to uniform sampling, although our experience is that uniform sampling often performs far worse than the norm-based sampling methods or JL.

Questions:

1- The reviewer is correct that, in some settings, communication is not a major concern – we simply need to communicate the indices of the rows sampled. However, in other settings, communication is clearly not useful, which is why we seek algorithms that do not require any. We will try to clarify this further in the introduction. Such settings are already discussed in Section 1.4 to an extent.

The issue with communication arises when we wish to compress multiple different matrices and to compute approximate matrix-matrix products between multiple pairs of sketches. For example, suppose we have a database of matrices $A_1,\ldots, A_q$ that need to be “pre-compressed” (as is the case in dataset search or multi-vector retrieval). At some later time, we receive a “query matrix” $B$, and wish to compute $A_1^TB, \ldots, A_q^TB$ based on a sketch of $B$. Because $A_1, \ldots, A_q$ could have very different row norms, a norm based sampling method will not sample the same set of rows from the matrices. So it is not clear what set of indices to communicate that will allow us to create a single sketch of $B$ that can be used when approximating each product $A_1^TB, \ldots, A_q^TB$. In fact, given how different $A_1, \ldots, A_q$ could be, our intuition is that, to achieve worst case guarantees, the sketch of $B$ should not depend on any individual $A_i$, necessitating communication-free sketching. We will think about if this claim can be made formal, but the main point is that naive “index sharing” approaches do not work in such settings. There might be alternative methods that use communication in a more careful way, but this would be surprising to us.

2- This is a good question. In the active learning setting, the goal is also to sample rows from $A$ and corresponding entries from $b$, although under different restrictions than our paper. In particular, in the active learning setting, $b$ is unknown, so the values of its entries cannot be taken into account when choosing which indices to select. However, the indices selected from $b$ can depend on which rows are selected from $A$ (in fact, the methods cited by the reviewer sample rows from $A$ using a probability distribution that depends only from $A$, and then select the corresponding entries from $b$ deterministically).

In contrast, in our setting, $b$ is known so the value of its entries can be taken into account when selecting which entries to subsample. However, we do not allow the sample taken from $b$ to depend on those taken from $A$, and vice versa. It seems that our setting is more challenging, in that it is not possible to obtain a relative error guarantee for regression using $o(n)$ samples, let alone $O(d/\epsilon)$ samples as in the active regression setting. This point is discussed further in our response to Reviewer St9U. Nevertheless, at least we are able to obtain some non-trivial additive approximation guarantee (Theorem 3).

We thank the reviewer for their thoughtful feedback. Indeed, the Threshold Sampling analysis is straightforward, although the method has a disadvantage in that it does not guarantee a sketch of a fixed size (the sketch size is random, albeit bounded in expectation).

Regarding Priority Sampling, the main difference between the prior analysis for inner products is as follows: The matrix-matrix product $A^TB$ can be viewed as a set of $d\times m$ inner products between the $d$ columns of $A$ and $m$ columns of $B$. Consider two fixed columns, $a$ from $A$ and $b$ from $B$. In contrast to the prior work, we do not sample the $i^\text{th}$ entries $a_i$ and $b_i$ with probability proportion to their squared magnitude – the sampling probability depends on the norm of the entire $i^\text{th}$ row in $A$ and $B$. As a result, we cannot prove our result by proving a “per inner product” guarantee – for some column pairs, the method will undersample heavy entries, and thus have high error. Instead, we must bound the total variance across all $d\times m$ inner products. Ultimately, doing so is not too difficult: the analysis of priority sampling is still < 2 pages, and all “from first principles”. However, we hope the simplicity of the methods and analysis can be an advantage of the result.

Thanks for the specific suggestions and comments to help improve clarity. We reply to the larger comments/questions below.

1- While we focus on the abstract problem of matrix product sketching, we acknowledge the importance of understanding how such approximations impact model performance in downstream tasks. To address this, we will cite relevant works, including [1], which uses projection matrices (linear compression) to reduce activation dimensionality, and [2], which applies a Quantized JL Transform {SimHash) to compress the key-value cache in a transformer-based autoregressive language model, where 'key' represents encoded past inputs and 'value' represents their associated outputs.

[1] Sakr et al. ("ESPACE: Dimensionality Reduction of Activations for Model Compression")

[2] Zandieh et al. ("QJL: 1-Bit Quantized JL Transform for KV Cache Quantization with Zero Overhead")

Additionally, we believe that empirically evaluating matrix sketching methods in the context of multi-vector retrieval is an important next step for future work. Compression schemes like product quantization have already been successfully applied in this setting (see, e.g. [3]).

[3] Dhulipala et al. ("MUVERA: Multi-Vector Retrieval via Fixed Dimensional Encodings")

2- The x-axis labeled "Sketch Size" in Figures 1–5 is the compression ratio of the sketch, specifically the total number of bits requires to store S(A) and S(B) divided by the total number of bits required to store $A$ and $B$. For a JL sketch with $k$ rows of length $d$, the storage complexity is $k \times d \times {64}$ bits, since numbers are stored using doubles. For a sampling based sketch with k rows that have s non-zeros each, the storage complexity is $k\times s \times (64+32)$ bits, since we used 64 bits to store the value of every entry and a $32$ bit int to store its index.

The term "Layer $k$" in Figures 4 and 5 refers to the individual layers (hidden layers) of the Transformer architecture in the LLaMA 2 model. In our experiments, we used a fine-tuned version of LLaMA 2 (meta-llama/Llama-2-7b-chat-hf), which consists of 32 layers. Each layer features multi-head self-attention and feedforward mechanisms. Notably, the $K, V$, and $Q$ matrices in the higher layers tend to have entries with greater differences between values compared to the lower layers, which impacts the behavior of the sketching techniques.

3- When referring to memory usage, do you mean sketch size? This is the main quantity we compare in the plots currently. In terms of memory usage during the computation of sketches, the difference between the methods is minimal. While one might expect JL to have higher memory requirements due to the multiplication with a dense $k \times n$ matrix, this matrix can be generated dynamically, column by column. As a result, both methods only incur a memory overhead of $O(k)$ when computing the sketch, where $k$ represents the final sketch size.

Questions:

1- We stated Theorem 3 as a constant probability guarantee for simplicity, but the dependence on the failure probability $\delta$ is identical to Theorem 2: we can obtain the result with probability $1-\delta$ if we pay an additional $\frac{1}{\delta}$ multiplicative factor in the sample complexity. In retrospect, this does not add to the theorem’s complexity by much, so we will update it to include this dependence.

This is a good question regarding the feasibility of obtaining a multiplicative $(1 + \epsilon)$ guarantee. Upon further thought, we believe that an additive error dependence is likely necessary for any sampling-based algorithm that sketches $b$ and $A$ independently. We will add a detailed discussion and formal argument to the next version of the paper.

To illustrate this, consider the case where $b$ exactly equals $Ax^*$ for some ground truth $x^*$, so achieving a multiplicative error guarantee would require zero error (since $||Ax^{\*} - b|| = 0$). Now, suppose $A$ is a matrix of height $d + m$, containing every row of the $d \times d$ identity matrix once, but with the first row (the first standard basis vector) appearing $m$ additional times. Finding $x^*$ in this case requires identifying all entries of $b$ that correspond to the $d$ identity rows of $A$. However, if $b$ is sampled independently of $A$, an algorithm has no way of identifying these specific entries, making a multiplicative guarantee unachievable in this setting. This limitation is in strong contrast to scenarios where $b$ is sampled with probabilities that are dependent on $A$.

2- With the above in mind, it is interesting to ask if there is a method that can both take advantage of sparsity (like our method) but achieve the stronger multiplicative regression guarantee of JL. It seems that doing so will require looking beyond direct “sample and solve” regression methods.