Large-Scale Riemannian Meta-Optimization via Subspace Adaptation

TL;DR

This work tackles the large memory footprint inherent in Riemannian meta-optimization when optimizing deep models with manifold constraints. It introduces a subspace adaptation strategy that separately learns row and column gradient subspaces via small LSTMs, allowing a single optimizer to be shared across Riemannian parameters of varying sizes. By refining gradients with structure-preserving matrices $R$ and $C$ derived from gradient covariance, the method achieves dramatic memory reductions (down to tens of kilobytes) while maintaining or improving convergence on tasks such as PCA on Grassmann, face recognition on Stiefel, and orthogonal deep networks. The approach is trained end-to-end through bi-level meta-learning, yielding a data-driven optimizer that generalizes across tasks and supports scalable, large-scale Riemannian optimization with practical impact for continual learning and real-world constrained problems.

Abstract

Riemannian meta-optimization provides a promising approach to solving non-linear constrained optimization problems, which trains neural networks as optimizers to perform optimization on Riemannian manifolds. However, existing Riemannian meta-optimization methods take up huge memory footprints in large-scale optimization settings, as the learned optimizer can only adapt gradients of a fixed size and thus cannot be shared across different Riemannian parameters. In this paper, we propose an efficient Riemannian meta-optimization method that significantly reduces the memory burden for large-scale optimization via a subspace adaptation scheme. Our method trains neural networks to individually adapt the row and column subspaces of Riemannian gradients, instead of directly adapting the full gradient matrices in existing Riemannian meta-optimization methods. In this case, our learned optimizer can be shared across Riemannian parameters with different sizes. Our method reduces the model memory consumption by six orders of magnitude when optimizing an orthogonal mainstream deep neural network (e.g., ResNet50). Experiments on multiple Riemannian tasks show that our method can not only reduce the memory consumption but also improve the performance of Riemannian meta-optimization.

Figures (7)

• Figure 1: The architecture of our optimizer. $\boldsymbol{Cov}_R^{\left ( t \right )}$ and $\boldsymbol{Cov}_C^{\left ( t \right )}$ are the row and column covariance matrices of Riemannian gradient $\boldsymbol{G}^{\left ( t \right )}$. The LSTMs take each coordinate scalar $\hat{\boldsymbol{r}}_{i}^{\left ( t \right )}$ or $\hat{\boldsymbol{c}}_{j}^{\left ( t \right )}$, its own previous hidden $\boldsymbol{h}^{\left ( t-1 \right )}$, and cell state $\boldsymbol{c}^{\left ( t-1 \right )}$ as input each time, and respectively produce a diagonal element of adaptive matrix $\boldsymbol{R}^{\left ( t \right )}$ or $\boldsymbol{C}^{\left ( t \right )}$. The LSTM parameters $\phi_{1}$ and $\phi_{2}$ are shared across coordinates of all Riemannian parameters.
• Figure 2: The illustration of the optimization procedure. The black dotted curve denotes a geodesic on manifold $\mathcal{M} _{k}$. The orthogonal projection and retraction operations are denoted as orange and red dotted curves, respectively. $\nabla _{\boldsymbol{W}_{k}}^{(t)}$ (orange solid line) is the Euclidean gradient and $\boldsymbol{G}^{(t)}$ (blue sold line) is the Riemannian gradient on the tangent space. $\boldsymbol{R}^{(t)}$ (yellow solid line) and $\boldsymbol{C}^{(t)}$ (cyan solid line) are the row and column subspace adaptive matrices. $\hat{\boldsymbol{G}}_{k}^{\left ( t \right )}$ (purple solid line) and $\boldsymbol{P}^{(t)}$ (red solid line) are the refined gradient and the update vector, respectively.
• Figure 3: Plots for the PCA and face recognition tasks (in the log scale).
• Figure 4: Testing loss curves of different Riemannian optimization algorithms on optimizing VGG16, ResNet18 and ResNet50.
• Figure 5: Testing loss vs. wallclock time (seconds) curves of different Riemannian optimization algorithms.

Theorems & Definitions (1)

• Remark 1