STORM: Sketch Toward Online Risk Minimization

1 INTRODUCTION . Empirical risk minimization ( ERM ) has been independently studied by the machine learning , statistics , and computer science communities . However , conventional ERM algorithms are becoming more expensive in terms of latency , memory , and energy due to the need for accurate inference over rapidly increasing data volumes ( He et al. , 2016 ; Vaswani et al. , 2017 ) . The net effect is an increase in storage costs of large data matrices ( Guhaniyogi & Dunson , 2015 ) , network latency for personalized inference/recommendations , and carbon emissions . There is a growing need for approximation algorithms that can minimize the cost , memory , and energy associated with training and inference ( Lin et al. , 2020 ; Motamedi et al. , 2016 ) . Additionally , there are distinct advantages of training on a small memory local/personalized device , as data is generally collected on these devices . This limits the data transfer to big servers that consumes substantial resources ( Bill Allcock , 2001 ; Xu , 2013 ) and minimizes the privacy vulnerabilities ( Mendez Mena et al. , 2018 ) . Existing compression methods generally fall into two categories : sketching and sampling . Linear Sketches : A large family of projection-based sketches uses ideas from compressed sensing , matrix factorization , and dimensionality reduction . A common theme is to approximate the large N × d data matrix with a smaller representation . For example , multiplication with a random ±1 matrix yields a sketch that attains the information-theoretic lower bound for compressed linear regression ( within a constant factor ) ( Clarkson & Woodruff , 2009 ) . Random Gaussian , Hadamard , Haar , and Bernoulli sign matrices have also been used in this framework with varying degrees of theoretical and practical success ( Dobriban & Liu , 2019 ) . The “ sketch-and-solve ” technique has also been extended to regression with the ` 1 objective and several other linear problems ( Sohler & Woodruff , 2011 ) . This framework can not accommodate regularized objectives because the model is computed directly from a set of precomputed sufficient statistics . Recent work by Shi & Phillips ( 2021 ) has shown that a matrix sketch based on the frequent directions method ( Liberty , 2013 ; Ghashami et al. , 2016 ) can be used for ridge regression . This algorithm uses a sketch to approximate the most costly part of the ridge regression objective , XTX + λI , in a streaming setting . The memory is O ( ld ) , where X ∈ Rn×d and l > 2 . That is , the memory is linear in the dimensions , with a minimum size of ( 2d+ 1 ) × floating precision that can be very large for high-dimensional data . Coresets and Sampling : A reasonable strategy is to extract a core subset ( or coreset ) of points that efficiently represent the complete dataset . The coreset framework supports the approximation of loss function sums within a ( 1± ) error margin . There are coresets for important problems such as ` p-regression ( Dasgupta et al. , 2009 ) and logistic regression ( Munteanu et al. , 2018 ) . Classification objectives also admit small coresets ( Har-Peled et al. , 2007 ; Tukan et al. , 2020 ) , as do many other regularized ERM problems ( Reddi et al. , 2015 ) . The problem is that many coresets require superlinear construction time . For example , many of the algorithms in the previous paragraph require k-means clustering , leverage scores , or multiple cycles through the dataset ( Tukan et al. , 2020 ; Munteanu et al. , 2018 ; Har-Peled et al. , 2007 ) . Coresets are also tricky to merge since it is not straightforward to decide which points to discard . As a result , composable and streaming coresets are an active area of research ( Indyk et al. , 2014 ; Karnin & Liberty , 2019 ) . Although leverage scores can be approximated online ( Cohen et al. , 2016 ) , the process is computationally expensive and not well suited to edge devices . In many practical problems , random sampling is the only feasible sampling baseline that supports regularized ERM . In this work , we propose a new sketching framework ( STORM ) that supports both regularization and streaming algorithms . Given an N -point dataset D = { ( x1 , y1 ) ... ( xN , yN ) } of d-dimensional examples observed in the streaming setting ( Fiat , 1998 ) , the goal is to construct a small structure S that can estimate a regression and classification loss over D. STORM consists of two main components , 1 ) STORM loss for regression and classification , 2 ) A tiny array of integer count values called STORM sketch . This array is sufficient to perform empirical risk minimization on STORM loss . In practice , our data structure allows for optimizations that are not possible with other methods . Specifically , our contributions are as follows : 1 . We propose an online sketching algorithm ( STORM ) that can estimate a special class of loss functions using only integer count values . Our sketch is appropriate for distributed data settings because it is small , embarrassingly parallel , and mergeable via addition . 2 . We characterize the set of STORM-approximable loss functions . The price of efficient loss estimation is a restriction on the kinds of losses we can use . While we can not directly estimate popular regression and classification losses , we construct surrogate losses for ` 2 regression and large margin classification . 3 . We show how to optimize models over STORM sketches . The count-based nature of our sketch makes gradient descent difficult , and black-box optimization techniques are of limited use in high-dimensional spaces . We propose an optimization algorithm that is custom-tailored to our sketch . 4 . We conduct a rigorous empirical evaluation with ridge regression and classification . We show that our approach has better train/test loss performance and lower variance than sampling baselines in the low-memory , resource-constrained computational regime . For some datasets , we can achieve strong performance with memory smaller than even one sample . 2 BACKGROUND . 2.1 EMPIRICAL RISK MINIMIZATION . In the standard statistical learning framework , we are given a training dataset D = { ( xi , yi ) ∈ D × Y } ni=1 of examples and asked to select a function h : X → Y that can predict y given x . In this paper , we will present our algorithms for the data space X = Rd and output space Y = R , but most of our results hold generally for other metric spaces too . The learning problem is to select a hypothesis hθ ( a function parameterized by θ ) that yields good predictions , as measured by a loss function loss ( hθ ( x ) , y ) . In empirical risk minimization ( ERM ) , we select θ to minimize the average loss on the training set . θ̂ = arg min θ n∑ i=1 loss ( hθ ( xi ) , yi ) In the case of Linear regression , hθ ( x ) = 〈θ , x〉 . For most applications , θ is found using the least-squares or ` 2 loss : loss ( hθ ( x ) ) = ||hθ ( x ) − y||22 . The unconstrained ` 2 loss is smooth and strongly convex , with favorable convergence criteria . The parameter θ is typically found using gradient descent or via a closed-form solution from the matrix formulation of the model y = Xθ . For our discussion , it will be important to express the loss in terms of the concatenated vector [ xi , yi ] : θ̂ = arg min θ n∑ i=1 ( 〈 [ xi , yi ] , [ θ , −1 ] 〉 ) 2 = arg min θ ‖y− Xθ‖22 2.2 LOCALITY SENSITIVE HASHING . A Locality sensitive hash ( LSH ) function ( l ( x ) ) provides the same hash value for similar points with high probability . The notion of similarity is usually based on the distance measure of the metric space X . For instance , there are LSH families for the Jaccard ( Broder , 1997 ) , Euclidean ( Datar et al. , 2004 ; Dasgupta et al. , 2011 ) and angular distances ( Charikar , 2002 ) . If we allow asymmetric hash constructions ( i.e . l1 ( x1 ) = l2 ( x2 ) , l1 6= l2 ) , then x1 and x2 may collide based on other properties such as their inner product ( Shrivastava & Li , 2014 ) . Asymmetric LSH functions use one hash function ( l1 ) for data and a different function ( l2 ) for queries . To accommodate asymmetric functions , we use the following definition of LSH , Our definition is a strict generalization of the original definition from Indyk ( 1998 ) Indyk & Motwani ( 1998 ) , which can be recovered by setting probability thresholds . Definition 1 . We say that a hash family F is locality-sensitive with collision probability k ( · , · ) if for any two points x1 and x2 in Rd , l ( x1 ) = l ( x2 ) with probability k ( x1 , x2 ) under a uniform random selection of l ( · ) ∈ F . The signed random projection ( SRP ) family for the angular distance ( Goemans & Williamson , 1995 ; Charikar , 2002 ) is an example of a symmetric LSH family . The SRP family is the set of functions l ( x ) = sign ( w > x ) , where w ∼ N ( 0 , Id ) . The SRP collision probability is k ( x1 , x2 ) = 1− 1 π cos−1 ( 〈x1 , x2〉 ‖x1‖2‖x2‖2 ) The inner product hash Shrivastava & Li ( 2014 ) is an example of an asymmetric LSH . To get this hash function , replace x1 with [ x1 , 0 , √ 1− ‖x1‖22 ] and x2 with [ x2 , √ 1− ‖x2‖22 , 0 ] in the SRP function . This procedure essentially uses different hash functions for x1 and x2 , since we augment x1 ( the data ) in one way and x2 ( the query ) in a different way . The collision probability k ( x1 , x2 ) is now a monotone function of the inner product 〈x1 , x2〉 . This construction implicitly assumes that x1 and x2 are inside the unit sphere , but other hash functions do not require this assumption ( Shrivastava & Li , 2015 ; Neyshabur & Srebro , 2015 ; Huang et al. , 2018 ) .1 RACE Sketch : The RACE sketch is a 2D integer array S ∈ ZR×B that we index using LSH functions ( Coleman & Shrivastava , 2020 ; Luo & Shrivastava , 2018 ) . Each row of S is called an Array-of-Counts ( ACE ) data structure ( Luo & Shrivastava , 2018 ) , and the entire sketch is a set of R independent ACE repetitions ( RACE ) . To construct a RACE sketch from a dataset D , we begin with an empty R×B integer array and a set of R independent LSH functions { l1 ( x ) ... lR ( x ) } . Each LSH function has the collision probability k ( θ , x ) . When an element x arrives from the data stream , we hash x to get R different hash values . We increment row i at location li ( x ) and repeat this process for all elements in the dataset . To query the sketch with a query θ , we return the average value of S [ i , li ( θ ) ] over the R rows . RACE sketches are interesting because they estimate the sum of LSH collision probabilities over the dataset with low variance for any query θ. Theorem 1 . Unbiased RACE EstimatorColeman & Shrivastava ( 2020 ) Let X = S [ i , li ( θ ) ] . Then E [ X ] = ∑ x∈D k ( x , θ ) var ( X ) ≤ ( ∑ x∈D √ k ( x , θ ) ) 2 1Regardless , we often scale and center the dataset when using inner product hash functions in practice .

Summary: The paper considers approximately computing ERM problems for real data, with an emphasis on regression and classification, in the *very low space* regime (most experiments are below 1Kb). This is done my maintaining a standard locality sensitive hashing table, but with specialized hash functions that allow for computing a surrogate for common regression and classification ERM objectives.   Due to the extreme space constraint, optimization is done with a heuristically characterized SGD. A significant amount of experiments is presented.