Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform

The Stiefel manifold $\text{St}(n,p)$, $n\geq p$, is a Riemannian manifold that is composed of all $n\times p$ orthonormal matrices $\{X\in\mathbb{R}^{n\times p}:X^{T}X=I\}$. In the rest of the paper, we will use notation $\mathcal{M}=\text{St}(n,p)$ to denote the Stiefel manifold. We regard $\mathcal{M}$ as an embeded submanifold of a Euclidean space. Hence, the Riemannian metric $\rho$ is the Euclidean metric as:$\rho_{X}(Z_{1},Z_{2})=tr(Z_{1}^{\top}Z_{2})$, where $Z_{1},Z_{2}$ are tangent vectors in $T_{X}\mathcal{M}$. The tangent space of $\mathcal{M}$ at $X$ is defined as

The projection of a matrix $Z\in\mathbb{R}^{n\times p}$ to $T_{X}\mathcal{M}$ can be computed as

Given a derivative of the objective function $\nabla f(X)$ at $X$ in the Euclidean space, we can compute the Riemannian gradient $\nabla_{\mathcal{M}}f(X)$ on the Stiefel manifold as a projection onto $T_{X}\mathcal{M}$ using $\pi_{T_{X}}(\nabla f(X))$ given by Eq.(3.2). It follows that optimization of $f$ on the Riemannian manifold can be computed as follows. First, compute $\nabla_{\mathcal{M}}f(X_{t})$ in $T_{X}\mathcal{M}$. Second, transport the momentum $M_{t}$ to the current tangent space $T_{X}\mathcal{M}$ and combine it linearly with the current Riemannian gradient $\nabla_{\mathcal{M}}f(X_{t})$ to update the momentum $M_{t+1}$. Finally, third, update the new parameter $X_{t+1}$ along a curve on the manifold with the initial direction as $M_{t+1}$.

While the exponential map and parallel transport can be used to update parameters and momentums in optimization on the Riemannian manifold, they are computationally infeasible on the Stiefel manifold. In the following section, we specify our computationally efficient alternatives.

We generalize the heavy ball (HB) momentum update (Ghadimi et al., 2015; Zavriev & Kostyuk, 1993) in the $k$th optimization step¹¹1Note that we use index $i$ in superscript to indicate our iterative steps in estimation of the Cayley transform, and index $k$ in subscript to indicate optimization steps. to the Stiefel manifold. Theoretically, it can be represented as:

where $Y(\alpha,X_{k},W_{k})$ is a curve that starts at $X_{k}$ with length $\alpha$ on the Stiefel manifold, specified by the Cayley transform in Eq.(3). In practice, we efficiently perform the updates in Eq.(7) by the proposed iterative Cayley transform and implicit vector transport on the Stiefel manifold. Specially, we first update the momentum as if it were in the Euclidean space. Then, we update the new parameters by iterative Cayley transform. Finally, we correct the momentum $M_{k+1}$ by projecting it to $T_{X_{k}}\mathcal{M}$. The details of our Cayley SGD algorithm are shown in Alg. 1.

ADAM is a recent first-order optimization method for stochastic objective functions. It estimates adaptive lower-order moments and uses adaptive learning rate. The algorithm is designed to combine the advantages of both AdaGrad, which performs well in sparse-gradient cases, and RMSProp, which performs well in non-stationary cases.

We generalize ADAM to the Stiefel manifold by making three modifications to the vanilla ADAM. First, we replace the standard gradient and momentum with the corresponding ones on the Stiefel manifold, as described in Section 4.1. Second, we use a manifold-wise adaptive learning rate that assign a same learning rate for all entries in a parameter matrix as in (Absil et al., 2009). Third, we use the Cayley transform to update the parameters. Cayley ADAM is summarized in Alg 2.

In CNNs, for a convolutional layer with kernel $\hat{K}\in\mathbb{R}^{c_{out}\times c_{in}\times h\times w}$, we first reshape $\hat{K}$ into a matrix ${K}$ of size $p\times n$, where $p=c_{out}$, $n=c_{in}\times h\times w$. In most cases, we have $p\leq n$. Then, we restrict the matrix ${K}$ on the Stiefel manifold using Cayley SGD or Cayley ADAM, while other parameters are optimized with SGD and ADAM.

Datasets: We evaluate Cayley SGD or Cayley ADAM in image classification on the CIFAR10 and CIFAR100 datasets (Krizhevsky & Hinton, 2009). CIFAR10 and CIFAR100 consist of of 50,000 training images and 10,000 test images, and have 10 and 100 mutually exclusive classes.

Models: We use two networks — VGG (Simonyan & Zisserman, 2014) and Wide ResNet (Zagoruyko & Komodakis, 2016) — that obtain state of the art performance on CIFAR10 and CIFAR100. For VGG, every convolutional layer is followed by a batch normalization layer and a ReLU. For Wide ResNet, we use basic blocks, where two consecutive 3-by-3 convolutional layers are followed by the batch normalization and ReLU activation, respectively.

Training Strategies: We use different learning rates $l_{e}$ and $l_{st}$ for weights on the Euclidean space and the Stiefel manifold, respectively. We set the weight decay as 0.0005, momentum as 0.9, and minibatch size as 128. The initial learning rates are set as $l_{e}=0.01$ and $l_{st}=0.2$ for Cayley SGD and $l_{e}=0.01$ and $l_{st}=0.4$ for Cayley ADAM. During training, we reduce the learning rates by a factor of $0.2$ at 60, 120, and 160 epochs. The total number of epochs in training is 200. In training, the data samples are normalized using the mean and variance of the training set, and augmented by randomly flipping training images.

Our baselines include SGD with momentum and ADAM. We follow the same training strategies as mentioned above, except for the initial learning rates set to 0.1 and 0.001, respectively.

Performance: Table 1 and Table 2 show classification errors on CIFAR10 and CIFAR100 respectively using different optimization algorithms. As shown in the tables, the proposed two algorithms achieve competitive performance, and for certain deep architectures, the best performance. Specifically, the network WRN-28-10 trained with Cayley ADAM achieves the best error rate of $3.57\%$ and $18.10\%$ on CIFAR10 and CIFAR100 respectively. Fig. 1 compares training curves of our algorithms and baselines in terms of epochs, and shows that both Cayley SGD and Cayley ADAM converge faster than the baselines. In particular, the training curves of the baselines tend to get stuck in a plateau before the learning rate drops, which is not the case with our algorithms. This might be because the baselines do not enforce orthonormality of network parameters. In training, the backpropagation of orthonormal weight vectors, in general, does not affect each other, and thus has greater chances to explore new parameter regions. Fig. 2 also compares the training loss curve in terms of time. Our Cayley SGD and Cayley ADAM converge the fastest among methods that also address orthonormality. Although the baselines SGD and ADAM converge faster at the beginning due to their training efficiency, our Cayley SGD and Cayley ADAM can catch up with the baseline after 12000 seconds, which corresponds to the 120th epoch of SGD and ADAM, and the 60th epoch of Cayley SGD and Cayley ADAM.

Comparison with State of the Art. We compare the proposed algorithms with two sets of state of the art. The first set of approaches are soft orthonormality regularization approaches (Bansal et al., 2018). Specially, for a weight matrix $K\in\mathbb{R}^{n\times p}$, SO penalizes $||KK^{\top}-I||^{2}_{F}$, DSO penalizes ($||KK^{\top}-I||^{2}_{F}+||K^{\top}K-I||^{2}_{F}$), the SRIP penalizes the spectral norm of $(KK^{\top}-I)$. The second set of approaches includes the following hard orthonormality methods: Polar decomposition(Absil et al., 2009), QR decomposition(Absil et al., 2009), closed-form Cayley transform, Wen&Yin (Wen & Yin, 2013), OMDSM(Huang et al., 2018a), DBN(Huang et al., 2018b). Note that we do not include momentum in Polar decomposition and QR decomposition as previous work does not specify the momentum. Also, we use the closed-form Cayley transform without momentum as an ablation study of the momentum effect. All experiments are evaluated on the benchmark network WRN-28-10. Table 3 shows that our algorithms achieve comparable error rates with state of the art. All algorithms are run on one TITAN Xp GPU, and their average training time are compared per epoch. Table 3 shows that our algorithms run much faster than existing algorithms, except for the baseline SGD and ADAM which do not impose orthonormality constraints.

In RNNs, the hidden-to-hidden transition matrix $K$ can be modeled as a unitary matrix (Arjovsky et al., 2016). Wisdom et al. (2016) model the hidden-to-hidden transition matrix as a full-capacity unitary matrix on the complex Stiefel manifold: $\text{St}(\mathbb{C}^{N})=\{K\in\mathbb{C}^{N\times N}:K^{H}K=I\}$.

Pixel-by-pixel MNIST: We evaluate the proposed algorithms on the challenging pixel-by-pixel MNIST task for long-term memory. The task was used to test the capacity of uRNNs (Arjovsky et al., 2016; Wisdom et al., 2016). Following the same setting as in Wisdom et al. (2016), we reshape MNIST images of $28\times 28$ pixels to a $T=784$ pixel-by-pixel sequences, and select 5,000 out of the 60,000 training examples for the early stopping validation.

Training: Wisdom et al. (2016) restricted the transition unitary matrix on the Stiefel manifold via a closed-form Cayley transform. On the contrary, we use Cayley SGD with momentum and Cayley ADAM to reduce the training time. Table 4 shows that the proposed algorithms reduce the training time by about 35% for all settings of the network, while maintaining the same level of accuracy. All experiments are performed on one TITAN Xp GPU.

Checking Unitariness: To show that the proposed algorithms are valid optimization algorithms on the Stiefel manifold, we check the unitariness of the hidden-to-hidden matrix $K$ by computing the error term $||K^{H}K-I||_{F}$ during training. Table 5 compares average errors for varying numbers of iterations $s$. As can be seen, the iterative Cayley transform can approximate the unitary matrix when $s=2$. The iterative Cayley transform performs even better than the closed form Cayley transform, which might be an effect of the rounding error of the matrix inversion as in Eq.(3).