High-Performance Row-Sampling for Khatri-Rao Products
Vivek Bharadwaj, University of California, Berkeley
The Khatri-Rao product is commonly defined as a column-wise Kronecker product of one or more matrices. This fundamental primitive appears in settings as diverse as cryptography, signal processing, compressed sensing, structured partial differential equation problems, and tensor decomposition. Unfortunately, a Khatri-Rao product of N matrices has a row count exponential in N, making it intractable to compute with or even materialize in computer memory in many cases. This poster addresses the following question: is there an efficiently computable map that can reduce the row count of a Khatri-Rao product while preserving the geometry of its column space? We will show that such a map exists by randomly sampling rows from the Khatri-Rao product according to a distribution of statistical leverage scores. We design a novel sampler that requires runtime quadratic in the column count, but only logarithmic in the row count of the Khatri-Rao product to draw each row sample. By applying our algorithm to alternating least-squares Candecomp / PARAFAC decomposition, we extract patterns from datasets with billions of entries. We also design a distributed version of our sampling algorithm; on thousands of LBNL Perlmutter CPU cores, we demonstrate speedups of up to 11x over SPLATT, a state-of-the-art library that does not use randomization.