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: The paper proposes two simple methods to learn the sketching matrix in low-rank matrix approximation. They serve as strong baselines to compare against other data-driven methods that learn the sketching matrix from training data. The first method (one-shot), 1Shot1Vec / 1Shot2Vec, sets the sketching matrix using the singular vectors of blocks of the matrices in the training data. The second method (few-shot), uses a different objective from the usual data-driven low-rank approximation that avoids costly calls to the SVD. The paper empirically validates these two approaches, showing that they are fast and are competitive with more complex methods.