Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform

1 INTRODUCTION . Orthonormality has recently gained much interest , as there are significant advantages of enforcing orthonormality on parameter matrices of deep neural networks . For CNNs , Bansal et al . ( 2018 ) show that orthonormality constraints improve accuracy and gives a faster empirical convergence rate , Huang et al . ( 2018a ) show that orthonormality stabilizes the distribution of neural activations in training , and Cogswell et al . ( 2015 ) show that orthonormality reduces overfitting and improves generalization . For RNNs , Arjovsky et al . ( 2016 ) ; Zhou et al . ( 2006 ) show that the orthogonal hidden-to-hidden matrix alleviates the vanishing and exploding-gradient problems . Riemannian optimization on the Stiefel manifold , which represents the set of all orthonormal matrices of the same size , is an elegant framework for optimization under orthonormality constraints . But its high computational cost is limiting its applications , including in deep learning . Recent efficient approaches incorporate orthogonality in deep learning only for square parameter matrices ( Arjovsky et al. , 2016 ; Dorobantu et al. , 2016 ) , or indirectly through regularization ( Bansal et al. , 2018 ) , which however does not guarantee the exact orthonormality of parameter matrices . To address the aforementioned limitations , our first main contribution is an efficient estimation of the retraction mapping based on the Cayley transform for updating large non-square matrices of parameters on the Stiefel manifold . We specify an efficient iterative algorithm for estimating the Cayley transform that consists of only a few matrix multiplications , while the closed-form Cayley transform requires costly matrix inverse operations ( Vorontsov et al. , 2017 ; Wisdom et al. , 2016 ) . The efficiency of the retraction mapping is both theoretically proved and empirically verified in the paper . Our second main contribution is aimed at improving convergence rates of training by taking into account the momentum in our optimization on the Stiefel manifold . We derive a new approach to move the tangent vector between tangent spaces of the manifold , instead of using the standard parallel transport . Specifically , we regard the Stiefel manifold as a submanifold of a Euclidean space . This allows for representing the vector transport ( Absil et al. , 2009 ) as a projection onto the tangent space . As we show , since the Cayley transform implicitly projects gradients onto the tangent space , the momentum updates result in an implicit vector transport . Thus , we first compute a linear combination of the momentum and the gradient in the Euclidean space , and then update the network parameters using the Cayley transform , without explicitly performing the vector transport . We apply the above two contributions to generalize the standard SGD with momentum and ADAM ( Kingma & Ba , 2014 ) to the Stiefel manifold , resulting in our two new optimization algorithms called Cayley SGD with momentum and Cayley ADAM . A theoretical analysis of the convergence of Cayley SGD is presented . Similar analysis for Cayley ADAM is omitted , since it is very similar to the analysis presented in ( Becigneul & Ganea , 2019 ) . Cayley SGD and Cayley ADAM are empirically evaluated on image classification using VGG and Wide ResNet on the CIFAR-10 and CIFAR-100 datasets ( Krizhevsky & Hinton , 2009 ) . The experiments show that Cayley SGD and Cayley ADAM achieve better classification performance and faster convergence rate than the baseline SGD with momentum and ADAM . While the baselines take less time per epoch since they do not enforce the orthonormality constraint , they take more epochs until convergence than Cayley SGD and Cayley ADAM . In comparison with existing optimization methods that also account for orthonormality – e.g. , Polar decomposition , QR decomposition , or closed-form Cayley transform – our Cayley SGD and Cayley ADAM run much faster and yield equally good or better performance in image classification . Finally , we apply the aforementioned two contributions to training of the unitary RNNs . Wisdom et al . ( 2016 ) proposes the full capacity unitary RNN that updates the hidden-to-hidden transition matrix with the closed-form Cayley transform . In contrast , for our RNN training , we use the more efficient Cayley SGD with momentum and Cayley ADAM . The results show that our RNN training takes less running time per iteration without compromising performance . 2 RELATED WORK . There is a host of literature on using orthonormality constraints in neural-network training . This section reviews closely related work , which can be broadly divided in two groups : orthonormality regularization and Riemannian optimization . Regularization approaches can be divided as hard , which strictly enforce orthonormality , and soft . Hard regularizations are computationally expensive . For example , Huang et al . ( 2018b ) extend Batch Normalization ( Ioffe & Szegedy , 2015 ) with ZCA , and hence require costly eigenvalue decomposition . Huang et al . ( 2018a ) derive a closed-form solution that also requires eigenvalue decomposition . Bansal et al . ( 2018 ) ; Cisse et al . ( 2017 ) ; Xiong et al . ( 2016 ) introduce mutual coherence regularization and spectral restricted isometry regularization ; however , these regularizations are soft in that they can not guarantee orthonormality . Riemannian optimization guarantees that the solution respects orthonormality constraints . For example , Cho & Lee ( 2017 ) replace Batch Normalization layers in a CNN with Riemannian optimization on the Grassmann manifold G ( n , 1 ) , where the parameters are normalized but not orthonormalized . Also , Vorontsov et al . ( 2017 ) ; Wisdom et al . ( 2016 ) ; Lezcano-Casado & Martı́nez-Rubio ( 2019 ) ; Helfrich et al . ( 2017 ) perform Riemannian optimization on the group of unitary matrices to stablize RNNs training , but their technique can not be applied to non-square parameter matrices . Becigneul & Ganea ( 2019 ) introduce a more general Riemannian optimization , but do not show how to efficiently perform this optimization on the Stiefel manifold . The key challenge of Riemannian optimization is that exponential mapping — the standard step for estimating the next update point — is computationally expensive on the Stiefel manifold . Some methods use an efficient pseudo-geodesic retraction instead of the exponential mapping . For example , Absil & Malick ( 2012 ) ; Gawlik & Leok ( 2018 ) ; Manton ( 2002 ) use a projection-based method to map the gradient back to the Stiefel manifold that rely on computational expensive SVD . Other approaches are based on the closed-form Cayley transform ( Fiori et al. , 2012 ; Jiang & Dai , 2015 ; Nishimori & Akaho , 2005 ; Zhu , 2017 ) , but require the costly matrix inversion . Also , Wen & Yin ( 2013 ) reduce the cost of the Cayley transform by making the restrictive assumption that the matrix size n× p is such that 2p n. Unfortunately , this algorithm is not efficient when 2p ≥ n . 3 PRELIMINARY . This section briefly reviews some well-known properties of the Riemannian and Stiefel manifolds . The interested reader is referred to Boothby ( 1986 ) ; Edelman et al . ( 1998 ) and references therein . 3.1 RIEMANNIAN MANIFOLD . Definition 1 . Riemannian Manifold : A Riemannian manifold ( M , ρ ) is a smooth manifold M equipped with a Riemannian metric ρ defined as the inner product on the tangent space TxM for each point x , ρx ( · , · ) : TxM× TxM→ R. Definition 2 . Geodesic , Exponential map and Retraction map : A geodesic is a locally shortest curve on a manifold . An exponential map , Expx ( · ) , maps a tangent vector , v ∈ Tx , to a manifold , M. Expx ( tv ) represents a geodesic γ ( t ) : t ∈ [ 0 , 1 ] on a manifold , s.t . γ ( 0 ) = x , γ̇ ( 0 ) = v. A retraction map is defined as a smooth mapping Rx : TxM → M on a manifold iff Rx ( 0 ) = x and DRx ( 0 ) = idTxM , where DRx denotes the derivative of Rx , idTxM denotes an identity map defined on TxM . It is easy to show that any exponential map is a retraction map . As computing an exponential map is usually expensive , a retraction map is often used as an efficient alternative . Definition 3 . Parallel transport and vector transport : Parallel transport is a method to translate a vector along a geodesic on a manifold while preserving norm . A vector transport τ is a smooth map defined on a retraction R of a manifoldM , τ : TxM× TxM→ TR ( ηx ) , ( ηx , ξx ) 7→ τηx ( ξx ) . τ satisfies the following properties : ( 1 ) Underlying retraction : τηx ( ξx ) ∈ TR ( ηx ) ; ( 2 ) Consistency : τ0xξx = ξx , ∀ξx ∈ TxM ; ( 3 ) Linearity : τηx ( aξx + bζx ) = aτηx ( ξx ) + bτηx ( ζx ) . Usually , a vector transport is a computationally efficient alternative to a parallel transport . 3.2 STIEFEL MANIFOLD . The Stiefel manifold St ( n , p ) , n ≥ p , is a Riemannian manifold that is composed of all n × p orthonormal matrices { X ∈ Rn×p : XTX = I } . In the rest of the paper , we will use notationM = St ( n , p ) to denote the Stiefel manifold . We regardM as an embeded submanifold of a Euclidean space . Hence , the Riemannian metric ρ is the Euclidean metric as : ρX ( Z1 , Z2 ) = tr ( Z > 1 Z2 ) , where Z1 , Z2 are tangent vectors in TXM . The tangent space ofM at X is defined as TXM = { Z : Z > X +X > Z = 0 } ( 1 ) The projection of a matrix Z ∈ Rn×p to TXM can be computed as πTX ( Z ) =WX where : W = Ŵ − Ŵ > , Ŵ = ZX > − 1 2 X ( X > ZX > ) . ( 2 ) Given a derivative of the objective function ∇f ( X ) at X in the Euclidean space , we can compute the Riemannian gradient ∇Mf ( X ) on the Stiefel manifold as a projection onto TXM using πTX ( ∇f ( X ) ) given by Eq. ( 2 ) . It follows that optimization of f on the Riemannian manifold can be computed as follows . First , compute ∇Mf ( Xt ) in TXM . Second , transport the momentum Mt to the current tangent space TXM and combine it linearly with the current Riemannian gradient ∇Mf ( Xt ) to update the momentum Mt+1 . Finally , third , update the new parameter Xt+1 along a curve on the manifold with the initial direction as Mt+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 . 3.2.1 PARAMETER UPDATES BY ITERATIVE CAYLEY TRANSFORM . The Cayley transform computes a parametric curve on the Stiefel manifold using a skew-symmetric matrix ( Nishimori & Akaho , 2005 ) . The closed-form of the Cayley transform is given by : Y ( α ) = ( I − α 2 W ) −1 ( I + α 2 W ) X , ( 3 ) where W is a skew-symmetric matrix , i.e . W > = −W , X is on the Stiefel manifold , and α ≥ 0 is a parameter that represents the length on the curve . It is straightforward to verify that Y ( 0 ) = X and Y ′ ( 0 ) =WX . ( 4 ) From Definition 2 and the definition of the tangent space of the Stiefel manifold given by Eq . ( 1 ) , the Cayley transform is a valid retraction map on the Stiefel manifold . By choosing W = Ŵ − Ŵ > , where Ŵ = ∇f ( X ) X > − 12X ( X > ∇f ( X ) X > ) as in Eq . ( 2 ) , we see that the Cayley transform implicitly projects gradient on the tangent space as its initial direction . Therefore , the Cayley transform can represent an update for the parameter matrices on the Stiefel manifold . However , the closed-form Cayley transform in Eq . ( 3 ) involves computing the expensive matrix inversion , which can not be efficiently performed in large deep neural networks . Our first contribution is an iterative estimation of the Cayley transform that efficiently uses only matrix multiplications , and thus is more efficient than the closed form in Eq. ( 3 ) . We represent the Cayley transform with the following fixed-point iteration : Y ( α ) = X + α 2 W ( X + Y ( α ) ) , ( 5 ) which can be proved by moving Y ( α ) on the right-hand side to the left-hand side in Eq. ( 3 ) . The expression in Eq . ( 5 ) is an efficient approximation of Eq. ( 3 ) . In Sec . 5 , we will analyze its convergence rate to the closed-form Eq . ( 3 ) .

Summary: This paper aims to improve upon current solutions for optimizing neural networks with orthonormal convolutional kernels/MLP layers. Optimizing neural networks while restricting the parameter matrices to remain orthonormal/on the Stiefel manifold is said to lead to faster convergence in terms of the number of epochs, and could reduce overfitting. One way to enforce orthonormality is to use the Cayley transform, which requires a matrix inversion that becomes expensive for large weight matrices in neural networks. The authors instead propose an iterative approximation to the Cayley transform that does not require a matrix inversion, and empirically only requires two iteration steps to lead to similar precision as a closed-form Cayley transform with numerical inverse. The authors combine this with SGD + momentum and ADAM by first performing the update step in Euclidean space and afterwards projecting the result back onto the stiefel manifold by using the iterative Cayley transform. They provide one assumption and two theorems on the convergence of the iterative approximation and its effect on optimization. The method is evaluated on classification for cifar10 and cifar100 and for modeling the hidden-to-hidden transition matrix in an RNN trained on MNIST.