On the Crucial Role of Initialization for Matrix Factorization

📄 arXiv: 2410.18965v3 📥 PDF

作者: Bingcong Li, Liang Zhang, Aryan Mokhtari, Niao He

分类: cs.LG, eess.SP, math.OC

发布日期: 2024-10-24 (更新: 2024-12-12)

💡 一句话要点

提出Nystrom初始化，加速非凸矩阵分解与LoRA微调收敛

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

关键词: 矩阵分解 Nystrom方法 初始化 低秩适配器 LoRA 大规模模型 优化算法

📋 核心要点

  1. 现有低秩矩阵分解方法收敛速度慢，依赖于良好的初始化，缺乏理论保证。
  2. 提出Nystrom初始化方法，加速ScaledGD在对称和非对称矩阵分解中的收敛速度。
  3. NoRA方法将Nystrom初始化应用于LoRA，在多种下游任务和模型规模上均表现出优越性能。

📝 摘要（中文）

本研究重新审视了经典的低秩矩阵分解问题，揭示了初始化在影响非凸非光滑优化收敛速度方面的关键作用。我们引入了Nystrom初始化，显著提高了缩放梯度下降法(ScaledGD)在对称和非对称矩阵分解任务中的全局收敛性。具体而言，我们证明了使用Nystrom初始化的ScaledGD可以实现二次收敛，而之前只知道线性收敛速度。此外，我们将这种初始化方法扩展到常用于微调基础模型的低秩适配器(LoRA)。我们的方法NoRA，即使用Nystrom初始化的LoRA，在大型语言模型和扩散模型中，从10亿到70亿参数的各种下游任务和模型规模上都表现出卓越的性能。

🔬 方法详解

问题定义：论文旨在解决低秩矩阵分解中，传统方法收敛速度慢且对初始化敏感的问题。现有的缩放梯度下降法(ScaledGD)虽然可以用于矩阵分解，但其收敛速度通常为线性，且对初始值的选择非常敏感，难以保证全局最优解。

核心思路：论文的核心思路是利用Nystrom方法进行初始化，为ScaledGD提供一个良好的起点，从而加速其收敛速度并提高全局收敛性。Nystrom方法通过对矩阵进行采样和近似，可以有效地估计矩阵的特征向量和特征值，从而提供一个更接近最优解的初始值。

技术框架：整体框架分为两个主要部分：首先，使用Nystrom方法对目标矩阵进行初始化；然后，使用ScaledGD进行迭代优化。对于LoRA，则将Nystrom初始化应用于LoRA的参数矩阵。具体流程为：1. 对原始矩阵进行Nystrom采样；2. 基于采样结果计算初始化矩阵；3. 使用该初始化矩阵作为ScaledGD或LoRA的初始值；4. 进行迭代优化。

关键创新：最重要的技术创新点在于Nystrom初始化方法。与随机初始化或其他启发式初始化方法相比，Nystrom初始化能够提供一个更接近最优解的初始值，从而显著加速收敛速度。此外，将Nystrom初始化应用于LoRA，并证明其在大型语言模型和扩散模型中的有效性，也是一个重要的创新。

关键设计：Nystrom初始化的关键在于采样策略和近似方法。论文可能采用了均匀采样或基于重要性的采样方法来选择矩阵的列或行。近似方法可能包括使用采样列/行构建低秩近似矩阵，并计算其特征向量和特征值。对于LoRA，关键在于如何将Nystrom初始化应用于LoRA的A和B矩阵，可能直接初始化A和B，或者初始化A@B的乘积。

📊 实验亮点

实验结果表明，使用Nystrom初始化的ScaledGD在矩阵分解任务中实现了二次收敛，显著优于之前的线性收敛速度。NoRA方法在各种下游任务和模型规模上都表现出卓越的性能，例如在大型语言模型和扩散模型中，从10亿到70亿参数的模型上都取得了显著的性能提升，具体提升幅度未知。

🎯 应用场景

该研究成果可广泛应用于推荐系统、图像处理、自然语言处理等领域，尤其是在需要进行大规模矩阵分解的场景下。通过使用Nystrom初始化，可以显著降低计算成本，提高模型训练效率。此外，NoRA方法在微调大型语言模型和扩散模型方面具有重要应用价值，可以帮助研究人员和工程师更高效地训练和部署这些模型。

📄 摘要（原文）

This work revisits the classical low-rank matrix factorization problem and unveils the critical role of initialization in shaping convergence rates for such nonconvex and nonsmooth optimization. We introduce Nystrom initialization, which significantly improves the global convergence of Scaled Gradient Descent (ScaledGD) in both symmetric and asymmetric matrix factorization tasks. Specifically, we prove that ScaledGD with Nystrom initialization achieves quadratic convergence in cases where only linear rates were previously known. Furthermore, we extend this initialization to low-rank adapters (LoRA) commonly used for finetuning foundation models. Our approach, NoRA, i.e., LoRA with Nystrom initialization, demonstrates superior performance across various downstream tasks and model scales, from 1B to 7B parameters, in large language and diffusion models.