Matrix Product Sketching via Coordinated Sampling

We revisit the well-studied problem of approximating a matrix product, ATB, based on small space sketches S(A) and S(B) of A ∈ Rn × d and B∈ Rn × 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 A and B are sparse, methods based on coordinated random sampling can outperform classical linear sketching approaches, like Johnson-Lindenstrauss Projection or CountSketch. For example, to obtain Frobenius norm error ε|A|F|B|F, coordinated sampling requires sketches of size O(s/ε2) when A and B have at most s ≤ d,m non-zeros per row. In contrast, linear sketching leads to sketches of size O(d/ε2) and O(m/ε2) for A and 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.