PoLoRA: A Preconditioned Orthogonalized LoRA Optimizer

Standard LoRA finetuning typically uses the Adam optimizer, which treats the low-rank factors as a flat vector of parameters, ignoring the underlying matrix structure and the product relationship between the two factors. This paper investigates whether a matrix-aware, product-aware optimizer can improve the training efficiency of LoRA compared to the standard Adam baseline.

The authors introduce PoLoRA (Preconditioned Orthogonalized LoRA), an optimizer designed specifically for the low-rank product structure. It incorporates three key innovations:

• Product-aware spectral update: Instead of updating factors independently, it computes the update direction by considering the linearized product of the factors, ensuring the update respects the merged weight matrix.
• Curvature preconditioning: It uses a Kronecker-factored approximation of the gradient curvature to normalize the update, effectively accounting for how different directions in the parameter space influence the loss.
• Magnitude control: It employs a specific rule to set the size of the factor updates, ensuring the spectral norm of the merged update is bounded and stable across different ranks.

PoLoRA consistently outperforms Adam across various models (1B to 8B parameters) and datasets (code and math). It reaches the final held-out loss of tuned Adam in 1.2 to 1.7 times fewer training steps. The authors demonstrate that the curvature preconditioning and magnitude control are both essential, with each contributing roughly half of the observed speedup. Furthermore, PoLoRA exhibits greater stability, as its optimal learning rate remains consistent across different LoRA ranks, unlike Adam, which requires re-tuning.