# Optimization on Stiefel Manifold

$$\text{Find}~~U, V$$ $$\text{to maximize}~~f(U,V)=\text{tr}(U^TAVN)$$ $$\text{subject}~~U^TU=I_p,V^TV=I_p$$ where $N=\text{diag}(\mu_1,\cdots,\mu_p)$ with $\mu_1>\mu_2>\cdots>\mu_p>0$.

I am reading the book Optimization Algorithms on Matrix Manifold, and the problem is on Page 11. It says if $(U,V)$ is a solution of this maximization problem, then the columns $u_i$ of $U$ and $v_i$ of $V$ are the $i$th dominant left and right singular vectors of $A$.

I am confused with results. How can we prove it? Is it obvious? Any advice is helpful. Thank you.

Here is most of the argument. There should be a shorter answer...

Part 1: show an upper bound on the optimal value

If $$(U, V)$$ is optimal, then certainly it is a critical point, that is: the Riemannian gradient of $$f$$ at $$(U, V)$$ is zero.

Let us first compute the Euclidean gradient with respect to the trace inner product:

$$\nabla f(U, V) = (AVN, A^TUN)$$.

Making the Stiefel manifold (that is, the set of matrices with orthonormal columns) into a Riemannian submanifold of the set of matrices, the book explains that the Riemannian gradient is the orthogonal projection of the Euclidean gradient to the tangent spaces. Working this out (it's all done in the book), we find that the Riemannian gradient of $$f$$ at $$(U, V)$$ vanishes exactly if

$$AVN = U \mathrm{sym}(U^T A V N)$$ and $$A^T U N = V \mathrm{sym}(V^T A^T U N)$$,

where $$\mathrm{sym}(M) = (M+M^T)/2$$ extracts the symmetric part of a matrix. Since $$U, V$$ have orthonormal columns, the above imply that $$U^T A V N = \mathrm{sym}(U^T A V N)$$ and $$V^T A^T U N = \mathrm{sym}(V^T A^T U N)$$, that is, they are symmetric. Therefore, the above conditions imply:

$$AVN = UU^T A V N$$ and $$A^T U N = VV^T A^T U N$$.

The matrix $$N$$ is invertible, hence:

$$AV = UU^T A V$$ and $$A^T U = VV^T A^T U$$.

The matrix $$U^T A V$$ has size $$p \times p$$; let us consider its SVD: $$U^T A V = \tilde U \tilde \Sigma \tilde{V}^T$$ ($$\tilde{U}$$, $$\tilde{V}$$ are $$p \times p$$ orthogonal matrices and $$\tilde \Sigma$$ is diagonal with nonnegative, decreasing entries). Plug this into the identities above:

$$AV = U \tilde U \tilde \Sigma \tilde V^T$$ and $$A^T U = V \tilde V \tilde \Sigma \tilde U^T$$.

Since $$\tilde U$$ and $$\tilde V$$ are orthogonal, the matrices $$\bar U = U\tilde U$$ and $$\bar V = V \tilde V$$ have orthonormal columns too, and the above reads:

$$A\bar V = \bar U \tilde \Sigma$$ and $$A^T\bar U = \bar V \tilde \Sigma$$.

These identities tell us that the diagonal entries of $$\tilde \Sigma$$ actually are singular values of $$A$$ ("distinct" singular values, counting multiplicities). Let us write $$\tilde \Sigma = diag(\sigma_{(1)}, \ldots, \sigma_{(p)})$$. The notation is meant to clarify that these are $$p$$ of the singular values of $$A$$, not necessarily the $$p$$ first ones; they are ordered though: $$\sigma_{(1)} \geq \cdots \geq \sigma_{(p)}$$.

Now consider the value of the cost function at that pair $$(U, V)$$:

$$f(U, V) = \mathrm{tr}(U^T A V N) = \mathrm{tr}(\tilde U \tilde \Sigma \tilde V^T N) = \mathrm{tr}(\tilde \Sigma \cdot \tilde V^T N \tilde U)$$

(I used the fact that inside a trace we can cyclically-shift matrices in a product.)

Let us define $$M = \tilde V^T N \tilde U$$. Then,

$$f(U, V) = \sum_{i = 1}^p M_{ii} \sigma_{(i)}$$

Here is how to finish the argument (I did not check this): notice that $$\mathrm{tr}(M) = \mathrm{tr}(\tilde U \tilde V^T N) = \mathrm{tr}(Q N)$$ for some orthogonal matrix $$Q = \tilde U \tilde V^T$$. Argue that this trace is at most equal to that of $$N$$, and this happens exactly if $$M$$ is a diagonal with the same entries as those of $$N$$ (possibly permuted). Using this, argue that $$\sum_{i = 1}^p M_{ii} \sigma_{(i)}$$ is at most equal to $$\mu_1 \sigma_{(1)} + \cdots + \mu_p \sigma_{(p)}$$. Finally, reason that this is at most equal to $$\mu_1 \sigma_{1} + \cdots + \mu_p \sigma_{p}$$, where $$\sigma_1 \geq \cdots \geq \sigma_p$$ are the $$p$$ largest singular values of $$A$$ (counting multiplicities).

Part 2: show the upper bound is attained

This long argument shows that $$f$$ is upper bounded by $$\mu_1 \sigma_{1} + \cdots + \mu_p \sigma_{p}$$. To conclude, it remains to show that this upper bound is attained for some choice of $$(U, V)$$; this is easy: just let the columns of $$U$$ and $$V$$ contain top singular vectors of $$A$$, in order: the computation is direct.