Large-Scale Riemannian Meta-Optimization via Subspace Adaptation
This addresses a memory bottleneck for researchers and practitioners using Riemannian meta-optimization in large-scale non-linear constrained optimization problems, offering a more efficient solution.
The paper tackles the high memory consumption of existing Riemannian meta-optimization methods in large-scale settings by proposing a subspace adaptation scheme that trains neural networks to adapt row and column subspaces of gradients, reducing memory by six orders of magnitude for models like ResNet50 while improving performance on Riemannian tasks.
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.