通过子空间适配，让优化器共享内存，大幅降低大模型优化的显存占用。
Large-Scale Riemannian Meta-Optimization via Subspace Adaptation
- 用神经网络分别适配梯度矩阵的行与列子空间，而非全量梯度。
- 在ResNet50上实现显存降低六数量级，同时提升优化性能。
- 适合需要高效大模型优化的科研与工程场景。
黎曼元优化为求解非线性约束优化问题提供了新思路，即训练神经网络作为优化器，在黎曼流形上执行优化。然而，现有方法在大规模优化场景中内存开销巨大，因学习到的优化器仅能适应固定大小的梯度，无法跨不同尺寸的黎曼参数共享。本文提出一种高效的黎曼元优化方法，通过子空间适配机制显著降低大规模优化的内存负担。该方法训练神经网络分别适配黎曼梯度的行和列子空间，而非直接处理完整的梯度矩阵。由此，学习到的优化器可跨不同尺寸的黎曼参数共享。在优化主流正交深度神经网络（如ResNet50）时，模型内存消耗降低六数量级。多组黎曼任务实验表明，本方法不仅大幅减少内存占用，还能提升黎曼元优化的性能。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。