A framework for learned CountSketch

1 INTRODUCTION . In recent years , we have seen the influence of machine learning extend far beyond the field of artificial intelligence . The underlying paradigm , which assumes that a given algorithm has an input distribution for which algorithm parameters can be optimized , has even been applied to classical algorithms . Examples of classical problems that have benefitted from ML include cache eviction strategies , online algorithms for job scheduling , frequency estimation of data stream elements , and indexing strategies for data structures ( Lykouris & Vassilvitskii , 2018 ; Purohit et al. , 2018 ; Hsu et al. , 2019 ; Kraska et al. , 2018 ) . This input distribution assumption is often realistic . For example , many real-world applications use data streaming to track things like product purchasing statistics in real time . Consecutively streamed datapoints are usually tightly correlated and closely fit certain distributions . We are interested in how this distributional paradigm can be applied to sketching , a data compression technique . With the dramatic increase in the dimensions of data collected in the past decade , compression methods are more important than ever . Thus , it is of practical interest to improve the accuracy and efficiency of sketching algorithms . We study a sketching scheme in which the input matrix is compressed by multiplying it with a “ sketch ” matrix with a small dimension . This smaller , sketched input is then used to compute an approximate solution . Typically , the sketch matrix and the approximation algorithm are designed to satisfy worst-case bounds on approximation error for arbitrary inputs . With the ML perspective in mind , we examine if it is possible to construct sketches which also have low error in expectation over an input distribution . Essentially , we aim for the best of both worlds : good performance in practice with theoretical worst-case guarantees . Further , we are interested in methods that work for multiple sketching applications . Typically , sketching is very application-specific . The sketch construction and approximation algorithm are tailored to individual applications , like robust regression or clustering ( Sarlos , 2006 ; Clarkson & Woodruff , 2009 ; 2014 ; 2017 ; Cohen et al. , 2015 ; Makarychev et al. , 2019 ) . Instead , we consider three applications at once ( regression , LRA , k-means ) and propose generalizable methods , as well as extending previous application-specific work . Our results . At a high level , our work ’ s aim is to make sketch learning more effective , general , and ultimately , practical . We propose a framework for constructing and using learned CountSketch . We chose CountSketch because it is a sparse , input-independent sketch ( Charikar et al. , 2002 ) . Specifically , it has one non-zero entry ( ±1 ) per column and does not need to be constructed anew for each input matrix it is applied to . These qualities enable CountSketch to be applied quickly , since sparse matrix multiplication is fast and we can reuse the same CountSketch for different inputs . Our “ learned ” CountSketch will retain this characteristic sparsity pattern and input-independence1 , but its non-zero entries will range in R. We list our main contributions and follow this with a discussion . • Two-stage sketch optimization : to first place the non-zero entries and then learn their values . • Theoretical worst-case guarantees , two ways : we derived a time-optimal method which applies to MRR , LRA , k-means , and more . We also proved a simpler method works for k-means . • SOTA experimental results : we showed the versatility of our method on 5 data sets with 3 types . Our method dominated on the majority of experiments . • Theoretical analysis on the necessity of two stages : we proved that including the first stage is strictly better for LRA and two common input distributions . • Empirical demonstration of the necessity of two stages : showed that including the first stage gives a 12 , 20 % boost for MRR , LRA . Our sketch learning algorithm first places the sparse non-zero entries using a greedy strategy , and then learns their values using gradient descent . The resulting learned CountSketch is very different from the classical CountSketch : the non-zero entries no longer have random positions and ±1 values . As a result , the usual worst-case guarantees do not hold . We sought a way to obtain worst-case guarantees that was fast and reasonably general . Our solution is a fast comparison step which performs an approximate evaluation of learned and classical sketches and takes the better of the two . Importantly , we can run this step before the approximation algorithm without increasing its overall time complexity . As such , this solution is time-optimal and applies to MRR , LRA , k-means , and more . An alternate method was proposed by a previous work , but it was only proved for LRA ( Indyk et al. , 2019 ) . This “ sketch concatenation ” method just involves sketching with the concatenation of a learned and a classical sketch . Since it is somewhat simpler , we wanted to extend its applicability . In a novel theoretical result , we proved this works for k-means as well . We also ran a diverse set of experiments to demonstrate the versatility and practicality of our approach . We chose five data sets spanning three categories ( image , text , and graph ) to test our method on three applications ( MRR , LRA , k-means ) . Importantly , these experiments have real-world counterparts . For example , LRA and k-means can be used to compress images , applying SVD ( LRA ) to text data is the basis of a natural language processing technique , and LRA can be used to compute approximate max cuts on graph adjacency matrices . Ultimately , our method dominated on the vast majority of tests , giving a 31 , 70 % improvement over classical CountSketch for MRR , LRA . Finally , we conducted ablation study of the components of our algorithm . In another novel theoretical result , we proved that including the time-consuming first optimization stage is strictly better than not to for LRA and two input distributions ( spiked covariance and Zipfian ) . Empirically , this is case for all 3 applications . Related work . In the last few years , there has been much work on leveraging ML to improve classical algorithms ; we only mention a few examples here . One related body of work is data-dependent 1While learned CountSketch is data-dependent ( it is optimized using sample input matrices ) , it is still considered input-independent because it is applied to unseen input matrices ( test samples ) . dimensionality reduction , such as an approach for pair-wise/multi-wise similarity preservation for indexing big data ( Wang et al. , 2017 ) and a method for learning linear projections for general applications ( Hegde et al. , 2015 ) . We note that multiplying an input matrix on the left with a sparse sketch is equivalent to hashing its rows to a small number of bins . Thus , we also find connections with the body of work on learned hashes , most of which addresses the nearest neighbor search problem ( see Wang et al . for a survey ) . However , in order to obtain approximation guarantees , our “ hash function ” ( sparse sketch ) must satisfy properties which these learned hashes usually do not , such as affine -embedding ( Def . A.1 ) . In particular , we build off of the work of Indyk et al . ( 2019 ) , which introduced gradient descent optimization for the LRA application . It also gave an LRA-specific method for worst-case guarantees . We have surpassed the sketching performance and breadth of this work . Namely , we introduce a sparsity pattern optimization step which is clearly crucial for sparse sketches . We also provide a more general method for worst-case guarantees and extend their method to k-means . 2 PRELIMINARIES . Our learned sketches have the sparsity pattern of the classical CountSketch . The construction of this sketch is described below . We also define affine -embeddings , which is a class of sketches that includes CountSketch . This -embedding property is desirable because it allows us to prove that certain sketching algorithms give ( 1 + ) -approximations . Definition 2.1 ( Classical CountSketch ) . The CountSketch ( abbreviated as CS ) matrix S has one non-zero entry in each column with a random location and random value in { ±1 } . Definition 2.2 ( Affine Embedding ) . Given a pair of matrices A and B , a matrix S is an affine -embedding if for all X of the appropriate shape , ‖S ( AX −B ) ‖2F = ( 1± ) ‖AX −B‖ 2 F . Notation . We denote the singular value decomposition ( SVD ) of A by A = UΣV > with orthogonal U , V and diagonal Σ. Relatedly , the Moore-Penrose pseudo-inverse of A is A† = V Σ−1U > , where Σ−1 is constructed by inverting the non-zero diagonal entries . 3 FRAMEWORK . We describe a framework for learned CountSketch that can be adopted by many different applications , including least-squares regression ( MRR ) , low-rank approximation ( LRA ) , and k-means clustering . We will return to these applications in the next section . In this section , we first describe how to optimize a CountSketch over a set of training samples . Then , we explain how to use the learned CountSketch to achieve good expected performance with worst-case guarantees . By running a “ fast comparison ” step before the approximation algorithm , we can do this in optimal time . SKETCH OPTIMIZATION . Most applications are optimization problems . That is , they are defined by an objective function , L ( · ) . For example , least-squares regression solves min X ‖AX −B‖2F given A ∈ Rn×d , B ∈ Rn×d′ . Of course , the optimal solution is a function of the inputs : X∗ = arg min X ‖AX −B‖2F = A †B However , in the sketching paradigm , we compute an approximately optimal solution as a function of the sketched ( compressed ) inputs . Taking CountSketch S ∈ Rm×n for m n , we have : X̂∗ = ( SA ) † ( SB ) Our goal is to minimize the expected approximation error with respect to S , which is constrained to the set of CountSketch-sparse matrices ( CS ) . However , this is simply equivalent to minimizing the objective value of the approximate solution . S∗ = arg min S∈CS E ( A , B ) ∼D [ LA , B ( ( SA ) † ( SB ) ) − LA , B ( X∗ ) ] = arg min S∈CS E ( A , B ) ∼D [ LA , B ( ( SA ) † ( SB ) ) ] For ease of notation , we will define G ( · ) as a function which maps a sketch and inputs to an approximate solution . G ( · ) is defined by the application-specific approximation algorithm . For MRR , G ( S , ( A , B ) ) = ( SA ) † ( SB ) . More generally , the sketch optimization objective is : S∗ = arg min S∈CS E A∼D [ LA ( G ( S , A ) ) ] ( 3.1 ) If the application is regression , we let A be ( A , B ) . We will solve this constrained optimization in two stages . For both stages , we approximate the expectation in empirical risk minimization ( ERM ) fashion . That is , we approximate the expectation over the true distribution by averaging over a sampled batch of the training set . Now , in the first optimization stage , we compute positions for the CountSketch-sparse nonzero entries . Then , in the second stage , we fix the positions and optimize the nonzero values . Stage 1 : Placing the nonzero entries . We want to maintain the sparsity pattern of CS ( one nonzero entry per column ) , but we are free to place that nonzero entry wherever we like for each column . A naı̈ve method would be to evaluate the objective for the exponential number of full placements . This is clearly intractable , so we consider a greedy alternative . In essence , we construct the sketch one nonzero entry at a time , and we choose the location of the next entry by minimizing ( 3.1 ) over the discrete set of possibilities . More precisely , we build the sketch S ∈ Rm×n iteratively , placing one nonzero entry at a time . For each nonzero entry , we consider m locations and 2 values for each location ( ±1 ) . We evaluate the sketch optimization objective ( 3.1 ) for all 2m incremental updates to S and choose the minimizing update . In the pseudo-code below , we iterate through the n columns of S , each of which contains a non-zero entry . Note that Sw , j = S + w ( ~ej ~ei > ) adds a single entry w to the i-th column , j-th row of the current , partially-constructed S. Algorithm 1 GREEDY STAGE Require : Atrain = { A1 , ... , AN } with Ai ∈ Rn×d ; sketch dimension m 1 : initialize S = Om×n 2 : for i = 1 to n do 3 : w∗ , j∗ = arg min w∈ { ±1 } , j∈ [ m ] ∑ Ai∈Atrain LAi ( G ( Sw , j , Ai ) ) where Sw , j = S + w ( ~ej ~ei > ) 4 : S [ j∗ , i ] = w∗ 5 : end for For some applications it can be inefficient to evaluate ( 3.1 ) , since it requires computing the approximate solution . For MRR and LRA , the approximate solution has a closed form , but for k-means , it must be computed iteratively . This is prohibitively expensive , since we perform many evaluations . In this case , we recommend finding a surrogate L ( · ) with a closed-form solution , as we illustrate in later sections . Stage 2 : Optimizing the nonzero values . We now fix the positions of the nonzero entries and optimize their values using gradient descent . To fix the positions , we represent S as just a vector of its nonzero entries , ~v ∈ Rn . We will denote H ( ~v ) : Rn → Rm×n as the function which maps this concise representation of S to the full matrix . H ( · ) depends on the positions computed in the last stage , which are fixed . Now , we simply differentiate E A∼D [ LA ( G ( H ( ~v ) , A ) ) ] ( 3.1 ) with respect to ~v . Algorithm 2 GRADIENT STAGE Require : Atrain = { A1 , ... , AN } with Ai ∈ Rn×d ; H ( · ) from Alg . 1 ; learning rate α 1 : for i = 1 to niter do 2 : S = Om×n 3 : sample Abatch from Atrain 4 : ~v ← ~v − α ( ∑ A∈Abatch ∂LA ( G ( H ( ~v ) ) , A ) ∂~v ) 5 : end for

Summary: The authors consider a specific sketching method, CountSketch, and three objective functions defined over the data design matrix: multiple-response regression (MMR), low-rank approximation (LRA), and k-means clustering. They compare the classical CountSketch with a random choices of the {-1,+1}-valued sketching matrix against: (1) Gradient descent optimization of the CountSketch weights. This was previously introduced for LRA and the authors extend it to MMR and k-means. (2) Greedy optimization of the positions of the CountSketch non-zero entries.