Large-Scale Riemannian Meta-Optimization via Subspace Adaptation

TL;DR

This paper introduces a memory-efficient Riemannian meta-optimization approach using subspace adaptation, enabling large-scale neural network training with significantly reduced memory and improved performance.

Contribution

It proposes a novel subspace adaptation scheme for Riemannian meta-optimization that shares the optimizer across parameters of different sizes, drastically reducing memory use.

Findings

Reduces memory consumption by six orders of magnitude on ResNet50.

Improves performance of Riemannian meta-optimization tasks.

Enables large-scale neural network training on Riemannian manifolds.

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.