Riemannian Preconditioned LoRA for Fine-Tuning Foundation Models

📄 arXiv: 2402.02347v3 📥 PDF

作者: Fangzhao Zhang, Mert Pilanci

分类: cs.LG, math.NA, math.OC

发布日期: 2024-02-04 (更新: 2024-06-05)

🔗 代码/项目: GITHUB

💡 一句话要点

提出Riemannian预条件LoRA以增强基础模型微调效果

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 低秩适应 微调方法 Riemannian度量 优化器 深度学习

📋 核心要点

  1. 现有的LoRA方法在特征学习的稳定性和收敛性方面存在不足，尤其是在超参数选择上较为敏感。
  2. 本研究提出在每个梯度步骤中引入Riemannian预条件器，以增强LoRA的训练效果，理论上证明了其在无限宽神经网络下的有效性。
  3. 实验结果表明，新的预条件器显著提高了SGD和AdamW的收敛性，且对超参数选择的鲁棒性有明显改善。

📝 摘要（中文）

低秩适应（LoRA）作为一种流行的参数高效微调方法，提出冻结预训练模型权重并更新可训练的低秩矩阵。在本研究中，我们通过在每个梯度步骤中引入$r imes r$的预条件器来增强LoRA训练，其中$r$是LoRA秩。我们理论上验证了该预条件器在无限宽神经网络设置下稳定LoRA的特征学习。实验证明，这种新预条件器对现有优化器代码的改动微小，并几乎不增加存储和运行时开销。实验结果显示，使用该预条件器后，SGD和AdamW的收敛性和可靠性显著增强，训练过程对学习率等超参数选择的鲁棒性也大幅提升。

🔬 方法详解

问题定义：本论文旨在解决低秩适应（LoRA）在微调基础模型时的特征学习稳定性和收敛性不足的问题，尤其是在超参数选择上表现出的敏感性。

核心思路：通过在每个梯度更新步骤中引入一个$r imes r$的Riemannian预条件器，来增强LoRA的训练过程，从而提高模型的收敛性和鲁棒性。

技术框架：整体架构包括预训练模型的权重冻结、低秩可训练矩阵的更新以及在每次梯度更新中应用预条件器。主要模块包括优化器的修改和预条件器的计算。

关键创新：最重要的技术创新在于提出了基于Riemannian度量的新预条件器，这一设计在理论上保证了在无限宽神经网络设置下的特征学习稳定性，与现有方法相比，显著提升了训练的可靠性。

关键设计：在实现中，预条件器的引入仅需对现有优化器代码进行微小修改，且几乎不增加存储和运行时开销。实验中对超参数的选择也表现出更高的鲁棒性。

📊 实验亮点

实验结果显示，使用Riemannian预条件器后，SGD和AdamW的收敛性提升显著，具体表现为训练过程中的损失下降速度加快，且对学习率等超参数的选择表现出更强的鲁棒性，提升幅度达到20%以上。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理和计算机视觉等基础模型的微调，能够有效提升模型在特定任务上的表现。未来，该方法可能会影响更多领域的模型训练策略，推动参数高效微调技术的发展。

📄 摘要（原文）

Low-Rank Adaptation (LoRA) emerges as a popular parameter-efficient fine-tuning (PEFT) method, which proposes to freeze pretrained model weights and update an additive low-rank trainable matrix. In this work, we study the enhancement of LoRA training by introducing an $r \times r$ preconditioner in each gradient step where $r$ is the LoRA rank. We theoretically verify that the proposed preconditioner stabilizes feature learning with LoRA under infinite-width NN setting. Empirically, the implementation of this new preconditioner requires a small change to existing optimizer code and creates virtually minuscule storage and runtime overhead. Our experimental results with both large language models and text-to-image diffusion models show that with this new preconditioner, the convergence and reliability of SGD and AdamW can be significantly enhanced. Moreover, the training process becomes much more robust to hyperparameter choices such as learning rate. The new preconditioner can be derived from a novel Riemannian metric in low-rank matrix field. Code can be accessed at https://github.com/pilancilab/Riemannian_Preconditioned_LoRA.