SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ziye Ma
Department of EECS
ziyema@berkeley.edu
Javad Lavaei
lavaei@berkeley.edu
Somayeh Sojoudi
Department of EECS, ME
sojoudi@berkeley.edu
Gradient descent (GD) is crucial for generalization in machine learning models, as it induces implicit regularization, promoting compact representations. In this work, we examine the role of GD in inducing implicit regularization for tensor optimization, particularly within the context of the lifted matrix sensing framework. This framework has been recently proposed to address the non-convex matrix sensing problem by transforming spurious solutions into strict saddles when optimizing over symmetric, rank-1 tensors. We show that, with sufficiently small initialization scale, GD applied to this lifted problem results in approximate rank-1 tensors and critical points with escape directions. Our findings underscore the significance of the tensor parametrization of matrix sensing, in combination with first-order methods, in achieving global optimality in such problems.
This paper is dedicated to addressing the non-convex problem of matrix sensing, which has numerous practical applications and is rich in theoretical implications. Its canonical form can be written as:
$$
\text {f i n d} \quad M \in \mathbb {R} ^ {n \times n} \tag {1}
$$
$$
\begin{array}{l} \text {s . t .} \quad \mathcal {A} (M) = \mathcal {A} (M ^ {*}) \quad \operatorname {r a n k} (M) \leq r, M \succeq 0. \end{array}
$$
$\mathcal{A}(\cdot):\mathbb{R}^{n\times n}\mapsto \mathbb{R}^m$  is a linear operating consisting of  $m$  sensing matrices  $\{A_i\}_{i = 1}^m\in \mathbb{R}^{n\times n}$  where  $\mathcal{A}(M) = [\langle A_1,M\rangle ,\dots ,\langle A_m,M\rangle ]^T$  . The sensing matrices and the measurements  $b = \mathcal{A}(M^{*})$  are given, while  $M^{*}$  is an unknown low-rank matrix to be recovered from the measurements. The true rank of  $M^{*}$  is bounded by  $r$  , usually much smaller than the problem size  $n$  . More importantly, since  $\mathcal{A}$  is linear, one can replace  $A_{i}$  with  $(A_{i} + A_{i}^{\top}) / 2$  without changing  $b$  , and therefore all sensing matrices can be assumed to be symmetric.
The aforementioned problem serves as an extension of both compressed sensing [1], which is widely applied in the field of medical imaging, and matrix completion [2, 3], which possesses an array of notable applications [4]. Additionally, this problem emerges in a variety of real-world situations such as phase retrieval [5-7], motion detection [8], and power system state estimation [9, 10]. A recent study by [11] established that any polynomial optimization problem can be converted into a series of problems following the structure of [1], thereby underscoring the significance of investigating this specific non-convex formulation. Within the realm of contemporary machine learning, [1] holds relevance as it is equivalent to the training problem for a two-layer neural network with quadratic activations [12]. In this context,  $m$  denotes the number of training samples,  $r$  is the size of the hidden layer, and the sensing matrices  $A_{i} = x_{i}x_{i}^{\top}$  are rank-1, with  $x_{i}$  representing the  $i^{th}$  datapoint.
To solve (1), an increasingly popular approach is the Burer-Monteiro (BM) factorization [13], in which the low-rank matrix  $M$  is factorized into  $M = XX^{\top}$  with  $X \in \mathbb{R}^{n \times r}$ , thereby omitting the
constraint, making it amenable to simple first-order methods such as gradient descent (GD), while scaling with  $\mathcal{O}(nr)$  instead of  $\mathcal{O}(n^2)$ . The formulation can be formally stated as follows:
$$
\min  _ {X \in \mathbb {R} ^ {n \times r}} f (X) := \frac {1}{2} \| \mathcal {A} \left(X X ^ {T}\right) - b \| ^ {2} = \frac {1}{2} \| \mathcal {A} \left(X X ^ {T} - Z Z ^ {\top}\right) \| ^ {2} \quad (\text {U n l i f t e d P r o b l e m}) \tag {2}
$$
with  $Z \in \mathbb{R}^{n \times r}$  being any ground truth representation such that  $M^* = ZZ^\top$ . Since (2) is a nonconvex problem, it can have spurious local minimal[1] making it difficult to recover  $M^*$  in general. The pivotal concept in solving (1) and (2) to optimality is the notion of Restricted Isometry Property (RIP), which measures the proximity between  $\|A(M)\|_F^2$  and  $\|M\|_F^2$  for all low-rank matrices  $M$ . This proximity is captured by a constant  $\delta_p$ , where  $\delta_p = 0$  means  $\mathcal{A}(M) = M$  for matrices up to rank  $p$ , leading to exact isometry case, and  $\delta_p \to 1$  implying a problematic scenario in which the proximity error is large. For a precise definition, please refer to Appendix A.3
Conventional wisdom suggests that there is a sharp bound on the RIP constant that controls the recoverability of  $M^{*}$ , with  $1/2$  being the bound for [2]. [14, 15] prove that if  $\delta_{2r} < 1/2$ , then all local minimizers are global minimizers, and conversely if  $\delta_{2r} \geq 1/2$ , counterexamples can be easily established. Similar bounds of  $1/3$  are also derived for general objectives [16, 17], demonstrating the importance of the notion of RIP. However, more recent studies reveal that the technique of over-parametrization (by using  $X \in \mathbb{R}^{n \times r_{\text{search}}}$  instead, with  $r_{\text{search}} > r$ ) can take the sharp RIP bound to higher values [18, 19]. Recently, it has also been shown that using a semidefinite programming (SDP) formulation (convex relaxation) can lead to guaranteed recovery with a larger RIP bound that approaches 1 in the transition to the high-rank regime when  $n \approx 2r$  [20]. These works all show the efficacy of over-parametrization, shedding light on a powerful way to find the global solution of complex non-convex problems. However, all of these techniques fail to handle real-world cases with  $\delta_{2r} \to 1$  in the low-rank regime. To this end, a recent work [21] drew on important concepts from the celebrated Lasserre's Hierarchy [22] and proposed a lifted framework based on tensor optimization that could convert spurious local minimizers of [2] into strict saddle points in the lifted space, for arbitrary RIP constants in the  $r = 1$  case. We state this lifted problem below:
$$
\min  _ {\mathbf {w} \in \mathbb {R} ^ {n o l}} \| \langle \mathbf {A} ^ {\otimes l}, \mathbf {w} \otimes \mathbf {w} \rangle - b ^ {\otimes l} \| _ {F} ^ {2} \quad (\text {L i f t e d P r o b l e m}, r = 1) \tag {3}
$$
where  $\mathbf{w}$  is an  $l$ -way,  $n$ -dimensional tensor, and  $\mathbf{A}^{\otimes l}$  and  $b^{\otimes l}$  are tensors "lifted" from  $\mathcal{A}$  and  $b$  via tensor outer product. We defer the precise definition of tensors and their products to Section 2. The main theorem of [21] states that when  $r = 1$ , for some appropriate  $l$ , the first-order points (FOP) of [2] will be converted to FOPs of [3] via lifting, and that spurious second order points (SOP) of [2] will be converted into strict saddles, under some technical conditions, provided that  $\mathbf{w}$  is symmetric, and rank-1. This rank-1 constraint on the decision variable  $\mathbf{w}$  is non-trivial, since finding the dominant rank-1 component of symmetric tensors is itself a non-convex problem in general, and requires a number of assumptions for it to be provably correct [23, 24]. This does not even account for the difficulties of maintaining the symmetric properties of tensors, which also has no natural guarantees. Therefore, although this lifted formulation may be promising in the pursuit of global minimum, there are still major questions to be answered. Most importantly, it is desirable to know whether the symmetric, rank-1 condition is necessary, and if so, how to achieve it without explicit constraints?
The necessity of the condition in question can be better understood through insights from [25]. The authors argue that over-parametrizing non-convex optimization problems can reshape the optimization landscape, with the effect being largely independent of the cost function and primarily determined by the parametrization. This notion is consistent with [21], which contends that over-parametrizing vectors into tensors can transform spurious local solutions into strict saddles. However, [25] specifically examines the parametrization from vectors/matrices to tensors, concluding that stationary points are not generally preserved under tensor parametrization, contradicting [21]. This implies that the symmetric, rank-1 constraints required in [3] are crucial for the conversion of spurious points.
It is essential to devise a method to encourage tensors to be near rank-1, with implicit regularization as a potential solution. There has been a recent surge in examining the implicit regularization effects in first-order optimization methods, such as gradient descent (GD) and stochastic gradient descent (SGD) [26], which has been well-studied in matrix sensing settings [27-29, 12]. This intriguing observation has prompted us to explore the possible presence of similar implicit regularization in tensor spaces. Our findings indicate that when applying GD to the tensor optimization problem
<sup>1</sup>A spurious point satisfies first-order and second-order necessary conditions but is not a global minimum.
(3), an implicit bias can be detected with sufficiently small initialization points. This finding does not directly extend from its matrix counterparts due to the intricate structures of tensors, resulting in a scarcity of useful identities and well-defined concepts for even fundamental properties such as eigenvalues. Furthermore, we show that when initialized at a symmetric tensor, the entire GD trajectory remains symmetric, completing the requirements.
In this paper, we demonstrate that over-parametrization alone does not inherently simplify non-convex problems. However, employing a suitable optimization algorithm offers a remarkably straightforward solution, as this specific algorithm implicitly constrains our search to occur within a compact representation of the over-parametrized space without necessitating manual embeddings or transformations. This insight further encourages the investigation of a (parametrization, algorithm) pair for solving non-convex problems, thereby enhancing our understanding of achieving global optimality in non-convex problems.
1.1 Related Works
Over-parametrization in matrix sensing. Except for the lifting formulation (3), there are two mainstream approaches to over-parametrization in matrix sensing. The first one is done via searching over  $Y \in \mathbb{R}^{n \times r_{\text{search}}}$  instead of  $X \in \mathbb{R}^{n \times r}$ , and using some distance metric to minimize the distance between  $A(YY^\top)$  and  $b$ . Using an  $l_2$  norm, [18, 19] established that if  $r_{\text{search}} > r[(1 + \delta_n) / (1 - \delta_n) - 1]^2 / 4$ , with  $r \leq r_{\text{search}} < n$ , then every second-order point  $\hat{Y} \in \mathbb{R}^{n \times r_{\text{search}}}$  satisfies that  $\hat{Y}\hat{Y}^\top = M^*$ . [30] showed that even in the over-parametrized regime, noise can only finitely influence the optimization landscape. [29] offered similar results for an  $l_1$  loss under good enough RIP constant. Another popular approach to over-parametrization is to use a convex SDP formulation, which is a convex relaxation of [1] [31]. It has been known for years that as long as  $\delta_{2r} < 1 / 2$ , then the global optimality of the SDP formulation corresponds to the ground truth  $M^*[32]$ . Recently [20] updated this bound to  $2r / (n + (n - 2r)(2l - 5))$ , which can approach 1 if  $n \approx 2r$ .
Algorithm regularization in over-parametrized matrix sensing. [12, 33] prove that the convergence to global solution via GD is agnostic of  $r_{\mathrm{search}}$ , in that it only depends on initialization scale, step-size, and RIP property. [29] demonstrates the same effect for an  $l_1$  norm, and further showed that a small initialization nullifies the effect of over-parametrization. Besides these works, [27] refined this analysis, showing that via a sufficiently small initialization, the GD trajectory will make the solution implicitly penalize towards rank-  $r$  matrices after a small number of steps. [28] took it even further by showing that the GD trajectory will first make the matrix rank-1, rank-2, all the way to rank-  $r$ , in a sequential way, thereby resembling incremental learning.
Implicit bias in tensor learning. The line of work [34-36] demonstrates that for a class of tensor factorization problems, as long as the initialization scale is small, the learned tensor via GD will be approximately rank-1 after an appropriate number of steps. Our paper differs from this line of work in three meaningful ways: 1) The problem considered in those works are optimization problems over vectors, not tensors, and therefore the goal is to learn the structure of a known tensor, rather than learning a tensor itself; 2) Our proof relies directly on tensor algebra instead of adopting a dynamical systems perspective, providing deeper insights into tensor training dynamics while dispensing with the impractical assumption of an infinitesimal step-size.
1.2 Main Contributions
1. We demonstrate that, beyond vector and matrix learning problems, optimization of differentiable objectives, such as the  $l_{2}$  norm, through Gradient Descent (GD) can encourage a more compact representation for tensors as decision variables. This results in tensors being approximately rank-1 after a number of gradient steps. To achieve this, we employ an innovative proof technique grounded in tensor algebra and introduce a novel tensor eigenvalue concept, the variational eigenvalue (v-eigenvalue), which may hold independent significance due to its ease of use in optimization contexts.
2. We show that if a tensor is a first-order point of the lifted objective (3) and is approximately rank-1, then its rank-1 component can be mapped to an FOP of (2), implying that all FOPs of (3) lie in a small sphere around the lifted FOPs of (2). Furthermore, these FOPs possess an escape direction when reasonably distant from the ground truth solution, irrespective of the Restricted Isometry Property (RIP) constants.
3. We present a novel lifted framework that optimizes over symmetric tensors to accommodate the over-parametrization of matrix sensing problems with arbitrary  $r$ . This approach is necessary because directly extending the work of [21] from  $r = 1$  to higher values may lead to non-cubical and, consequently, non-symmetric tensors.
Please refer to Appendix A.1 and A.2 for the notations and definitions of first-order and second-order conditions. Here, we introduce two concepts that are critical in understanding our main results.
Definition 1 (Tensors and Products). We define an  $l$ -way tensor as:
$$
\mathbf {a} = \left\{a _ {i _ {1} i _ {2} \dots i _ {l}} \mid 1 \leq i _ {k} \leq n _ {k}, 1 \leq k \leq l \right\} \in \mathbb {R} ^ {n _ {1} \times \dots \times n _ {l}}
$$
Moreover, if  $n_1 = \dots = n_l$ , then we call this tensor an  $l$ -order (or  $l$ -way),  $n$ -dimensional tensor.  $\mathbb{R}^{n\circ l}$  is an abbreviated notion for  $n \circ l \coloneqq n \times \dots \times n$ . In this work, tensors are denoted with bold letters unless specified otherwise. The tensor outer product, denoted as  $\otimes$ , of 2 tensors  $\mathbf{a}$  and  $\mathbf{b}$ , respectively of orders  $l$  and  $p$ , is a tensor of order  $l + p$ , namely  $\mathbf{c} = \mathbf{a} \otimes \mathbf{b}$  with  $c_{i_1\dots i_lj_1\dots j_p} = a_{i_1\dots i_l}b_{j_1\dots j_p}$ . We also use the shorthand  $\mathbf{a}^{\otimes l}$  for repeated outer product of  $l$  times for arbitrary tensor/matrix/vector  $\mathbf{a}$ .  $\langle \mathbf{a},\mathbf{b}\rangle_{i_1,\ldots ,i_d}$  denotes tensor inner product along dimensions  $i_1,\ldots ,i_d$  (with respect to the first tensor), in which we simply sum over the specified dimensions after the outer product  $\mathbf{a} \otimes \mathbf{b}$  is calculated. This means that the inner product is of  $l + p - 2d$  orders. Please refer to Appendix A.3 for a more in-depth review on tensors, especially on its symmetry and rank.
Definition 2 (Restricted Strong Smoothness (RSS) and Restricted Strong Convexity (RSC)). The linear operator  $\mathcal{A}:\mathbb{R}^{n\times n}\mapsto \mathbb{R}^m$  satisfies the  $(L_s,r)$ -RSS property and the  $(\alpha_{s},r)$ -RSC property if
$$
f (M) - f (N) \leq \langle M - N, \nabla f (N) \rangle + \frac {L _ {s}}{2} \| M - N \| _ {F} ^ {2}
$$
$$
f (M) - f (N) \geq \langle M - N, \nabla f (N) \rangle + \frac {\alpha_ {s}}{2} \| M - N \| _ {F} ^ {2}
$$
are satisfied, respectively for all  $M,N\in \mathbb{R}^n$  with  $\mathrm{rank}(M),\mathrm{rank}(N)\leq r$ . Note that RSS and RSC provide a more expressible way to represent the RIP property, with  $\delta_r = (L_s - \alpha_s) / (L_s + \alpha_s)$ .
3 The Lifted Formulation for General  $r$
A natural extension of (3) to general  $r$  requires that instead of optimizing over  $X \in \mathbb{R}^{n \times r}$ , we optimize over  $\mathbb{R}^{[n \times r] \circ l}$  tensors, and simply making tensor outer products between  $\mathbf{w}$  to be inner products. However, such a tensor space is non-cubical, and subsequently not symmetric. This is the higher-dimensional analogy of non-square matrices, which lacks a number of desirable properties, as per the matrix scenario. In particular, it is necessary for our approach to optimize over a cubical, symmetric tensor space since in the next section we prove that there exists an implicit bias of the gradient descent algorithm under that setting.
In order to do so, we simply vectorize  $X \in \mathbb{R}^{n \times r}$  into  $\operatorname{vec}(X) \in \mathbb{R}^{nr}$ , and optimize over the tensor space of  $\mathbb{R}^{nr \circ l}$ , which again is a cubical space. In order to convert a tensor  $\mathbf{w} \in \mathbb{R}^{nr \circ l}$  back to  $\mathbb{R}^{[n \times r] \circ l}$  to use a meaningful objective, we introduce a new 3-way permutation tensor  $\mathbf{P} \in \mathbb{R}^{n \times r \times nr}$  that "unstacks" vectorized matrices. Specifically,
$$
\langle \mathbf {P}, \operatorname {v e c} (X) \rangle_ {3} = X \quad \forall X \in \mathbb {R} ^ {n \times r}, n, r \in \mathbb {Z} ^ {+}
$$
Such  $\mathbf{P}$  can be easily constructed via filling appropriate scalar "1"s in the tensor. Via Lemma A.4 we also know that
$$
\langle \mathbf {P} ^ {\otimes l}, \operatorname {v e c} (X) ^ {\otimes l} \rangle_ {3 * [ l ]} = \left(\langle \mathbf {P}, \operatorname {v e c} (X) \rangle_ {3}\right) ^ {\otimes l} = X ^ {\otimes l} \tag {4}
$$
where  $[l]$  denotes the integer set  $[1, \ldots, l]$ , and  $c * [l]$  denotes  $[c, 2c, \ldots, c * l]$  for some  $c \in \mathbb{Z}^+$ . For notational convenience, we abbreviate  $\langle \mathbf{P}^{\otimes l}, \mathbf{w} \rangle_{3*[l]}$  as  $\mathbf{P}(\mathbf{w})$  for any arbitrary  $z$ -dimensional tensor  $\mathbf{w}$  where  $z$  can be broken down into the product of two positive integers. Thus, using (4), we can extend (3) to a problem of general  $r$ , yet still defined over a cubical tensor space:
$$
\min  _ {\mathbf {w} \in \mathbb {R} ^ {n r o l}} \| \langle \mathbf {A} ^ {\otimes l}, \langle \mathbf {P} (\mathbf {w}), \mathbf {P} (\mathbf {w}) \rangle_ {2 * [ l ]} \rangle - b ^ {\otimes l} \| _ {F} ^ {2} \quad \text {(L i f t e d f o r m u l a t i o n , g e n e r a l} r) \tag {5}
$$
Let us define a 3-way tensor  $\mathbf{A} \in \mathbb{R}^{m \times n \times n}$  so that  $\mathbf{A}_{kij} = (A_k)_{ij} \forall k \in [m], (i,j) \in [n] \times [n]$ . Define  $f^l(\cdot): \mathbb{R}^{n \circ 2l} \mapsto \mathbb{R}$  and  $h^l(\cdot): \mathbb{R}^{[n \times r] \circ l} \mapsto \mathbb{R}$  as  $f^l(\mathbf{M}) := \| \langle \mathbf{A}^{\otimes l}, \mathbf{M} \rangle - b^{\otimes l} \|_F^2$  and  $h^l(\mathbf{w}) = f^l(\langle \mathbf{w}, \mathbf{w} \rangle_{2*[l]})$ , with  $\nabla f^l(\cdot) = \nabla_{\mathbf{M}} f^l(\cdot)$  and  $\nabla h^l(\cdot) = \nabla_{\mathbf{w}} h^l(\cdot)$ .
We prove that (5) has all the good properties detailed in [21] for (3). In particular, we prove that the symmetric, rank-1 FOPs of (5) have a one-to-one correspondence with those of (2), and that those FOPs that are reasonably separated from  $M^{*}$  or have a small  $r^{th}$  singular value can be converted to strict saddle points via some level of lifting. For the detailed theorems and proofs, please refer to Appendix B
4 Implicit Bias of Gradient Descent in Tensor Space
In this section, we study why and how applying gradient descent to (5) will result in an implicit bias towards to rank-1 tensors. Prior to presenting the proofs, we shall elucidate the primary intuition behind how GD contributes to the implicit regularization of (2). This will aid in comprehending the impact of implicit bias on (5), as they share several crucial observations, albeit encountering greater technical hurdles. Consider the first gradient step of (2), initialized at a random point  $X_0 \in \mathbb{R}^{n \times r_{\mathrm{search}}} = \epsilon X$  with  $\|X\|_F^2 = 1$  and  $r_{\mathrm{search}} \geq r$ :
$$
\begin{array}{l} X _ {1} = X _ {0} - \eta \nabla h (X _ {0}) = (I + \eta [ \mathcal {A} ^ {*} \mathcal {A} (M ^ {*}) ]) X _ {0} - [ \mathcal {A} ^ {*} \mathcal {A} (X _ {0} X _ {0} ^ {\top}) ] X _ {0} \\ = \left(I + \eta \left[ \mathcal {A} ^ {*} \mathcal {A} \left(M ^ {*}\right) \right]\right) X _ {0} - \epsilon^ {2} \left[ \mathcal {A} ^ {*} \mathcal {A} \left(X X ^ {\top}\right) \right] X _ {0} \\ = (I + \eta [ \mathcal {A} ^ {*} \mathcal {A} (M ^ {*}) ]) X _ {0} + \mathcal {O} (\epsilon^ {3}) \\ \end{array}
$$
where  $\eta$  is the step-size. Therefore, if  $\epsilon$  is chosen to be small enough, we have that
$$
X _ {t} \approx (I + \eta \mathcal {A} ^ {*} \mathcal {A} (M ^ {*})) ^ {t} X _ {0} \quad \text {a s} \epsilon \rightarrow 0
$$
Again, according to the symmetric assumptions on  $\mathcal{A}$ , we can apply spectral theorem on  $\mathcal{A}^*\mathcal{A}(M^*) = \sum_{i=1}^{n}\lambda_i v_i v_i^\top$  for which the eigenvectors are orthogonal to each other. It follows that  $X_t \approx \left(\sum_{i=1}^{n}(1 + \eta \lambda_i)^t v_i v_i^\top\right) X_0$ .
In many papers surveyed above on making an argument of implicit bias, it is assumed that there is very strong geometric uniformity, or under the context of this paper, it means that  $L_{s} / \alpha_{s} \approx 1$ . Under this assumption, we have  $f(M) \approx f(N) + \langle M - N, \nabla f(M) \rangle + \| M - N \|_{F}^{2} / 2$ , leading to the fact that  $\nabla^2 f(M) = \mathcal{A}^* \mathcal{A} \approx I$ . This immediately gives us  $\mathcal{A}^* \mathcal{A}(M^*) \approx M^*$  so that  $\lambda_{r+1}, \ldots, \lambda_n \approx 0$  as  $M^*$  is by assumption a rank- $r$  matrix. This further implies that  $X_t \approx \left( \sum_{i=1}^{r} (1 + \eta \lambda_i)^t v_i v_i^\top \right) X_0$ , which will become a rank- $r$  matrix, achieving the effect of implicit regularization, as  $X$  is now over-parametrized by having  $r_{\text{search}} \geq r$ .
However, when tackling the implicit regularization problem in tensor space, one key deviation from the aforementioned procedure is that  $\bar{L}_s / \alpha_s$  will be relatively large, as otherwise there will be no spurious solutions, even in the noisy case [14, 30], which is also the motivation for using a lifted framework in the first place. Therefore, instead of saying that  $A^*A(M^*) \approx M^*$ , we aim to show that the gap between the eigenvalues of a comparable tensor term will enlarge as we increase  $l$ , making the tensor predominantly rank-1. This observation demonstrates the power of the lifting technique, while at the same time eliminates the critical dependence on a small  $L_s / \alpha_s$  factor that is in practice often unachievable due to requiring sample numbers  $m$  in the asymptotic regime [37].
Therefore, in order to establish an implicit regularization result for (5), there are four major steps that need to be taken: