PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform

Jun Li, Li Fuxin, Sinisa Todorovic

arXiv:2002.01113v1cs.LGstat.ML

TL;DR

Strict orthonormality constraints are useful in deep learning, but Stiefel-manifold optimization is computationally expensive. The paper introduces iterative Cayley retractions and implicit vector transport to build Cayley SGD and Cayley ADAM; experiments report faster convergence and lower per-iteration runtime than relevant baselines and orthonormality methods, without compromising performance.

  • Problem

    Stiefel-manifold optimization offers exact orthonormality constraints, but its high computational cost limits deep-learning use, while prior efficient approaches are restricted to square matrices or do not guarantee exact orthonormality.

    Source: §1 · p. 1

  • Method

    The paper combines an iterative Cayley-transform retraction with implicit vector transport to extend SGD with momentum and ADAM to the Stiefel manifold.

    Source: §1 · pp. 1–2

  • Results

    Cayley SGD and Cayley ADAM achieve faster convergence than baseline SGD and ADAM and run faster per iteration than existing hard and soft orthonormal methods, with comparable performance across CNN and RNN experiments.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    The algorithms provide an exact-orthonormality optimization approach that applies to non-square parameter matrices with lower per-iteration runtime than existing hard and soft orthonormal methods.

    Source: §7 · p. 9

    The convergence analysis assumes a Lipschitz-continuous objective gradient; for some ReLU models, the derivative meets this condition almost everywhere except near zero.

    Source: §5 · p. 6

Abstract

from arXiv · show

Strictly enforcing orthonormality constraints on parameter matrices has been shown advantageous in deep learning. This amounts to Riemannian optimization on the Stiefel manifold, which, however, is computationally expensive. To address this challenge, we present two main contributions: (1) A new efficient retraction map based on an iterative Cayley transform for optimization updates, and (2) An implicit vector transport mechanism based on the combination of a projection of the momentum and the Cayley transform on the Stiefel manifold. We specify two new optimization algorithms: Cayley SGD with momentum, and Cayley ADAM on the Stiefel manifold. Convergence of Cayley SGD is theoretically analyzed. Our experiments for CNN training demonstrate that both algorithms: (a) Use less running time per iteration relative to existing approaches that enforce orthonormality of CNN parameters; and (b) Achieve faster convergence rates than the baseline SGD and ADAM algorithms without compromising the performance of the CNN. Cayley SGD and Cayley ADAM are also shown to reduce the training time for optimizing the unitary transition matrices in RNNs.

1 INTRODUCTION

The introduction frames Stiefel-manifold optimization as a way to enforce orthonormality, but its computational cost limits use in deep learning. It presents an iterative Cayley retraction and implicit vector transport, then evaluates the resulting optimizers on CNNs and RNNs.

1 INTRODUCTION

  • Riemannian optimization on the Stiefel manifold enforces orthonormality, but its computational cost limits deep-learning applications; prior efficient approaches address only square matrices or use regularization without exact guarantees.Orthonormality has been associated with improved CNN accuracy, activation stability, and generalization, and with alleviating vanishing and exploding gradients in RNNs.Source: §1 · p. 1
  • An iterative Cayley-transform retraction updates large non-square parameter matrices using a few matrix multiplications instead of the costly matrix inverses required by the closed-form transform.The paper reports theoretical proof and empirical verification of the retraction's efficiency.Source: §1 · p. 1
  • The implicit vector transport combines momentum and gradient in Euclidean space before the Cayley update, avoiding explicit vector transport because the transform implicitly projects onto the tangent space.This momentum-based contribution is intended to improve training convergence rates.Source: §1 · pp. 1–2
  • The two contributions yield Cayley SGD with momentum and Cayley ADAM on the Stiefel manifold; the paper analyzes convergence theoretically for Cayley SGD.Similar analysis for Cayley ADAM is omitted because it is very similar to analysis in Becigneul & Ganea (2019).Source: §1 · p. 2
  • Cayley SGD and Cayley ADAM achieve better classification performance and faster convergence than baseline SGD with momentum and ADAM, while running faster than existing orthonormality-aware methods with equally good or better image-classification performance.The baselines take less time per epoch because they do not enforce orthonormality, but require more epochs to converge.Source: §1 · p. 2
  • For unitary RNN training, Cayley SGD with momentum and Cayley ADAM reduce running time per iteration without compromising performance relative to the closed-form Cayley-transform approach.The comparison is framed against the full-capacity unitary RNN of Wisdom et al. (2016).Source: §1 · p. 2

2 RELATED WORK

Related work spans regularization and Riemannian optimization, with existing approaches trading exact orthonormality against computational cost or applicability to non-square matrices.

2 RELATED WORK

  • Hard regularization enforces orthonormality but can require costly eigenvalue decomposition, while soft mutual-coherence and spectral-restricted-isometry regularizers cannot guarantee orthonormality.Examples include Huang et al. (2018b) and Huang et al. (2018a) for eigenvalue-decomposition methods, and Bansal et al. (2018); Cisse et al. (2017); Xiong et al. (2016) for soft regularization.Source: §2 · p. 2
  • Riemannian approaches guarantee orthonormality, but cited unitary-matrix methods cannot handle non-square parameters, while the general method of Becigneul & Ganea (2019) lacks an efficient Stiefel-manifold procedure.Cho & Lee (2017) use the Grassmann manifold for normalized but not orthonormalized CNN parameters.Source: §2 · p. 2
  • Projection-based pseudo-geodesic retractions rely on computationally expensive SVD, closed-form Cayley methods require matrix inversion, and Wen & Yin (2013) is inefficient when 2p ≥ n.Wen & Yin (2013) assumes 2p ≪ n for its reduced-cost Cayley transform.Source: §2 · p. 2

3 PRELIMINARY

The preliminary section establishes the manifold concepts behind constrained optimization, then introduces a Cayley-transform retraction and an implicit transport for momentum updates.

3.1 RIEMANNIAN MANIFOLD

  • A retraction maps tangent-space steps back onto the manifold; it can provide an efficient alternative to the usually costly exponential map.Vector transport moves tangent vectors to a new tangent space and is often a computationally efficient alternative to parallel transport.Source: §3.1 · p. 3

3.2 STIEFEL MANIFOLD

  • A Stiefel matrix has orthonormal columns, and optimization combines a projected Riemannian gradient, transported momentum, and a step along the manifold.Exponential maps and parallel transport are computationally infeasible on the Stiefel manifold, motivating efficient alternatives.Source: §3.2 · p. 3

3.2.1 PARAMETER UPDATES BY ITERATIVE CAYLEY TRANSFORM

  • The iterative Cayley transform estimates a valid retraction using matrix multiplications, avoiding the costly matrix inversion in its closed form.The iteration approximates the closed-form Cayley transform.Source: §3.2.1 · p. 4

3.2.2 MOMENTUM UPDATES BY THE IMPLICIT VECTOR TRANSPORT

  • Implicit vector transport combines momentum and the gradient in Euclidean space, then uses the Cayley transform to update parameters without explicitly transporting the vector.The transport is a projection onto the tangent space, and the Cayley transform implicitly projects the combined direction.Source: §3.2.2 · p. 4

4 ALGORITHMS

The algorithms generalize SGD with momentum and ADAM to Stiefel-manifold optimization by combining momentum with the stochastic gradient and applying an iterative Cayley update. Cayley SGD and Cayley ADAM implement this approach with manifold-specific momentum and learning-rate handling.

4 ALGORITHMS

  • Both algorithms combine the previous momentum with the current stochastic gradient, then use the iterative Cayley transform to estimate the next Stiefel-manifold point.This generalizes conventional SGD with momentum and ADAM to Riemannian optimization on the Stiefel manifold.Source: §4 · p. 4

4.1 CAYLEY SGD WITH MOMENTUM

  • Cayley SGD generalizes the heavy-ball update by projecting momentum onto the current tangent space and updating parameters along a Cayley-transform curve.Its practical update uses an iterative Cayley transform and implicit vector transport.Source: §4.1 · p. 5

4.2 ADAM ON THE STIEFEL MANIFOLD

  • Cayley ADAM replaces the Euclidean gradient and momentum with their Stiefel-manifold counterparts, uses one adaptive learning rate per parameter matrix, and updates with the Cayley transform.The shared rate across matrix entries is described as a manifold-wise adaptive learning rate.Source: §4.2 · pp. 5–6

5 CONVERGENCE ANALYSIS

Under a Lipschitz-gradient assumption, the paper proves convergence of the iterative Cayley transform and an asymptotic stationarity result for Cayley SGD with momentum. It also compares the Cayley iteration's convergence rate with two other approximation algorithms.

5 CONVERGENCE ANALYSIS

  • The analysis assumes a Lipschitz-continuous objective gradient; the paper notes that ReLU derivatives satisfy this almost everywhere except near zero.The stated exception is a small neighborhood around zero whose measure tends to zero.Source: §5 · p. 6
  • For α ∈ (0, min{1, 2∥W∥}), the iterative Cayley transform converges to closed-form Y(α), with iteration-i error o(α^(2+i)), faster than Newton iterative and Neumann Series bounds of o(α^i).Source: §5 · p. 6
  • For an objective f satisfying Assumption 1, Cayley SGD with momentum has min_{k=0,...,t} E[||∇_M f(X_k)||²] = o(1/√(t+1)), tending to zero as t → ∞.Source: §5 · p. 6

6 EXPERIMENTS

The experiments evaluate Cayley SGD and Cayley ADAM on CIFAR image classification and pixel-by-pixel MNIST, comparing performance, convergence, runtime, and unitarity. The methods attain competitive classification accuracy, improve runtime or convergence in the reported comparisons, and maintain unitary constraints with few iterative Cayley steps.

6.1 ORTHONORMALITY IN CNNS

  • For the CNN experiments, convolution kernels are reshaped as p × n matrices and constrained on the Stiefel manifold; other parameters use SGD or ADAM.The evaluation uses VGG and Wide ResNet on CIFAR10 and CIFAR100.Source: §6.1 · p. 6
  • 3.57% and 18.10% error are achieved by WRN-28-10 with Cayley ADAM on CIFAR10 and CIFAR100, respectively.The paper identifies these as the best error rates in the reported experiments.Source: §6.1 · p. 7, Table 1 · p. 7, Table 2 · p. 7
  • Cayley SGD and Cayley ADAM are among the two fastest methods by training-loss convergence across epochs on CIFAR10 and CIFAR100.Baselines may plateau before a learning-rate drop, whereas the Cayley methods do not; the authors suggest orthonormality may explain this difference.Source: Figure 1 · p. 7, §6.1 · p. 7
  • With training time measured in seconds, the Cayley methods converge fastest among methods that also enforce orthonormality, while SGD and ADAM are faster initially.The Cayley methods catch up with the baselines after 12,000 seconds.Source: §6.1 · p. 7, Figure 2 · p. 8
  • Cayley methods achieve error rates comparable to state of the art and run faster per epoch than orthonormality methods, except unconstrained SGD and ADAM.The comparison uses WRN-28-10 on CIFAR10 and CIFAR100.Source: §6.1 · p. 8, Table 3 · p. 8

6.2 ORTHONORMALITY IN RNNS

  • About 35% less training time per iteration is reported for Cayley SGD and Cayley ADAM across Full-uRNN settings, with the same accuracy.The comparison is against the closed-form Cayley transform on pixel-by-pixel MNIST.Source: §6.2 · pp. 8–9, Table 4 · p. 9
  • The iterative Cayley transform approximates a unitary matrix with s = 2 and has lower reported unitarity error than the closed-form transform.The authors suggest rounding error from matrix inversion as a possible explanation for the comparison.Source: Table 5 · p. 9, §6.2 · p. 9

7 CONCLUSION

The conclusion presents iterative Cayley optimization as an efficient way to enforce exact orthonormal constraints, with competitive performance and faster convergence across the evaluated CNN and RNN tasks.

7 CONCLUSION

  • The proposed approach enforces exact orthonormal constraints through optimization on the Stiefel manifold.It applies the iterative Cayley transform to conventional SGD and ADAM, yielding Cayley SGD with momentum and Cayley ADAM.Source: §7 · p. 9
  • The paper theoretically analyzes convergence for Cayley SGD with momentum.Source: §7 · p. 9
  • Across the evaluated CIFAR10, CIFAR100, and pixel-by-pixel MNIST tasks, the algorithms achieve comparable performance and faster convergence than baseline SGD and ADAM.They also take less runtime per iteration than existing hard- and soft-orthonormal methods and apply to non-square parameter matrices.Source: §7 · p. 9

A PRELIMINARY

The preliminary analysis establishes boundedness of gradients and the momentum sequence under the stated manifold and gradient assumptions.

A PRELIMINARY

  • The bounded Stiefel manifold and Lipschitz gradient assumption imply bounded Euclidean and projected manifold gradients.Source: A PRELIMINARY · p. 12
  • For a finite dataset, the stochastic gradients and their Riemannian projections are also bounded by the same upper bound G.Source: A PRELIMINARY · p. 12
  • The momentum recursion is bounded, and the paper consequently establishes that Wk is bounded.Source: A PRELIMINARY · p. 12

B PROOF OF THEOREM 1

The proof establishes that Eq.(5) is a contraction under a bound on α and ∥W∥, so its iterates converge to the closed-form solution. It uses derivative bounds and a Taylor expansion.

B PROOF OF THEOREM 1

  • The cited passage states that the iteration in Eq.(5) is a contraction mapping and converges to the closed-form solution Y(α), but gives no complete condition in the fragment.Source: B PROOF OF THEOREM 1 · p. 12
  • The proof differentiates Eq.(5) and bounds the resulting derivatives by a positive constant C.Source: B PROOF OF THEOREM 1 · p. 12
  • The proof uses a Taylor expansion of Y from Eq.(3), with the initialization Y⁰=Xᵏ+αMᵏ⁺¹.Source: B PROOF OF THEOREM 1 · p. 12

C PROOF OF THEOREM 2

The proof concludes that the minimum expected squared Riemannian gradient norm decays as o(1/√(t+1)) and tends to zero. Its analysis uses Taylor expansion, auxiliary variables, smoothness and Lipschitz assumptions, and expectations.

C PROOF OF THEOREM 2

Related papers

Loading 2002.01113v1…