On the Benefits of Weight Normalization for Overparameterized Matrix Sensing
Authors: Yudong Wei
Liang Zhang
Bingcong Li
Niao He
Venue: ICLR 2026 Poster (2026)
Keywords: Weight normalization
Overparameterization
Matrix sensing
Non-convex optimization
Primary area: optimization

Abstract: While normalization techniques are widely used in deep learning, their theoretical understanding remains relatively limited. In this work, we establish the benefits of (generalized) weight normalization (WN) applied to the overparameterized matrix sensing problem. We prove that WN with Riemannian optimization achieves linear convergence, yielding an $\textit{exponential}$ speedup over standard methods that do not use WN. Our analysis further demonstrates that both iteration and sample complexity improve polynomially as the level of overparameterization increases. To the best of our knowledge, this work provides the first characterization of how WN leverages overparameterization for faster convergence in matrix sensing.

Published as a conference paper at ICLR 2026
ON THE BENEFITS OF WEIGHT NORMALIZATION FOR OVERPARAMETERIZER MATRIX SENSING
Yudong Wei Liang Zhang Bingcong Li Niao He
ETH Zurich
yudwe @ethz.ch
(liang. zhang, bingcong.li, 1i, niao. hlao.hee@inf.ethz:
ABSTRACT
While normalization techniques are widely used in deep learning, their theoretical understanding remains relatively limited. In this work, we establish the benefits of (generalized) weight normalization (WN) applied to the overparameterized ma- trix sensing problem. We prove that WN with Riemannian optimization achieves linear convergence, yielding an exponential speedup over standard methods that do not use WN. Our analysis further demonstrates that both iteration and sample complexity improve polynomially as the level of overparameterization increases. To the best of our knowledge, this work provides the first characterization of how WN leverages overparameterization for faster convergence in matrix sensing.
INTRODUCTION
Normalization schemes, such as layer, batch, and weight normalization, are essential in modern deep networks and have proven highly effective for stabilizing training in both vision and language models (loffe & Szegedy, 2015; Ba et al., 2016; Salimans & Kingma, 2016). Despite their practical success, theoretical explanations of why they work remain elusive, even for relatively simple problems.
This work focuses weight normalization (WN), which decouples parameters (i.e., variables) into directions and magnitudes, and then optimizes them separately. It has recently regained consider- able attention because of the seamless integration with LoRA (Hu et al., 2022), leading to several powerful strategies for parameter-efficient fine-tuning of large language models; see c.g., (Liu et al., 2024; Lion et al., 2025). Yet, theoretical support for WN remains relatively limited. Prior results in (Wu et al., 2020) show that WN applied to overparameterized least squares induces implicit regu- larization towards the minimum & fo-no solution. The implicit regularization of WN on diagonal linear neural networks is studied in (Chou et al., 2024). WN is also observed to reduce Hessian spectral norm and improve generalization in deep networks (Cisneros-Vellan
al., 2025).
Our work broadens the understanding of WN by establishing its merits for the overparameterized matrix sensing problem. The goal is recover low-rank positive semi-definite (PSD) matrix A E STh from linear measurements. In the vanilla formulation without WN, one can exploit the low-rankness of ground-trune -ruth matrix, i.c., FA = rank(A) < for efficient parameterization. Specifically, we can optimize for Y R such that YY = A (Burer & Monteiro, 2005). The overparameterized regime r > FA of interest due to the need of exact recovery without knowing r' priori. This problem has wide applications in machine learning and signal processing (Candès al., 2013), and serves as a popular testbed for theoretical deep learning given non-convexity and rich loss landscape; see e.g.. (Li et al., 2018; Jin al., 2023; Arora et al., 2019).
Without WN, prior work (Xiong et al., 2024) establishes sublinear lower bound on the convergence rate when the above sensing problem is optimized via gradient descent (GD), even with infinite data samples. We circumvent this lower bound by i) extending WN for coping with matrix variables; and, ii) proving that applying this generalized WN with Riemannian gradient descent (RGD) ables a linear convergence rate the finite sample regime, leading to an exponential improvement. Remarkably, WN leverages higher level of overparameterization achieve both faster convergence
"Equal supervision.

---

Published as a conference paper at ICLR 2026
Table 1: Comparison with existing algorithms for overparamettrrize matrix sensing. WN gives ex- act convergence with linear rate. "E.C." is the abbreviation of "exact convergence", that is, whether the reconstruction error bound will go to zero when the iteration number + co. UB and LB are short for upper and lower bound, respectively. OP stands for overparameterization.
Algorithm
WN
E.C.
Initialization
Convergence Rate
Faster with OP
GD (UB) (Stöger& & Soltanolkotabl, 2021)
X
X
Small random
N/A 2 inge
GD (LB) (Xiong al,, 2024)
Small random
x
RGD (Theorem: 3.2)
Random
(r-r x Spplr
and lower sample complexity. To the best of knowledge, this is the first theoretical result demon- strating that normalization benefits from overparameterization.
More concretely, our contributions are summarized follows:
+ Exponentially faster rate. For overparameterized matrix sensing problems, we prove that ran- (m-ra)"
domly initialized WN achieves linear convergence rate of exp(- ( m ()), where the condition number of the ground-truth matrix A. This linear rate exponentially faster than the sublinear lower bound ( Te obtained without WN. Moreover, additional overparameteriza- tion in WN provides quantifiable benefits: the iteration complexity scales down polynomially the overparameterization level r increases; see Table for a summary.
22
Two-phase convergence behavior: We further investigate the optimization trajectory and reveal a two-phase behavior in WN. The iterates first move from a random initialization to a neighborhood around global optimum, potentially traversing and escaping from several saddle points in polynomial time. Our results demonstrate that this phase ends faster with additional overparameterization Once iterates approach the global optimum, a local phase begins. With the benign loss landscape shaped by WN, we prove that linear convergence rate can be obtained.
+ Empirical validation. We conduct experiments on overparameterized matrix sensing using both synthetic and real-world datasets. The numerical results corroborate our theoretical findings.
1.1 RELATED WORK
Overparameterized matrix sensing. Overparameterized matrix sensing arises from many ma- chine learning and signal processing applications such as collaborative filtering and phase retrieval (Schafer al., 2007; Srebro Salakhutdinov, 2010; Candès et al., 2013; Duchi et al., 2020). The problem is now canonical benchmark in theoretical deep learning, mainly because the loss land- scape riddled with saddle points and lacks global smoothness or global PL condition. Conver- its population loss, i.e., matrix factorization, can be found
gence analyses for various algorithms in (Ward & Kolda, 2023; Li al., 2025; Kawakami & Sugiyama, 2021). ). Small random initialization in overparameterized matrix sensing has been studied in (Stöger & Soltanolkotabi, 2021; Jin al., 2023; Xiong et al., 2024; Xu et al., 2023), while (Ma et al., 2023; Zhuo et al., 2024; Cheng & Zhao, 2024) are based on spectral initialization. Besides saddle escaping under small initialization, an- other intriguing phenomenon is that overparameterization can exponentially slow the convergence of GD compared to the exactly parameterized case (Zhuo al., 2024; Xiong al., 2024). Our work proves that WN avoids this slowdown and achieves an improved rate. Moreover, additional overparameterizztion leads to faster convergence and lower sample complexity.
Riemannian optimization. Riemannian optimization is naturally connected to WN for learning the direction variables, which are constrained on a smooth manifold, e.g., a sphere. Existing literature has extended gradient-based methods to problems with smooth manifold constraints; et al., 2008; Smith, 2014; Mishra et al., 2012; Boumal, 2023). This work follows standard notions of Riemannian gradient descent (RGD). In its simplest form, RGD iteratively moves along the negative direction of the Riemannian gradient, obtained by projecting the Euclidean gradient onto the tangent space, and then maps the iterate back to the manifold via retraction.
e.g., (Absil
2

---

Published as a conference paper at ICLR 2026
Notational conventions. Bold capital (lowercase) letters denote matrices (column vectors); (-) and || | refer to transpose and Frobenius norm of matrix; || denotes the spectral (f2) norm for matrices (vectors); (A, B) Tr(A BB represents the standard matrix inner product; and, (A) denotes the i-th largest singular value of matrix A. Moreover, S and $ denote symmetric and positive semi-definite (PSD) matrices of size m, respectively.
PROBLEM FORMULATION
We foeus on applying WN the symmetric low-rank matrix sensing problem. The objective is to recover a low-rank and positive semi-definite (PSD) matrix A E S from collection of n data {(M, 34) [=) where each feature matrix M, 3th is symmetric and the corresponding label is Y/ Tr A). For notational coneiseness, we let y [y1..... ................................ 3n]" Iº and define linear mapping M ~R" with [M(A)]; Tr(M, A). Given that r A rank(A) m, a parameter economical formulation is based on the Burer-Monteiro factorization (Burer & Monteiro, 2005) that introduces matrix RRx such that YY approximates A accurately. This leads
[IM(YYT) - -yl?
min YERMXr
(1)
Despite its seemingly simple formulation, loss landscape contains saddle points, hence achieving global optimum from random initialization is nontrivial. Moreover, overparameterization. i.e.,
r > TA. is often considered in practice to ensure exact recovery of A without prior knowledge of its rank. It is established in (Xiong et al., 2024) that such overparameetrrizatin induces optimization challenges even in the population setting (n -> oo). In particular, a lower bound of GD shows that |YY,Y, A || converges faster than 211/t), where & is the iteration number. This rate is exponentially slower than the linear one when A is known to employ r = (Ye & Du, 2021).
Applying WN problem (1). For a vector variable, WN decouples it into direction and magnitude, and optimizes them separately. Extending this idea to matrix variables in (1), we leverage polar decomposition to write Y = Xê, where X E St(m,, r) lies in a Stiefel manifold and e E S Here, the Stiefel manifold St(m,r) r) is defined as {X E Rm = I,). One can geometrically interpret as orthonormal bases for an r-dimensional subspace, thus representing "directions", and Ô captures the "magnitude" of a matrix. Substituting in (1), we arrive at
1 1 M x yl1 s.t. X St(m,))) r), Ô
min x,6
The above problem can be further simplified by merging θθ into single matrix O e S and ii) relaxing the PSD constraint on to only symmetry, i.c., O S". This relaxation achieves the same global objective in the overparameterized regime, yet significantly improves computational efficiency by avoiding SVDs or matrix exponentials needed for optimizing over PSD cones (Van- denberghe & Boyd, 1996; Todd, 2001). In sum, applying WN gives the objective
min f fX, +) = } M - 1 s.t. X E St(m,,) r), OES.. X.6
(2)
For convenience, we continue to refer to this generalized variant as WN, since aligns with the direction-magnitue decomposition paradigm. Similar reformulations of (1) have appeared in (Mishra et al., 2014; Levin et al., 2025). The former empirically studies the faster convergence on matrix completion problems, while the latter tackles local geometry around stationary points. Our work, on the other hand, characterizes the hehaviors of WN along the entire trajectory and clarifies its interaction with overparameterizatton.
2.1 SOLVING WN RIEMANNIAN OPTIMIZATION
Generalizing the vector WN! on matrix problems, Riemannian optimization is adopted for coping with the manifold constraint
rt(m,r). We simply treat the manifold as an embedded one
While the practical update rule of WN (Salimans & Kingma, 2016, eq. (4)) lies between Riemannian and Fuclidean optimization, (Wu al., 2020, Lemma 2.2) shows that the limiting flow Riemannian flow.
3

---

Published as a conference paper at ICLR 2026
Algorithm Riemannian gradient descent (RGD) for solving WN (2)
1: Input: Initial point X, St(m,r), r), Ou Sr stepsizes 17,
2: for 0, ,, do
3: Calculate G22 the Riemannian gradient of X,, via (3)
4: Update X,41 via (4)
I/ The direction variable
5: Calculate K, Ve (X O₂) via (9)
6: Update O,11 via (5)
The magnitude variable
7: end for
8: Output: 01+1
in Euclidean space. Extensions to other geometry are straightforward. To optimize the direction variable X,, let G, := (X,, ④,) denote the Euclidean gradient on X; (a detailed expression is given (8) of Appendix C). The Riemannian gradient for X, can written as
X x G. X.).
:= -
(3)
2
Further applying the polar retraction" to ensure feasibility, the update for X, is given by
X (X, (), G G
where 1) >> 0 is the stepsize. Detailed derivations of (3) and (4) are deferred to Appendix C. Note that polar retraction is used here for theoretical simplicity. Shown later sections, other popular retractions for Stiefel manifolds such as QR and Cayley' share almost identical performance numerically.
An alternative update method is adopted for the magnitude variable O₁ Denote its gradient as K, Ve (XX+1, a₂), whose expression can be found from (9) in Appendix C. We use GD with a stepsize > 0 to optimize θ i.e.
= O ,K,,
(5)
This update ensures feasibility the symmetric constraint on O,, VI 0, whenever initialized with On E Sr; see a proof in Lemma F.9. The step-by-step procedure for solving (2) is summarized in Algorithm 1, and it is termed as RGD for future reference.
3 ON THE BENEFITS OF WN
This section demonstrates that WN delivers exact convergence at a linear rate for overparameterized matrix sensing (2) and leverages additional overparameterization to yield faster optimization and lower sample complexity. Recall that the rank of A is denoted by rA- Let the compact SVD of A be A U where U E R and E S Without loss of generality, we assume (Σ{) and or (Σ)= 1/k with x 2 denoting the condition number. We will use the restricted isometry property (RIP) (Recht al., 2010), a standard assumption matrix sensing, in our proofs; see more in, e.g., (Zhang et al., 2021; Stöger & Soltanolkotabi, 2021; Xu et al., 2023; Xiong et al., 2024).
Definition 3.1 (Restricted Isometry Property (RIP)) The linear mapping M(-) is (r, 8)-P with 8 e [0, 1). for all matrices A 5tm of rank most satisfies
(1 5) < | < (1 + ()
RIP ensures that the linear measurement approximately preserves the Frobenius norm of low-rank matrices. This property has been shown to hold for wide range of measurement operators. For ex- ample, when M, is symmetric Gaussian, sample size of n Omrr/8' suffices to guarantee that (r, RII holds with high probability. A detailed discussion with illustrative examples is provided in Appendix A.3. With these preparations, we are ready to uncover the merits WN.
LLt X E St(m, r) and point in its tangent space G E TxSt(m, r). The polar retraction for X + G given by Rx(G) = (X + GT, + G G) -1/12
"See e.g., (Absil al., 2008), for detailed discussions retractions.
4

---

Published as a conference paper at ICLR 2026
3.1 MAIN RESULTS
We consider WN under random initialization, meaning that Xo is chosen uniformly at random from the manifold St(m, r). One possible approach is to set X Zo (Z) Zn)-1/2, where the entries of Zo E TRxx are i.i.d. Gaussian random variables (0, 1) (Chikuse, 2012).
Theorem 3.2 Consider solving the WN-aided sensing problem (2) initialized with random Xu E St(m,r) r) and O E S satisfying ||ool| S 2. Assume that TA S m and M(-) is (r+ra+1, RIPP ())R with 8 = ( (r-ra) Algorithm I using stepsizes Ll O kmmArtra) and μ 2 generates sequence (Xt, O )= With high probability over the initialization, this sequence proceeds in two distinct phases, separated by burn-in time to with an upper bound O r2
(r-TA))
Ty-ra)
i) Saddle phase. For some universul constant E E00, 1), it follows that
c2(r ra) |X,,,] AlE S 2 V/T KAmaprr
+ 1, 1 1565
i) Linearly convergent phase. For some universal constant Cg E (0, 1), it is guaranteed that
ca(r ca(r-ra) x+m2r2rA
toto
X - AlIF 3
Vt to +1.
This theorem showcases the two regimes of convergence behavior using randomly initialized RGD. In the first phase, the upper bound of reconstruction error |X AllF is proved to monoton- ically decrease over iterations. This seemingly slow convergence arises from the potential saddle escape that will be discussed in the next section. Eventually, RGD achieves linear rate until exact convergence, i.e., lim X, - A IF 0. Next, we break down Theorem 3.2 to demon- strate the benefits of WN for the overparameterized matrix sensing from two different perspectives.
Optimization benefits of WN include i) faster convergence rate, and ii) less stringent initialization requirements. Theorem 3.2 shows that WN achieves exact convergence to the ground-truth truth matrix A with a linear rate. In contrast, without WN, the convergence behavior of randomly initialized GD on (1) weaker. Specifically, (Stöger Soltanolkotabi, 2021) shows that GD can only attain a constant reconstruction error with early stopping, but not guarantee last-iteration convergence. On the other hand, (Xiong et al., 2024) establishes a lower bound for exact recovery of GD, giving a sublinear dependence on /; see a detailed comparison in Table 1. In addition, our guarantee of this linear rate is obtained without striet requirements on initialization, which stands in stark contrast to the WN setting, where the magnitude of random initialization must be carefully controlled, often inversely proportional to x (Stöger& & Soltanolkotabi, 2021; Jin et al., 2023; Xu et al., 2023).
WN makes overparameterization friend. Because the additional parameters induce computation and memory overheads, it is natural to expect more gains from overparameterization. It can be seen from Table 1 that GD does not benefit from overparamettrizzaton, while the benefits of overparame- terization for WN are twofold. Setting r pr for some p > 1, one can rewrite the upper bound of which decreases polynomially with p. In the linearly convergent
the burn-in time to as O mm 2
phase, WN achieves convergence rate of exp (- O (p- (r- t)). which is also faster with larger k+m
In terms of iteration complexity, this translates into polynomial improvement with the level of overparameterization To quantitatively understand the merits of overparameterization,, we consider two cases. In the mildly overparameterized regime, where ra + c for some constant e O(1). the convergence rate reads exp
O ( ( )). When the level of overparameteriiztion increases to Cr a the rate improves to exp O ( )). Through comparison, it is readily seen that additional overparameterizzaion yields up factor of (r4)) improvement in the exponent. On the statistical side, the sample complexity of WN determined by the RIP assumption on M(.). Under
the Gaussian design, as detailed in Appendix A.3, the RIP holds w.h.p. when O (r-ra)
Notably, the sample complexity n reduces polynomially as r increases. In particular, following the same analysis for the convergence rate, this reduction reach up factor of (r)
5

---

Published as a conference paper at ICLR 2026
10
grad of
grad
10
10
10.4
10
10
10
1000
1500
500
1000
1500
1000
1500
2000
(a) Squared reconstruction error
(b) Gradient norm
(c) |X,,,X; Apll
10
1.00
0.75
0.50
500 1000
1500
2000
500
1000
1500
2000
0.00
500
1000
1500
2000
(d) 0r (.)
(c) Tr( V7 (T)
(f) (O.)
The saddeee (i.e., sequential learning) behaviors in WN. The x-axis corresponds to the iteration number, and the y-axis follows the subfigure title. (a) Each plateau signifies saddle; (b) gradient norm at saddles drops by orders; (c) saddles strongly relate to the best rank-p pan-pppppon approximation of A; (d) sequential learning in the alignment between X, and U; (e) sequential learning in the alignment between X, and U1; and, (f) sequential pattern in the magnitude variable O
Figure 1:
3.2 GEOMETRY PROOF SKETCH
The proof of Theorem 3.2, while involved, admits clear geometrical interpretation originated from the direction-magniuee decoupling of WN. Here, we only focus on the "direction" X; to gain more intuition. Given that both X, and are bases of linear subspace, it is desirable that span(U) C span(X,) at convergence. Equivalently, the principle angles between span (U) and (Xan at optimal should all be This can be depicted via the alignment matrix 4 singular values coincide with the cosine of these principle angles (Björck & Golub, 1973). Our proof builds upon this and shows that Tr( (( + T', i.e., the two subspaces align.
X, whose
The convergence unfolds in two phases. In the first phase, Tr. grows from near 0 (due to random initialization) to near optimal r' 0.5. Through consecutive lemmas, it can be shown
that Tr( 1- 4 TT,) Tr( Tr ( 2 O (r-ra) This monotonic increase in alignment sures polynomial time escape (potential) saddles, and also translates to the decreasing bound of | - A/F Theorem 3.2. The second phase starts after Tr( / > TA 0.5, where the alignment error A Tr( Φ Tr(I, 0 2, decreases linearly to 0.
DIVING DEEPER INTO THE SADDLE PHASE
Next, take closer look at the convergence of RGD WN in the saddle phase, that is t to. or equivalently Tr < тA 0.5. Our numerical experiments in Figure indicate that RGD traverse sequence of saddles. The saddle-to-saddil behavior is known for GD (1) (Li et al., 2021; Jin et al., 2023). This section shows that this behavior persists for (2), yet can be faster with a higher level of overparameterization. To bypass the randomness associated with M, we begin by pinpointing the saddles for the population loss, i.c., problem (2) in the infinite data limit n + co. More precisely, the objective is given by foo X, ① = XOXT All

---

Published as a conference paper at ICLR 2026
109
-GD
109
RGD
109
-930m
RGD
RGD
10
-00
100
10
500
RGD 00)
10
Gay 75
10"
-G).-75)
-GD)-75)
100,=100
.GDDy1000
10
10
10
10
10"4
19
10"2
10
10
30
16
100
330
(a) Different K
(b) Different r
(c) Full rank r m
Figure Convergence comparison RGD on WN and GD on (1) under varying problem conditions
(squared reconstruction error vs. iteration). (a): WN enables RGD to converge linearly regardless of K; (b): with WN larger r leads to shorter saddle phase and a faster convergence rate; (c): WN converges remarkably fast in the full rank case
Lemma 4.1 For given p {0, 1,..,....77 17a-1, 1], let Ap be the best rank-p approximation of A, i.e., A, arg rank(A) rank Asso ||A AlE In particular, let Ao 0. A point (X, +) is saddle the population loss foo XOX Ap and Tr(X X)
Lemma 4.1 indicates that the saddles of foo are closely related to the best rank-p approximation of A. It further suggests that a baddle-oosaddle dynamic is aligned with incremental learning": the algorithm successively learns A, for increasing p until the ground-truth matrix is recovered. Lemma 4.2 below shows that in the finite-sample regime, the saddles of foe also have small gradient norm on f, i.e., no larger than O (r-ra) under the parameter choices of Theorem 3.2.
Lemma 4.2 Assume that M(-) is + and || ||| < 2, the finite sample loss in (2) satisfies & θ) VfX FF S 12mó and Vef oz (X, θ) Vef(X, O)Il S mm. Here, V denotes the Riemannian gradient with respect to X.
Having characterized the saddles, we now turn to the saddle-to saddle-to-saddle trajectory in Figure 1. This figure traces the optimization trajectory of Algorithm 1 on WN with 300, r = and k 3, with more details shown in Appendix G.I. Figure la plots the squared reconstruction error across iterations. Each plateau marks escape from a saddle, as confirmed by the small gradient norm shown in Figure lb. Figure le further shows that these saddles exactly those characterized in Lemma 4.1, where |, Al In other words, each saddle escape corresponds leaving the neighborhood of Ap.
10,
2 for E {0, 1..... F 1} stays close to sequentially.
In addition, the optimization variables, geometrically interpretable as direction and magnitude, also exhibit a sequential learning behavior. For the direction variable X,, the singular values of Φ T (which characterize the squared cosine of the principle angles between X, and U) are visualized in Figure 1d. Further, let U1 E I m (m -ra) be an orthonormal basis for the orthogonal complement of span (U) The alignment of X, and U j is plotted in Figure le, with the alignment matrix defined as V, == X.X, The singular values of the magnitude variable θ are plotted Figure 1f. A clear sequential learning pattern observed among all these figures.
Lastly, we highlight that polynomial time is needed to escape all saddles: Theorem 3.2 bounds the duration of this phase to be most O ( (m"mm)"²² 2,2
(r-ra)² iterations. This bound decreases with larger r,
indicating that overparameteriizaioon facilitates saddle escape under WN.
5 NUMERICAL EXPERIMENTS
Numerical experiments using both synthetic and real-world data are conducted in this section to validate our theoretical findings for WN on overparameterized matrix sensing problems. In the ex- periments with synthetic data, the target matrix is generated as A = UΣU E gmxm where
Also known deflation; see e.g, (Ge al., 2021; Anandkumar al., 2014; Seddik al., 2023)

---

Published as a conference paper at ICLR 2026
10
100
RGD )
Cayley QR
RODK=10
1
10
RODK=10 9
10
10
GDK=10")
GD(=-10
10
10
GD(2-109)
10
10
10
10
10
10
10-14
10-14
0
250
500
750
1000
0
400
800
1200
(a) Various retractions
(b) Noisy measurement
Figure 3: Additional numerical results of WN (squared reconstruction error vs. iteration).
U TR is random matrix with orthonormal columns, and Σ STA diagonal matrix with condition number K. In the image reconstruction experiments, the target matrix A is directly con- structed from the underlying image. For the measurements, we use independent random Gaussian feature matrices (M; h to ensure RIP. More details on setups are deferred to Appendix G.2.
5.1 FASTER CONVERGENCE OF WN
Under various choices of the condition number K, we compare the convergence behavior of RGD on problem (2) with random initialization, against that of GD (1) with small random initialization.
In this experiment, we consider target matrices with large condition numbers, i.e., x: E (50, 75, 100). F. = 3, and 60000. The squared reconstruction
We set 10, of iterations is plotted in Figure 2a. We observe that WN enables RGD to converge linearly to zero after a saddle phase, regardless of the condition number k. This is consistent with our theoretical result in Theorem 3.2. In contrast, GD slows down to sublinear rate after its initial phase, yielding substantially larger errors at the same iteration count.
versus the number
5.2 ON THE BENEFIT OF OVERPARAMETERIIZATOO
Next, demonstrate that WN leverages overparameterrization for faster convergence. To this end, we consider randomly initialized problem instances of (1) and (2) under different
In this experiment, focus setting with overparameteriizatonn is chosen from E (50, 75, 100), and the number of measurements is set to 50000. RGD with random initialization and GD is with small random initialization. The squared reconstruction eтrоr versus the number of iterations is plotted in Figure 2b.
300, F. 5, and
10. The level of
The results show that under WN, RGD converges faster as increases. This behavior consistent with our analysis in Theorem 3.2. In comparison, although the theoretical convergence rate given by (Xiong al., 2024) is independent rf our empirical results indicate that larger leads slightly slower convergence of GD. Moreover, Figure 2b clearly shows that saddle escape becomes faster with larger as reflected in shorter plateaus earlier onset of linear convergence. Figure 2b also shows that larger in WN leads to steeper slope in the linearly convergent phase, demonstrating that additional overparameteriiztion prompts faster rate. This aligns well with our theoretical observations and discussions in Sections 3.1 and 4.
We also demonstrate that WN is remarkably effective in the full rank setting with r = m in Figure 2c, where the convergence three instances with m = E (50, 75, 100), rA {10, 15, 20), , E (1, 15, 50), and n = 30000 is plotted. The faster convergence arises from the fact that at initializa- tion. X, E St(m,, m) already aligns with the target subspace spanned by U, ie.. Tr(I, 2 Φ 0. Equivalently, this is the case where only the magnitude is optimized. This faster convergence implies that learning the correct direction (i.e., U) is more challenging than magnitude.
8

---

Published as a conference paper at ICLR 2026
(a) Ground truth
(b) RGD reconstruction
(c) GD reconstruction
(d) Ground truth
(e) RGD reconstruction
(l) GD reconstruction
Figure 4: The advantages of WN on image reconstruction.
5.3 ADDITIONAL EXPERIMENTS
Moreover, additional experiments reveal other interesting behaviors of WN.
Alternative manners of refraction. Although algorithm for WN tackles only the polar retrac- tion, other popular retractions share similar performance. In Figure 3a, we plot the performance of Algorithm with different manners for retraction, such Cayley and QR,
instance of (2) with m 10, 5, F = 3, 6 = 2, = 1000. The three curves of squared reconstruction errors nearly scale the errors of Cayley and QR by and 1/3, respectively.
coincide. For better visualization,
Noisy measurements. To examine the robustness of WN, we consider setting with corrupted labels, i.e., yi Tr(M, A) by for i.i.d. Gaussian noise by NO, g²). Figure 3b compares WN with the vanilla problem (I) under the choices of & 10
= 10 and & 10 10-5, It can be seen that RGD holds a linear rate under all choices of &, and the final squared reconstruction error stabilizes around O( On the other hand, the error of GD is mainly confined by its slow convergence rate. This demonstrates that the power of WN carries to noisy settings as well.
5.4 IMAGE RECONSTRUCTION EXPERIMENTS
Lastly, we evaluate the advantages of WN on two image reconstruction problems.
The first experiment follows (Duchi et al., 2020) to consider generalized phase retrieval problem on 32 x 32 horse image from the CIFAR- 10 dataset (Krizhevsky & Hinton, 2009). The image is converted to grayscale and vectorized as a E 1024 Standard lifting reformulation converts this problem to a sensing problem on rank-one ground-truth matrix A aa E 11024, see (Candès & Li, 2014). The second considers direct matrix sensing of a structured image given by A E $128 with " = In both cases, we set the overparametcrization level to r = 100 and use n 50000 feature matrices. RGD and GD randomly initialized and run for ÍRGD 100, faD 200 iterations in both experiments to make the overall runtime comparable; see Appendix G.2.2 for details.
The reconstructions from the two experiments presented in Figure 4. As shown, WN enables RGD to achieve more accurate recovery of the ground truth compared to GD. These results demon- strate that WN provides significant improvement for image reconstruction problems.

---

Published as a conference paper at ICLR 2026
6 CONCLUSION
This work provides new theoretical insights into the role of weight normalization (WN) in over- parameterized matrix sensing. We prove that randomly initialized WN with proper Riemannian optimization guarantees linear rate, yielding an exponential improvement on overparameterized sensing problems without WN. Moreover, we show that overparameterizatoon can be exploited un- der WN to achieve faster optimization and lower sample complexity. Our analysis also reveals a two-phase convergence behavior, with detailed characterizations of faster convergence in both phases. Numerical experiments on both synthetic and real-world data further validate our findings. Future work includes extending these results broader non-convex learning settings, such as tensor problems (Tong et al., 2022), and developing new algorithms that build on WN.
REPRODUCIBILITY STATEMENT
We have taken several steps ensure the reproducibility of our work. For the theoretical results, we provide complete proofs of all theorems and lemmas in Appendices E and F. For the empirical results, Section 5 contains detailed descriptions of the numerical experiments, and the experimental setups are presented in Appendix G.
REFERENCES
P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization Algorithms on Matrix Manifolds. Princeton University Press, 2008
Animashree Anandkumar, Rong Ge, Daniel Hsu, Sham M Kakade, and Matus Telgarsky. Tensor decompositions for learning latent variable models. Journal of Machine Learning Research, 15 (1):2773-2832, 2014.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06455 2016.
Åke Björck and Gene H. Golub. Numerical methods for computing angles between linear subspaces. Mathematics of Computation, 27(12)):579-594 1973.
Nicolas Boumal. An Introduction to Optimization on Smooth Manifolds. Cambridge University Press. 2023.
Samuel Burer and Renato D. C. Monteiro. Local minima and convergence in low-rank semidefinite programming. Mathematical Programming, 103(3):427-444, 2005.
Emmanuel J Candès and Xiaodong Li. Solving quadratic equations via PhaseLift when there are about as many equations as unknowns. Foundations of Computational Mathematics, 1015):1017-
1026, 2014.
Emmanuel J Candès, Thomas Strohmer, and Vladislav Voroninski. PhaseLift: Exact and stable signal recovery from magnitude measurements via convex programming. Communications on Pure and Applied Mathematics, 66(8):1241-1274 2013.
Emmanuel J. Candès and Yaniv Plan. Tight oracle inequalities for low-rank matrix recovery from minimal number of noisy random measurements. IEEE Transactions on Information Theory, 57 (4):2342-2359, 2011.
Cheng Cheng and Ziping Zhao. Accelerating gradient descent for over-parameterrieed asymmetric low-rank matrix sensing via preconditioning. In IEEE International Conference Acousties, Speech and Signal Processing, pp. 7705-7709, 2024.
Yasuko Chikuse. Statistics on Special Manifolds, volume 174. Springer Science & Business Media, 2012.
10

---

Published as a conference paper at ICLR 2026
Hung-Hsu Hsu Chou, Holger Rauhut, and Rachel Ward. Robust implicit regularization via weight nor- malization. Information and Inference: A Journal of the IMA, 13(3)iaac022, 2024.
Pedro Cisneros- Velarde, Zhijie Chen, Sanmi Koyejo, and Arindam Banerjee. Optimization and generalization guarantees for weight normalization. Transactions on Machine Learning Research, 2025. ISSN 2835-8856.
John C Duchi, Oliver Hinder, Andrew Naber, and Yinyu Ye. Conic descent and its application to memory-efficiet optimization over positive semidefinite matrices. Advances in Neural Informa- fion Processing Systems, 33:8308 8317, 2020.
Rong Ge. Yunwei Ren, Xiang Wang, and Mo Zhou. Understanding deflation process in over- parametrized tensor decomposition. Advances in Neural Information Processing Systems, 34: 1299-1311, 2021.
Gene H. Golub and Charles F. Van Loan. Matrix Computations 4th Edition. Johns Hopkins University Press, 2013.
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen -hun, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation large language models. In International Con- ference on Learning Representations, 2022.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pp. 448-456. PMLR, 2015.
Jikai Jin, Zhiyuan Li, Kaifeng Lyu, Simon Shaolei Du, and Jason D Lee. Understanding incre- mental learning of gradient descent: A fine-grained analysis of matrix sensing. In International Conference on Machine Learning, pp. 15200-15238. PMLR, 2023.
Yuhi Kawakami and Mahito Sugiyama. Investigating overparameterization for non-negative matrix factorization in collaborative filtering. In ACM Conference on Recommender Systems, pp. 685- 690. 2021.
Alex Krizhevsky and Geoffrey Hinton, Learning multiple layers of features from tiny images. Tech- nical report, University of Toronto, 2009.
Eitan Levin, Joe Kileel, and Nicolas Boumal. The effect of smooth parametrizations on nonconvex optimization landscapes. Mathematical Programming, 209(1):63-111) 2025
Bingcong Li, Liang Zhang, Aryan Mokhtari, and Niao He. On the crucial role of initialization for
matrix factorization. International Conference on Learning Representations, 2025.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over- over- wer-parameterized matrix sensing and neural networks with quadratic activations. In Conference on Learning Theory, pp. 2-47. PMLR, 2018.
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference Learning Representations, 2021.
Kai Lion, Liang Zhang, Bingcong Li, and Niao He. PoLAR: Polar-decomposed low-rank adapter representation. arXiv preprint arXiv:2506.03733, 2025.
Shih- Shih-Yang Liu, Chien-Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang- Ting Cheng, and Min-Hung Chen. DoRA: Weight-decomposed low-rank adaptation. In Interna- tional Conference on Machine Learning, 2024.
Cong Ma, Xingyu Xu, Tian Tong, and Yuejie Chi. Provably accelerating ill-conditioned low- rank estimation via scaled gradient descent, even with overparameterizatton. arXiv preprint arXiv:2310.06159, 2023.
Jianhao Ma and Salar Fattahi. Convergence of gradient descent with small initialization for unregu- larized matrix completion. In Conference Learning Theory, pp. 3683-3742. PMLR, 2024.
11

---

Published as a conference paper at ICLR 2026
Bamdev Mishra, K Adithya Apuroop. and Rodolphe Sepulchre. A Riemannian geometry for low- rank matrix completion. arXiv preprint arXiv:/211.1550, 2012.
Bamdev Mishra, Gilles Meyer, Silvere Bonnabel, and Rodolphe Sepulchre. Fixed-rank matrix fac- torizations and Riemannian low-rank optimization. Computational Statistics, 29(3):591-621, 2014.
Benjamin Recht. A simpler approach to matrix completion. Journal of Machine Learning Research, 12(12), 2011.
Benjamin Recht, Maryam Fazel, and Pablo A. Parrilo. Guaranteed minimum-rank solutions of linear matrix equations via nuclear norm minimization. SIAM Review, 52(3):471-50,, 2010.
Mark Rudelson and Roman Vershynin. Smallest singular value of a random rectangular matrix. Communiccatonns on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences, 62(12):1707-1755 2009.
Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in Neural Information Processing Systems, 29, 2016.
J Ben Schafer, Dan Frankowski, Jon Herlocker, and Shilad Sen. Collaborative filtering recommender systems. In The Adaptive Web: Methods and Strategies of Web Personalization, pp. 291-324. Springer, 2007.
Mohamed El Amine Seddik, Mohammed Mahfoud, and Merouane Debbah. Optimizing orthogo- nalized tensor deflation via random tensor theory. arXiv preprint arXiv:2302.05799, 2023.
Steven Thomas Smith. Optimization techniques on Riemannian manifolds. arXiv preprint 1r7 2014.
Nathan Srebro and Russ Salakhutdinov. Collaborative filtering in a non-uniform world: Learning with the weighted trace norm. Advances in Neural Information Processing Systems, 23, 2010.
Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruc- tion. Advances in Neural Information Processing Systems, 34:23831-2388 2021.
M. Todd. Semidefinite optimization. Acta Numerica, pp. 515-560, 2001.
Tian Tong, Cong Ma, and Yuejie Chi. Accelerating ill-conditioned low-rank matrix estimation via scaled gradient descent. Journal of Machine Learning Research, 22(150):1-63, 2021.
Tian Tong, Cong Ma. Ashley Prater-Bennette, Erin Tripp. and Yuejie Chi. Scaling and scalabil- ity: Provable nonconvex low-rank tensor estimation from incomplete measurements. Journal of Machine Learning Research, 23(163): 2022.
Lieven Vandenberghe and Stephen Boyd. Semidefinite programming. SIAM review, 38(1):49-95,
1996.
Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint 101 2010.
Rachel Ward and Tamara Kolda. Convergence of Alternating Gradient Descent for Matrix Factor- ization. Advances in Neural Information Processing Systems, 01-22369-22382 2023.
Xiaoxia Wu, Edgar Dobriban, Tongzheng Ren, Shanshan Wu, Zhiyuan Li, Suriya Gunasekar, Rachel Ward, and Qiang Liu. Implicit regularization and convergence for weight normalization. Ad- vances Neural Information Processing Systems, 33:2835-2847, 2020.
Nuoya Xiong, Lijun Ding, and Simon Shaolei Du. How over over- per-parramterizzaionn slows down gradient descent in matrix sensing: The curses of symmetry and initialization. In International Conference on Learning Representations, 2024.
Weihang Xu and Simon Du Over paramttttttiinnnnoo exponentially slows down gradient descent for learning single neuron. In Conference Learning Theorr, pp. 1155-1198. PMLR, 2023.
12

---

Published as a conference paper at ICLR 2026
Xingyu Xu, Yandi Shen, Yuejie Chi, and Cong Ma. The power of preconditioning in overparame- terized low-rank matrix sensing. In International Conference on Machine Learning, pp. 38611- 38654. PMLR, 2023.
Baturalp Yalçı,, Haixiang Zhang, Javad Lavaei, and Somayeh Sojoudi. Factorization approach for low-complexity matrix completion problems: Exponential number of spurious solutions and failure of gradient methods. In International Conference on Artificial Intelligence and Statisties, 319-341. PMLR, 2022.
Tian Ye and Simon S Du. Global convergence of gradient descent for asymmetric low-rank matrix factorization. Advances in Neural Information Prucessing Systems, 14:1429-1439, 2021.
Jialun Zhang, Salar Fattahi, and Richard Y Zhang. Preconditioned gradient descent for over- parameterized nonconvex matrix factorization. In Advances in Neural Information Processing Systems, volume 34, pp. 5985-5996, 2021.
Mo Zhou, Weihang Xu, Maryam Fazel, and Simon S Du. Global convergence of gradient EM for over-parameterized gaussian mixtures. arXiv preprint arXiv:2506. 06584, 2025.
Jiacheng Zhuo, Jeongyeol Kwon, Nhat Ho, and Constantine Caramanis. On the computational and statistical complexity of over-parameterized matrix sensing. Journal of Machine Learning Research, 25(169):1-47, 2024.
13

---

Published as a conference paper at ICLR 2026
USAGE OF LLMS
The authors conducted all aspects the research, including conception, thcoretical proofs, experi- mentation, analysis, and writing of the manuscripts. Large language models (LLMM)) were employed exclusively for the purpose language refinement.
A MORE ON BACKGROUNDS
A.1 POLAR DECOMPOSITION
The definition of the polar decomposition is provided below; see (Golub & Van Loan, 2013, Section 9.4.3) for a detailed discussion and theoretical background.
Definition A.1 The polar decomposition of matrix X E R with 2 is defined as
X -UPP
where E mxxr has orthonormal columns and E S4 a positive semi-defiitte matrix.
This decomposition can be interpreted as expressing X as the product of directions (U) and mag- nitude part (P). unique when X has full column rank.
A.2 RIEMANNIAN OPTIMIZATION
Riemannian optimization provides principled framework for optimization problems whose vari- ables are constrained on a smooth manifold, such as spheres, Stiefel and Grassmann manifolds.
Let M be smooth manifold and f M R be differentiable objective function. At any point X E M, the feasible directions form the tangent space T.M. The Riemannian gradient, denoted fX)), is defined as the orthogonal projection of the Fuclidean gradient V onto Tx M. Intuitively, it is the direction of steepest descent that remains compatible with the manifold geometry.
basic Riemannian gradient descent (RGD) iteration consists of steps:
G2 V (Xx) E Tx, M,
X = Rx.((G)
where Rx, Tx, M + M is retraction, namely smooth mapping satisfying Rx, (0) X, and whose curve c(8) Rx, (sG,) has initial velocity ((0)- = G, Such mapping brings a tangent step back to the manifold while approximating the true geodesic. Retractions admit simple closed forms on many manifolds, such as normalization on the sphere.
This framework generalizes standard gradient methods to curved spaces while preserving their in- tuitive interpretation. As result, Riemannian optimization has become a popular tool for problems with geometric constraints, and is supported by rich theoretical foundation and efficient algorithms; see, e.g. (Absil et al., 2008; Smith, 2014; Mishra et al., 2012; Boumal, 2023).
A.3 RESTRICTED ISOMETRY PROPERTY (RIP)
The RIP condition (Recht al., 2010) in Definition 3.1 is standard assumption in matrix sensing, ensuring that the linear measurement operator approximately preserves the Frobenius norm of low- rank matrices. This property has been verified hold with high probability for a wide variety of measurement operators. The following lemma establishes RIP for Gaussian design measurements.
Lemma A.2 (Candès & Plan, 2011) If M(-) is a Gaussian random measurement ensemble, i.e., the entries of M C 3th are independent up to symmetry with diagonal elements sampled from (0, 1/n) and off-diagonal elements from N0, 1/2n), then with high probability, M(-) is (r, RRIP, as long > Cmr/&8 for some sufficiently large universal constant C >> 0.
A.4 OVERPARA IN OTHER NONCONVEX ESTIMATION PROBLEMS
Beyond matrix sensing, the role of overparameterizzaton has also been examined a range of non- convex estimation problems. For matrix completion, (Ma & Fattahi, 2024) proves that the vanilla
14

---

Published as a conference paper at ICLR 2026
gradient descent with small initialization converges to the ground truth matrix without requiring any explicit regularization, even in the overparameterized scenario. In Gaussian mixture learning. (Zhou et al., 2025) establishes that Gradient EM achieves global convergence at polynomial rate with polynomial samples, when the model is mildly overparameterized. For neural network training, (Xu & Du, 2023) shows that in the problem of learning single neuron with ReLU activation, randomly initialized gradient descent can suffer from an exponential slowdown when the model is overparam- eterized. These studies illustrate that overparameterization appears across diverse problem settings, while its precise influence on the convergence behavior is problem-dependent
A.5 PRECONDITIONED ALGORITHMS
Preconditioning is a popular tool for BM-based matrix sensing to improve the convergence rate. For example, considering the problem
min 9Y = 1:00==1110000 [lIM(YYT)) - y11?,
Y
preconditioned gradient descent (PrecGD) (Zhang et al., 2021) and scaled gradient descent (ScaledGD) (Tong et al., 2021) adopt the following updates:
PrecGD: Y = Y,,
ScaledGD: Y = Y, - VV)
Since (Y, Y,) may be singular in the overparameterized regime, ScaledGD cannot be directly ap- plied. The variant ScaledGD(A) proposed in (Xu al., 2023) addresses this by using similar update to PrecGD with a particular choice of X. Next, we provide a detailed comparison of pro- posed approach with these preconditioned methods.
Comparison with PrecGD. PrecGD establishes only local convergence guarantee, requiring initialization sufficiently close to the ground truth. Although (Zhang et al., 2021) also dis- cusses globally convergent variants of PrecGD, they rely on gradient perturbations of the form Y(+1 Y, n/Vg)(Y,) Y, + AI)-1 + Ge] with some random noise Ç to escape potential saddles. In addition, they require a multi-stage switching mechanism that monitors several quanti- ties, including [V)((,)) Amin (
and Amin (Y, Y,). Notably, V²)(Y,) E m= is fourth-order tensor, which is meiery-intensivv. In fact, one key motivation for adopting the -urer-Monteiro factorization is to reduce the parameter dimension to mr by exploiting the low-rank structure, whereas forming such tensor negates this advantage. Moreover, computing Amin (vn(()) is especially expensive in large-scale matrix sensing problems. In contrast, our algorithm has a global convergence guarantee from random initialization without requiring pertur- bations or multi-stage switching rules.
Comparison with ScaledGD(A). ScaledGD(A) requires carefully controlled small initialization with magnitude a. To reach accuracy E, the method must satisfy a < ), implying that exact convergence (e 0) not be guaranteed. Moreover, ScaledGD(A) requires (rA + 0) RIP condition with < Okx-Ca) Co) for sufficiently large constant Cg. In contrast, we just need 8 < (xx 2). As a result, our sample complexity is significantly smaller, especially when the condition number k is large, i.e., in ill-conditioned settings.
Comparison of the benefits of overparameterization. A major advantage of our approach that higher level of overparameterizztion can not only improve the convergence rate, but also reduce the required sample complexity. In contrast, ScaledGD(A) does not show explicit benefits from increasing PrecGD's local convergence improves only with square-root dependence on which much weaker than the polynomial improvement achieved by our algorithm. In addition, PrecGD does not gain reduction in sample complexity from additional overparameterizzttin..
Comparison of potential extensions. A further benefit is the generality of our weight normalization formulation. This way of factorization
be directly applied to arbitrary low-rank PSD optimiza- tion problems. In contrast, PrecGD and ScaledGD(A) rely on second-order information of the loss function g, restricting their applicability beyond matrix sensing.
Comparison of iteration complexity. For the iteration complexity, PrecGD and ScaledGD(A) achieve better x-dependence than algorithm. However, the faster rates partially arise from the
15

---

Published as a conference paper at ICLR 2026
quasi-newton nature of their update, where (Y, Y Y. AI) is estimation to Hessian. On the other hand, our algorithm is a purely first-order method, and we believe that incorporating second-ore order information improve the convergence of our algorithm well. To further validate this point, initialize a preconditioned version proposed approach as follows.
WN with preconditioner. Motivated by the designs of PrecGD and (caledGD(A), we also explore incorporating second-order information to improve convergence empirically. To this end, we first derive a preconditioner for RGD.
For any direction H E Rmxr, the Hessian with respect to X takes the form
V (H, θ) = M(HOX + XOH x M A)) = -M'M XO + M X6 + M" M(xex AOHO.
When the RIP constant å < 1, we can approximate M R I, which yields
(H, v HO² + XOH X + (XeX -)
Near the optimum, the residual term satisfies (XOX AOHO O. If we further ignore the term XOH xe, the Hessian is well approximated by HO? Vectorizing both sides gives vec( V fHH )=(([ @ (e)-vee(H), which implies the approximated Hessian structure V s I @ Motivated by this approximation, we design a preconditioner replacing the Euclidean gradient of X, by (x,, ( AI)-1, where is regulariza- tion parameter that may be changed from iteration to iteration. We call this variant preconditioned Riemannian gradient descent (PrecRGD).
( + + for RGD, i.e.,
We conduct experiments on an instance of m 50, r 40, rA 3, 10, with sensing matrices as well as second instance with larger overparameterization level r while keeping all other parameters fixed; see Appendix G.3 for more details on setups. shown in Figure 5, PrecRGD exhibits higher convergence rate than RGD, demonstrating that preconditioner design is highly effective for faster convergence under the WN formulation. Moreover, PrecRGD outperforms PrecGD and achieves a convergence behavior that is compa- rable to ScaledGD(A). This suggests that WN is fully compatible with preconditioning tech- niques, and we believe that this promising direction for further improving the convergence rate. Furthermore, when the level of overparameterization r increases, PrecRGD converges even faster, while PrecGD and ScaledGD(A) do not show explicit improvements when increases. This again highlights the benefits of overparameteriiztion for WN.
8000 10
GU
PracGo PrecRop
101
de Preetraagge
10
10
10-12
10"10
Figure 5: Comparison with preconditioned algorithms
(squared reconstruction iteration).
B OTHER EXTENSIONS
B.1 EXTENSION TO ASYMMETRIC PROBLEMS
We have shown the high effectiveness of WN in overparameterizee matrix sensing problems, and its underlying parameterization reveals that is broadly applicable to a wide range of low- rank optimization tasks, even when the target matrix is asymmetric. Consider a general matrix A E Rmxa with Burer-Monteiro factorization Y Y,Y,, where Y, E Rexr and Y₂ Raxr, We can apply the Polar decomposition to each factor, i.e., Y, X Y2 X2O2, with X, E St(m, X2 E St(n,r) r) and O O E Sr. By combining the two magnitude matrices into θ = O E Rx",, we obtain the representation A =x,ex,
This generality suggests that WN has substantial potential in variety of applications, including collaborative filtering (Schafer et al., 2007), compressed sensing (Candès et al., 2013), matrix com- pletion (Recht, 201 1), and other related problems.
16

---

Published as a conference paper at ICLR 2026
5076
40%
BAIDIN RARRIE
1054
09
(a) Rank-1 completion
(b) Rank-r completion
Figure 6: Comparison of RGD with WN and GD on the challenging matrix complction problems
(% successful convergence vs. perturbation size).
B.2 EXTENSION TO CON-BENIGN LOSS LANDSCAPE
We further evaluate WN on the challenging matrix completion tasks proposed in (Yalçın et al., 2022), where the loss is constructed to have exponentially many spurious local minima, leading to the failure of most gradient-baut based methods. Using Burer-Monteiro factorization, the matrix completion objective can be written
min Y Ye
'IYY M.
where M is low-rank ground truth matrix and the measurement operator (+)a is constructed from specially designed combinatorial structures.
Applying WN, the problem becomes
min X.e
|Xxex
MoalE, s.t. X St(m,, r), O Sr.
We use an update rule similar to Algorithm (1), with the operator M(.) replaced by ())a- Follow- ing the experimental setup (Yalçın al., 2022) (see details in Appendix G.4), we evaluate the success rate across range of ranks = 1, 3, maximum independent set sizes || = 2, 4, 6 and perturbation levels g E [0.05, 0.5). Our findings are summarized as follows:
Rank 1: Our algorithm performs particularly well. For |S| 2 and |S| 4, the success rates are over 90% under almost all the perturbation levels, substantially higher than that of GD. Even for the more difficult case |S| 6, method still achieves successful rate around 40%, again significantly outperforming GD.
Rank 3: these regimes, both our algorithm and GD exhibit similarly low success rates, consistent with the intrinsic difficulty of the problem reported in (Yalçın al., 2022).
These experiments on this challenging setting show that our approach has clear advantages over GD. The results indicate that WN remains effective when the objective involves specially designed combinatorial structures or exhibits highly nontrivial optimization landscapes. This further high- lights the potential of WN broadly applicable framework for low-rank optimization problems.
C ALGORITHM 1 DERIVATION
We consider the overparameterized setting ra and apply joint update both X, and O in an alternating manner. Let M T + gm denote the adjoint of M with explicit form My Em M, The Stiefel manifold St(m,r) is embedded in the Euclidean space, then first compute the Euclidean gradient of X, as
G = [ A) X,,
(8)
[X, - O) + (MMMMI))XOOA - -A) X,),
17

---

Published as a conference paper at ICLR 2026
Projecting it onto the tangent space of St(m, r)) yields the Riemannian gradient
G. = (Im X,X, X [T x X.)
Using polar retraction, the update of X, along the direction G, with stepsize vi is given by
X.+1 - G G.)-1/2.
For the magnitude variable , the Euclidean gradient is
K, MM
A))X:+1
Denoting the identity mapping by I, the update of O with stepsize y becomes
θ O. 1. x |MM " 2 "I) -
X++1
(10)
2 XX I -
X++1
Note that we do not impose diagonal or nonneg- ative constraints on during the updates. In fact, forcing @ to be diagonal and nonnegative often worsen the loss landscape and may lead non-convergence Levin al. (2025). We il- lustrate this phenomenon with a simple experi- ment in Figure 7, where we constrain + to be diagonal and nonnegative via SVD followed by hard-thresholding: see Appendix G.5 for details of the experimental setup. As shown by the curve labeled "Diagonalized", such constraints indeed hinder the algorithm from converging to the ground truth.
108
OOiginal Diagonsized
10
10
0
200
400
800
800
Figure 7: Lack of convergence diagonal
(squared reconstruction error vs. iteration).
D PROOF STRATEGIES AND SUPPORTING LEMMAS
D.1 PROOF STRATEGIES
To establish convergence of Theorem 3.2, analyze the evolution of the principle angles between span(U) (U) and sX)n(X.). Specifically, we track the quantity Tr(I, Φ This term reflects the subspace alignment error between su) and span(X.). For notational convenience, we set 2, which is consistent with our choice in Theorem 3.2.
Our proof is structured into two phases:
Phase I (Saddle phase): When the alignment emror is large, i.e., Tr(I,, -., 2 0.5", we rely on the fact that o (4;)) remains bounded away from zero. This property guarantees that the alignment error decreases by least constant amount each iteration.
Phase II (Linearly convergent phase): Once Tr(I. - 0.5, we enter contraction regime. this regime, we establish that the reconstruction emror and the alignment error decrease jointly, governed by coupled inequality system.
Throughout both phases, two error terms caused by the limited number of measurements must be carefully controlled. Formally, introduce the following definitions:
A (M" M (X x O
41 A),
E, = I - A).
$0.5 is chosen for simplicity, any constant E 1) is valid; see proof E.8 for detailed analysis.
18

---

Cimr Published as a conference paper at ICLR 2026
Incorporating these two error terms, rewrite G, and @++1 as follows:
G. (X,O,X A),,O E,
O+1 AX,+1 A,X,++
These two terms will be used repeatedly throughout the proofs in the following sections.
D.2 SUPPORTING LEMMAS
Since Theorem 3.2 considers random initialization, is conditioned on the following high- probability event F, which gives a lower bound on the smallest singular value of 4o U Xo:
(r (r-ra) rA) }.
F = = (U X.)
where 61 > (1, a universal constant, with C₁ given Lemma E7.
Lemma D.1 With respect to the randommess in Xo. event F occurs with probability at least
1 - exp(-m/2) - C3 exp(-Xxp
where C2 > O and C3 = 6C: E (0, 1) are universal constants.
Vei
This lemma ensures that the smallest singular value of the initial alignment between U and Xo is bounded away from zero with high probability, which critical initialize Phase
Lemma D.2 Suppose that at iteration the alignment error satisfies that
Tr(I, $ S P.
then the reconstruction envor at iteration t satisfies that
|,,,X, A S 2 + 11,-11|e.
The lemma above connects the reconstruction error |X,O,X A|IF with the alignment error Tr(I,, D2 and the measurement error || 1 It means that the reconstruction error is small once X, and U are sufficiently aligned and the measurement error small.
Lemma D.3 Assuming 300 M(-) is tr 1, 1R with 5 Vmr E E [0, 1), and
|Or||| s Then, the measurement errors satisfy that
A X Allr,
| S X AllF-
This provides upper bounds the norm of the measurement error terms A, E, by the reconstruc- tion error X, Al|r, which is guaranteed by the RIP property of M(.).
Lemma D.4 Let Xt = (1l A (-1 + ( + Tr(I,, $, )(1 (|) I.)
8, = 01 (I, $.
H; = (Im X,X, (AX,X 1,-XX
X,X 0, XOX
2
X
A X X,X, XXX, AX.)
2
Assuming A |F, |E 1, 17 10 and | < 2, then the following inequality holds:
Tr(I,a - $ TT Tr(I, -
(11)
< (8, 16x Tr( 27(1 - 13 xx)ox))022__ (1.) Tr((I,, 2, )b.
R²
+ +277 Tr(I,, Φ )(1 A, FF + 1F
+ 227 Tr(I,, 2.0,) H.
19

---

Published as a conference paper at ICLR 2026
This lemma quantifies how the alignment error Tr(I, Φ evolves between iterations. This the key lemma that drives the reduction of the alignment error.
Lemma D.5 Assuming M() is (r + TA + 1,6)-RIP with 5 guaranteed that
S 3/m If | S 2, then it is
0 S
As shown in Lemmas D.3 and Lemma D.4, the analyses require that ||O,]| upper bounded by 2. This condition has already been guaranteed at initialization. Moreover, based on the update rule of + given in (10), we observe that [Oll remains close to | AX in each iteration.
Lemma D.6 Assuming 7 S L.M is (r + +,,) and |0: 0 S 2, we have that
D Wit 3 (1+6)) 2 4. (0111/7---2200--2222((111
m-rA-
Moreover, il is also guaranteed that
V I 3 (1+2:++2(( -ODT Im-ra
This lemma establishes an upper bound on the growth of v V. Together with Lemma F2, we can ensure that o²2 (P, remains adequately large throughout Phase
Lemma D.7 Assuming S 500r Then for any 2 0, the alignment error satisfies that
MM() is (r TA + RRIP with 8 S Vm and |+rl| 2.
Tr(I. T p + [+1)
Tr(I, $.D.)) 0.1.
This guarantees that the alignment error Tr(I,, Φ (b,)) does increase too much in step when we choose suitable stepsize which crucial for bridging Phase I and Phase
E PROOFS
E.1
PROOF OF LEMMA D.
Proof. Since the initialization X. satisfies the conditions stated in Lemma F.8, we can apply the lemma directly. In particular, substituting yields the desired result.
E.2 PROOF OF LEMMA D.2
Proof. Directly substituting the expression of O into the Frobenius norm term, have that
X O, AlIF
X,X AX,X A - X,X, 1 llF
< X,X, AX,X
A |X, A IFF
< X,X AX
AXXX, ||F + |AXXX] AllF+ | A |F
2| (Im X,X | + ||IFF
= - V Tr(I,,
TT ||.-I|FF
< 2 Σ VP A
= 2/p + A-11/FF
where (a) is by the inequality AB F < ||||| that is valid for any conformable matrices.
20

---

Published as a conference paper at ICLR 2026
E.3 PROOF OF LEMMA D.3
Proof. We first prove that ||G:||= S 2//G..... Indeed,
|Gal|= = II( - X,X G. + - [x] G - (.),
< | - |1| + |X |||G.||F
$ 2/
We now proceed to estimate the update distance |X++1 - X,I/F.
X X,IFF -,) X,IlF
$ [|([-++((0,00,)""---"000+-+11110 -/)31/
|1 G.) -1/2 G||
=11 G.llF
$ -1)
( G.)) + 2/| G.IlF
2 VFll 1 (172 1/
1
2m G |F
S Vrnl + 2nl G.llF.
where (a) by VI+ Vz for any x 2 Since |Ge||F S 2||Gel|.. we arrive at
|x+1 X 2n((F 1)
= -2/( - +
3
IX +
(4)
< 4n(vr+ XO,X AllF)
< 4n(r++ (MM X + IX -AllF)
19 4 1 - AlIF,
where (b) is from |B < B||] (c) is due to |0,|| 2, |X,|| S 1; and (d) follows from
Lemma F.11 and O A) I rank O rank A) Er+ra.
21

---

Published as a conference paper at ICLR 2026
Finally, we turn to estimating
||
| and |
|A |F M (X FF
1. II(
A)ll
W X - ALLF
x Vm T - A + |X - X - X,,,,XXII (FF)
2 X A
XxI(F)
S vm5(( + 16n F 1)(vm + 1)) X, AlIF
EVm (1+ 64
X A
Vmr 300
S XX AlF
|. M OX A)]lF
W MM O - A)lI
M
SV 1A3
( - ALF
XXX,, AllF,
where (e) is by Lemma F.1l and 1x A) S rank(X:+1 O,X,11 +rank(A) < r + Tas (f) is from |X E and ||O,| < 2: (g) due to tn S 300K-Fx and 6 1 _ /mr
and (h)
follows from Lemma E.1l and ran A) S xx) rank(A) A) S r rr+ra +TA.
E.4 PROOF OF LEMMA D.4
Proof. Noting that |X, |1 S 1, || < 1, |, 2,, X,X; || S obtain
H IF [m
X,X; AXX A Im X,X, ,,
X
E X 0 | X
+ |
AX,
A X + A XXXXXXXXXXXXXXX
$ A, +
(12)
In the same way, it follows that H - 2-1l| + =|
From the update of X,, we have X = - (I, + G. - -G) Premultiplying by UT and postmultiplying by U, follows that
P T
++1
= =( U GL) G. G.)
-G.
(Ira + (IIA 2, 1.. " u
(IL, (I,, Φ E. Σ o. n UTU, H.)
(b)
[Tra + 1/Ir - (..0))40000 Σ x. - (I, 11 G. G.)
(I.,, + (I,, कΦ EΦ Σ $ n H.)
where (a) is from directly expanding G and (b) is by Lemma F.1.
22

---

Published as a conference paper at ICLR 2026
We next derive an upper bound for G Gg. Substituting the expression of G2, obtain
G G. = X AX,XX - XXXXXXXXXXXXX AX, + HHH H.
X, A A(Im X,X, H.
H. (Im X, | AX,
(Ir, - $ I, + | 2 + (I - XXI H I,
01TA - 0, ), 16x
where (c) follows from Lemma F.12, X, Σ 1 and |||| x,X, U S [I X,X, U
I and (d) is due to II v Tr(I,, Φ ). and H s 2 A + S 6.
Combining the lower bound 1++1 1 1+ the upper bound G, G, derived above, and the inequality 1 17 8. 167 X 2 100 (1 96) > we derive
1
$.+1 1+1
- 12 - 16172xx Xt
Σ [Ir + n((rr D E Σ/ n7 rHHH,
(13)
(II, + nIr,
1 10 7--UHH H,
Let the compact SVD of 4, be Q.A,P.. where Q RrAXEA, A, E R'AXYA, and P, RrxxA
Denote S, = Q, EQ, is a positive definite matrix. This gives that
Tr [ + (I,, 1. EΦ,
Σ क [L, + ,T,,, Φ. )E Σ
(Tr(Qu(LrA + + A A? [L SSAAS, A3]] Q7
2 TT(Q:(A? - nI,,, - A.)S ASS A3))
= =T(O.AA))))) TT
A.S. - -A7)
L T((Q.AQT)) 2010A (AP)
Tr (Irr A7)AA)
x2
2 (A?) Tr([IrA (..
TT( +
x²
where (c) follows from the fact that 17 Q A/S - is PSD: and (f) is by Lemma F.3 and Lemma F.4. More precisely, o,,, (S,A7S,) of (S,Jr, (A?) TrA (A²)
2
23

---

Published as a conference paper at ICLR 2026
Taking trace on both sides of (13), arrive at Tr( Tr( 1+1 (1+1) (+1)
(14)
1 - 3/ Xi
21 (A
2 Tr(
-, )000)
x²
- -2nTr (Irr + n(Tra 2. 20
+ H,H,, U)
2n0rg (A Tr (I,,, - b, -..0)
2 Tr(
n²
-2 H U + nI,,
D T 0 H U)
@ Tr(29 2 (A})
TT (.
x²
21/T((U (Im - X,X, AX,X, A x.
2n/Tr(U" (Im - X,X, ,
=X,,
-,TT(( X A X 4,X, A X (x
21 2 D E U)
TT( 2 (A1)
Tr (T((I D 1
k²
20/Tr(U" (Im - X,X, JUEU X,X, A 1- x.
2n)Tr(U (m,-X,X) )x 0
-, - 2. H U)
where (g) is by substituting H, in; and (h) arises from Tr(M) Tr(M for any M E RRAATA By the Cauchy-Schwarz inequality, we can upper bound the three trace terms as follows.
For the first term, have that
Tr(UT(T, (Im X,X, JEU X,X, (
(15)
(Im X,X, JUEU
|| X,X, 1 I|
1m - X,X, |
Tr(I,A क ΦT) 1.-1|
For the second term, we can obtaín that Tr(UT (Im
x,X, X,,,,,
(16)
U (Im xXXx)))| (x IF
2 (Im - X,x, |
= $ E.
For the third term, it holds that
Tr (
D 2 Σ U)
(17)
Ir. $ || 1. E Ull:
(Im X,X. U E , U||
[Im X,X, | H
V Tr(Ira 2 H |F.
24

---

Published as a conference paper at ICLR 2026
Here (t) is from ||]| < 1, 11 (14), (15), (16), and (17). follows that
1 X S /
1, and | I 2. Combining inequalities
2nor, (At)
Tr( ( (1) 2 _T(
Tr Tr((Tr, $, )
1 1- 1- 3 16n/xx Xt
n²2
2TV Tr(I, - 2 (11 1-1FFF + =
201² Tr(Ira D.
Reorganizing the terms, we arrive at
Tr(L..- P+t1 2 [++
Tr(Lra (.
211(1 1 16) (A?))
$ 18e 16xx Trxx)
Tr Tr((IrA $.
k²
+ (2n 21 32n x1)\ Tr(I,A $. ((1 A - 1||F =.
+ (2r²² 21 Br -32n" xilV Tr(IrA $ HellF
$ (B( 16x( Tr(b,,0)) 211(1 -1² 161/ xr) (A Tr((I,A -D 1..
x²
+ 2n/ Tr(I, 0 (11 (11 | + 2
+ +2n² Tr(IrA- D. HHl|F
Together with or (A?) .r (Q D
T Q+) =07, (D.), we conclude the proof.
E.5 PROOF OF LEMMA D.5
Proof. From the update formula of O, we obtain
o 10+1ll = X 1+ AX 141 X X,Il
X 14 AX |
(a)
x 1+ +||l
+ [[M M - IX A)l
(6)
1
< 1
X O.X
AlI
3/m
+ Vm
[X,+ O,X 11 All
3/m
< 1 + (||O,| + |A|)
2,
where is) by [XII |||| S 1; and (b) follows from Lemma F.11 and rank 1,,, - A) < rank(X,,1 x (1) + rank(A) < t
E.6 PROOF OF LEMMA D.6
Proof. Let L, = X, AXXXX A X XX + -
+ ( X AX, X. AXXX. A.)X)
25

---

Published as a conference paper at ICLR 2026
Applying the triangular inequality, obtain
| = | X A
X x + 2 +xxxI
+ 2 [|X]
AX AXX
(18)
(a) 2 A -
||
+ 4 | ||
where (a) is from |
1 1, ||||| < 1 and |0,||| < 2. Multiplying the update formula (4) on the
left by U.. we have that
V e+1 = U (X, (T, + G. G.) -1/2
I X A AX, + m N - - G G,)
| (I,- x AX.X, AX + nb - - _ (I, Gf G.)
where (b) is by expanding G, directly. Consequently, we have the following upper bound for L 2+1 v +t1
241 W = (I, (X AX, AX,+
-U (AAAAAAAA--HHA----EEEe/--------
(v, (L. nX AX.X AX +.) -U[=
(c)
1 (I, - x AX,X, AX AX, L.) -
(v n/ A AX AX, L - -U[ T
= 1 ( x AX,X, AX + L.)
(I, X AX, AX, + nL,) /
-, (I, - X AX,X, AX, + L,) V
(19)
n (1, X AX,X, AX, + +L) O,X, 3.U1
+ + EU1
where (c) is from that (I, G - is PSD and all of its eigenvalues are smaller than 1. Since YYT || holds for any symmetric matrix Y Rexr and by Lemma F.12, we can upper bound the three terms as follows.
For the first term, we can obtain that
V. (I. x AX.X, AX, AX,+
II - -X A AX, + nL)) T
(20)
3 II,- AX,X, AX, + L P.
For the second term, it holds that
U (I, 7 AX,X AX, (L.) T V..
+ V. "|, AX, AX, + LL) O,X, E,U,
(21)
= v (I, - AX,X, AX, LL ||| L
For the third term, we have that
T 3 U. R-TA-
(22)
26

---

Published as a conference paper at ICLR 2026
Combining inequalities (19), (20), (21) and (22), it follows that
L 11112/41 1+1 = ||, x X AX, + P.
2m v (I. x AX, AX; nL))
, E,U E, L
+ 1 U X, 0
IM-rA
(III, - x x A || + +
+ 1, x X AX,I n [L)) E. | m-
+ U U E. I-FA
(1= L+ V 4 + 7 L. D) E. I + +4 1
where (d) is by triangular inequality: and (e) is from that all the eigenvalues of the PSD matrix
X, A AX, are smaller than I, along with |/ S 1, |X.I| S 1, || || S 1 and |0:| 2.
From (18), we obtain
|L,||
< 2|| A || + =| ||. Then, we can further simplify the inequality as
(1 ( || 2 I)) V.
4+1
+ Ex| + (5 1| 2/| - E. I-
(23)
From Lemma F.11 and our assumption of the RIP property of M(-), obtain upper bounds for the two error terms.
A
A)ll
< O ALF
$ 1X [[A|))
(2Vr + vrajo,
|. = - X - AIII
$ ,, - AlIF
$ x +
2 (2Vr+ Vrajo,
where (f) is from | - 1, e, ||F S v S 2/5 1 S 0i S 2/r, and
|||| S /ra A AA Plugging these two upper bounds into (23), we arrive
P TT b+ 3 (1 + (olebnnl T. +
m-
We now consider the relationship between V T and VoV.
Let L. = (x. AX - x, A X E O + x,)
Multiplying the update formula (4) at t = 0 on the left by U.. we have that
V U (X, - G)) + G -1/2
1/2
Consequently, we derive the following upper bound on v WT
V M T = U - G/ + G Go) - G/ TU
(x) 2 U[[Xo (X,
(k)
(Vo([, nLi)) - - ( Vo vLo) - -U
=ol,, - (L,) - nLo) T V Hol, - -Lo) E.U1
(24)
-,, pLo) T Vo U U.
27

---

Published as a conference paper at ICLR 2026
where (g) is from that (I, + G Go) Gal-1in is PSD and all of its eigenvalues smaller than 1; and (h) is by expanding the expression of Go directly. Since YY" N I holds for any symmetric matrix Y R2x* and by Lemma F.12, upper bound the three terms follows.
For the first term, holds that
o TL,,(T,, / 3 Ir, nTLoll V.
(25)
For the second term, have that
nLLn) TuO
z 2 o - - U | Im--A-
(26)
For the third term, we can obtain that
X < X U I--
(27)
Combining inequalities (24), (25), (26) and (27), it follows that
V V I, n/on| 2 V 1 2n o [ Im-ra
+ // |
X E I-mra
(1 n LoR - 4n + 1/ اo Boll mm + 10l Im-ra
(28)
where (i) is by
Xool || s
U_l < 1, and |10ol| s 2.
From Lemma F.11 and our assumption of the RIP property M(-), we have that
| M Xx -A)ll
S X AlIF
S 0, + ||A|)
s (2VF+ vTalo.
Then, we can bound |Loll as follows:
1
|
X 0x + 1,X, AX, O xxx,,
2
2
< + 2(2 + VTA)ō.
Plugging theses two upper bounds into incquality (28), we finally arrive at
V 3 21--2--- (1
E.7 PROOF OF LEMMA D.7
Proof. We first estimate ||G,]| and |Ge|l. From the expression of Go we have that
||G. = M[X, 0),
x M2 A)lI
< 2(1 M XXX AIII+ XX - Al))
(6) 2 IX,O,XXXT A + 2 -AlI
Vm
x X AIl
< 4( +|AAI)
12,
28

---

Published as a conference paper at ICLR 2026
where (a) is due we can upper bound [Gel as follows:
|X, I 1. | 2, and
||||| S 1; and (b) is from Lemma F.11. Analogously,
II X,X, G X [ G. G x,)
||
S Im
XI
G+ll
+ X G X,I
2 2||
S 24,
where (c) follows from X < 1 and the fact that all the cigenvalues of the PSD matrix Im XX,X less than 1. Multiplying the update formula (4) on the left by UT, obtain
① - $
=U X X U
GG = ( -C
( U +
(d)
= -6,0, -B
+ U.
GUUGGGGG
-
4. - (17 (,2 G U
2 o G G.G. U| tn u G.G. G.G,G,
(
10r Im
where (d) is from Lemma F.1; (e) is by Lemma 12 and (f) is due to || < 1, |
$ G, s
24 and r) S 500rA By subtracting the inequality from I follows that
I P 1 3 Ira $. + 10r IPA
Taking trace on both sides yields
Tr(I.
4 11+1 (T S Tr(Ir, - DD + 0.1.
E.8 PROOF OF THEOREM 3.2
Proof. For the proof, we take ") 97bof and <4(r-r)) kmm where 6₁ O From
(
Lemma D.5, we have that ||O,|| < 2 holds for all f 2 0 by mathematical induction. For later use, we define the following three terms in the same way as in Lemma D.4:
B. (I.. F. 1
X: (1 - + 1) + V Tr(l,, 2, [(1 A 1 1..1)
H, = (Im X,X, (AX A -XX,,)
+ X.X.E.XX - X,OXX
XXXXX A XXX, XXX. AX.).
Lemma D.3 with the RIP property of M(-) implies that |A,1| E < 1 for all t 2 1. Thus, the assumptions of Lemma D.4 are met, guaranteeing that inequality (11) holds for all iterations. Building on inequality (11), we divide the convergence analysis into phases.
29

---

Published as a conference paper at ICLR 2026
Phase I (Saddle phase). Tr(I,, + 2 0.5.
We assume for now that o (t) 2 (r-rz))/(2c holds in Phase I, which will be proved later.
Let the compact SVD of Φ be Q Q,P.P. where Q₁ E R' A, E R'ax'a, and P, E R We simplify (11) follows:
Tr(I 2 2:41 2 T
Tr(L. - (2
27y11 1/2 167 xelo? ($t) Tr((Ira - (
$ (r
n²
+ 22 (1)
+ [1 + 27j² /ra HellF
(2)
27 - -/ 16 (1)))
(1 + 16xx)
Tr(L..- $
x²
+ 2 A
F + 2 2 + 21/² /rA HellF
77 17² 16))2x) rT)" Tr(IrA ..
(1+ 16X
(29)
mm²22
+ +2n/rA((1A + 27²² Helle
where (a) is by Lemma F.3 and Tr((IrA - 4 $ = Tr((I, - A7) and (b) is from our assumption that of (.)) 2 (r - (2)2/(2
Using Lemma F.11 and the RIP property of MM+), we can control the quantities of the ror terms. In particular, following inequalities imply that both || and ||=,||F are uniformly
bounded by constant that depends only on m, and va but is independent of t.
Expanding the expression of A and applying Lemma F.11, we have that
11 S
- A)l
air r6)
| AlI
12
ca(r- r 2m2 (2V++ (TA)
(r -
3/37/277/27
(r FA)2
mint 18cFn2m2-224 48
(30)
Applying the same reasoning to S., it follows that
E || <
A)I
r ca(r r
X AlIF
km3/22444 A
- 6
(2Vr + Vra)
xm33/2/47
3ca(r - -TA)"
2233/277/2A
(r-rA)²
min( m8 48
(31)
30

---

Published as a conference paper at ICLR 2026
Here, (c) is from Co ( and rA < m. Since Tr(I,, + T.A together
with (30) and (31), we can upper bound Xr as follows:
Xt = (1) A || +
11) + V Tr(Lra - 2 2 ( ( (.)
< (I) - | | 2 + (Ta(lA--1ll) (.11F 1E:1F
+ 1
$
+
)a
Vral 48 TT
48
48
48VTA
$ 16'
From incqualities (12), (30), and (31), we obtain the following upper bound on [H:]|e.
H |F $ 2 A
A-111F + 3
x
+ +4x
48
48
2VTA
With these upper bounds, incquality (29) can be simplified as follows:
Tr(I,
D.+1 1 (+1 - Tr(I,A DD)
I rT 711 (r -ra)"
Tr(Ina- 21 + (r- + +12
207 22²
8c2
(4) -(r-rA) r2)*
+ T((-rA)
2 m2? + +61jrA+ + rr(r-ra)² Tr(IIA- - $.
S
2rqr:m22p2
Gr2m2?p
'/( ra)*
TA + (/ - ra)*
Tr(InA - 2
m22
GRR2m222
+ G + (r
2 Re2
GR3m2,2
where (d) is by Tr(IrA $ +. 2 0.5; and (e) holds if the expression in bracket is less than zero
(r-ra) Recall that 17 975cfx²m
The summation of the terms in bracket is negative, which implies
(r-ra)* that at each step, Tr(I,, $ decreases at least by A= == 70 Consequently, after at most (AA 0.5) A < 7000c n 'mmm 2 iterations, RGD leaves Phase I.
(F-
Let 02 7000c7
(0, Denote to 2 as the last iteration this phase. The analysis above
implies that Tr(I, 2 S rA cx(r-ra)+t for all S S to and < 7000 1c m 2r2 kmm","r
(r-ra))
From Lemma D.2 and inequality (30), we obtain the following bound for 1 S lo:
|X,,. AllF 2
Tr(I, 2. + A 1IF
a(r-)
2
TA
cr -(r--r)) *m
2
PA
+ 1.
We now prove that ,2, (Φ.) 2 (r r / holds Phase By Lemma D.1, it holds w.h.p.
02 ( 0z, (U X (r-ra)²
31

---

Published as a conference paper at ICLR 2026
Moreover, by Lemma F.2 and the assumption
rA S m it follows that
o² (oo) S 1 (r (r=
of (v))
C170r
Since 17
(r-rA)) 2760fm2m33000 and 6 c((r-ra) we can deduce that "'m²"n'ra
c1(r cr(r-ra)'0
nVTA + 2VT)8 325c
From Lemma D.6 and the upper bound on 17( VTA + 2V)), we obtain the following inequality
L V 3 (1 + 4(r 97 )² L v
4c1(r-rA) 10
6
Using Weyl's inequality and c4 = ( ) we have the following upper bound on of (V))
975ck 10022 (oo) + 4r - rA) 10
4(r TA)*
of 11 S (1
65c 6
4(r-rA) )² (r (r-r))
+ 4ca(r - rA) 10
< (1 +
m
cimr
6
5
16(r 16(r-TA) T*)*
(r 2
195
cymr
5 1 2(r - -ra)?
3cmmr
TA S S and CA, 2 1. Applying Lemma D.6 with the upper bound on TVTA + 22TT) we obtain
where (f) is by r
Ve V 6+1 3 (1+ ((r-ra)) 10
c((-rr)) r 10
2 V.
40c7mtmbro1, Im-ra: 121. 1.
40c7
Using Weyl's inequality, we have the following relationship between of ( (P and o² (.))) ca(r - 220 40c7no'm
Ora (V,.D, + ca(r - rA) 10 40c7
2 ra (+++1) = T V V 1 ( S (1
cr(r TA) 10 (1))
<4 - TA) 10
= =(1+
10c
40c
Denote { := " 70
By iterating the recursive inequality, the following upper bound holds
1-2
o² (V.) S ( 2(t-
+
c
Σ (1+
S 5 () 0 + Σ
(1 11+00 () 2
1=0
= -1++) + () 21 - []/[11+)2= -1]
5(1+ 2t 22 (v))+ +1++11111 (1+5)²* -1,
for all 1 St 7000cf"' (1-ra)*
Invoking Lemma F.13 and noting that ( S 2. which is ensured by
C4 = 0( 0() ), we obtain
1 < (1 + 6tt)o²a (L1) + +tς
(9)
7 (v,) 1) 2100cfca - 22)
mr
< 2(r T2)
2100cfc, - r2)2
3c, mer
mr
(N)
(r-rA) -2
201 mer
32

---

Published as a conference paper at ICLR 2026
where (g) is by IÇ v (1 and o7 (V)) S 1; and (h) holds by c4 0 By Lemma
mr
F.2 and the assumption TA S 2 it can be seen that o72 (Dr) o7 (v)) 2 all & < lo
(r-ra) 20rm
holds for
7000cf 2
i.e., throughout Phase I.
(M-ra)s
Phase (Linearly convergent phase). Tr(Ira P. 0.5.
This corresponds to near-optimal regime. An immediate implication of this phase is that Tr( + A 0.5 2 AΛ 0.6. Recall that to Z the last iteration the first phase.
We assume that Tb, 2 0.6 for all 2 to + 1, and we will prove this later.
Given that the singular values of 4; T lie [0, 1], we have
0.4 - o, (1)) = Tra (2. S [f((2)) S 1.
Moreover, since 3 = 01 Φ s Tr 0 and 3 S 1, it follows that
3, Tr STA Tr(La 2, xTTr((,0)) TAXI
In addition, it can be derived that
Tr(L,,
S , Tr((I,, -, 1..
25
I (P:) Tr(Lra -
$ Tr(I.A $
With the inequalities above, we can simplify (11) as follows:
8n)
271²
32n
Tr(I..
+.+1 TTT) S
1rr
Xx)
Tr(TTA +.
(32)
25
x2
x2
+
+617 + 20y
Tr(Ir, $
A
+ 3 + H, IF).
Recall that X² (|] |||||
|. ))) Tr(,, $. [[[ A
|x) and H, |F <
2|| 111F + E, F. Since
A < 1 and |3 < 1, (32) can be written as
81
194n²
Tr(I, P.+1 2 TT.)
< (1
+ +
]r 2
25
R²
+ 16n TA A + ||)
+ Tr(In, - +.
(167)rr A + |E,
+ Tr(I,, - +. (4n) 2n) A + + lFF)
(r-ra)
Substituting = 975cfx"m
into the inequality above, it follows that
1 1(11 A ||,
Tr(I.A P..1 (. TT( +D
(11
TT(IrA - .1 kATA (1 (| |B.
(33)
Tr(TrA - -.
+
((1 A(1 + El)r,
ntt2
R²rA
where 9 := 6 "rA is a constant in 1).
(r-ra)*
33

---

Published as a conference paper at ICLR 2026
From Lemma D.3 and the RIP property of M(.) with 8 ca(r-ra)
guarantees that
cr(
A S 1
| AlLF
62m5/2,7/2r
cir rA)
1x AlIF,
13mpp²
crrr TA) A x2mb/227/2rA
| S |E
|X,0,X| ALF
cr(r TA)) 22222
X ALIF.
Together with = ) we can rewrite inequality (33) as
Tr(I. 11+1 T
S TT(rI - D 1-9
O A
XX -X,,
AIF
180
+ +2 x X Al|F||X,,000. A
AllF
(34)
+ V Tr(Ir, - $ X A + |X,0, - All=))
Denote b, : Tr(L, D. ). a, a= = X ALF. Inequality (34) can be expressed as
be+1 S qb: + 1-9 (a at ia, + Vbl + az)).
(35)
180
Combining Lemma D.2, Lemma D.3 and the RIP property of M(.) with ca(r-ra)
"m" we obtain
a, < b +
6 a-
(36)
Since to + 1 the first iteration in Phase II, we have Tr(I., Poot1 1 to+1) S 0.5. From Lemma D.7, follows that Tr(I,, P1o22 (o+2) S 0.6. Hence, bio biot2 E [0, 0.6].
From Lemma F.14, to establish the linear convergence rate of o,, it suffices to analyze the following equality system of (b, t and 2
1 1- (ap àg 22 + Vollār-++ + āz)).
b b = qd. + 180
0 = 2 by + - t= to 2, lo + 3,...
ūto-1 = ato11, beot1 = 0.6, Ino = 0.6.
By Lemma D.2 and the RIP property of M(-) with 8 ca(r-ra) kmm we derive
( Ūtg+1 = aa S 22 + | S 2 brott + 48 3 biot1
3/2
VI+ V Bhot22
where (1) is from inequality (30). From the update of ā, at l = ln + 2, it follows that
uoot2 = 2//inot boot2 + 2ta+1 S 2/bortt + 1 Vbboti+ +
S 3 3/brot
288
Therefore, applying Lemma F.14 and Lemma F.15, we arrive at
0tg+144 S 00o+1+4 S broti 1 1+9 q
1/2
33 1 1 1+1
ca(1 ca(r rA)) 6+1
2
2
for all t 2 0, with = 3250067 E (0, 1). This establishes the linear convergence rate of
34

---

Published as a conference paper at ICLR 2026
We prove that Tr( Φ rA 0.6 for all t to + 1. This amounts to proving that Tr(IrA 2 = b, I 0.6 for all 2 to + 1.
Since 2 0.6, inequality (35) holds for t to 2. Hence,
1 1-9 (4
bto+3 S qbro+22 + 180
42ot2 + 20 +2 + 2 + aso+2))
1- Q (9 x 0.6 +
0.6g + 180
9 0.6 18 0.6 6x 0.6)
1 1+9 4
x 0.6
2
K 0.6,
where (j) is from the fact that 0ta+1 S 33 boott 3v0.6, atot2 S 0to+2 S 3 boot2 = 3.6.6
Therefore, inequality (35) holds for to + 3. From inequality (36), we have ato+3 S 2 biots atot2 I 3.0.6. By recursion, follows that Tr(Ira
$ = b, I 0.6 for all t 2 to + 1.
(e-ra)* To conclude, by choosing stepsizes 97 and μ = 2, we have that X AAIF S
ca("-ra)" t-to ComprrA
3
for all t 2 to 1, with high probability over the initialization.
E.9 PROOF OF LEMMA 4.1
Proof. Let U the compact SVD of A, where In, u2..... Ural and E diag(A) A1, λ2, ; )r with A 2 X2 2 , the diagonal matrix whose diagonal entries are λι, Xr > 0. Here, diag(A1, 22 .......... )ra) denotes
Ar
We first consider p 2 1. From the Eckart- Yount- Mirsky theorem, have that the best rank approximation of A under the Frobenius norm is A, = U, Σ, UJ, where U, = [n), U2, Up and Σ diaa , without considering the ordering of the eigenvalues.
and range( Ap) )) S ran it follows that range(A,) range(U,). Together with range(A,) range( X x C range(XX (X), we can obtain (ange(U,) C range(X( (X). Therefore, there exsits matrix Q E Rrxp, such that U, = XQ. By the definition of U,, we derive that
We begin by analyzing the form of X and e Since tank(Ap) = rank(U,)
UTu, = Q XQ - Q Q -1,
x
which implies that column-orthonormal matrix.
We extend Q to an orthogonal matrix Q []. Let V, XP, then [U1, Vi] [XQ, XP XQ Since x Tx Q T Q I,,, then UU, V1] also a column-orthonormal
matrix, which means that
V, = [V1, V2. Vr-pl, with V1, V2,..., Vr-p E U. VIV, I.
Let U2 -[up+1 ... u. and then U = UU, U2]. By substituting U and X, we obtain
X
U U2]
i P VTU,
0
where (a) is from V1, Va Vr-p UF.
35

---

Published as a conference paper at ICLR 2026
Since Tr(X X) p, it follows that
Tr(X UU
X)
I
I o
0 V Φ-)
Tr Q
0
V
-r (B. VTU.U.V.) ))
p ++ U
After cancelling the term p on both sides, we obtain Tr(V, U2U, V1) = 1 0. Hence, we have that U. V1 O. which implies that V1. V2. V--p E U.. Moreover, since V1, V Vr-p U as well, we conclude that V1, V2,...
U+
Substituting X = U1. VijQT into XeX = Ap- p. we can obtain
U1, TiQTeQQU, V1]"= = Ap
U Ujdiag(Au \1, U
[U1, Vildiag A1 A2... ................................ Ap. 0 O........ O Vi]
Expanding both sides of the equation, together with v1, V2, V....., v. E U can obtain
eQ diag (1, A2...., Ap. 0, .......0 0).
This implies that
- Qdiag(A1 ......, Ap,
Q) 0)Q
To proceed, we first verify that (X ([U,, V1l. diag(A (,, A2...., Apt 0, .......)) is indeed
saddle point and then prove (X, θ) is also a saddle point.
We compute the Euclidean gradients of fox with respect X and O as follows:
Vxfoo (X, -) (XxX A)XŌ XO² - AXO.
Volx X, θ = x (X6X - - 2
1 -
By plugging the expression of (X, Q in, we obtain
Vxfxo (X, θ) = x A
= [,fu,, uuu
up ...... 0] - - A3up, 0, U O....,
= 0,
1 ② -
Vofx(XX
2 1
(diag(An (1, Azy..., :p, Op......,00 diag(A1, ................................ ... O,..0....
2
0.
Then, the Riemannian gradient is
X X Vx/a (X, -xx 0.
- XX XX-XX/VX// X θ)
Therefore, (X, ②) is stationary point the Riemannian sense.
We now show that (X, θ neither local minimum nor local maximum of the objective function.
For any 0 <V < Ar. we will construct a pair (X+. @+). such that foo (x+, θ+) > for (X. θ),
(x+, 6)). = V X
O 6 < v and X[ I,-
36

---

Published as a conference paper at ICLR 2026
Let x =X = 1U1. Vi and O+ diag(Ar A1 -v, -2,....... O, ............ 0). By construction, X 0,.....))))
I, and d ((++, θ+), (X, 6) V2² Σ hold. The value of the objective function is
Jnx [xx 6+) - +1x θ - Ali
1 1l(1
Vunu +
Σ
Σ uu l
AA |nnnal
Σ Nuu.u!
impt1
(b) 1
1 (v² | T1 |XOX AlF)
loo Xx +)
where (b) is by the orthogonality of (u1, ug..... J.
We now try to construct a pair (X_, )_), such that foo fx0( fx)(X,
((X é (x, - VIX X1 + 1 6 S and X_X = Ip.
Since V, E U- for any i E (1, (1,2...., p}. it follows that v, E spann Ur J Accordingly, we consider
X_ U1, sup+1, V2. ................................ Vr-pl:
θ = diag(Ai, A Aps vo, 0, 0,...,0),
where k s, 50 > 02 = and k; 9,20 vo will be given later. We can easily verify that X_X_
I, holds. The distance is
6( (x,, 6)) 1 X|2 || 611
[[k [[k--)v + sup+111 +
V(k- - 1)2 52 5² v/
V 2k + v8.
2² Let h: 1
k2
and up 5. then d ((x, 0-), S /V oh _²
+4 v. The
T
value of the objective function is
fao (x_,
1 x AlE
1
| VIIV + Isu + ksviu + (uos² - -p))up+1up4
FA Σ
nu 1|2
i=pt
( 1 (1* + +k2s² 11 IF
1/2
2=3 v (1 + (uDs² - (p) foo X, θ)
1
1
4 0 (k+ + 222 + +
2 _opps
Soo [X ²):
where (c) is from the orthogonality of {in ug...., un v1}. Let vo > 0 be sufficiently small.
Then 1u (R+ + 21:23² 22 + 54) 1 2p < 0. This ensures that the perturbed pair leads to a strictly smaller objective value, i.c., foo [__, ②_) < foo (X.
Therefore, have verified that (X, θ) is saddle point. Building upon this result, now proceed to show that (X, θ) (XQT Q T) is also saddle point.
37

---

Published as a conference paper at ICLR 2026
Plugging in the expression of (X, θ), we obtain the Euclidean gradients as follows:
Vxfx (X, )) (Xox* - A)XO
=XQ QeQ QX" AAXQ'QQQQ
(X6² AX
(X, = X (Xex - AAx
2
QX (XQ QOQ QX A|XQ
2
Q(6 X X'AXI
Then, the Riemannian gradient is
X Vx (X, x) 0.
(Im Vx
(X,
Therefore, (X, θ is stationary point in the Riemannian sense.
Let [X +. θ+) [ QT, Q++QT),
QT, (x
The distance is
0+), 0)) V VIX X + O OIR
VI(, X)Q 1 QQ Q
VI - 1 + |+ - O1
=d d(X,, (x),,+).(X.6))
In the same manner, we can obtain that d O(), θ)) orthogonality of Q, the following three identities hold:
By the
xex XQ QÔQ QXT = XOXT,
X _ Q QX = -
x o X x x QÔ Q QX_ x X
Then, we have /(x,0)= θ) /X, 0.) = /( θ_ and J(_ --)= ,(X_,__ --). Thus, obtain the strict inequality ,(X_, /X, θ) (XX 2 Therefore, (X, θ) also a saddle point.
We now turn to the case p 0, i.c., XOXT Ao 0. Consequently, θ A Let be expressed as X = [x1, X2, ..... x,]], where each x; is a column vector. Since Tr(X UU | = follows that U X Hence, each x, lies in U for f1.2....,r))
We compute the Euclidean gradient of the objective function for with respect to X and θ
Vx/a X θ) = (Xex A)XO = X AXO = 0,
Vel X, θ) X (XxxT- A)X
1 ( 2
- AX)
x
Then, the Riemannian gradient is
X ( [ Vx
(Im XX 0)
Therefore, ()) is stationary point in the Riemannian sense.
38

---

Published as a conference paper at ICLR 2026
For any 0 < < A construct the pair (X+, (+) follows:
X [xx + su1, X2 ......,,,. x,),
θ = diag( 111 0.. 0),
where l: distance is
1. 8.8 -1- R2, and 1 S 3- We can easily verify that X_X, X
and the
0+), d 04), (X, 0)) VIX, X | O 1
V |1k - 1)x1 + su 1 + us
V (k: - 1)2 + $2 + 127
122 V2
$
The value of the objective function is
04) - x 1x,, AlE
7 1s (k2x1x] + uu
A Nun
ra
4 11n | s + Σ Au
( (ugk* + +1 + 221 + XOX -
foo (),
where (d) due the orthogonality of (u,, ug..... ur X1). Now consider the pair (x_, A_) defined as:
X [Ax] + sul, X2, X2, Xrl,
θ = diag(12 (22 .,....),. 0),
where I 1 8.8 vI - b²_ and
V/2 I z. It can be verified that X X I, and the
distance is
((x,,0))((,0))) = VIX X O θ
|lkk: 1)x1 + su 12 2
V (k 112 32 +
122
4
$
39

---

Published as a conference paper at ICLR 2026
The value of the objective function is
foo
X + 1_ - All AlI
4 11us xxxx + +
u
4 |1nk x u
u
(e) (ught + 1220 - 2222110 + |XXX - All
4 (k² + + - 2222d103) + - (xx))))))(((,,))))
where (e) is by the orthogonality of {u1 ug,.... , Let 12 > 0 be sufficiently small. Then (v2 (K+ + s*) s²) 0. This guarantees that fx (,_, .-) fxo(X, θ). Therefore,
(X, )) is also saddle point when p = 0.
E.10
PROOF OF LEMMA 4.2
Proof. We begin by computing the Euclidean gradients of foo and f with respect to X and O:
Vx/x = (XOX - A)XO,
Volx(X,0) = XXXxX - AJX,
Vxf 9) = =M.M(XOX - A)XO,
Vol(X, = (XXT=)))XXMMMXXXXX X.M(XXXXXXXX AAX
Then, we can obtain that the gap between population gradicnt and sensing gradient is
VVx/ V = - A)XO|F
$ - A)]
3 MM(M -
A)IIF
$ (M' M (XXX -AIII
2 2 1 T AlLf
2md TOXT AlI
+ [[A])
K9 6md,
40

---

Published as a conference paper at ICLR 2026
(efx(XX)) Vef(X, O)||=
2 X
A]X||
[XII(MMM-I(XO) - I
2
_
IIM'M I - A)lEF
2
5.
1 A)l
V
Vmd|XxX - ALI
2
XXx All
X + I|A))
32 2 2
where (f) is by |X| S ||| < 2, |||
1; and (g) is from Lemma F.11. Then, the difference
between the two Riemannian gradients can be bounded as
| (X, (X) O)|| 1[
XX
- -
(Vx Vx/(X,
X(Vxf 9) Vx)(X, )) |X|F
2
< 11 mm - XX (VX (X, - Vx/(X O))
XX (Vx/ X - Vx/(X O))(F
X(Vxfoo X, 0) - Vxf(X, θ) |
2
< ( - XX
X
XII
2
(
12m8,
where (h) is due to [[
XX
|| 1
F OTHER USEFUL LEMMAS
Lemma E1 Given a PSD matrix A, we have that (I A)= > I-A A.
Proof. Diagonalizing both sides and using 1/(1 A) 2 X, AA 0 yields the result.
Lemma F2 Let X E St(m,r) and U E St(m,rr). Let U. E (m=((m-rA) be an orthonormal
basis for the orthogonal complement of SU). Denote &- T E and V = U X e (((m-FA) It is guaranteed that (T) + ( holds for i 2
Proof. Since X lies in the Stiefel manifold, we have that
U
I,=X = x I x x
U.
x
(37)
U
$ Φ 4 TV.
41

---

Published as a conference paper at ICLR 2026
Equation (37) shows that v
"T and + & commute, i...
v TTu v = ( ))
②) ②- ① ② T Φ
= (I, + ② 4) = I ( $).
The commutativity shows that the eigenspaces T and v T coincide. As a result, we have again from (37) that ( (V) for i E 21,
Lemma F.3 Suppose that P and Q x diagonal matrices, with non-negative diagonal entries. Let S E S be a positive definite matrix with smallest eigenvalue Amin. then we have that Tr(PSQ) 2 Amin Tr(PQ).
Proof. Let pi and qi be the (i, entry of P and Q. respectively, Then have that Tr(PSQ) DiStun 2 Amin Epi:: Amin Tr(PQ),
where the last inequality comes from S being positive definite, i.e., S e, Se, Amin.
Lemma F.4 Let A mxxx be matrix with full columa rank and B Rexp be non-zero matrix, Let Jmin (-) denote the smallest non-zero singular value. Then it holds that a (AB) 2 Jmin (A) Jiin (B).
Proof. Using the min-max principle for singular values,
Min (AB)
min
||ABx||
||x|| |x||=11xx XXColSSpn B3)
Bx
min
A
|Bx||
|x/l=en B)
Bx
min
Ayll
[x
|xll== B). (ECoSpannBB)
min
Ayll
min
Bx
1yl|=1.yEColSppi
min
Ayyl (x))ssCC)))
min
[x
|
min (A)o A min (B),
where (a) is by changing of variables, i.e., y Bx/| Bx|]
Lemma F.5 (Theorem 2.2.1 of (Chikuse, 2012)) If Z E gmxr has entries drawn i.i.d. from Gaus- 1/2 sian distribution (0, 1), then X Z(Z Z) 11/2 is random matrix uniformly distributed on St(m,r) r).
Lemma F6 (Vershynin, 2010) If E |gxr is a matrix whose entries are independently drawn from (0, 10. 1). Then for every 0, with probability at least cxp(-7² 22/2), we have
0, (Z) /m + +T+
Lemma E7 (Rudelson Vershynin, 2009) If mxxr is matrix whose entries are indepen- dently drawn from (0, 1). Suppose that m 2 r. Then for every 7 2 0, we have for two universal constants C₁ and C2 > that
ar (Z) S --Vmm-V---- - Vr- S (CIr)" + exp(-C2m))
Lemma F88 If U St(m, TA) )) is a fixed matrix, X St(m,r) r) is uniformly sampled from St(m, r) using methods described in Lemma F.5, and > TA, then we have that with probability at least
- --exp--m//1 - (Cir)- -r++1 exp(-C2r).
Tr (UT) T(r - TA +A+1)
6 mr
42

---

Published as a conference paper at ICLR 2026
Proof. Since X E St(m, r) is uniformly sampled from St(m,r) using methods described in Lemma F.5, we write X Z(Z 1/2 where Z IR has entries ii.d. sampled from W/0, 1). We thus have
rra (U x Tr (U z(Z Z) (U))
We now consider U Z E R It is clear that the entries of U U are also i.,. random variables. As consequence of Lemma F.7, have that with probability at least 1-astt-CC17 (C11)) exp(-C2r),
or (U'Z) 2 T(V[ VrA-
We also have from Lemma F.6 that with probability at least - exp(-m/2).
q (Z 2) (Z) < (2Vm+ vF)2
Taking union bound, we have with probability least 1 exp( wxpp-mpppmpmmmmmm (() exp(-C2rr,
Tr (UX) X (a) U Z) (Z)
T(Vī VTA
7(r rA +1) -(r-ra+1)
2/m + VT
3 3 2/T 6/mr
where (a) comes from Lemma F.4.
Lemma F9 Suppose θ E S". Then the update rule (10) guarantees that 0,-1 also belongs to S".
Proof. From the update rule, we have that
0 AX X
1 -A) X1+1++
2
Since O, E and A E S" it follows that X,+1 O,X t+1 E Sm and 1+1 AX,+1 E By definition of M and M", the composition M M defines self-adjoint operator in Sm Hence,
X
x [ M M (I) AA]] ES.
Thus, O+t1 E S", which completes the proof.
Lemma F.10 Let M(-) 9th + RR be a linear mapping that is
with ō E [0, 1). Then for any symmetric matrix Z of rank most and any symmetric mairix Y of rank most r". we have that
(IM'M-I)(Z) YX / ZlI
Proof. Denote A(Z, Y) == ((M' - (Z). Y) = (M()),M(Y)) (Z, Y). The above
inequality trivially holds when / or Y 0. Without loss of generality, assume that |Z]]= * and || 3 0. Define
z and Y = TYIE. It then follows that
|
A Y) A Y) ||||F|/|||
Using the polarization identity, we obtain
M(( 1 (IM + 8)12 1 - 8)1
1
(
(11 + i - -1 r112)
1
43

---

Published as a conference paper at ICLR 2026
Substituting the two equalities into the expression of AZZ Y), we have that
A Y) ((M(2), (z, Y)
1I Y) - |M Y) (1) Y - 7 Y
4
1
(IIM(( Y)12² 2 1I + M 2)12 1 Y
4
(2)
(l|Z
YI
|| YIA
4
8
(1
||
2
8,
where (a) is from the facts that M(-) is (r + RIP with constant 6, rank (Z Y) S rank () +
rank(Y) I r and rank(Z Y) - rank () + rank YY) I Artr Therefore, we have that
K(M-M - [))), Y = (Z, Z, Y = A Y) | F Y - 6 |
which completes the proof.
Lemma F.11 (Lemma 7.3 of (Stöger & Soltanolkotabi, 2021)) Let (M.:::):.. 8 + R be a linear mapping that is (r + TA + 1, RRIP with [0, 1), then [[MM - [))ll S A|| for all
matrices A E 8m of rank at most r TA.
Proof. By Lemma F.10, if A E 9t has rank at most TA and Y E 9 has rank at most 1, then it holds that
M - ()A), Y A
Hence, it suffices to prove that there exists matrix Y of rank 1, such that ((M" M 1--MM-D--- M-II(A),, IIM - (A)l and || || s 1. Since (M* M (A( is symmetric matrix, follows that
( = max u ||||=1
(A)uu
max ||u|==1
max MMMMMMII(C un ||||=|
Let Y Íü where ú arg max (M'' M - (A),, uu T) We then have that rank(Y) 1, ||]|==1
Y YMM K(M (A) Y) = IKM'M-I(C - and |Y||F 1
Lemma F.12 Let A E Raxm B E Rmx be two real matrices, then the following inequality holds
AB A < A B n
Proof. For any unit vector E R. with ||=||| 1, we can obtain that
TT + B x x AB B A (a) 2x ABx.
where (a) is from the fact that x B A æ is scalar. By the Cauchy-Schwee inequality and the definition of the spectral norm, we have that
x ABa < |Aax|
|||||| S |A||
BB|
||
|||||||]
Hence, we obtain the following inequality:
a (AB + BB A )m < 2|
Since this holds for any unit vector it follows that
AB BB+BAAT 3 ||| B L.
44

---

Published as a conference paper at ICLR 2026
Lemma F.13 Let t 2 1 be positive integer. For all real numbers x, satisfying 0 S S 7. the
following inequality holds:
(1 + x² < + 3tz.
Proof. Let f(x) + 31x + x) x E [0, | Then, for all x E [0, Hl obtain
f'(x) 3t t(1 + 2) 1-1 > 3t t(1 { ) 2 (3 e)t > 0.
Therefore, for all x E [0, Hl. f(z) 2 f(0) = 0, which means (1 +1+ x)) < + 3tz for all x E [0, 2
Lemma F.14 Let & R21, 1 E (2) 1). Suppose that sequences {at) (0 Rzo satisfy
1- 4
bi-1 qb. +
(ak 1 + at + 22 + (b((0-1+a;)).
(38)
180k2
1
a1 S 2k
+ 0--,, t= 1.2.....
(39)
6
and another pair of sequences 1 {b:]
C Ryo satisfy
b = qb: +
1-4 q
(
+ at + 20t- Tā + V ba + -z)),
(40)
180
1
ā.- 2k b
+
t 1,2....
(41)
6
If the initial conditions satisfy then a, S āg and by S by hold for all
ao x āo. bo S bu, b1 S bi.
2 0.
Proof. We proceed by mathematical induction. From inequality (38), we obtain
1
a1 S 2k bi
a0
(0) 2k b1
ao
σ1.
where (a) is by initial conditions; and (b) is from equality (40). Analogously, inequality (39) implies
b2 I qbi + 1804
q
(ag + of + 2apa1 + vbi + a1))
1 1
(ag + aî 2ū0ā v (1 + +ā))
9o1 + 180
b2.
where (c) is due to initial conditions and 01 S 0σ: and (d) by equality (41). By induction, we conclude that a, < ã and b, < b, for all t 2 0, which completes the proof.
Lemma E.15 Let k R21. 8 (2 1). Suppose that sequences (ar fbc Rzo satisfy:
b =qb; + 180
ạ (ak + at + 2û, 1 + V (/ (ā)))
(42)
à 22 b + - 11
(43)
If the initial conditions satisfy
3/2h
b.. < i,.
āo. 21. bo. by E Rzo. 0o S 3 <
VI+ q
then have that ā. < 3o 1+g)
1/2 for all 1 112 0.
45

---

Published as a conference paper at ICLR 2026
Proof. We proceed by mathematical induction. We first consider the following auxiliary system:
bo+1 = max{qò: + 1- (az + a² + 2ā + V + +o))) 1 1+ b.).
(44)
18042
2
å max v b, - sky/bb,, t= 1=1,2....
1
(45)
6
Let 0g do, bo bo and b, bb. It holds that ān S 33kbbn Vitq
3r vb, and thus we have
2k in +
1 :0 s 2 + 2/1+9 Vin S 3k 34
V2k
From equality (45) at = 1, we obtain âi 3/. Since åo S follows that
3v24 vI+9
Vb, and â, = 3k / /i,
qdi + 1-4
(o) +
aî + 2ā,â1 + 6 k1))
180k-2
$ <qb +
9
18h²
by + +9.26 + 18
v1 + b1 + 3/2k in + 3kbi)
180
1+0 9
V1+4
qbi + 1 1 9
(186² + 9
+ 18822k2 + +322k2 - 3b3)bb
180
1 1-91 ii
< qoو +
2
bb
2
From equality (44) at = we have b2 1t.. Using the same reasoning for t yields
2k b2 + Ô1 = 2k be + b
k 2 2 b, = -2 be 2 1+ V b2 s 3k b2.
Equality (45) at / = implies that جة = 3k 3k Vbz. Since ài = 3k / b, and â2 = 3k / bz. we obtain
abz + 1-q q (of + a3 + 0ū + //2(a + a2))
180122
qb2 + 180
9 (9
+ 2 82 + 18h2 V 6,82 + +3k( b,b2 + b2))
1 q 18k
9k 18
3v2k
=gb +
V1+9 + v1+q + +3k)02
180
1+0 g
S qbz + 180
(18h?
+ 91.2
+ 182 + 3 3v2222 + b
1 bz
< qbz +
2
1 9 irz-
Applying equality (44) at 2, 8s 3kVbb and by 1 Assume that à:-1 3k 3 b 33 / br. and 61p1 4 claim that âr+l 3k
14462 is derived. Therefore, we have that â 3k 33k//1,,0 0z
2
3k / B841 and B442 = 2 1-
From equality (45), we obtain
âet1 max (2k bett
1
6
max/2kc
bis1
1
be. 3ky betal
2
k
max(2k
b1+1 +
1+9
/beti, 3k berit
2
= =3kybo41
46

---

Published as a conference paper at ICLR 2026
Analogously, equality (44) implies that
2142 = max + 180 (a7 + +ū²+1 + 22gā + Vetta + ( âc-1)); 2 bs+1}
1 1-9 9
1+ q
1-4 y (9 + b + 18k2 1 /bbbet + (eet b, + 35 36. b(41)),
max +
180
1 9 b
2
+max
1-9,
2
+ 1 2(
2
1/2
1
2
1/2
+
1
1 4 bbt11)
20
1
1+9
1 1+0
3k
1+
q
3h
2
1 1+ be+1-
2
Therefore, we have that (6:) decreases linear rate and that à; = 3k / b2 in the system (44)
and (45), which means that å, s /b ( f/2 ) 3k 6o ( f/2 for all 2
We now prove that ā; S â. b, < b, for all t 2 0. Obviously, āo S âo ã1 S â1. bo S bu. and
b & b, hold. Applying equality (42) at f and equality (43) at t = 2, we obtain
1 9 (a) + aî 2āgā Vbn(āo + +ā))
b2 = مàب + 180
1 9 (aß + at + +2n 6á + +())
$ وà1 + 180k2
4 (a + a + 2ōg + vbi + =a))), 2
1+
< max{qbi + 180
= bz,
ā2 2k 02 + a
s 2k v 02 + à1
0d
I /
02 +
1 aai 3k / 02)
= 02
Hence, ăy < 0z. by S b2 and recursively, we can obtain ā S ar, b, < by for all t 2 0. Consequently, (5o achieves at least a linear convergence rate in the system (42) and (43). which means that ā. 3bv (122)/2 for all 2
G EXPERIMENTAL SETUPS
In this section, we provide experimental setups in detail.
G.1 SETUP FOR FIGURE
We apply Algorithm to problem (2) and study the trajectory generated by the algorithm.
In this experiment, set 300, r 5, and 3. The ground-truth matrix A E pmxm is constructed as A = U where U E guxra is random orthonormal matrix and Σ E ST is diagonal with entries generated by power spacing scheme. Specifically, the j-th entry of Σ is given by o, = k - for = j I,...,
TA where we set 0.6.
50000 independent feature matrices (M.) C 8 in the following manner. For each E {1, ... n.). we sample R, E Rm with i...d standard Gaussian entries and define M, 2/7 R, + R,),, which ensures the symmetry of M.
We generate
47

---

Published as a conference paper at ICLR 2026
We initialize Xo Gaussian entries. This initialization ensures that X. lies on the manifold St(m, and Og E S". For this experiment, we set the stepsizes to r7 0.2 and μ = 2.
Zo Zz) Zo)-1/2 and O₂ 0.5I,, where Zo R has ii.d. standard
G.2 SETUP FOR SECTION 5
G.2.1 SETUP FOR THE EXPERIMENTS WITH SYNTHETICC DATA
We apply Algorithm to problem (2) and compare its convergence with standard GD applied to (1).
In the noisy measurement experiment, we set m = 10, = TA = 3, and n = 1000, while for other experiments, the corresponding values are given in the main text. In these settings, the ground truth matrix A E Tmmx is formed as A = UEUT, where U E mR=x*A a random orthonormal matrix and Σ S is diagonal matrix with entries evenly distributed on a logarithmic scale in the interval [1/M, 1I. Feature matrices {M are generated as described in G.I.
For other experiments, we initialize Xo = Zo Z -n)-1// and O = I, for Algorithm 1, where Zo E _ consists of i.i.d. standard Gaussian entries, and X 0.1Zg for GD.
0.1 and μ 2 for RGD and v = 0.1 for GD, except for the first experiment in Subsection 5.2. In that case, we set 1) 0.1, 0.12, 0.14 and μ = 2 for RGD, and 7 0.1, 0.12, 0.14 for GD, corresponding to r = 50, 75, 100, respectively,
Throughout the experiments, we use stepsizes
G.2.2 SETUP FOR IMAGE RECONSTRUCTION EXPERIMENTS
For the image reconstruction experiments, conduct two setups: one based recovering a CIFAR- 10 image from linear measurements and the other on direct matrix sensing of a structured image.
For the CIFAR-I0 experiment, we take the first horse image from CIFAR- 10 dataset, convert it to grayscale, and vectorize it as a E I 1024 The ground-truth matrix is set as A = aa E $S124, The overparameterizzttoo level set r 100, with 50000 feature matrices generated as in G.1. RGD is initialized with X,- Zo Z Zo) - and O I,, where Zo IR m has i.i.d. standard Gaussian entries. GD uses small random initialization: X = 0.1 Zo. We run RGD for RRGD 100 and GD for toD 200 iterations. For RGD, we adopt stepsizes of r 0.01 for updating X and for updating O For GD, we apply stepsize = 0.01 to update X.
After optimization, following the approach of (Duchi et al., 2020, Section 4.1), we perform a rank- one truncated SVD on the recovered matrix A, and the estimate the original signal constructed as the leading singular vector multiplied by the square root of its corresponding singular value. The resulting vector is then reshaped into 32 x 32 reconstruction image.
For the structured image experiment, we generate a grayscale matrix A E S of rank A = 2 using bloek-wave basis functions. Specifically, we construct T' one-dimensional signals of length 128, where each signal is normalized block wave taking values in +1 with random period. Stacking these signals forms a matrix U E R 1288xr The ground- truth image is defined as A = UAUT, where A is 2 x 2 diagonal matrix with diagonal entries 1 and 0.9. This diagonal matrix assigns geometrically decaying weights to different block-wave modes.
We again fix 100 and use are randomly initialized as above. We run RGD for fRGD = 100 and GD for fGD = 200 iterations. We adopt stepsizes of 7 = 0.03 and μ 2 in RGD and a stepsize of 1 = 0.03 GD.
50000 feature matrices generated in G.1. Both RGD and GD
The per-iteration computational complexity of both RGD and GD is Onm²²r) which is dominated by the operation of sensing. Since each RGD iteration requires performing two sensing operations while GD requires only one, we set the number of iterations as top = 25RGD to make the overall runtime roughly comparable between the methods.
G.3 SETUP FOR THE EXPERIMENTS OF PRECONDITIONED ALGORITIIMS
We apply Algorithm and PrecRGD to problem (2) and compare their convergence behavior with GD, PrecGD and ScaledGD(A) applied to problem (1).
48

---

Published as a conference paper at ICLR 2026
In this experiment, set 50, 40, =3, 10, and 8000 for the first instance and increase r to 45 in the second instance while keeping all other parameters fixed. The ground-truth matrix A IR=x is constructed A UEU where U mmx is a random orthonormal matrix and Σ E S is diagonal matrix with entries evenly distributed on a logarithmic scale in the interval [1 /Kx, 1]. Feature matrices (M, generated as described in G.I.
The regularization parameter À used in PrecRGD is chosen adaptively as X ||M* MM, A)llee. where I/ denotes the matrix infinity norm. For PrecGD and ScaledGD (A), we set the [I Y. a) and À 0.1, respectively.
regularization parameters
We initialize X. Zo(Za (Z Zo) -1/2 and O L, for RGD and PrecRGD, where Zo E R consists of i.i.d. standard Gaussian entries. For GD and PrecGD, we use initialization X = Zu- For ScaledGD(A),
set X.calld 10-6 6 Ze. For RGD and PrecRGD, adopt stepsizes of 0.02 for updating X and μ = 2 for updating @. For GD, PrecGD and ScaledGD(A), we apply stepsize v 0.02.
G.4 SETUP THE MATRIX COMPLETION PROBLEMS
We apply RGD to problem (7) and compare its convergence behavior with that of GD applied to (6).
We directly use the code provided in the supplementary material of (Yalçın et al., 2022), replacing only the objective function with our WN-based formulation (7). The update of X, follows the same step-size strategy as that of Y, in their implementation, and O is updated using a fixed step size of 2. Except for these modifications, other settings are kept identical to the original code.
G.5 SETUP FOR THE EXPERIMENT OF DIAGONAL AND NONNEGATIVE +
We apply algorithm (1) to WN and compare its convergence behavior with that of constraining 2 to be diagonal and nonnegative.
In this experiment, we set m = 10, r = 5, 3, = 3, and n = 1000. The ground-truth matrix A mm is constructed as A UΣUT. where U mmxr is a random orthonormal matrix and Σ E S is diagonal matrix with entries evenly distributed on a logarithmic scale in the interval [1/x, 1]. Feature matrices [M] are generated as described in G.I.
We initialize X (o -1/2 and O I, for both settings, where Zo R consists of i.i.d. standard Gaussian entries. We adopt stepsizes of = 0.1 for updating X and μ 2 for updating @. To enforce O to be diagonal and nonnegative in the "Diagonalized" setting, we perform an SVD on θ after each update, extract the diagonal matrix, and apply hard thresholding to ensure that all the diagonal entries are nonnegative.
49


Reviews:
[{"review_id": "tXiIxtP5dG", "reviewer": "Reviewer_mDtz", "summary": "The authors address the problem of overparameterized matrix sensing and present an analysis of the benefits of applying weight normalization in this context. They reformulate the classical matrix sensing problem by decoupling the magnitude and direction components of the matrix, and they demonstrate that using Riemannian gradient descent with weight normalization improves the convergence rate, particularly when the level of overparameterization is large. Experiments on both synthetic and real-world datasets are provided to support the analysis.", "questions": "N/A", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is well written, clearly structured, and the ideas are well explained. The theoretical motivation and experimental design are both sound.", "weaknesses": "I have a major concern regarding the positioning of the paper within the existing literature. While the paper emphasizes the proposed weight normalization approach, its main benefit—improved convergence rate—falls within a broader line of research on optimization techniques for matrix sensing. In particular, related approaches such as preconditioned gradient descent and their follow-up works have already been proposed in this domain. Although some of the related (e.g. preconditioned gradient descent) papers are cited, they are not discussed thoroughly. These works should be more deeply reviewed in the related work section, and, ideally, comparisons should be included in the experimental evaluation to better situate the contribution and highlight the advantages of the proposed approach over existing methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DeB3ZEbMnE", "reviewer": "Reviewer_5U1W", "summary": "This paper studies a matrix factorization similar to the weight normalization that can be helpful for the matrix sensing problem. The author proves that this matrix factorization with Riemannian optimization can achieve a linear convergence rate, which is an exponential improvement over the previous lower bound for symmetric matrix sensing.", "questions": "Lemma 4.1 states that a point $(X, \\Theta)$ is a saddle point if certain conditions are satisfied. However, this does not necessarily imply that all saddle points satisfy these conditions. Could the authors clarify how this lemma connects to the saddle-to-saddle dynamics and the subsequent incremental learning process?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The paper is well-organized and easy to follow. The contributions of this paper is proposed in a clear way, and the main technical idea is also clear.  The author also provides solid theoretical proof and some discussions about the main result. \n\n2. The authors also further study the initial increment learning phase in the optimization process, which is good for understanding the optimization dynamics of WN with Riemannian optimization. \n\n3. The empirical part is comprehensive. Beyond the simulated optimization problems, the authors also provide experiments on image reconstruction problem. The performance of WN with Riemannian optimization is better compared to GD, which also matches the theoretical discovery of this paper.", "weaknesses": "1. My main concern is that after relaxing the PSD constraint, the problem no longer corresponds to symmetric matrix sensing. Specifically, when the constraint $\\Theta \\in S^r_{+}$ is relaxed to $\\Theta \\in S^r$, the term $X^T \\Theta X$ may no longer be written as $YY^T$. Thus, the proposed method is not a true solution to symmetric matrix sensing but rather an approach that accelerates matrix sensing in general. This weakens the paper’s contribution, as several existing methods [1,2] also focus on improving the efficiency of matrix sensing.\n\n2. The paper introduces Riemannian optimization equations without providing sufficient background or basic explanations. It would be helpful if the authors could include a high-level introduction and some intuitive discussion of Riemannian optimization to improve readability.\n\n[1]. Xu et al. 2023. The power of preconditioning in overparameterized low-rank matrix sensing\n\n[2]. Xiong et al. 2024. How over-parameterization slows down gradient descent in matrix sensing: The curses of symmetry and initialization", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pLBrxlVPA6", "reviewer": "Reviewer_BeU2", "summary": "This paper provides a theoretical characterization of how Weight Normalization (WN) accelerates convergence in overparameterized matrix sensing. The authors prove that Riemannian gradient descent (RGD) with WN achieves linear convergence and improved sample complexity, while standard gradient descent (GD) may exhibit exponential slowdown. Theoretical findings are supported by synthetic experiments that align well with the analysis.", "questions": "1.See Weakness. 1.\n\n2.Since $A$ is PSD, one might consider reformulating (1) as \n\n$\\min_{X, \\Theta} f(X, \\Theta) = \\frac{1}{4}\\|M(X\\Theta X^{\\top})-y\\|^2$\n\nwith $\\Theta$  being diagonal and nonnegative. Would this restriction simplify the analysis or improve interpretability?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.Provides the first clear theoretical explanation of how WN accelerates convergence under overparameterization.\n\n2.Solid and rigorous analysis using Riemannian optimization tools.\n\n3.Well-aligned experiments that validate the theory.\n\n4.Overall writing and presentation are clean and accessible.", "weaknesses": "1.The paper identifies a two-stage convergence pattern with a transition at $r_A-\\frac{1}{2}$ yet this boundary is not theoretically justified or intuitively explained.\n\n2.The “Full-rank case (r = m)” in Sec. 5.3 is conceptually an extension of Sec. 5.2 (“ON THE BENEFIT OF OVERPARAMETERIZATION”) and could be merged there for better logical flow.\n\n3.The title seems somewhat broader than the actual technical scope, which is limited to PSD matrix sensing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5N7K4wAqSD", "reviewer": "Reviewer_zbqV", "summary": "This paper investigates the advantages of a generalized weight normalization (WN) technique in the context of overparameterized matrix sensing, where the goal is to recover a low-rank PSD matrix from linear measurements. The authors prove that WN, combined with Riemannian gradient descent (RGD), achieves linear convergence with random initialization, offering an exponential speedup over standard gradient descent (GD) methods without WN. They demonstrate that increased over-parameterization reduces both iteration and sample complexity polynomially. The analysis reveals a two-phase convergence behavior (initial saddle-escape phase followed by linear convergence) and is supported by experiments on synthetic data and image reconstruction tasks.", "questions": "Overall, I don't see major issues that undermine the core contributions—the theory is sound, and the claims are well-supported. \nThat said, here are a few constructive suggestions:\n1. I am curious about how the proposed algorithm would perform in more challenging matrix sensing settings, such as those with specially designed structures or nontrivial optimization landscapes (e.g., as discussed in arXiv:2110.10279).\n2. A short note on potential applications (e.g., in signal processing) could emphasize the work's relevance to ICLR's audience.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper's primary strength lies in its novel theoretical insights, providing the first characterization of how WN leverages overparameterization for faster convergence in matrix sensing, with a rigorous proof of linear convergence that exponentially outperforms sublinear rates in non-WN methods. It quantifies the benefits clearly, showing polynomial improvements in iteration and sample complexity as overparameterization increases, which contrasts positively with prior work where overparameterization hinders performance. Empirically, the experiments robustly validate the theory, including comparisons under varying conditions like condition numbers, overparameterization levels, and noise, with real-world image reconstruction adding practical value. The manuscript is well-organized and clear, featuring helpful tables, figures, and a reproducibility statement with full proofs and setups in appendices.", "weaknesses": "1. While the analysis is focused and insightful, its scope is limited to symmetric PSD matrix sensing. It would be helpful if the authors could briefly explain why it is reasonable to consider only the symmetric PSD case.\n2. Experimentally, while solid, the studies are somewhat constrained in scale (e.g., small matrix dimensions in synthetic tests).\n3. The RIP condition required for the main theorem scales as $\\delta = O((r - r_A)^6 / (\\kappa^2 m^3 r^4 r_A))$, which appears to be rather stringent. A brief analysis or discussion of this condition's implications would be valuable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)