Accelerated Optimization With Orthogonality Constraints

TL;DR

This paper introduces a generalized accelerated gradient method tailored for optimization problems with orthogonality constraints, demonstrating improved iteration scaling and performance on large, ill-conditioned problems.

Contribution

It presents a novel accelerated gradient method specifically designed for orthogonality-constrained optimization, outperforming existing quasi-Newton methods.

Findings

Iteration count scales with the square root of the condition number.

The method outperforms state-of-the-art quasi-Newton methods on large, ill-conditioned problems.

Numerical experiments validate the efficiency and effectiveness of the proposed approach.

Abstract

We develop a generalization of Nesterov's accelerated gradient descent method which is designed to deal with orthogonality constraints. To demonstrate the effectiveness of our method, we perform numerical experiments which demonstrate that the number of iterations scales with the square root of the condition number, and also compare with existing state-of-the-art quasi-Newton methods on the Stiefel manifold. Our experiments show that our method outperforms existing state-of-the-art quasi-Newton methods on some large, ill-conditioned problems.

Tables

Table 1. Table 1: Comparison of the methods on a problem of size n = 1000 𝑛 1000 n=1000 , k = 10 𝑘 10 k=10 .

Method	Iterations	Function Evals	Gradient Evals	Time (s)
Quasi-Newton	31027.4	32528.9	31028.4	26.3
Accelerated Gradient	17266.2	43513.4	17267.2	17.8

Table 2. Table 2: Comparison of the methods on a problem of size n = 2000 𝑛 2000 n=2000 , k = 20 𝑘 20 k=20 .

Method	Iterations	Function Evals	Gradient Evals	Time (s)
Quasi-Newton	84767.2	89013.8	84768.2	169.4
Accelerated Gradient	28758.8	93747.8	28759.8	64.5

Equations

$$\underset{X^{T}X=I_k}{\mathrm{argmin}}f(X),$$

$$Hf(x)(v)={\frac{d^{2}f({\exp }_x(tv))}{d{t}^2}|}_{t=0}.$$

$$\kappa (Hf(x))=\frac{{\sup }_{\Vert v{\Vert }_g=1}Hf(x)(v)}{{\inf }_{\Vert v{\Vert }_g=1}Hf(x)(v)}.$$

$$R(x,0)=x$$

$${\frac{d}{dt}|}_{t=0}R(x,tv)=v$$

$$S_{n,k}=\{X\in {\mathbb{R}}^{n\times k}:X^{T}X=I_k\}.$$

$$\left[\begin{array}{cc}{I}_{k\times k}& 0_{k\times n}\\ 0_{n\times k}& O_{(n-k)\times (n-k)}\end{array}\right],$$

$$g(Y,Z)=\text{Tr}\left(Y^T\left(I-\frac{1}{2}X{X}^T\right)Z\right)$$

$$g^{\ast }(Y,Z)=\text{Tr}\left(Y^T\left(I+X{X}^T\right)Z\right).$$

$${\varphi }_g(W)=\left(I+X{X}^T\right)W$$

$${\varphi }_g^{-1}(V)=\left(I-\frac{1}{2}X{X}^T\right)V.$$

$$X(t)=\exp \left(t(V{X}^T-X{V}^T+X{V}^{T}X{X}^T)\right)X.$$

$$X(t)=\exp \left(t(W{X}^T-X{W}^T)\right)X.$$

$$R_C(X,{\varphi }_g(W))={\left(I-\frac{1}{2}(W{X}^T-X^{T}W)\right)}^{-1}\left(I+\frac{1}{2}(W{X}^T-X^{T}W)\right)X,$$

$$R_1(X,{\varphi }_g(W))=X+2U(I-Z^{T}U{)}^{-1}{Z}^{T}X.$$

$$f(y)\ge f(x)+⟨\nabla f(x),y-x⟩+\frac{\mu }{2}\Vert x-y{\Vert }_2^2.$$

$$f(y)\le f(x)+⟨\nabla f(x),y-x⟩+\frac{L}{2}\Vert x-y{\Vert }_2^2.$$

$$x_0=y_0,x_{t+1}=y_t-{\gamma }_t\nabla f(y_t),y_{t+1}=x_{t+1}+{\alpha }_t(x_{t+1}-x_t).$$

$$f(x_t)-f(x^{\ast })\le 2{\left(1-\sqrt{\frac{\mu }{L}}\right)}^t(f(x_0)-f(x^{\ast })).$$

$$f(x_t)-f(x^{\ast })\le 2({\gamma }_t{q}_t(q_t+2){)}^{-1}\Vert x_0-x^{\ast }{\Vert }_2^2.$$

$$f(x_t)-f(x^{\ast })\le 2L{t}^{-2}\Vert x_0-x^{\ast }{\Vert }_2^2.$$

$$\Vert x_t-x^{\ast }{\Vert }_2^2\le 4(L/\mu )t^{-2}\Vert x_0-x^{\ast }{\Vert }_2^2=4\kappa t^{-2}\Vert x_0-x^{\ast }{\Vert }_2^2.$$

$$\Vert x_t-x^{\ast }{\Vert }_2^2\le \frac{\Vert x_0-x^{\ast }{\Vert }_2^2}{2}.$$

$$y_{t+1}=x_{t+1}+{\alpha }_t(x_{t+1}-x_t)$$

$$f(x_{t+1})>f(x_t)-c_R{\gamma }_t\Vert \nabla f(y_t){\Vert }^2,$$

$$\underset{t\to \infty }{\lim }\Vert \nabla f(x_t)\Vert \to 0.$$

$$f(x_{t+1})\le f(x_t)-c_R{\gamma }_t\Vert \nabla f(y_t){\Vert }_2^2.$$

$$c_R\sum _{t=0}^N{\gamma }_t\Vert \nabla f(y_t){\Vert }_2^2\le f(x_0)-f(x_N).$$

$$\sum _{t=0}^{\infty }\Vert \nabla f(y_t){\Vert }_2^2\le \frac{L(f(x_0)-M)}{c_{R}c}<\infty .$$

$$\Vert \nabla f(x_t){\Vert }_2\le (1+L{\gamma }_t)\Vert \nabla f(y_{t-1}){\Vert }_2\le (1+L{\gamma }_0)\Vert \nabla f(y_{t-1}){\Vert }_2$$

Full text

{AMS}

65K05, 65N25, 90C30, 90C48

1 Introduction

Optimization problems over the set of orthonormal matrices arise naturally in many scientific and engineering problems. Most notably, eigenfunction and electronic structure calculations involve minimizing functions over the set of orthonormal matrices [1, 3, 9, 20]. In these applications, the objective functions are smooth but often ill-conditioned. There are also more recent applications which involve non-smooth objectives, most notably the calculation of compressed modes [13], which involve an $L^{1}$ penalization of variational problems arising in physics.

In this paper, we consider optimization problems with orthogonality constraints, i.e. problems of the form

[TABLE]

where $X$ is an $n\times k$ matrix, $I_{k}$ is the identity matrix, and $f$ is a smooth function. The manifold of orthonormal matrices over which we are optimizing is referred to as the Stiefel manifold in the literature. Many methods have been proposed for solving (1), including variants of gradient descent, Newton’s method, quasi-Newton methods, and non-linear conjugate gradient methods [1, 3, 17, 21, 4]. However, existing methods can suffer from slow convergence when the problem is ill-conditioned, by which we mean that the Hessian of $f$ at (or near) the minimizer is ill-conditioned [3]. Such problems are of particular interest, since they arise when doing electronic structure calculations, or when solving non-smooth problems by smoothing the objective, for instance. Moreover, preconditioning such problems can be very difficult due to the manifold constraint.

In an attempt to solve ill-conditioned problems more efficiently, we develop an extension of the well-known Nesterov’s accelerated gradient descent algorithm [10] designed for optimizing functions on the Stiefel manifold. For the class of smooth, strongly convex functions on $\mathbb{R}^{n}$, accelerated gradient descent obtains an asymptotically optimal iteration complexity of $O(\sqrt{\kappa})$, compared to $O(\kappa)$ for gradient descent with optimal step size selection (see [2], section 3.7, here $\kappa$ is the condition number of the problem). Our method extends this convergence behavior from $\mathbb{R}^{n}$ to the Stiefel manifold, thus providing an efficient method for solving ill-conditioned optimization problems with orthogonality constraints.

Other work on accelerated gradient methods on manifolds includes [8] and [19]. In [8] an accelerated gradient method on general manifolds is presented. However, their algorithm involves solving a non-linear equation involving both the metric on the manifold and the objective function $f$. Unfortunately, solving this equation is only feasible for the special type of model problem which they consider and cannot be generally applied to arbitrary optimization problems on the Stiefel manifold. In [19], a theory is developed which shows that a certain type of accelerated method can achieve accelerated convergence locally. However, their method involves calculating a geodesic logarithm in every iteration and has not yet been implemented, although an iterative method for calculating the geodesic logarithm has been developed in [22]. In constrast, our method only involves very simple linear algebra calculations in each iteration and can be run efficiently on large problems.

The paper is organized as follows. In section 2, we briefly introduce the necessary notation and ideas from differential geometry. In section 3, we discuss accelerated gradient descent on $\mathbb{R}^{n}$. We recall results which are relevant to our work. In section 4, we detail the design of our method. One of the key ingredients is an efficient procedure for performing approximate extrapolation and interpolation on the manifold, which we believe could be useful in developing other optimization methods. In section 5, we show numerical results which provide evidence that our method achieves the desired iteration complexity. Finally, in section 6, we present comparisons with other optimization methods on the Stiefel manifold. We show that our method outperforms existing state of the art methods on some large, ill-conditioned problems.

2 Riemannian Manifolds

In this section, we briefly introduce the notation we will use in the rest of the paper concerning the Stiefel manifold and differential geometry in general. We also collect some formulas for calculating on the Stiefel manifold which will be used later. Some references for differential geometry include [15, 6] and for the geometry of the Stiefel manifold, see [3].

Let $M$ be a smooth manifold and $x\in M$. We denote the tangent space of $M$ at $x$ by $T_{x}M$ and the dual tangent space by $(T_{x}M)^{*}$. We denote the tangent bundle of $M$by $TM$, and likewise the dual tangent bundle by $(TM)^{*}$.

Suppose $f$ is a $C^{1}$ function on $M$. We denote the derivative of $f$, by $\nabla f(x)\in(T_{x}M)^{*}$. Notice that the derivative of $f$ is naturally an element of the dual tangent space $(T_{x}M)^{*}$. If $M$ is a Riemannian manifold, then each tangent space $T_{x}M$ is equipped with a positive definite inner product $g:T_{x}M\times T_{x}M\rightarrow\mathbb{R}$. We denote the norm induced by $g$ as $\|v\|^{2}_{g}$ and the norm induced by $g$ on the dual space as $\|w\|_{g*}$.

Additionally, the inner product $g$ provides an isometry between the tangent space and its dual, which we denote by $\phi_{g}:(T_{x}M)^{*}\rightarrow T_{x}M$ and $\phi^{-1}_{g}:T_{x}M\rightarrow(T_{x}M)^{*}$, and which are also known as raising and lowering indices. Given a $C^{1}$ function $f$ on a Riemannian manifold $M$, an object which is often considered is the Riemannian gradient, obtained by raising the indices of the gradient $\nabla f(x)\in(T_{x}M)^{*}$ to obtain a tangent vector (instead of a dual tangent vector). In this work, we will not work with the Riemannian gradient explicitly, but will rather work with the true gradient $\nabla f(x)\in(T_{x}M)^{*}$, which we consider to be more natural. Whenever we need to raise or lower indices this will be explicitly stated.

Assuming that the geodesic equations on $M$ can be solved globally in time (which is true for the Stiefel manifold that we are interested in) we denote the exponential map based at $x\in M$ by $\exp_{x}:T_{x}M\rightarrow M$.

Let $x\in M$ and $f\in C^{2}(M)$. We denote the Hessian of $f$ at $x$ (viewed as a quadratic form on $T_{x}M$) by

[TABLE]

The condition number of $Hf(x)$ will be important in what follows and is given by

[TABLE]

In Riemannian optimization, geodesics are often expensive to compute exactly, which leads to the concept of a retraction. A retraction associates to each point and tangent vector a curve on the manifold which approximates a geodesic to (at least) first order. This means that if we only move a small distance, a retraction will be very close to the geodesic, even though a retraction can be vastly different globally (see [1] for examples and more discussion).

Definition 2.1**.**

Let $M$ be a (smooth) manifold. A retraction on $M$ is a (smooth) map $R:TM\rightarrow M$ (here $TM$ denotes the tangent bundle of $M$) satisfying for all $x\in M$ and $v\in T_{x}M$

[TABLE]

[TABLE]

*(Here I write $R(x,v)$ for the image of the point $(x,v)\in TM$ under $R$.) *

We proceed to collect explicit formulas for each of these quantities on the Stiefel manifold.

2.1 The Stiefel Manifold

The Stiefel manifold $S_{n,k}$ is the set of $n\times k$ orthonormal matrices, i.e.

[TABLE]

There are two metrics commonly put on the Stiefel manifold in the literature. One is obtained by viewing $S_{n,k}\subset\mathbb{R}^{n\times k}$ and considering the metric induced by the ambient space $\mathbb{R}^{n\times k}$. The other, called the canonical metric and which we will be considering for the remainder of this paper, is obtained by viewing $S_{n,k}=O(n)/O(n-k)$ as the quotient of the orthogonal group $O(n)$ by the right action of $O(n-k)$. Specifically, the action is given by right multiplication by

[TABLE]

where $O_{(n-k)\times(n-k)}\in O(n-k)$. This induces a quotient metric on $S_{n,k}$. For more details on the former metric and the differences between these two viewpoints, see [3].

We fix the following representation of the elements of $S_{n,k}$, its tangent and dual tangent spaces. The elements of $S_{n,k}$ will be represented by $n\times k$ orthonormal matrices (even though our metric is induced by viewing $S_{n,k}$ as a quotient $O(n)/O(n-k)$).

The tangent space at a point $X\in S_{n,k}$ is identified with the set $T_{X}=\{V\in\mathbb{R}^{n\times k}:V^{T}X+X^{T}V=0\}$.

We represent the dual space by $(T_{X})^{*}=\{W\in\mathbb{R}^{n\times k}:W^{T}X+X^{T}W=0\}$, with the pairing between $T_{X}$ and $(T_{X})^{*}$ given by $\langle V,W\rangle=\text{Tr}(V^{T}W)$ (i.e. the usual inner product on $\mathbb{R}^{n\times k}$). Utilizing a slight abuse of notation, we can also interpret any matrix $W\in\mathbb{R}^{n\times k}$ as an element of the dual space using this inner product. This corresponds to identifying the matrix with its orthogonal projection onto the dual space $(T_{X})^{*}$.

Using these representations, the quotient metric on $S_{n,k}$ is given by (see [3])

[TABLE]

and the inner product on the dual space is given by

[TABLE]

The maps corresponding to raising and lowering the indices are

[TABLE]

and

[TABLE]

The advantage of using canonical metric, i.e. the metric induced by the quotient structure of $S_{n,k}$, is that geodesics can be computed using the matrix exponential. In fact, the constant-speed geodesic starting at $X\in S_{n,k}$ and moving initially in the direction $V\in T_{X}S_{n,k}$ is given by (see [3], equation 2.42, details in appendix (B))

[TABLE]

If we give our direction via a dual vector $W$ and raise indices first, we obtain the simpler expression

[TABLE]

This formula can be applied to any $W\in\mathbb{R}^{n\times k}$. This is equivalent to viewing $W$ as an element of the dual tangent space via the Frobenius inner product, i.e. orthogonally projecting $W$ onto the dual tangent space. Notice, however, that we do not need to perform this projection explicitly, we can simply insert $W$ into the above formula (see [3, 17]). The matrix $WX^{T}-XW^{T}$ has rank at most $2k$, and so this exponential can be calculated by diagonalizing a $2k\times 2k$ antisymmetric matrix, as shown in [3].

In [17], a different retraction is introduced, which can be viewed as a Padé approximation of the above exponential. Their retraction, called the Cayley retraction, is defined on a dual vector (or any arbitrary $W\in\mathbb{R}^{n\times k}$ viewed as an element of the dual tangent space as mentioned above) by the formula

[TABLE]

which can be calculated using the Sherman-Morrison-Woodbury formula [14] as (with $U=[\frac{1}{2}W,X]$ and $Z=[X,-\frac{1}{2}W]$)

[TABLE]

This retraction avoids the need to calculate a matrix exponential and reduces the number of floating point operations required by a significant constant factor over the geodesic retraction (although both retractions have the same asymptotic complexity, see [17]).

3 Accelerated Gradient Descent

Let $f:\mathbb{R}^{n}\rightarrow\mathbb{R}$ be a differentiable convex function. We say that $f$ is $\mu$-strongly convex if

[TABLE]

We also say that $f$ is $L$-smooth if

[TABLE]

One way of thinking about these definitions is that $\mu$-strong convexity implies that the eigenvalues of the Hessian of $f$ at every point are greater than $\mu$ and $L$-smoothness implies that the eigenvalues are less than $L$.

In his seminal paper [10], Nesterov introduced first-order methods which achieves the asymptotically optimal objective error for the class of $L$-smooth convex functions and for the class of $L$-smooth and $\mu$-strongly convex functions. These methods take the form

[TABLE]

The choice of $\gamma_{t}$ and $\alpha_{t}$ depend on whether the function $f$ is strongly convex (as opposed to only convex and $L$-smooth), and also on the precise parameters $\mu$ and $L$.

If $f$ is $\mu$-strongly convex and $L$-smooth, then setting $\alpha_{t}=\frac{\sqrt{L}-\sqrt{\mu}}{\sqrt{L}+\sqrt{\mu}}$ and $\gamma_{t}=1/L$ produces the asymptotically optimal objective error of $O((1-\sqrt{\frac{\mu}{L}})^{-t})$ (compared with $O((1-\frac{\mu}{L})^{-t})$ for gradient descent), as the following theorem shows.

Theorem 3.1**.**

Assume that $f$ is $\mu$-strongly convex and $L$-smooth. Let $x^{*}$ be the minimizer of $f$. If we let $\alpha_{t}=\frac{\sqrt{L}-\sqrt{\mu}}{\sqrt{L}+\sqrt{\mu}}$ and $\gamma_{t}=1/L$ in (17), then we have that

[TABLE]

Proof 3.2**.**

*See, for instance, section 3.7 in [2]. *

One disadvantage of the method analyzed in Theorem (3.1) is that setting the proper step size and momentum parameter requires knowing the smoothness parameter $L$ and the strong convexity parameter $\alpha$.

The optimal method for $L$-smooth functions is more flexible. In particular, no knowledge about the smoothness parameter is needed. One can use a line search to determine the correct step size and still obtain the optimal objective error of $O(t^{-2})$ (compared with $O(t^{-1})$ for gradient descent). In particular, we have the following result (which generalizes the results in [16] to obtain a larger family of accelerated schemes).

Theorem 3.3**.**

Assume that $f$ is convex and differentiable with minimizer $x^{*}$.

Let $q_{t}$ be any sequence of non-negative real numbers satisfying $q_{0}=0$ and $(q_{t+1}+1)^{2}\leq(q_{t}+2)^{2}+1$ (in particular $q_{t+1}\leq q_{t}+1$ works).

Then, if in iteration (17), $\gamma_{t}$ is chosen so that $\gamma_{t}\leq\gamma_{t-1}$ and $f(x_{t+1})\leq f(y_{t})-(\gamma_{t}/2)\|\nabla f(y_{t})\|_{2}^{2}$, and $\alpha_{t}=\frac{q_{t}}{2+q_{t+1}}$, we have

[TABLE]

Proof 3.4**.**

*See the appendix (A). *

Note that in the above theorem we made no assumption that $f$ was $L$-smooth. This emphasizes that the method is independent of the particular value of $L$. We choose the step size $\gamma_{t}$ to provide a sufficient decrease in the objective. Such a $\gamma_{t}$ can be found using a line search and will be about $1/L$ in the worst case (within a constant depending on the precise line search scheme).

Also, setting $q_{t}=\alpha t$ and $\gamma_{t}=1/L$ for $\alpha\leq 1$ recovers the result from [16] (with $r=1+2/\alpha$). In particular, the special case $\alpha=1$ gives $f(x_{t})-f(x^{*})\leq 2Lt^{-2}\|x_{0}-x^{*}\|_{2}^{2}$.

3.1 Adaptive Restart for Strongly Convex Functions

It is often the case in practice that an objective is strongly convex and smooth, but the strong convexity parameter $\mu$ and the smoothness parameter $L$ are unknown. This creates a problem because the correct momentum and step size in Theorem (3.1) requires knowledge of the strong convexity and smoothness parameters. Many researchers have proposed methods for adaptively estimating the parameters $\mu$ and $L$ (see, for instance, [11], [7] and [12]).

These ideas are relevant to our work because we will in general not have knowledge of the local (i.e. near the minimizer) smoothness and strong convexity parameters for objectives on the Stiefel manifold, and we build upon the work presented in [12]. The methods introduced there are based on the following observation.

Suppose we are given a $\mu$-strongly convex, $L$-smooth function $f$. Then since $f$ is convex and $L$-smooth, we can run iteration (17) with the parameters given in Theorem (3.3) (setting $q_{t}=t$) and obtain the following objective error

[TABLE]

The strong convexity of $f$ now allows us to bound the iterate error by the objective error, since strong convexity implies that $(\mu/2)\|x_{t}-x^{*}\|_{2}^{2}\leq f(x_{t})-f(x^{*})$. Combining this with equation (20) we see that

[TABLE]

where $\kappa=(L/\mu)$ is the condition number of $f$. This implies that after $t=\sqrt{8\kappa}$ iterations, we will have

[TABLE]

So by restarting the method (i.e. setting $x_{0}=x_{t}$ and resetting the momentum parameter) every $\sqrt{8\kappa}$ iterations, we halve the iterate error every time we restart. This means that it takes $O(\sqrt{\kappa}\log(\epsilon))$ iterations to attain an $\epsilon$-accurate solution and thus restarting the method at this frequency recovers the asymptotically optimal convergence rate (for $\mu$-strongly convex $L$-smooth functions).

Of course, in order to apply this scheme, we must know the condition number $\kappa$ in order to determine the correct restart frequency. To get around this, the method proposed in [12] adaptively chooses when to restart based on an observable condition on the iterates. Specifically, they consider two restart conditions

• •

Function Restart Scheme: Restart when $f(x_{t})>f(x_{t-1})$

• •

Gradient Restart Scheme: Restart when $\nabla f(y_{t-1})\cdot(x_{t}-x_{t-1})>0$

Both of these restart conditions are based upon the analysis of a quadratic objective and it is an open problem to fully analyze their behavior when applied to an arbitrary strongly convex, smooth function. However, experimental results in [12] show empirically that the adaptively restarted methods perform well in practice.

4 Acceleration on the Stiefel Manifold

We now come to the heart of the paper. In this section, we develop a version of Nesterov’s accelerated gradient descent which is designed for efficiently optimizing functions on the Stiefel manifold. We do this by generalizing the adaptively restarted methods of the previous section to the Stiefel manifold. There are two main difficulties which we must overcome in doing this, the first of which results from the non-convexity of our objectives and is thus not specific to optimization on the Stiefel manifold.

Our first difficulty is that the functions which we will be minimizing are non-convex. This is due to the fact that all globally convex functions on the Stiefel manifold are constant [18] (since the manifold is compact). Because of this, we cannot hope to obtain a global convergence rate. Our goal is to construct a method which is guaranteed to converge and which will achieve an accelerated rate once it is close enough to the (local) minimizer. Notice that this difficulty applies equally to non-convex optimization problems on other manifolds, including Euclidean space $\mathbb{R}^{n}$.

The second difficulty, which is specific to the Stiefel manifold, is that we must find an efficient way of generalizing the momentum step

[TABLE]

of (17) to the manifold. We will develop a very efficient method for averaging and extrapolating on the Stiefel manifold, which can potentially be used to design a variety of other optimization methods as well.

In the following subsections, we will describe how each of these difficulties is overcome.

4.1 Restart for Non-convex Functions

When adapting accelerated gradient methods to the Stiefel manifold, we are faced with the issue that the manifold is compact and so the only convex functions are constant. Consequently, the functions which we are optimizing are necessarily non-convex. In this case the convergence results of Theorems (3.1) and (3.3) can not be applied and in fact we cannot hope for a ‘global’ convergence rate.

Instead, what we note is that in a small neighborhood of a local optimum $X^{*}$ the function will be strongly convex and smooth, provided that the Hessian at $X^{*}$ is positive definite. Moreover, the ratio of the strong convexity and smoothness parameters in this neighborhood will be the close to the condition number of $\nabla^{2}f(X^{*})$, which we denote by $\kappa(X^{*})$.

Thus the accelerated gradient method analyzed in Theorem (3.1) suggests that we should be able to find a method which achieves a convergence rate of $O((1-\kappa(X^{*})^{-1/2})^{t})$ once it is close enough to the local minimum $X^{*}$. But since we have to deal with functions which are not globally convex, we hope to design a method which is guaranteed to converge to a local minimum even for non-convex functions, and which achieves the optimal convergence rate once it is close enough to the local minimum.

Our approach is to modify the function restart scheme considered in [12] and described in the previous section. We introduce the following restart condition, which forces a sufficient decrease in the objective.

• •

Modified Function Restart Scheme: Restart when

[TABLE]

where $c_{R}$ is a parameter we take to be a small constant (recall that $\gamma_{t}$ is the step size at step $t$).

The norm above is the norm of the gradient in the dual tangent space if we are on a Riemannian manifold and is taken to be the Euclidean norm if we are considering $\mathbb{R}^{n}$. We will apply this condition to optimization problems on the Stiefel manifold, but we analyze its convergence for non-convex functions on Euclidean space, a situation in which it applies equally well.

Theorem 4.1**.**

Let $f$ be a differentiable, $L$-smooth function on $\mathbb{R}^{n}$, i.e. $\nabla f$ is Lipschitz with constant $L$. Assume also that $f$ is bounded below.

Consider the iteration (17) with step size $\gamma_{t}$ chosen to satisfy $c/L\leq\gamma_{t}\leq\gamma_{t-1}$ for some $c\leq 1$ and $f(x_{t+1})\leq f(y_{t})-(\gamma_{t}/2)\|\nabla f(y_{t})\|_{2}^{2}$ (which can be done by choosing $\gamma_{t}\leq\frac{1}{L}$).

If this iteration is restarted whenever (23) holds (with the new $\gamma_{0}$ chosen to be $\leq\gamma_{t}$), then we have

[TABLE]

Proof 4.2**.**

Note first that our condition on the step size $\gamma_{t}$ can always be satisfied, since by the $L$-smoothness of $f$ we have that $\gamma_{t}=c/L$ will always work.

Also note that since $x_{0}=y_{0}$, the condition on the step size always guarantees that $f(x_{1})\leq f(x_{0})-c_{R}\gamma_{0}\|\nabla f(y_{0})\|_{2}^{2}$.

So we can always run the algorithm (17) in a way which satisfies the conditions of the theorem.

To complete the proof, we note that the restart condition combined with the observation that we always take at least one step implies that

[TABLE]

Summing this, we obtain

[TABLE]

Since $f$ is bounded below, say by $M$ and $\gamma_{n}\geq c/L$ we see that

[TABLE]

This implies that $\|\nabla f(y_{t})\|\rightarrow 0$. Now we simply note that since $f$ is $L$-smooth and $x_{t}=y_{t-1}-\gamma_{t}\nabla f(y_{t-1})$, we have that

[TABLE]

*where the last inequality is because $\gamma_{t}\leq\gamma_{0}$ by assumption. Thus, $\|\nabla f(x_{t})\|\rightarrow 0$ as desired. *

4.2 Extrapolation and Interpolation on the Stiefel Manifold

We have now seen how to get around knowing the strong convexity and smoothness parameters and how to deal with non-convex functions in the process. We proceed to address the second difficulty mentioned at the beginning of the section. Namely, we consider the problem of generalizing the momentum step of (17)

[TABLE]

to the manifold setting.

More generally, we will consider the problem of efficiently extrapolating and interpolating on the Stiefel manifold, i.e. given two points $X,Y\in S_{n,k}$ and $\alpha\in\mathbb{R}$, we want to calculate points $(1-\alpha)X+\alpha Y$ on a curve through $X$ and $Y$. By setting $\alpha\in(0,1)$ this gives a way of averaging points on the manifold and by setting $\alpha>1$ or $\alpha<0$ we can extrapolate as in (29).

A possible approach would be to perform the extrapolation or interpolation in Euclidean space and then project back onto the Stiefel manifold. However, this projection step, which consists of taking the orthonormal part from the polar decomposition of the matrix, is more expensive than the method we propose, especially for large matrices. One could also replace the projection by a reorthogonalization procedure such as Gram-Schmidt (or a QR factorization). However, this is quite inaccurate if $k$ (the number of vectors) is large and is also not as cheap as our method, which only involves matrix products and inversions (but no factorizations).

The approach we take is both simpler and more efficient. What we propose for generalizing

[TABLE]

is to solve for a $V\in(T_{X}S_{n,k})^{*}$ which satisfies (here $R$ is a retraction which we have fixed in the course of designing our method, we present the corresponding equations in Euclidean space on the right to clarify the method)

[TABLE]

and to then extrapolate or average by setting

[TABLE]

Note that the use of $\phi_{g}$ simply allows us to work in the dual tangent space.

The obvious difficulty with this is solving equation (31) for $V$, i.e. finding a $V$ such that $R(X,\phi_{g}(V))=Y$ for some given $X$ and $Y$. However, if we take our retraction to be $R_{1}$ from the previous section (this is the Cayley retraction introduced in [17]), then this boils down to solving

[TABLE]

for $V$. Since $X^{T}X=Y^{T}Y=I$, one can now easily check that $V=2Y(I+X^{T}Y)^{-1}$ solves this equation. Here we are viewing $V\in\mathbb{R}^{nk}$ as an element of the dual tangent space via the Frobenius inner product, as mentioned in section (2). In addition, we can check that if we replace $V$ by $V^{\prime}=V+XS$ for any symmetric $k\times k$ matrix $S$, then $VX^{T}-XV^{T}=V^{\prime}X^{T}-X(V^{\prime})^{T}$. This means that $V^{\prime}$ also satisfies equation (33). In particular, we can replace $V$ by its orthogonal projection onto the dual tangent space $(T_{X})^{*}$, which then gives us the desired vector.

One potential issue with this approach is the possibility that the matrix $(I+X^{T}Y)$ could be singular or ill-conditioned. To address this issue, we have the following lemma showing that the matrix $(I+X^{T}Y)$ is well-conditioned as long as $X$ and $Y$ are not too far apart on $S_{n,k}$. We do not explicitly check this condition in our algorithms, but we have not experienced any numerical issues with the inversion of $I+X^{T}Y$ in our experiments.

Before giving the result, we introduce some notation. Let $X,Y\in S_{n,k}$. We write the geodesic distance between $X$ and $Y$ with respect to the quotient metric as

[TABLE]

where the infemum is taken over all paths $C(t)$ which connect $X$ and $Y$, i.e. for which $C(0)=X$ and $C(1)=Y$. Similarly, we write the geodesic distance with respect to the embedding metric as

[TABLE]

where the infemum is taken over the same set where $C(0)=X$ and $C(1)=Y$.

We now have the following result.

Lemma 4.3**.**

Suppose that $X,Y\in S_{n,k}$ and $d_{S_{n,k}^{Q}}(X,Y)<\sqrt{3/2}$. Then we have

[TABLE]

*Here $d_{S_{n,k}^{Q}}$ is the distance on the Stiefel manifold with the quotient metric. *

Proof 4.4**.**

We first note that

[TABLE]

where $\|\cdot\|_{F}$ is the Frobenius norm. The first inequality above is clear since the Frobenius norm is defined by the same infemum as in equation (35), but without the restriction that the path $C(t)$ needs to lie on the Stiefel manifold. The second inequality follows trivially from equations (34) and (35), as noted in [3], sections 2.2.1 and 2.3.1.

Now consider the condition number of the matrix $I+X^{T}Y$

[TABLE]

Since $X$ and $Y$ are orthonormal, we clearly have $\|X^{T}Yz\|_{2}\leq 1$ and so the numerator above is bounded by $2$. For the denominator, we note that

[TABLE]

We combine this with the fact that

[TABLE]

Since $X$ and $Y$ are orthonormal and $\|z\|_{2}=1$, we have that $\|Xz\|_{2}^{2}=\|Yz\|_{2}^{2}=1$. Additionally, since the Frobenius norm bounds the operator norm, we have that $\|(X-Y)z\|_{2}^{2}\leq\|X-Y\|_{F}^{2}$. Plugging this into equation (39) we see that

[TABLE]

Utilizing equation (37) and (38), we finally get

[TABLE]

*if $d_{S_{n,k}^{Q}}^{2}(X,Y)<\frac{3}{2}$. *

This gives us a computationally efficient procedure for averaging and extrapolating on the Stiefel manifold. We have already described how this can be used to generalize accelerated gradient methods to the Stiefel manifold. We also propose that this averaging and extrapolation procedure could potentially be a building block in other novel optimization algorithms on the manifold.

4.3 Gradient Restart Scheme

We can use the idea of the previous subsection to generalize the gradient restart scheme to the Stiefel manifold. Recall that the gradient restart scheme restarts iteration (17) whenever

[TABLE]

We begin by noting that $X_{t}=Y_{t-1}-\gamma_{t-1}\nabla f(Y_{t-1})$ and so we can rewrite this condition as

[TABLE]

Now it is clear that on the manifold $\|\nabla f(Y_{t-1})\|_{2}^{2}$ should become $\|\nabla f(Y_{t-1})\|_{g*}^{2}$. Here we are viewing $\nabla f(Y_{t-1})$ as an element of the dual tangent space which it naturally is an element of, not as the Riemannian gradient which is obtained by raising the indices. The tricky part is generalizing $\nabla f(Y_{t-1})\cdot(Y_{t-1}-X_{t-1})$. What we propose is to solve for a $V\in(S_{Y_{t-1}})^{*}$ such that

[TABLE]

This element $V$ then serves as $X_{t-1}-Y_{t-1}$ and the analogue of the gradient restart condition becomes

[TABLE]

As in the previous subsection, we see that equation (44) can be efficiently solved for $V$ if the retraction we are using is $R_{1}$ (the Cayley retraction introduced in [17]).

4.4 Accelerated Gradient Descent on the Stiefel Manifold

We now put together all of the the ideas presented in this section to obtain two versions of accelerated gradient descent on the Stiefel manifold; the function restart variant, presented in algorithm (1), and the gradient restart scheme, presented in algorithm (2). In the next section we will present numerical results which demonstrate empirically that our methods achieve the desired iteration complexity, i.e. that the number of iterations scales as $O(\sqrt{\kappa})$, where $\kappa$ is the condition number of the objective function.

5 Numerical Results

In this section, we provide the results of numerical experiments which test the convergence properties and robustness of the two algorithms described in the previous section. We test the algorithms on a sequence of eigenvector calculations with increasing condition numbers. This allows us to investigate how the iteration count scales with the condition number of the problem. The reason we do eigenvector calculations is that the condition number of the corresponding objectives can be evaluated with relative ease.

5.1 Single Eigenvector Calculations

We begin by testing our algorithms on the sphere (which is a special case of the Stiefel manifold $S_{n,k}$ with $k=1$). The problem we solve is the eigenvector calculation

[TABLE]

where $A$ is a symmetric matrix. The solution to this problem is the eigenvector corresponding to the smallest eigenvalue of $A$.

In order to evaluate the performance of our algorithm, we must investigate how the number of iterations scales with the condition number of (46) (not to be confused with the condition number of $A$). A trivial calculation shows that this condition number is given by

[TABLE]

where $\lambda_{1},...,\lambda_{n}$ are the eigenvalues of $A$.

We choose a random initial point drawn uniformly at random on the sphere and apply both versions of our accelerated gradient descent algorithm to optimize the objective (46). We take the matrix $A$ to be a diagonal matrix with diagonal entries (and eigenvalues) $d_{ii}=\lambda_{i}=i$. Our stopping criterion is based on the relative gradient norm, i.e. we stop when $\|\nabla f(X_{n})\|_{g^{*}}\leq\epsilon\|\nabla f(X_{0})\|_{g^{*}}$, with tolerance $\epsilon=10^{-10}$. For both algorithms (1) and (2) the initial step size is taken to be $\gamma_{0}=0.1$ and the line search parameters are taken as $\lambda_{d}=1.7$ and $c_{L}=0.7$. The restart parameter for algorithm (1) is taken as $c_{R}=0.01$. We remark that the performance of the algorithm is not particularly sensitive to these parameter values.

We run this experiment for $21$ values of $n$ evenly spaced in log-space from $10^{2}$ to $10^{4}$ and plot the logarithm of the number of iteration vs the logarithm of the condition number in figure (1). To reduce random fluctuations, we take the average of the logarithm of the number of iterations over $50$ trials for each $n$. We also compare our results with a gradient descent scheme (using the same line search to determine the step size as our methods) and give the coefficients of a log-linear fit to the iteration data.

We see that our method empirically achieves the desired convergence behavior. Indeed, we can see that the number of iterations scales slightly better than with the square root of the condition number, whereas the number of iterations for gradient descent scales about with the condition number. In addition, our accelerated method significantly outperforms the gradient descent method even for relatively small condition numbers.

5.2 Multiple Eigenvector Calculations

We now test our algorithms on the Stiefel manifold $S_{n,k}$ with $k>1$. The problem we consider is that of calculating the smallest $k$ eigenvectors of a symmetric linear operator $A$ by minimizing the Brockett cost

[TABLE]

where $X_{k}$ denotes the $k$-th column of $X$ and $0<\alpha_{1}<\alpha_{2}<...<\alpha_{k}$ are coefficients which force the minimizer to consist of eigenvectors of $A$ rather than eigenvectors up to an orthogonal transformation.

As before, we want to investigate how the number of iterations depends upon the condition number of (48). A simple calculation shows that the condition number of the Brockett cost (48) is

[TABLE]

We briefly note that if one knew the eigenvalues $\lambda_{1},...,\lambda_{n}$, then the coefficients $\alpha_{1},...,\alpha_{n}$ which minimize the condition number (49) produce a minimial condition number of

[TABLE]

Indeed, this minimum condition number is achieved for

[TABLE]

or any multiple of this choice.

As with the calculations on the sphere, we choose a random initial point on $S_{n,k}$ and utilize the same relative gradient norm stopping criterion with a tolerance $\epsilon=10^{-10}$. We again let $A$ be a diagonal matrix with $d_{ii}=\lambda_{i}=i$. We take for the weights $\alpha_{i}$ the condition number minimizing choice corresponding to these eigenvalues, which is simply $\alpha_{i}=i$. We use the same value for the initial step size $\gamma_{0}$ and the line search and restart parameters $\lambda_{d},c_{L},c_{R}$ as for the single eigenvector calculations. We remark that as before the performance of the algorithm is not particularly sensitive to these parameter values.

As before, we run this experiment for $21$ values of $n$ evenly spaced in log-space from $10^{2}$ to $10^{4}$ and plot the logarithm of the number of iteration vs the logarithm of the condition number in figure (2). To reduce random fluctuations, we take the average of the logarithm of the number of iterations over $50$ trials for each $n$. We also compare our results with a gradient descent scheme (using the same line search to determine the step size as our methods) and give the coefficients of a log-linear fit to the iteration data.

We again see that the gradient restart scheme empirically achieves the desired convergence behavior. Indeed, we see that the number of iterations scales better than with the $\sqrt{\kappa}$, but that gradient descent achieves a scaling close to $\kappa$. As before, our methods also significantly outperform gradient descent for all problems tested.

5.3 Comparison with State-of-the-Art Quasi-Newton Methods

In this section, we compare our algorithm with an existing state of the art method. We only consider our function restart scheme since it performed better than the gradient restart scheme in our previous tests. We compare with the state of the art L-RBFGS quasi-Newton method implemented in the ROTLIB library [5] (see also [4]).

We run both methods against each other on an ill-conditioned Brockett cost function (48), where the matrix $A$ is taken to be diagonal with spectrum $\{\frac{k^{2}}{n}\}_{k=1,...,n}$ and the weights are $\alpha_{i}=i$ (note that this is not the optimal choice of $\alpha_{i}$ given in equation (51), however, this is not particularly important as the main point is that this objective is ill-conditioned). In order to obtain reliable results, we test both methods with the same objective and the same random initial points. We repeat each test $10$ times and record the average number of iterations, number of gradient evaluations, number of function evaluations, and total computational time. We use a relative gradient norm stopping condition, i.e. we stop when $\|\nabla f(X_{n})\|_{g^{*}}\leq\epsilon\|\nabla f(X_{0})\|_{g^{*}}$, with tolerance $\epsilon=10^{-9}$ (we use the slightly larger tolerance of $10^{-9}$ because the line search implemented in ROPTLIB failed when optimizing further). The parameters we use for the function restart scheme (1) are initial step size $\gamma_{0}=0.1$, line search parameters $\lambda_{d}=1.7$ and $c_{L}=0.9$, and restart parameters $c_{R}=0.01$.

In our tests we consider two different values of $n$ and $k$. First, we test with $n=1000$, $k=10$, whose results are shown in table (1). Then we test with the larger values $n=2000$, $k=20$. The results of this test are shown in table (2). We see that our method outperforms the method in [4], using about a third as many gradient evaluations, about the same number of functions evaluations, and less than half the time as the quasi-newton method for the larger problem $n=2000$, $k=20$.

6 Conclusion

In this paper, we developed novel accelerated first-order optimization methods designed to handle orthogonality constraints. The algorithms developed are a generalization of Nesterov’s gradient descent to the Stiefel manifold. In the process, we constructed an efficient way of averaging and extrapolating points on the manifold, which we believe can be useful in developing other novel optimization algorithms. Numerical experiments indicate that our methods not only achieve the desired scaling with the condition number of the problem, but also outperform state of the art quasi-Newton methods on some large, ill-conditioned problems.

We would also like to note that although the algorithms we have constructed make explicit use of formulas specific to the Stiefel manifold, we believe the ideas presented in this paper can be generalized to other manifolds as well. This would depend upon generalizing the momentum step (29) efficiently to the manifold of interest. In fact, for quotients of the Stiefel manifold, for example the Grassmann manifold, the algorithms developed in this paper are already immediately applicable.

7 Acknowledgements

We are grateful to Russel Caflisch, Stanley Osher, and Vidvuds Ozolins for suggesting this project and providing continued guidance. This work was supported by AFOSR grant FA9550-15-1-0073.

Appendix A Proof of Theorem (3.3)

Proof A.1** (Proof of Theorem (3.3)).**

Consider the Lyapunov function

[TABLE]

We will show that $J_{t+1}\leq J_{t}$ which proves the theorem since $\gamma_{t}q_{t}(q_{t}+2)(f(x_{t})-f(x^{*}))\leq J_{t}$ and $J_{0}=2\|y_{0}-x^{*}\|_{2}^{2}=2\|x_{0}-x^{*}\|_{2}^{2}$. To this end, we denote

[TABLE]

and

[TABLE]

Then we see that

[TABLE]

Since by assumption $\gamma_{t+1}\leq\gamma_{t}$ and $q_{t+1}(q_{t+1}+2)=(q_{t+1}+1)^{2}-1\leq(q_{t}+2)^{2}$, we see that $\gamma_{t+1}q_{t+1}(q_{t+1}+2)\leq\gamma_{t}(q_{t}+2)^{2}$, and the bottom line in the above equation is bounded by

[TABLE]

Thus we see that

[TABLE]

The step sizes $\gamma_{t}$ are chosen so that $f(x_{t+1})-f(y_{t})\leq-(1/2)\gamma_{t}\|\nabla f(y_{t})\|_{2}^{2}$, and so we can rewrite the above to obtain

[TABLE]

The convexity of $f$ implies that $f(y_{t})-f(x^{*})\leq\nabla f(y_{t})\cdot(y_{t}-x^{*})$ and $f(y_{t})-f(x_{t})\leq\nabla f(y_{t})\cdot(y_{t}-x_{t})$, so we get

[TABLE]

Now we consider $J^{2}_{t+1}-J^{2}_{t}$. Note that $J^{2}_{t}=(1/2)\|s_{t}\|_{2}^{2}$ with

[TABLE]

Thus

[TABLE]

Considering that

[TABLE]

we will be done if we can show that $s_{t+1}-s_{t}=-\gamma_{t}(q_{t}+2)\nabla f(y_{t})$. To this end we compute

[TABLE]

Recalling the update formulas for the iterates $x_{n}$ and $y_{n}$, we see that

[TABLE]

and

[TABLE]

Thus equation (62) simplifies to

[TABLE]

*which is equal to $-(q_{t}+2)\gamma_{t}\nabla f(y_{t})$ by our choice of $\alpha_{t}$. *

Appendix B Derivation of equation (11)

We begin from equation 2.42 in [3], which, using their notation is

[TABLE]

where

[TABLE]

for some $A\in\mathbb{R}^{p\times p}$ skew-symmetrix and some $B\in\mathbb{R}^{n-p\times p}$. We rewrite this using the fact that $Q$ is an orthonormal matrix as

[TABLE]

Now, translating into the notation we have used, we see that $Q=[X,X^{\prime}]$, $Y(t)=X(t)$, and $QXI_{n,p}=V$ (here $X^{\prime}$ is a basis for the orthogonal complement of $X$). Equation (11) now boils down to checking that (in our notation)

[TABLE]

where $V=XA+X^{\prime}B$ and $X^{T}X^{\prime}=0$ since $X^{\prime}$ is an orthogonal complement of $X$.