Few-Shot Data-Driven Algorithms for Low Rank Approximation

1 introduction . Low-rank approximation ( LRA ) of matrices is a fundamental operation in data analysis and machine learning , forming the basis for many efficient algorithms that involve storing and analyzing large matrices . Its importance has led to a large body of work on algorithms for computing LRA . While it is well-known how to find an optimal LRA of a given matrix in polynomial time using the Singular Value Decomposition ( SVD ) , the degree of that polynomial makes the algorithm too slow for large matrices , defeating the purpose of leveraging it to design fast and scalable algorithms . With this motivation in mind , a long and prominent line of research has discovered fast approximate LRA algorithms , laying the foundations for a rich and expansive theory of fast algorithms for numerical linear algebra ( e.g. , [ 26 , 51 , 15 , 16 , 18 , 24 ] ) . Due to their strong practical performance , these algorithms have become standard and widely used tools in machine learning over massive datasets ( e.g. , [ 57 , 30 , 43 , 27 , 13 , 53 ] ) . See the surveys [ 41 , 56 , 42 ] for a comprehensive overview of this area . These algorithms are mostly based on sketching . In its most basic form , the input matrix A ∈ Rn×d is left-multiplied by a carefully chosen sketching matrix S ∈ Rm×n , whose row dimension is much smaller than that of A ( i.e. , m n ) . Then , a good low rank approximation to A can be computed from SA . Reducing the row dimension from n to m improves the algorithmic efficiency in multiple ways . First , SA can be stored using less space than A , which is important in the context of space-efficient streaming algorithms . Furthermore , the dimensionality reduction step can be combined with other steps to obtain LRA algorithms with running time nearly linear in the size of A. Identifying suitable sketching matrices is a highly non-trivial task , and by now there are several 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . known randomized constructions , including sparse sign matrices [ 15 , 16 , 43 , 46 ] , random projection matrices [ 51 , 2 ] , and subsampled Hadamard matrices [ 21 , 13 ] . Unfortunately , a notable downside of sketching-based algorithms is that they lead to a larger approximation error than the optimal LRA , due to the lossy nature of sketching . Since this error propagates downstream to all tasks that use LRA as a subroutine , minimizing this error is highly important . This raises the question — can we design better sketching matrices , that offer a better tradeoff between the sketch width m and the approximation quality ? Such an improvement would directly translate into a more accurate ( or alternatively , more space efficient ) streaming algorithm for LRA ( see Section 1.2 for more details ) . Furthermore , it would also be an important step towards obtaining more time-efficient algorithms as well . Data-driven approach to LRA . A natural approach to the above challenge is to rely on the fact that in many cases , we have access to large quantities of past data ( seen as a training set ) , which is related to the future input matrices ( seen as a test set ) . In such a scenario , we could potentially use a preprocessing ( or learning ) stage on the past data in order to better handle future data . This approach , often called data-driven or learning-augmented algorithm design , has gained increased popularity lately , and led to improved algorithms for a host of classical problems ( see Section 1.3 ) . For low rank approximation , this line of work was recently pioneered by [ 33 ] , who proposed a datadriven LRA algorithm that we refer to henceforth as IVY . Specifically , IVY learns the values of the non-zero entries in a sparse sketching matrix S , instead of choosing them at random ( as in classical LRA , e.g. , [ 16 ] ) . Learning is done using stochastic gradient descent ( SGD ) , by minimizing the empirical ( normed ) difference between matrices from the training set and their corresponding LRAs computed using S. Computing gradients for this loss function is a highly non-trivial operation , but nonetheless , [ 33 ] show it can be made tractable by exploiting modern auto-gradient software and GPU hardware , which are typically used nowadays to train deep neural networks . Challenges and limitations . Despite its improved accuracy , the IVY algorithm has several drawbacks related to its computational cost . First , it requires a differentiable implementation of the SVD , which is slower than standard implementations , leading to long training times . Second , it relies on specialized software and hardware ( automatic gradient computation in PyTorch , and access to a GPU ) . Third , in order to gain significant improvement over the oblivious ( non-data-driven ) baseline , it requires training over dozens of existing matrices , which may not necessarily be available . Fourth , it chooses the sketching matrix in a black-box , non-interpretable way , using black-box learning . All these limitations are also shared by [ 3 , 38 ] . This leads us to consider whether much more efficient data-driven algorithms can attain a comparable advantage in their LRA accuracy . 1.1 Our Results . In this work , we propose a new approach to data-driven LRA , motivated by a different and more analytic-minded viewpoint . As mentioned above , previous work takes a “ black-box ” approach to the problem , directly minimizing the empirical loss via SGD . However , since LRA is a well-studied topic with a rich underlying mathematical structure , it is worth to consider alternative approaches that make direct use of our analytic understanding . This may allow us to simplify the learning algorithm — for example , by devising more effective loss functions — and perhaps even eliminate the heavy learning machinery altogether , in favor of analytic techniques . This approach leads us to data-driven algorithms that are considerably faster to train than IVY , do not necessarily require access to specialized hardware and software , and use only one or few training matrices ( thus , they perform few-shot learning ) . Specifically , we give the following two algorithms : One-shot closed-form algorithm : Our first algorithm does not require learning it all ; it only picks a random matrix from the dataset , and extracts a sketching matrix by an analytic computation . It requires neither GPU access nor autograd functionality . Few-shot SGD algorithm : For our second algorithm , we use a new loss function for learned LRA ( given in eq . ( 3 ) in Section 2 ) . It is faster to compute and differentiate than the empirical loss used in [ 33 ] , being a low degree polynomial , and can be ( locally ) optimized with only 2 or 3 steps of SGD , requiring only that many randomly chosen matrices from the dataset . Both algorithms achieve significantly better sketch length to accuracy tradeoffs than data-oblivious LRA . When compared to IVY and its variants , they do attain similar ( and eventually better ) accuracy than our algorithms if allowed to run for sufficiently many SGD iterations , but with much longer training times and many more training matrices . In addition , our closed-form algorithm is also interpretable , as it computes a sparse sketching matrix S with a clear mathematical meaning . Specifically , the non-zero entries of S are set to be singular vectors of a block-partitioned randomly chosen training matrix , with a block partition determined by the sparsity pattern of S. The singular vectors in each block are chosen either greedily , or at random proportionally to their squared singular values . Finally , we experiment with combining our approach with IVY , by initializing the weights of the sketching matrix according to our one-shot algorithm , instead of at random as in [ 33 ] . The result is an algorithm that uses the same memory , time and training samples as IVY , with significantly improved accuracy . One takeaway message for designing data-driven algorithms , is that while it may seem tempting to invoke heavy-weight machine learning technology , there are possibly much simpler algorithms that attain similar performance with a significantly reduced computational cost , and an analytic understanding of the problem may help us arrive at those algorithms . Limitations of our approach . Like most previous methods [ 33 , 2 ] , we focus on optimizing the nonzero entries of a sparse sketching matrix with a fixed sparsity pattern ( chosen at random ) . Learning the sparsity pattern from data during training may improve performance even further . However , since the search space is discrete and very large , it is not clear how to do so without a very large increase in training time ( more in section 3 ) . We leave this as an important direction for future work . Broader impact . Incorporating data into algorithms entails risks related to bias and fairness , as well as environmental effects of massive processing . Our methods carry benefits of interpretability and of significantly reduced GPU utilization , which may help alleviate these concerns . However , our work is focused on a generic algorithmic operation , and any application to a specific data domain would require further consideration of possible bias in the data as well as of privacy concerns . 1.2 Preliminaries . Low rank matrix approximation . In the LRA problem , the input is a matrix A ∈ Rn×d , say with n ≥ d , and an integer 0 < k d. The goal is to output a matrix A′ ∈ Rn×d of rank k which minimizes the approximation error ‖A − A′‖2F , where ‖·‖F denotes the Frobenius matrix norm . It is a standard fact ( known as the Eckhart-Young-Mirsky Theorem ) that an optimal solution , often denoted as Ak , can be computed deterministically in time O ( nd2 ) via the singular value decomposition . A long line of work has focused on algorithms that incur slightly larger approximation — namely ‖A−Ak‖2F + ‖A‖2F , or better yet ( 1 + ) ‖A−Ak‖2F , for an arbitrarily small > 0 — and run in time nearly linear in the size of A , and polynomial in k and −1 . See [ 56 ] for a survey . The SCW algorithm . The data-oblivious LRA algorithm used as the basis for [ 33 ] , henceforth called SCW , is due to [ 16 ] , building on [ 51 , 15 ] . It uses a random sketching matrix S ∈ Rm×n ( where m n is called the sketching dimension ) , chosen such that each column contains one nonzero entry at a uniformly random row , with a value which is uniformly random in { −1 , 1 } . Given an input matrix A and the sketching matrix S , SCW computes the best LRA of A in the rowspace of A . This can be computed in several equivalent ways , for example , by Algorithm 1 below . Algorithm 1 : Sketch-based matrix low-rank approximation Input : Input matrix A ∈ Rn×d , sketching matrix S ∈ Rm×n . Output : Rank-k approximation Ã ∈ Rn×d of A . 1 Procedure SCW ( S , A ) : 2 Ũ Σ̃Ṽ > ← SVD of SA 3 [ AṼ ] k ← Optimal rank-k approximation of AṼ ( computed by full SVD ) 4 return [ AṼ ] kṼ > It is proven in [ 16 ] that a sketching dimension of m = poly ( k , −1 ) suffices for the output rank-k matrix A′ to satisfy ‖A − A′‖2F ≤ ( 1 + ) ‖A − Ak‖2F with high probability . Furthermore , the algorithm can be implemented as a two-pass streaming algorithm , using space O ( m ( n+ d ) ) . Thus , reducing the sketch length m directly reduces the storage requirements of the algorithm . Data-driven LRA . The data-driven version of the LRA problem , introduced in [ 33 ] , has two phases . In the training phase , we are given a training set of matrices , Atrain ⊂ Rn×d . Our goal is to construct , or learn , a sketching matrix S. In the testing phase , we are given a test set of matrices , Atest ⊂ Rn×d . We compute a rank-k approximation for each of them using the SCW algorithm ( Algorithm 1 ) , but with the learned sketching matrix from the training phase instead of a random sketching matrix . The goal is to choose S so as to minimize the testing error , Err ( S , Atest ) : = ∑ A∈Atest ‖A− SCW ( S , A ) ‖2F . ( 1 ) The IVY algorithm . In [ 33 ] , S is initialized by sampling it from the same distribution as SCW – sparse random signs . The sparsity pattern ( i.e. , the locations of the non-zero entries in S ) remains fixed , while the values of those entries are learned by SGD over the training loss , S = argminS′Err ( S ′ , Atest ) = argminS′ ∑ A∈Atest ‖A− SCW ( S′ , A ) ‖2F . ( 2 ) This requires computing gradients over the steps of SCW , which in turn uses two invocations of the SVD ( in steps 2 and 3 of Algorithm 1 ) . While not all implementations of the SVD are differentiable , IVY employs a differentiable one which is based on the power method . Safeguard guarantees . A common desideratum of data-driven algorithms is to guarantee they never perform much worse than their data-oblivious counterparts ( see , e.g. , [ 47 , 44 ] ) , which could happen for example by over-tailoring to an irrelevant training set . For IVY , [ 33 ] note this can be achieved by simply concatenating an oblivious sketching matrix to the learned sketching matrix . This property is shared by subsequent work on data-driven LRA [ 3 , 38 ] as well as by our algorithms .

Summary: Problem: Sketching based techniques (such as randomized algorithms, etc) is the go to approach for deriving a low rank approximation. Despite their power in terms of space and speed to deduce the solution, these algorithms induce large approximation errors, which regresses at the downstream tasks and need to be addressed Solution: A data-driven approach for low rank approximation -- A low rank representation is learned by optimizing a frobenius loss function over the input training and approximation. There are some challenges with this approach. The proposed algorithm is less data hungry with a two step analytical approach. Results: Significant gains in terms of accurate low rank approximations with less resource usage.