引言

像 Llama-2 和 RoBERTa 这样的大型语言模型 (LLM) 的快速演进彻底改变了自然语言处理领域。然而，将这些庞大的模型适配到特定任务 (即微调过程) 面临着巨大的计算障碍。随着模型规模激增至数十亿参数，通过标准方法训练它们所需的 GPU 显存变得极其昂贵。

罪魁祸首就是反向传播 。 在传统的一阶 (FO) 优化 (如 SGD 或 AdamW) 中，训练过程需要存储每一层的中间激活值来计算梯度链式法则。对于一个数十亿参数的模型，这道“显存墙”往往需要企业级 GPU 集群才能逾越。

最近，研究人员转向了零阶 (ZO) 优化方法，如 MeZO，试图绕过这堵墙。ZO 方法仅利用前向传播来估计梯度，完全消除了存储反向传播图的需求。这可以将显存使用量减少高达 8 倍。然而，天下没有免费的午餐: ZO 方法以不稳定著称。它们受困于高方差、收敛缓慢，并且在应用于超大模型时经常发散 (即无法学习) 。

AdaZeta 登场。

在一篇题为 AdaZeta: Adaptive Zeroth-Order Tensor-Train Adaptation for Memory-Efficient Large Language Models Fine-Tuning 的新论文中，来自加州大学圣塔芭芭拉分校和 Amazon AGI 的研究人员提出了一个新的框架，旨在稳定 ZO 微调。通过结合超低参数的张量化适配器 (tensorized adapters) 和智能的自适应梯度估计调度，AdaZeta 在实现 ZO 方法显存优势的同时，达到甚至超越了传统微调方法的性能。

在这次深度解读中，我们将探索 AdaZeta 的机制，了解它如何驯服显存高效训练中的不稳定性。

背景: 高效微调的挑战

要理解 AdaZeta，我们首先需要了解高效 LLM 训练的现状以及当前方法的不足之处。

参数高效微调 (PEFT)

标准微调会更新模型中的每一个权重。PEFT 方法，像 LoRA (低秩自适应) , 通过冻结预训练模型并仅训练注入的小型适配器层来工作。虽然这减少了可训练参数的数量，但标准实现仍然依赖反向传播，这意味着显存的节省受限于存储计算图的需求。

零阶 (ZO) 优化

ZO 优化将神经网络视为一个“黑盒”。它不是通过微积分 (反向传播) 计算精确梯度，而是进行估计。

想象你在黑暗中身处一座崎岖的山上 (损失景观) 。你想找到底部 (最小损失) 。

  • 一阶 (反向传播) : 你有地图和手电筒。你计算出脚下确切的坡度，然后向下迈出一步。
  • 零阶: 你看不到坡度。你向随机方向迈出一小步，看看是向上还是向下，退回来，然后利用这个信息来猜测坡度。

这种“猜测”的数学公式通常通过随机梯度估计 (RGE) 来实现。虽然显存效率高，但这种猜测游戏引入了噪声 (方差) 。随着维度 (参数) 的增加，这些猜测的准确性会急剧下降，导致模型难以学习。

张量列 (Tensor-Train) 分解

为了对抗高维问题，AdaZeta 利用了张量列 (TT) 分解 。 这是一种数学技术，用于将大矩阵表示为一系列较小的低秩张量。

Illustration for tensorized linear layer and tensorized adapters.

如上图 Figure 2 所示:

  • (a) 张量化线性层: 一个巨大的权重矩阵 \(\mathbf{W}\) 被分解为一连串较小的张量 \(\mathcal{G}\)。
  • (b) 张量化适配器的结构: 这些轻量级的张量层被插入到 Transformer 块 (前馈和注意力机制) 中。

通过使用 TT 分解，可训练参数的数量大幅下降——甚至比 LoRA 还要少。这对于 ZO 优化至关重要，因为在低维空间中估计梯度要准确得多。

AdaZeta 框架

该论文的核心贡献是 AdaZeta 框架，它解决了以前 ZO 方法的两个主要痛点: 因高维度导致的估计精度差以及因梯度方差导致的训练发散。

该框架包含两大支柱:

  1. 快速前向张量化适配器 (用于降维) 。
  2. 自适应查询调度 (用于减少方差) 。

1. 快速前向张量化适配器

标准 ZO 方法之所以举步维艰，是因为它们试图同时估计数百万个参数的梯度。AdaZeta 通过使用张量化适配器压缩可训练空间来缓解这一问题。

在标准线性层中，权重存储在矩阵 \(\mathbf{W}\) 中。在 Tensor-Train 格式中，该矩阵被重塑为张量并进行分解。然而，计算 TT 层的输出涉及缩并 (相乘) 一系列张量，如果按顺序进行，速度可能会很慢。

由于 ZO 优化每一步梯度更新都需要两次前向传播，推理速度至关重要。AdaZeta 引入了一种并行缩并方法 。 它不再是一个接一个地相乘张量 (\(1 \to 2 \to 3 \dots\)) ，而是将张量因子分组以便同时处理。

重塑和缩并过程可以数学表示为:

Equation for weight reshaping

在这里，\(\mathcal{G}\) 代表张量因子。通过分组 (例如，分离索引 \(1\) 到 \(o\) 和 \(o+1\) 到 \(2o\)) ，该框架避免了高维中间张量，并加快了前向传播速度。这使得 AdaZeta 能够使用超压缩适配器，而不会在训练期间遭受速度惩罚。

2. 自适应查询调度

第二个，或许也是最具影响力的创新，是 AdaZeta 处理梯度估计的方式。

在 ZO 优化中，“查询 (Query) ”指的是一次扰动权重并检查损失的过程。

  • 1 次查询: 扰动一次权重，检查损失，估计梯度。噪声大，速度快。
  • N 次查询: 用不同的随机种子扰动权重 \(N\) 次，取结果的平均值。噪声小，速度慢。

以前的方法如 MeZO 通常使用固定数量的查询 (通常仅为 1) 。研究人员发现，对于大模型 (如 Llama-2-7B) ，单次查询噪声太大，导致发散 (模型停止学习，损失激增) 。相反，使用大量固定的查询 (例如 10 次) 会让训练慢得令人发指。

AdaZeta 提出了自适应查询调度 。 其核心思想是，在模型学习的早期阶段使用少量的查询，随着训练的进行，逐步增加查询数量以细化梯度估计。

该调度遵循基于训练 epoch \(e_k\) 的次线性增长:

Adaptive Query Equation

其中:

  • \(Q_k\) 是步骤 \(k\) 的查询数量。
  • \(\alpha\) 和 \(\beta\) 是缩放因子 (0 到 1 之间) 。
  • \(Q_{max}\) 是上限，防止过程变得太慢。

AdaZeta 算法

让我们看看它是如何在训练循环中发挥作用的。算法在每次更新时执行以下步骤:

  1. 确定查询数量: 根据当前 epoch 计算 \(Q_k\)。
  2. 估计梯度: 循环 \(Q_k\) 次以累积梯度估计。

在循环内部，对于每个查询 \(q\)，算法执行经典的两点估计:

Forward and Backward Perturbation Loop

在上述片段中:

  • \(z_q\) 是一个随机高斯扰动向量。
  • \(\ell_+^q\) 是权重偏移 \(+\epsilon z_q\) 后的损失。
  • \(\ell_-^q\) 是权重偏移 \(-\epsilon z_q\) 后的损失。

最后，估计的梯度是这 \(Q_k\) 次扰动的平均值，用于更新权重 \(w\):

Gradient Averaging and Update

通过对多次查询取平均值，梯度估计的方差显著降低。通过自适应地增加 \(Q_k\)，AdaZeta 确保训练在开始时保持快速，并在需要时变得更加精确——从而防止了困扰标准 MeZO 的发散问题。

理论收敛性

这为什么有效？作者提供了理论分析，证明该方法保证收敛。

算法的收敛界推导如下:

Convergence Bound Equation

不纠结于符号细节，关键结论是: 不等式右边的项包含 \(\sum \frac{1}{Q_k}\)。

  • 如果 \(Q_k\) 是常数 (例如 1) ，误差项保持较大。
  • 如果 \(Q_k\) 随时间增加 (如 AdaZeta) ，总和会收缩，收紧界限并保证梯度趋近于零 (收敛) 。
  • 此外，项 \(C(d, \epsilon)\) 依赖于维度 \(d\)。因为张量化适配器保持 \(d\) 非常小，总体误差界限远低于将 ZO 应用于完整模型。

实验与结果

研究人员在 RoBERTa-Large (中等规模) 和 Llama-2-7B (大规模) 上进行了评估，涵盖了情感分析、推理和问答等各种任务。

1. 收敛稳定性

最引人注目的直观结果是 AdaZeta 与其他方法相比的稳定性。

Loss curves for SST-2, WiC, and CB tasks

在 Figure 1 中，观察 Llama-2-7B 的训练曲线:

  • 蓝/橙线 (MeZO-LoRA): 显现出明显的波动。注意损失是如何波动或几乎没有下降的。
  • 绿线 (Sparse-MeZO): 表现较好，但仍表现出方差。
  • 红线 (AdaZeta): 显示出平滑、陡峭且持续的评估损失下降，收敛点比所有基线都低。

这证实了自适应查询调度成功减轻了发散风险。

2. 大模型上的准确率

当在低数据资源 (SuperGLUE 和 SQuAD 数据集) 的 Llama-2-7B 上测试时，AdaZeta 显著优于基线。

Table 2: Llama-2-7B Results

Table 2 重点展示了:

  • AdaZeta vs. MeZO-LoRA: AdaZeta 在几乎所有任务中都取得了更高的分数。例如，在 RTE 任务上，AdaZeta 得分为 74.0 , 而 MeZO-LoRA 仅为 59.6 。
  • AdaZeta vs. 一阶 (FT) : 值得注意的是，在几个任务中 (如 COPA 和 CB) ，AdaZeta 甚至超越了标准的一阶微调 (FT) ，尽管它使用的显存只有后者的一小部分。这表明张量适配器的正则化效果结合 ZO 噪声实际上可能有助于低数据环境下的泛化。

3. 显存效率

ZO 方法的首要目标是减少显存。AdaZeta 在这方面表现出色。

Figure 3: Accuracy vs Memory Trade-off

Figure 3 展示了这种权衡。理想情况下，你希望处于左上角 (高准确率，低显存) 。

  • 蓝色圆圈 (FT - 一阶): 高准确率，但显存使用量巨大 (最右侧) 。
  • 红/蓝三角形 (AdaZeta): 位于最左侧 (低显存) ，但在准确率轴上很高，超越了 Sparse-MeZO 和 MeZO-LoRA 等其他显存高效方法。

详细的显存分析揭示了节省的程度:

Table 10: Quantitative Memory Profiling

如 Table 10 所示，使用标准一阶方法在 SST-2 任务上微调 Llama-2-7B 需要 118.65 GB 的显存。AdaZeta 仅需 14.73 GB 。 这实现了 8 倍的显存减少 , 使得在单个消费级 GPU (如 RTX 3090 或 4090) 上微调 7B 模型成为可能。

4. 与 LoRA 的比较

一个常见的问题是: “为什么不直接使用 Batch Size 为 1 的 LoRA 呢？”

Table 4: Comparison with LoRA

Table 4 回答了这个问题。即使 Batch Size 为 1 且 Rank 为 1 (LoRA 的绝对最小值) ，标准一阶 LoRA 也需要 35.60 GB 。 Rank 为 8 的 AdaZeta 仅需 14.05 GB 。

因为一阶 LoRA 仍然需要存储反向传播所需的激活图 (即使图更小) ，它无法与 AdaZeta “仅前向传播”的本质相竞争。

结论

AdaZeta 框架代表了让大型语言模型微调变得触手可及的重要一步。通过解决高维度和梯度估计方差这两个双重挑战，它将零阶优化从一种不稳定的尝试转变为一种强大、实用的工具。

核心要点:

  1. 张量化适配器大幅减少了搜索空间，使梯度估计更容易、更准确。
  2. 自适应查询调度消除了困扰大规模 ZO 训练的发散问题，提供了平滑的收敛，且没有巨大的速度惩罚。
  3. 效率: AdaZeta 使得在显存小于 16GB 的硬件上微调 7B+ 参数模型成为可能，且性能与全显存方法相当或更好。

对于硬件资源有限的学生和研究人员来说，AdaZeta 为如何在适配庞大的基础模型时不仅仅是更努力地工作，而是更聪明地工作提供了蓝图。