Paper: Reweighted Solutions for Weighted Low Rank Approximation (ICLR-2024)

Reviewer comment (Theory):
The output of the main algorithm is not guaranteed to be of low rank, which is inconsistent with the original goal of WLRA. It may be better to discuss more about this type of output, specifically, why it can replace a low-rank matrix in practice.

Author response:
We have indeed departed from outputting a low rank matrix for the WLRA problem, and instead opted for a generalized solution. However, this solution can still be stored using few parameters, and can also be applied quickly if we use our practical variation of Algorithm 2. Our extensive experiments show that these relaxed solutions are indeed useful in practice, and demonstrates that it allows for solutions that can be stored and applied efficiently, allowing for state of the art performance for tasks such as neural network compression with a weighted objective.

Reviewer comment (Theory):
There is no theoretical guarantee for Algorithm 2, which makes it less interesting.

Author response:
While we have not been able to analyze Algorithm 2, our experiments demonstrate its effectiveness in practice. We believe this is already compelling evidence for practitioners to adopt our approach.

Reviewer comment (Theory):
On page 2 before Algorithm 1, the authors claim that the storage space $O((n+d)rk)$ is nearly optimal for $r=O(1)$. Does there a matching lower bound of space for WLRA when $r=O(1)$?

Author response:
When $r = O(1)$, the storage space is $O((n+d)k)$. Note that the size of the output requires this many bits of space, since the output is a factorization of the $n\times d$ matrix into a $n\times k$ matrix and a $k\times d$ matrix.  
We thank the reviewer for the question. If $r = 1$, then $\mathbf W$ can be set to the all ones matrix, in which case WLRA becomes the standard low rank approximation problem. A formal information theoretic lower bound for the fact that low rank approximation requires $\Omega((n+d)k)$ bits is standard, and follows from, e.g., applying our lower bound of Theorem 1.5 with $r = 1$. A stronger lower bound against any algorithm which recovers even just an approximate right factor of the SVD is shown in Theorem 82 in Section 10.2 of Boutsidis-Woodruff-Zhong 2016 "Optimal Principal Component Analysis in Distributed and Streaming Models"

Reviewer comment (Theory):
Is rank $r$ known in advance? If not, suppose we run Algorithm 1 with $r > rank(W)$, what happens?

Author response:
We assume that $r$ is known in advance. Note that $r$ can be readily computed with access to the weight matrix $\mathbf W$. As we do in our experiments, one can also approximate $\mathbf W$ by a low rank matrix via an approximate SVD. We do not prove guarantees when $r > \mathrm{rank}(\mathbf W)$.

Reviewer comment (Novelty):
This paper relaxes the classical WLRA problem with two conditions: 1) removing the low-rank requirement for matrix $\widetilde{\mathbf{A}}$; 2) assuming the weight matrix $\mathbf{W}$ is low rank. Given these two conditions, the problem becomes much easier, and the proposed algorithm is kind of trivial. Furthermore, if one discarded the low-rank requirement for matrix $\widetilde{\mathbf{A}}$, the relaxed WLRA problem would be kind of insignificant.

Author response:
Please note that even if we allow $\tilde{\mathbf A}$ to violate the low rank condition, we can still study a non-trivial problem if $\tilde{\mathbf A}$ is structured, i.e. can be stored using much fewer bits of memory or can be applied much quicker than $\mathbf A$. We agree that our proposed algorithm is simple, yet the analysis is nontrivial and our experiments show that our idea leads to great practical benefits over prior work, as demonstrated by our experiments.

Reviewer comment (Novelty):
The unsupervised feature selection with a weighted $F$-norm objective directly follows Theorem 1.2. Also, although the bounds for communication complexity of WLRA problem are almost tight, the results and processes are kind of straightforward. Therefore, the contributions in this paper are quite limited.

Author response:
The fact that our Theorem 1.2 immediately gives the unsupervised feature selection result shows that Theorem 1.2 is a powerful result. We ask the reviewer to consider the importance of the problem of unsupervised feature selection with a weighted objective, which is a commonly encountered problem in practical data analysis, and the fact that such results were not possible prior to our result.

Our communication complexity lower bound indeed uses standard tools, but we consider the conceptual message resulting from this theorem to be our more important contribution. Our result shows that when considering the communication complexity of the WLRA problem, the rank of the weight matrix must necessarily factor into the communication complexity. That is, without the low-rank assumption on the weight matrix, the WLRA problem is difficult in the communication complexity model. We believe this perhaps surprising message to be one of our most important contributions, as a priori it is not clear the rank would parameterize the communication complexity.

Reviewer comment (Novelty):
In Related Work, the following reference is missing. $\textit{Recovery guarantee of weighted low-rank approximation via alternating minimization}$.

Author response:
We have included this work in the Related Work section.

Reviewer comment (Theory):
The main motivation of the algorithm as well as the theoretical results are tailored to the case where the weight matrix W is low-rank. However, the fundamental problem in this setting is that there is no reason why the entrywise inverse matrix W^{\circ -1} is low-rank, which leads to the necessity of computing a dense matrix in Algorithm 1, as the authors state, making the algorithm rather impractical in a large-scale setting.

Author response:
(3) We cannot prove that Algorithm 2 works for every low rank plus sparse matrix - however, we believe we make a significant theoretical advance by applying it to a large class of popular weight matrices used in practice and giving the first theoretical bounds for such matrices. In Appendix A, for each weight matrix described there, either $\mathbf W^{\circ -1}$ already has low rank, or one can apply $\mathbf W^{\circ -1}$  quickly as we describe there. Note that if $\mathbf W = u \cdot v^T$, then $\mathbf W^{\circ -1} = u^{\circ -1} \cdot (v^T)^{\circ -1}$, where zero entries of $u$ or $v$ can be replaced with an arbitrary number (e.g., $1$ for concreteness) for the reason described in (2) above. Thus, $\mathbf W^{\circ -1}$ does partially inherit low rank structure from $\mathbf W$. Also, as noted in Appendix A, this inheritance of low rank structure holds for larger-rank matrices $\mathbf W$ provided the combinatorial rectangles defined by the supports of the low rank factors are disjoint, which together with our ability to handle sparsity, solves five different previously studied and practical weight patterns - in fact for every motivating weight matrix studied in Musco-Musco-Woodruff we give an efficient algorithm. Besides our theoretical results, which at least to us are surprising since a priori $\mathbf W^{\circ -1}$ may not look easy to apply, we give experimental results for an even wider class of weight matrices $\mathbf W$ by heuristically computing a low rank approximation of $\mathbf W^{\circ -1}$.

Reviewer comment (Reproducibility):
With the practical variant of Algorithm 2, only the storage issue of the resulting approximation is mitigated, but not the fact that W^{\circ -1} needs to be computed in a dense manner, which requires O(nd) of storage in intermediate calculations.

Author response:
In Appendix A, we have added a discussion of how a broad family of structured matrices formed by the sum of support-disjoint rank 1 matrices and a sparse matrix can be stored in a way such that $\mathbf W^{\circ -1}$ can be applied and stored efficiently, without forming the matrix $\mathbf W$. This family captures, for example, all of the classes of structured matrices of interest discussed in the prior work of Musco-Musco-Woodruff "Simple Heuristics Yield Provable Algorithms for Masked Low-Rank Approximation" including Low-Rank Plus Sparse, Low-Rank Plus Diagonal, Low-Rank Plus Block Diagonal, Monotone Missing Data Pattern, and Low-Rank Plus Banded matrices.

Reviewer comment (Reproducibility):
While the communication complexity discussion of Section 1.1.2 is interesting, it is unclear whether and how the algorithm can be implemented efficiently in a distributed manner.

Author response:
Our algorithm essentially only relies on the primitive of an SVD, which is known to be efficiently implementable in a distributed setting (see, e.g., Boutsidis-Woodruff-Zhong 2016 "Optimal principal component analysis in distributed and streaming models"). Our work shows that this primitive is indeed useful for weighted low rank approximation, especially for structured weight matrices (see Response 1)x, and that this algorithm is nearly optimal for certain distributed formulations of the weighted low rank approximation problem.

Reviewer comment (Novelty):
While the authors claim that their result is the first about communication complexity for weighted low-rank approximation problems, it seems that communication complexity for such results was already previously discussed, e.g., in Musco, Cameron, Christopher Musco, and David Woodruff. "Simple Heuristics Yield Provable Algorithms for Masked Low-Rank Approximation." _Innovations in Theoretical Computer Science Conference (ITCS 2021)_. 2021.

Author response:
(4) We will stress that in the Musco-Musco-Woodruff paper, there is no communication or distributed setting at all. The notion of ``communication complexity" is a combinatorial notion about a matrix, and that is used to describe the rank of their bicriteria approximations. In our submission, however, we describe a distributed setting and use communication complexity to mean the maximum length of a protocol transcript, over all inputs and random choices.

Reviewer comment (Theory):
Can you discuss the time complexity of the resulting algorithms and intermediate space complexity of Algorithm 1 and Algorithm 2?

Author response:
Both Algorithms 1 and 2 involve computing rank $rk$ approximations to matrices of size $n\times d$. Exact SVD can be computed in $O(nd^{\omega-1})$ operations using $O(nd)$ words of space, whereas randomized approximation algorithms can approximate the Frobenius objective up to a $(1+\varepsilon)$ factor in time and space $O(\mathsf{nnz}(\mathbf A)) + \tilde O(n) \mathrm{poly}(k/\varepsilon)$, where $\mathsf{nnz}(\mathbf A)$ denotes the number of nonzero entries of $\mathbf A$ (Clarkson-Woodruff 2013). Thus for dense matrices where $\mathsf{nnz}(\mathbf A) = O(nd)$, the intermediate calculations do not affect the space complexity of the algorithm. We note that in the setting of machine learning applications, dense linear algebra is often used due to improved hardware utilization, and thus $\mathsf{nnz}(\mathbf A) = nd$ is a common setting.

Reviewer comment (Novelty):
Please discuss your result in the context of previous communication complexity results.

Author response:
The study of communication complexity of WLRA is, to the best of our knowledge, new in this work. The work by Musco-Musco-Woodruff mentioned by the reviewer indeed discusses communication complexity and WLRA together. However, in that work, communication complexity is used as a tool to analyze the "complexity" of the weight matrix $\mathbf W$, and the analysis results in an offline algorithm for WLRA. That is analogous to our use of the rank of $\mathbf W$ as a tool to analyze the ``complexity'' of $\mathbf W$. On the other hand, that work does not address the problem of how a solution to the WLRA problem can be communicated between servers, which is the topic of our work.

Reviewer comment (Reproducibility):
I did not find the code for the timing experiments in the provided code submission. Furthermore, it would be good if the hyper parameter choices of the reference algorithms are provided in the manuscript, e.g. of adam and em.

Author response:
The code for the timing experiments was provided in `weighted_lra.py` in the provided code submission. The adam algorithm was run with 100 epochs, with an initial learning rate of 1.0 that is decayed by a factor of 0.7 for every 10 steps. The em algorithm was run for 25 epochs. These details have been included in the manuscript.