PoLoRA: A Preconditioned Orthogonalized LoRA Optimizer

2026-07-20

AI Summary

Existing LoRA optimization methods overlook the intrinsic low-rank matrix product structure, limiting training efficiency and stability. This work proposes a structure-aware optimizer that effectively exploits LoRA’s inherent matrix geometry through three key innovations: spectral-direction updates, curvature preconditioning based on per-sample loss variations, and an orthogonalization-based magnitude control mechanism. Evaluated across models ranging from 1B to 8B parameters, the proposed method achieves comparable final loss to Adam with 1.2–1.7× fewer training steps, while incurring no more than a 3% increase in per-step computational overhead. Furthermore, it demonstrates strong robustness to learning rate selection, significantly easing hyperparameter tuning.

Abstract

Low-rank adaptation (LoRA) makes finetuning large language models cheaper by adding to each weight matrix a trainable low-rank update parameterized as the product of two matrices. These matrices are usually trained with Adam, which treats them as a single flat vector of parameters and ignores both the matrix and product structure of LoRA. Applying a matrix-aware optimizer such as Muon to each factor does not consistently improve over Adam, and neither do the product-aware Muon variants proposed in concurrent works. To realize consistent gains, we introduce PoLoRA, a Preconditioned Orthogonalized LoRA optimizer built from three ingredients: a product-aware spectral update direction, curvature preconditioning derived from controlling the per-sample loss change, and a magnitude rule that controls the sizes of both the factor and merged updates. We evaluate PoLoRA on instruction-tuning datasets for code and math across models from 1B to 8B parameters, and find that it reaches the final held-out loss achieved by tuned Adam in 1.2-1.7 times fewer steps, while adding at most 3% per-step overhead. Compared to Adam, PoLoRA is also less sensitive to the learning rate, and its optimal learning rate is stable across ranks.

Problem

• LoRA
• optimizer
• matrix structure
• product-aware
• large language models

Innovation

• LoRA
• optimizer
• preconditioning
• low-rank adaptation
• spectral update