Implicit Regularization in Matrix Factorization

Introduction

When optimizing underdetermined problems with multiple global minima, the choice of optimization algorithm can play a crucial role in biasing us toward a specific global minima, even though this bias is not explicitly specified in the objective or problem formulation. For example, using gradient descent to optimize an unregularized, underdetermined least squares problem would yield the minimum Euclidean norm solution, while using coordinate descent or preconditioned gradient descent might yield a different solution. Such implicit bias, which can also be viewed as a form of regularization, can play an important role in learning.

In particular, implicit regularization has been shown to play a crucial role in training deep models (Neyshabur et al., 2015, 2017; Zhang et al., 2017; Keskar et al., 2017): deep models often generalize well even when trained purely by minimizing the training error without any explicit regularization, and when there are more parameters than samples and the optimization problem is underdetermined. Consequently, there are many zero training error solutions, all global minima of the training objective, some of which my generalize horribly. Nevertheless, our choice of optimization algorithm, typically a variant of gradient descent, seems to prefer solutions that do generalize well. This generalization ability cannot be explained by the capacity of the explicitly specified model class (namely, the functions representable in the chosen architecture). Instead, it seems that the optimization algorithm biases us toward a “simple" model, minimizing some implicit “regularization measure”, and that generalization is linked to this measure. But what are the regularization measures that are implicitly minimized by different optimization procedures?

As a first step toward understanding implicit regularization in complex models, in this paper we carefully analyze implicit regularization in matrix factorization models, which can be viewed as two-layer networks with linear transfer. We consider gradient descent on the entries of the factor matrices, which is analogous to gradient descent on the weights of a multilayer network. We show how such an optimization approach can indeed yield good generalization properties even when the problem is underdetermined. We identify the implicit regularizer as the nuclear norm, and show that even when we use a full dimensional factorization, imposing no constraints on the factored matrix, optimization by gradient descent on the factorization biases us toward the minimum nuclear norm solution. Our empirical study leads us to conjecture that with small step sizes and initialization close to zero, gradient descent converges to the minimum nuclear norm solution, and we provide empirical and theoretical evidence for this conjecture, proving it in certain restricted settings.

Factorized Gradient Descent for Matrix Regression

We are particularly interested in the regime where $m\ll n^{2}$, in which case (1) is an underdetermined system with many global minima satisfying $\mathcal{A}(X)=y$. For such underdetermined problems, merely minimizing (1) cannot ensure recovery (in matrix completion or recovery problems) or generalization (in prediction problems). For example, in a matrix completion problem (without diagonal observations), we can minimize (1) by setting all non-diagonal unobserved entries to zero, or to any other arbitrary value.

Instead of working on $X$ directly, we will study a factorization $X=UU^{\top}$. We can write (1) equivalently as optimization over $U$ as,

When $d<n$, this imposes a constraint on the rank of $X$, but we will be mostly interested in the case $d=n$, under which no additional constraint is imposed on $X$ (beyond being p.s.d.) and (2) is equivalent to (1). Thus, if $m\ll n^{2}$, then (2) with $d=n$ is similarly underdetermined and can be optimized in many ways — estimating a global optima cannot ensure generalization (e.g. imputing zeros in a matrix completion objective). Let us investigate what happens when we optimize (2) by gradient descent on $U$.

To simulate such a matrix reconstruction problem, we generated $m\ll n^{2}$ random measurement matrices and set $y=\mathcal{A}({X^{*}})$ according to some planted $X^{*}\succeq 0$. We minimized (2) by performing gradient descent on $U$ to convergence, and then measured the relative reconstruction error $\norm{X-X^{*}}_{F}$. Figure 1 shows the normalized training objective and reconstruction error as a function of the dimensionality $d$ of the factorization, for different initialization and step-size policies, and three different planted $X^{*}$.

First, we see that (for sufficiently large $d$) gradient descent indeed finds a global optimum, as evidenced by the training error (the optimization objective) being zero. This is not surprising since with large enough $d$ this non-convex problem has no spurious local minima (Burer and Monteiro, 2003; Journée et al., 2010) and gradient descent converges almost surely to a global optima (Lee et al., 2016); there has also been recent work establishing conditions for global convergence for low $d$ (Bhojanapalli et al., 2016; Ge et al., 2016).

The more surprising observation is that in panels $(a)$ and $(b)$, even when $d>m/n$, indeed even for $d=n$, we still get good reconstructions from the solution of gradient descent with initialization $U_{0}$ close to zero and small step size. In this regime, (2) is underdetermined and minimizing it does not ensure generalization. To emphasize this, we plot the reference behavior of a rank unconstrained global minimizer $X_{gd}$ obtained via projected gradient descent for (1) on the $X$ space. For $d<n$ we also plot an example of an alternate “bad" rank $d$ global optima obtained with an initialization based on SVD of $X_{gd}$ (‘SVD Initialization’).

When $d<m/n$, we understand how the low-rank structure can guarantee generalization (Srebro et al., 2005) and reconstruction (Keshavan, 2012; Bhojanapalli et al., 2016; Ge et al., 2016). What ensures generalization when $d\gg m/n$? Is there a strong implicit regularization at play for the case of gradient descent on factor space and initialization close to zero?

Observing the nuclear norm of the resulting solutions plotted in Figure 2 suggests that gradient descent implicitly induces a low nuclear norm solution. This is the case even for $d=n$ when the factorization imposes no explicit constraints. Furthermore, we do not include any explicit regularization and optimization is run to convergence without any early stopping. In fact, we can see a clear bias toward low nuclear norm even in problems where reconstruction is not possible: in panel (c) of Figure 2 the number of samples $m=nr/4$ is much smaller than those required to reconstruct a rank $r$ ground truth matrix $X^{*}$. The optimization in (2) is highly underdetermined and there are many possible zero-error global minima, but gradient descent still prefers a lower nuclear norm solution. The emerging story is that gradient descent biases us to a low nuclear norm solution, and we already know how having low nuclear norm can ensure generalization (Srebro and Shraibman, 2005; Foygel and Srebro, 2011) and minimizing the nuclear norm ensures reconstruction (Recht et al., 2010; Candès and Recht, 2009).

Can we more explicitly characterize this bias? We see that we do not always converge precisely to the minimum nuclear norm solution. In particular, the choice of step size and initialization affects which solution gradient descent converges to. Nevertheless, as we formalize in Section 3, we argue that when $U$ is full dimensional, the step size becomes small enough, and the initialization approaches zero, gradient descent will converge precisely to a minimum nuclear norm solution, i.e. to $\operatorname*{argmin}_{X\succeq 0}\norm{X}_{*}\text{ s.t. }\mathcal{A}(X)=y$.

Gradient Flow and Main Conjecture

The behavior of gradient descent with infinitesimally small step size is captured by the differential equation $\dot{U}_{t}:=\derivative{U_{t}}{t}=-\nabla f(U_{t})$ with an initial condition for $U_{0}$. For the optimization in (2) this is

The dynamics (3) define the behavior of the solution $X_{t}=U_{t}U_{t}^{\top}$ and using the chain rule we can verify that $\dot{X}_{t}=\dot{U}_{t}U_{t}^{\top}+U_{t}\dot{U}_{t}^{\top}=-\mathcal{A}^{*}(r_{t})X_{t}-X_{t}\mathcal{A}^{*}(r_{t})$, where $r_{t}=\mathcal{A}(X_{t})-y$ is a vector of the residual. That is, even though the dynamics are defined in terms of specific factorization $X_{t}=U_{t}U_{t}^{\top}$, they are actually independent of the factorization and can be equivalently characterized as

We can now define the limit point $X_{\infty}(X_{\textrm{init}}):=\lim_{t\to\infty}X_{t}$ for the factorized gradient flow (4) initialized at $X_{0}=X_{\textrm{init}}$. We emphasize that these dynamics are very different from the standard gradient flow dynamics of (1) on $X$, corresponding to gradient descent on $X$, which take the form $\dot{X}_{t}=-\nabla F(X_{t})=-\mathcal{A}^{*}(r_{t})$.

Based on the preliminary experiments in Section 2 and a more comprehensive numerical study discussed in Section 5, we state our main conjecture as follows:

For any full rank $X_{\textrm{init}}$, if $\widehat{X}=\lim_{\alpha\to 0}X_{\infty}(\alpha X_{\textrm{init}})$ exists and is a global optima for (1) with $\mathcal{A}(\widehat{X})=y$, then $\widehat{X}\in\operatorname*{argmin}_{X\succeq 0}\ \norm{X}_{*}\ \textrm{s.t.}\ \mathcal{A}(X)=y$.

Requiring a full-rank initial point demands a full dimensional $d=n$ factorization in (2). The assumption of global optimality in the conjecture is generally satisfied: for almost all initializations, gradient flow will converge to a local minimizer (Lee et al., 2016), and when $d=n$ any such local minimizer is also global minimum (Journée et al., 2010). Since we are primarily concerned with underdetermined problems, we expect the global optimum to achieve zero error, i.e. satisfy $\mathcal{A}(X)=y$. We already know from these existing literature that gradient descent (or gradient flow) will generally converge to a solution satisfying $\mathcal{A}(X)=y$; the question we address here is which of those solutions will it converge to.

The conjecture implies the same behavior for asymmetric problems factorized as $X=UV^{\top}$ with gradient flow on $(U,V)$, since this is equivalent to gradient flow on the p.s.d. factorization of $\left[\begin{smallmatrix}W&X\\ X^{\top}&Z\end{smallmatrix}\right]$.

Theoretical Analysis

We will prove our conjecture for the special case where the matrices $A_{i}$ commute, and discuss the more challenging non-commutative case. But first, let us begin by reviewing the behavior of straight-forward gradient descent on $X$ for the convex problem in (1).

Getting started: 𝐦=𝟏𝐦1\mathbf{m=1}

Consider the simplest case of the factorized problem when $m=1$ with $A_{1}=A$ and $y_{1}=y$. The dynamics of (4) are given by $\dot{X}_{t}=-r_{t}(AX_{t}+X_{t}A)$, where $r_{t}$ is simply a scalar, and the solution for $X_{t}$ is given by, $X_{t}=\exp\left(s_{t}A\right)X_{0}\exp\left(s_{t}A\right)$ where $s_{T}=-\int_{0}^{T}r_{t}dt$. Assuming $\widehat{X}=\lim_{\alpha\to 0}X_{\infty}(\alpha X_{0})$ exists and $\mathcal{A}(\widehat{X})=y$, we want to show $\widehat{X}$ is an optimum for the following problem

The KKT optimality conditions for (5) are:

We already know that the first condition holds, and the p.s.d. condition is guaranteed by the factorization of $X$. The remaining complementary slackness and dual feasibility conditions effectively require that $\widehat{X}$ is spanned by the top eigenvector(s) of $A$. Informally, looking to the gradient flow path above, for any non-zero $y$, as $\alpha\to 0$ it is necessary that $|s_{\infty}|\to\infty$ in order to converge to a global optima, thus eigenvectors corresponding to the top eigenvalues of $A$ will dominate the span of $X_{\infty}(\alpha X_{\textrm{init}})$.

The characterization of the the gradient flow path from the previous section can be extended to arbitrary $m$ in the case that the matrices $A_{i}$ commute, i.e. $A_{i}A_{j}=A_{j}A_{i}$ for all $i,j$. Defining $s_{T}=-\int_{0}^{T}r_{t}dt$ – a vector integral, we can verify by differentiating that solution of (4) is

In the case where matrices $\left\{A_{i}\right\}_{i=1}^{m}$ commute, if $\widehat{X}=\lim_{\alpha\to 0}X_{\infty}(\alpha I)$ exists and is a global optimum for (1) with $\mathcal{A}(\widehat{X})=y$, then $\widehat{X}\in\operatorname*{argmin}_{X\succeq 0}\norm{X}_{*}\ \textrm{s.t.}\ \mathcal{A}(X)=y$.

Let $\beta=-\log\alpha$, then $\lambda_{k}(X_{\infty}(\alpha I))=\exp(2\lambda_{k}(\mathcal{A}^{*}(s_{\infty}(\beta)))-2\beta)$. For all $k$ such that $\lambda_{k}(\widehat{X})>0$, by the continuity of $\log$, we have

Defining $\nu(\beta)=\nicefrac{{s_{\infty}(\beta)}}{{\beta}}$, we conclude that for all $k$ such that $\lambda_{k}(\widehat{X})\neq 0$, $\lim_{\beta\to\infty}\lambda_{k}(\mathcal{A}^{*}(\nu(\beta)))=1$. Similarly, for each $k$ such that $\lambda_{k}(\widehat{X})=0$,

Thus, for every $\epsilon\in(0,1]$, for sufficiently large $\beta$

Therefore, we have shown that $\lim_{\beta\to\infty}\mathcal{A}^{*}(\nu(\beta))\preceq I$ and $\lim_{\beta\to\infty}\mathcal{A}^{*}(\nu(\beta))\widehat{X}=\widehat{X}$ establishing the optimality of $\widehat{X}$ for (5). ∎

Interestingly, and similarly to gradient descent on $X$, this proof does not exploit the particular form of the “control" $r_{t}$ and only relies on the fact that the gradient flow path stays within the manifold

Since the $A_{i}$’s commute, we can verify that the tangent space of $\mathcal{M}$ at a point $X$ is given by $T_{X}\mathcal{M}=\textrm{Span}\left\{A_{i}X+XA_{i}\right\}_{i\in[m]}$, thus gradient flow will always remain in $\mathcal{M}$. For any control $r_{t}$ such that following $\dot{X}_{t}=-\mathcal{A}^{*}(r_{t})X_{t}-X_{t}\mathcal{A}^{*}(r_{t})$ leads to a zero error global optimum, that optimum will be a minimum nuclear norm solution. This implies in particular that the conjecture extends to gradient flow on (2) even when the Euclidean norm is replaced by certain other norms, or when only a subset of measurements are used for each step (such as in stochastic gradient descent).

However, unlike gradient descent on $X$, the manifold $\mathcal{M}$ is not flat, and the tangent space at each point is different. Taking finite length steps, as in gradient descent, would cause us to “fall off" of the manifold. To avoid this, we must take infinitesimal steps, as in the gradient flow dynamics.

In the case that $X_{\textrm{init}}$ and the measurements $A_{i}$ are diagonal matrices, gradient descent on (2) is equivalent to a vector least squares problem, parametrized in terms of the square root of entries:

Unfortunately, in the case that the matrices $A_{i}$ do not commute, analysis is much more difficult. For a matrix-valued function $F$, $\derivative{t}\exp(F_{t})$ is equal to $\dot{F_{t}}\exp(F_{t})$ only when $\dot{F_{t}}$ and $F_{t}$ commute. Therefore, (7) is no longer a valid solution for (4). Discretizing the solution path, we can express the solution as the “time ordered exponential":

where the order in the products is important. If $A_{i}$ commute, the product of exponentials is equal to an exponential of sums, which in the limit evaluates to the solution in (7). However, since in general $\exp(A_{1})\exp(A_{2})\neq\exp(A_{1}+A_{2})$, the path (12) is not contained in the manifold $\mathcal{M}$ defined in (11).

Nevertheless, in order to make finite progress moving along a commutator direction like $[A_{i},A_{j}]X_{t}+X_{t}[A_{i},A_{j}]^{\top}$, it is necessary to use an extremely non-smooth control, e.g., looping $1/\epsilon^{2}$ times between $\epsilon$ steps in the directions $A_{i},A_{j},-A_{i},-A_{j}$, each such loop making an $\epsilon^{2}$ step in the desired direction. We expect the actual residuals $r_{t}$ to behave much more smoothly and that for smooth control the non-commutative terms in the expansion of the time ordered exponential (12) are asymptotically lower order then the direct term $\mathcal{A}^{*}(s)$ (as $X_{\textrm{init}}\rightarrow 0$). This is indeed confirmed numerically, both for the actual residual controls of the gradient flow path, and for other random controls.

Empirical Evidence

Beyond the matrix reconstruction experiments of Section 2, we also conducted experiments with similarly simulated matrix completion problems, including problems where entries are sampled from power-law distributions (thus not satisfying incoherence), as well as matrix completion problem on non-simulated Movielens data. In addition to gradient descent, we also looked more directly at the gradient flow ODE (3) and used a numerical ODE solver provided as part of SciPy (Jones et al., 2001). But we still uses a finite (non-zero) initialization. We also emulated staying on a valid “steering path" by numerically approximating the time ordered exponential of 12 — for a finite discretization $\eta$, instead of moving linearly in the direction of the gradient $\nabla f(U)$ (like in gradient descent), we multiply $X_{t}$ on right and left by $e^{-\eta\mathcal{A}^{*}(r_{t})}$. The results of these experiments are summarized in Figure 3.

In these experiments, we again observe trends similar to those in Section 2. In some panels in Figure 3, we do see a discernible gap between the minimum nuclear norm global optima and the nuclear norm of the gradient flow solution with $\|U_{0}\|_{F}=10^{-4}$. This discrepancy could either be due to starting at a non-limit point of $U_{0}$, or numerical issue arising from approximations to the ODE, or it could potentially suggest a weakening of the conjecture. Even if the later case were true, the experiments so far provide strong evidence for atleast approximate versions of our conjecture being true under a wide range of problems.

Finally, we also did experiments on an exhaustive grid search over small problems, capturing essentially all possible problems of this size. We performed an exhaustive grid search for matrix completion problem instances in symmetric p.s.d. ${3\times 3}$ matrices. With $m=4$, there are $15$ unique masks or $\{A_{i}\}_{i\in}$’s that are valid symmetric matrix completion observations. For each mask, we fill the $m=4$ observations with all possible combinations of $10$ uniformly spaced values in the interval $$. This gives us a total of$15\times 10^{4}$problem instances. Of these problems instances, we discard the ones that do not have a valid PSD completion and run the ODE solver on every remaining instance with a random$U_{0}$such that$\|U_{0}\|_{F}=\bar{\alpha}$, for different values of$\bar{\alpha}$. Results on the deviation from the minimum nuclear norm are reported in Figure 4. For small$\bar{\alpha}=10^{-5},10^{-3}$, most of instances of our grid search algorithm returned solutions with near minimal nuclear norms, and the maximum deviation is within the possibility of numerical error. This behavior also decays for$\bar{\alpha}=1$.

Discussion

It is becoming increasingly apparent that biases introduced by optimization procedures, especially for under-determined problems, are playing a key role in learning. Yet, so far we have very little understanding of the implicit biases associated with different non-convex optimization methods. In this paper we carefully study such an implicit bias in a two-layer non-convex problem, identify it, and show how even though there is no difference in the model class (problems (1) and (2) are equivalent when $d=n$, both with very high capacity), the non-convex modeling induces a potentially much more useful implicit bias.

We also discuss how the bias in the non-convex case is much more delicate then in convex gradient descent: since we are not restricted to a flat manifold, the bias introduced by optimization depends on the step sizes taken. Furthermore, for linear least square problems (i.e. methods based on the gradients w.r.t. $X$ in our formulation), any global optimization method that uses linear combination of gradients, including conjugate gradient descent, Nesterov acceleration and momentum methods, remains on the manifold spanned by the gradients, and so leads to the same minimum norm solution. This is not true if the manifold is curved, as using momentum or passed gradients will lead us to “shoot off” the manifold.

Much of the recent work on non-convex optimization, and matrix factorization in particular, has focused on global convergence: whether, and how quickly, we converge to a global minima. In contrast, we address the complimentary question of which global minima we converge to. There has also been much work on methods ensuring good matrix reconstruction or generalization based on structural and statistical properties. We do not assume any such properties, nor that reconstruction is possible or even that there is anything to reconstruct—for any problem of the form (1) we conjecture that (4) leads to the minimum nuclear norm solution. Whether such a minimum nuclear norm solution is good for reconstruction or learning is a separate issue already well addressed by the above literature.

We based our conjecture on extensive numerical simulations, with random, skewed, reconstructible, non-reconstructible, incoherent, non-incoherent, and and exhaustively enumerated problems, some of which is reported in Section 5. We believe our conjecture holds, perhaps with some additional technical conditions or corrections. We explain how the conjecture is related to control on manifolds and the time ordered exponential and discuss a possible approach for proving it.

References