Convergence of Alternating Gradient Descent for Matrix Factorization

Summary

This problem considers a matrix factorization problem: it seeks to minimize a function of the form $f(X,Y) = 1/2 \Vert A - XY^T \Vert_F^2$. In general, matrix factorization problems have applications to matrix sensing, phase retrieval, and are seen as prototypical non-convex optimization problems. The special case of the matrix factorization problem provided in this paper is very simple, and should be seen as a first theoretical step to understand other matrix factorization problems with more practical applications. The contribution of this paper is to propose a particular initialization of the optimization problem that induces an improved convergence rate of alternating gradient descent to recover the matrix $A$. This is supported by both rigorous results and simulations.

Strengths

The paper is clear and well-written. I checked the proofs of the main text and believe they are correct. I think the paper addresses an interesting problem in matrix factorization, but I am not convinced by the specific initialization that they propose. This is detailed below.

Weaknesses

I would like the authors to discuss more the motivations underlying their paper. My understanding is that factorizing a matrix $A$ in a product $XY^T$ is not hard, for instance one can choose $X=A$ and $Y = Id$. Of course, it is interesting to see if gradient descent is able to recover a factorization of the matrix $A$, given that there is non-convexity and non-smoothness. However, it seems to me that the initialization that you are proposing is similar to the solution $X=A$ and $Y = Id$. In this sense, the papers [YD21] and [JCD22] are stronger as they prove convergence from a Gaussian initialization, that seems further away from the solution. Of course, a contribution of your paper is to improve the convergence rates. However, it does not seem to me that obtaining optimal convergence rates is an identified challenge in the literature. For instance, you compare your convergence rate to the rate of [JCD22], but again these authors use a more challenging initialization, and prove that incremental dynamics appear, which is much harder than proving a simple convergence rate. In light of the above points, I think that the comparison with related works in Section 3.1 is a bit unfair. Given that this model is toyish, it is important to focus on behaviors that we believe to generalize on more sophisticated (matrix multiplication, non-convex) models. While incremental learning is one of them, I didn't understand whether an improved convergence with an atypical initialization satisfies this criteria.

Questions

- Given that the convergence rates are exponential, wouldn't logarithmic y-axes be more suitable to plot the evolution of the errors? - I think it is a bit tough to skip the proof of Lemma 4.1. The least would be to explain that X -> f(X,Y) is a convex ||Y||^2-smooth function and thus you can use classical (but not really "direct"!) manipulations from convex optimization to derive the bounds. - Could you comment on why you are using alternating gradient descent rather that full gradient descent? Is it related to using only marginal smoothness in Lemma 4.1? Could your techniques be adapted for full gradient descent? - l.312: Are you using different Gaussian matrices for X_0 and Y_0? - l.319: Can you comment on the very small value for \nu? - Figure 3: It seems that the benefits of the proposed initialization reduces when using large stepsizes. Could you try even larger stepsizes? Can you push simulations to the limit of stability of the gradient descent equations? Typos - l.407: "non-asympototic" -> "non-asymptotic" - l.80: "and" -> "an" - l.263 and footnote same page: "PL-equality" -> "PL inequality"? - l.280: "interation" -> "iteration" - l.300: inconsistent notation for the transpose

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

1 poor

Limitations

See weaknesses section.

Summary

This paper provides a new analysis for the convergence of alternating gradient descent on the low-rank matrix factorization problem $min_{X,Y} f(X,Y) = ||XY - A||_F^2$. The authors show that, by warm-starting the solution using the target matrix $A$, and appropriate step size scaling, they can achieve $\epsilon$ error in $f$ in $O(\kappa^2 \log \frac{||A||_F}{\epsilon})$ iterations of alternating gradient descent, where $\kappa$ is the pseudo-condition number of $A$, as long as the rank of $XY$ is larger than the rank of $A$ by a constant factor. This improved the best known dependence of $\kappa$ from $\kappa^3$ to $\kappa^2$. Also, in contrast to some previous works, the initialization step does not require computing an SVD of $A$.

Strengths

- The main technical contribution of this paper is the introduction and analysis of an asymmetric warm starting rule for matrix factorization. To the best of my knowledge this is an original technical contribution. - The theoretical analysis is interesting and advances the state of the art. Matrix factorization is a fundamental optimization task that among others, can be viewed as a subproblem of neural network training. As a result, progress in theoretical understanding of optimization algorithms is highly relevant. - The experimental results are interesting and show a surprising quality gap between different initialization methods. - The paper is very well written and technically sound.

Weaknesses

- It seems to me like the assumption in line 454 that $V^\top \Phi_1$ has i.i.d entries is incorrect, and it only has i.i.d columns. In this case, a different lemma should be used in place of Proposition A.1. Does this change anything in the bounds? - I have some comments about the experiments in Section 6. It seems like the authors compare different algorithms using the same step size, however in my opinion it would be fairer to compare using the best fixed step size for each algorithm. This is especially so because of the special imbalanced step size setup in the authors' algorithm. For example, in Figure 3, Random and ColSpan significantly improve when increasing the step size from $0.001$ to $0.1$. What happens if the step size is increased even more? Can we rule out that the improvement in convergence is because of the step size magnitude? Some typos: - decent -> descent in multiple theorem/lemma statements. - Line 231: it should not be a strict subset

Questions

1. Could the authors give some intuition behind the importance of imbalanced scaling? I understand that it is used to balance the norms of $X$ and $Y$, but the fact that it is so significant in the experiments is somewhat surprising so I am wondering if there is something deeper. 2. Could this approach be used in the noisy setting, i.e. when $A$ is not exactly low rank? What are the difficulties?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

NA

Summary

The authors explore the use of alternating gradient descent (AGD) with a fixed step size for asymmetric matrix factorization. They demonstrate that a finite number of AGD iterations can achieve a high-probability -optimal factorization, even when starting from an asymmetrical random initialization. Empirical evidence supports the significant convergence rate improvements resulting from their proposed initialization. The authors' proof leverages a uniform Polyak-Lojasiewicz (PL) inequality and Lipschitz smoothness constant, providing a simpler approach for analyzing nonconvex low-rank factorization problems. The findings in this paper challenge the conventional belief that symmetrical randomization is necessary when the underlying data structure allows for symmetry. This novel result, supported by solid theoretical justification, showcases promising improvements and highlights the potential for broader impacts in machine learning and optimization.

Strengths

- Excellent novelty - Theoretical proof is elegant - Promising numerical results

Weaknesses

It is worth noting that there are numerous competitors beyond first-order methods in the field to solve an exact rank-r matrix decomposition problem. To further enrich the discussion, it would be valuable if the authors explored potential avenues for generalizing their results. For example, investigating the application of their findings in matrix completion or the singular value decomposition (SVD) of a low-rank matrix with added noise could provide insights into the broader applicability of their approach.

Questions

There are two parts to the initialization. One is the $A\Phi_1$ vs $\Phi_2$, another is $\sqrt{\eta}$ vs $\frac{1}{\sqrt{\eta}}$. I am curious whether $\sqrt{\eta}$ vs $\frac{1}{\sqrt{\eta}}$ is the major factor for bringing improvements. For example, what are the results if the authors use symmetrical random initialization with $\sqrt{\eta}$ on one side and $\frac{1}{\sqrt{\eta}}$ on another side?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

I don't foresee any potential negative societal impact resulting from this research.

Summary

This paper proves an improved convergence bound on alternating gradient descent for asymmetric matrix factorization problem, i.e. given $A \in R^{m \times n}$, finding $X \in R^{n \times d}$ and $Y \in R^{m \times d}$ that minimize $||XY^{\top} - A||_F^2.$ This paper establishes that if $A$ is rank $r$ for some $r < d,$ (so the optimal solution is 0) then the algorithm converges to a solution with error at most $\varepsilon ||A||_F^2$ in roughly $\frac{d}{(\sqrt{d}-\sqrt{r})^2}\kappa^2 \log \frac{1}{\varepsilon}$ iterations, where $\kappa = \frac{\sigma_1(A)}{\sigma_r(A)}$. Thus, if the problem is constantly over-parameterized, i.e. $d-r = \Omega(r),$ this gives a dimension independent convergence bound. A crucial piece in their algorithm and analysis is starting from an unconventional asymmetric initialization, where $X$ is initialized as $\frac{1}{\sqrt{\eta}}A\Phi_1$ and $Y$ as $\sqrt{\eta}\Phi_2,$ where $\Phi_1, \Phi_2$ are (normalized) iid Gaussian matrices, and $\eta$ is the step size. The paper also reports several numerical simulations demonstrating the significant advantage of their initialization.

Strengths

The (asymmetric) matrix factorization problem is both an important problem in itself, and an important test bed for developing techniques to analyze non-convex optimization. I think this paper takes a significant step towards understanding (alternating) gradient descent in this setting. It improves on the previous best known bound for GD which was ~ $\kappa^3.$ I find the improvement due to the asymmetric initialization quite illuminating, and I think it will inspire other advantageous (yet quick) initializations for descent algorithms.

Weaknesses

I don't see any major weaknesses.

Questions

- I think the bound on line 30 is incorrect? there should be a square-root on d and r.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

N/A