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 the problem of ``"sketching" – a popular compression technique in machine learning - used for reducing the size of the data enabling one to quickly compute an approximate solution using this compressed input. This paper introduces a general framework for learning and applying sparse sketching matrices. A two-stage procedure for learning sketches with the same sparsity pattern as CountSketch is proposed which involves first learning the sketch’s non-zero entries, and then optimizing their values. They then show how to apply the obtained sketch so that it has worst case approximation error guarantees. This procedure is applied to three applications, namely least squares regression, low rank approximation (LRA) and k-means clustering. Experimental results demonstrate a substantial reduction in the approximation error compared to other baseline approaches (including classically learned sketches). On the theoretical front, it is shown for regression and LRA that the proposed approach obtains improved error guarantees for fixed time complexity.  Additionally, it is shown for LRA (under certain input distributions) that including the first stage is strictly better than not including it. Finally, a more straightforward way of retaining worst case approximation guarantees for k-means is shown.