TeZO: Empowering the Low-Rankness on the Temporal Dimension in the Zeroth-Order Optimization for Fine-tuning LLMs

📄 arXiv: 2501.19057v1 📥 PDF

作者: Yan Sun, Tiansheng Huang, Liang Ding, Li Shen, Dacheng Tao

分类: cs.LG

发布日期: 2025-01-31

💡 一句话要点

TeZO：利用时序维度低秩性，提升大语言模型零阶优化微调效率

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

关键词: 零阶优化 大语言模型 低秩分解 模型微调 张量分解 计算效率 内存优化

📋 核心要点

  1. 现有基于零阶优化的大语言模型微调方法，忽略了训练过程中梯度在时序上的关联性，未能充分利用梯度的低秩特性。
  2. TeZO将零阶优化扰动在时序维度上建模为3D张量，并利用CPD分解提取低秩2D矩阵，同时考虑模型和时序维度上的低秩性。
  3. 实验结果表明，TeZO在降低内存消耗的同时，能够达到与SOTA方法相当的性能，并可扩展到Adam优化器。

📝 摘要（中文）

零阶优化(ZO)已在大语言模型(LLM)高效微调任务中展现出显著潜力。特别是，最近的研究将梯度的低秩性纳入考虑，引入低秩ZO估计器以进一步降低GPU内存消耗。然而，现有工作大多只关注每个单独梯度的低秩性，忽略了训练过程中所有梯度共享的一个更广泛的属性，即所有梯度近似位于一个相似的子空间内。本文同时考虑这两个因素，提出了一种新的低秩ZO估计器TeZO，它捕捉了模型和时间维度上的低秩性。具体来说，我们将时间维度上的ZO扰动表示为一个3D张量，并采用Canonical Polyadic Decomposition (CPD)来提取每个低秩2D矩阵，从而显著降低训练成本。TeZO也可以很容易地扩展到Adam变体，同时消耗比MeZO-SGD更少的内存，并且只需要MeZO-Adam约35%的内存。全面的理论分析和广泛的实验研究都验证了其效率，以更低的时间和内存开销实现了与SOTA相当的结果。

🔬 方法详解

问题定义：现有基于零阶优化的LLM微调方法，虽然利用了梯度的低秩性来降低计算和存储成本，但主要关注单个梯度的低秩性。忽略了训练过程中，所有梯度在时间维度上存在关联，近似位于一个相似的子空间中。这种时序上的关联性没有被充分利用，限制了优化效率的进一步提升。

核心思路：TeZO的核心思路是同时利用模型参数和时间维度上的低秩性。具体来说，将训练过程中产生的梯度扰动视为一个3D张量，其中两个维度对应模型参数，一个维度对应时间步。通过对这个3D张量进行低秩分解，可以提取出在模型参数和时间维度上都具有低秩结构的表示，从而更有效地进行优化。

技术框架：TeZO的整体框架包括以下几个步骤：1）在每个时间步，对模型参数进行零阶扰动，得到扰动后的模型参数。2）使用扰动后的模型参数计算损失函数。3）利用损失函数的差分来估计梯度。4）将所有时间步的梯度扰动组合成一个3D张量。5）使用Canonical Polyadic Decomposition (CPD)对3D张量进行低秩分解，提取低秩表示。6）使用低秩表示来更新模型参数。

关键创新：TeZO的关键创新在于同时考虑了模型参数和时间维度上的低秩性，并将梯度扰动建模为一个3D张量。通过对3D张量进行低秩分解，可以更有效地提取梯度中的低秩结构，从而降低计算和存储成本，并提高优化效率。与现有方法相比，TeZO能够更好地捕捉梯度在时间维度上的关联性，从而实现更高效的微调。

关键设计：TeZO的关键设计包括：1）使用Canonical Polyadic Decomposition (CPD)进行低秩分解。CPD是一种常用的张量分解方法，可以将一个高阶张量分解为多个低秩矩阵的乘积。2）将时间维度上的梯度扰动组合成一个3D张量。这种表示方式能够有效地捕捉梯度在时间维度上的关联性。3）将TeZO扩展到Adam优化器。通过将TeZO与Adam优化器相结合，可以进一步提高优化效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TeZO在多个数据集上取得了与SOTA方法相当的性能，同时显著降低了内存消耗。例如，TeZO在使用Adam优化器时，内存消耗仅为MeZO-Adam的35%。此外，TeZO在训练时间上也具有优势，能够更快地收敛到最优解。

🎯 应用场景

TeZO可应用于各种大语言模型的微调任务，尤其适用于资源受限的场景，如边缘设备或低GPU内存环境。该方法能够降低微调过程中的内存消耗和计算成本，加速模型部署，并促进大语言模型在更广泛领域的应用，例如移动设备上的个性化助手和智能客服。

📄 摘要（原文）

Zeroth-order optimization (ZO) has demonstrated remarkable promise in efficient fine-tuning tasks for Large Language Models (LLMs). In particular, recent advances incorporate the low-rankness of gradients, introducing low-rank ZO estimators to further reduce GPU memory consumption. However, most existing works focus solely on the low-rankness of each individual gradient, overlooking a broader property shared by all gradients throughout the training, i.e., all gradients approximately reside within a similar subspace. In this paper, we consider two factors together and propose a novel low-rank ZO estimator, TeZO, which captures the low-rankness across both the model and temporal dimension. Specifically, we represent ZO perturbations along the temporal dimension as a 3D tensor and employ Canonical Polyadic Decomposition (CPD) to extract each low-rank 2D matrix, significantly reducing the training cost. TeZO can also be easily extended to the Adam variant while consuming less memory than MeZO-SGD, and requiring about only 35% memory of MeZO-Adam. Both comprehensive theoretical analysis and extensive experimental research have validated its efficiency, achieving SOTA-comparable results with lower overhead of time and memory.