LoRA 与权重衰减 (2023)

# irhum.github.io - LoRA 与权重衰减 来源：https://irhum.github.io/blog/lorawd/ LoRA（Hu 等人，2021 (https://irhum.github.io/blog/lorawd/#ref-hu2021lora)）是目前流行的替代大型语言模型（LLM）全量微调的方法：我们不调整整个模型的数十亿权重，而是添加小型的“适配器”权重矩阵来*修改*原始权重矩阵，并只调整这些适配器。 本篇博文深入探讨一个有趣的行为：虽然 LoRA 通常被视为全量微调的即插即用替代方案，但它与权重衰减的相互作用意味着它求解的*优化问题*与全量微调不同。具体来说，在 LoRA 中，解权重被正则化到冻结的基模型\\\(\(W \\rightarrow W\_\{\\text\{init\}\}\)\\\)，而不是像全量微调中那样\\\(W \\rightarrow 0\\\)。 这意味着，即使资源越来越多（甚至等同于全量微调的资源），LoRA 也不会*越来越*接近全量微调，因为其目标函数与全量微调隐含不同。根据使用场景，这既可以视为**缺陷或特性**，但从业者应明确考虑这一点。 ## 回顾：微调 对于 LLM，我们通常微调一个初始模型（在广泛的文本到文本任务上“表现良好”），以提升在*特定*目标任务（例如从自然语言生成数据库查询）上的性能。我们通过两步完成： - 首先，创建一个微调训练数据集\\\(\{\(x\_i, y\_i\)\_n\}\\\), 包含输入\\\(x\\\)和目标\\\(y\\\)的配对。1 (https://irhum.github.io/blog/lorawd/#fn1) - 优化初始模型的权重，使得我们的微调训练数据集\\\(\{\(x\_i, y\_i\)\_n\}\\\)变得更“可能”。这里的想法是：一个能更大概率在训练集上对\\\(x\\\)生成正确答案\\\(y\\\)的模型，也会泛化到在*新*的\\\(x\\\)上更可能生成\\\(y\\\)。 ### 全量微调 全量微调意味着调整模型中*所有*的权重。对于像 GPT-3 175B（Brown 等人，2020 (https://irhum.github.io/blog/lorawd/#ref-brown2020language)）这样的模型，这意味着给我们的优化算法 1750 亿个可以按需上下“调节”的数字，以使微调训练数据更“可能”。让我们深入一点，更具体地定义这里的权重含义。 Transformer 中的每一层主要由两个组件组成：一个多头注意力网络，后跟一个前馈网络。这意味着构成每一层的大部分“权重”存储在六个矩阵中2 (https://irhum.github.io/blog/lorawd/#fn2)，如图所示。然后，\\\(\\theta\\\) 用作所有存储在模型所有层所有矩阵中的权重的简写。 在全量微调中，\\\(\\theta\\\) 中的每一个权重都可以进行更新。我们的目标是生成更新后的权重，使左侧所示的负对数似然（NLL）最小化3 (https://irhum.github.io/blog/lorawd/#fn3)。没有封闭形式的方法来获得“最优”权重，因此我们通过反复应用许多步梯度下降来求解优化问题，如右侧所示。 现在，直接这样进行梯度下降会很快导致过拟合4 (https://irhum.github.io/blog/lorawd/#fn4)，因此我们通常对问题进行正则化。对于 LLM，通常选择的正则化工具是权重衰减。具体来说，当使用普通 SGD5 (https://irhum.github.io/blog/lorawd/#fn5) 时，权重衰减等价于在损失中添加一项，即权重的平方和： \\\[R\(\\theta\)=\\sum\_i \\sum\_j\[W\_\{\{\\color\{RoyalBlue\}q\}\}^\{\\color\{PineGreen\}\{1\}\}\]\_\{ij\}^2\+\\cdots\\\] 因此，总目标如下（其中\\\(\\lambda\\\)是控制权重衰减强度的超参数）： \\\[\\min\_\{\\color\{YellowOrange\}\{\\theta\}\} \\biggl\[\\underbrace\{\-\\log P\_\{\\color\{YellowOrange\}\{\\theta\}\}\(\{\\color\{PineGreen\}\{y\}\} \\mid \{\\color\{RoyalBlue\}\{x\}\}\)\}\_\{\\color\{BrickRed\}\{L\}\} \+ \\frac\{\\lambda\}\{2\} R\(\{\\color\{YellowOrange\}\{\\theta\}\}\)\\biggr\]\\\] 对该目标求导得到梯度，我们注意到梯度更新有两个不同的项6 (https://irhum.github.io/blog/lorawd/#fn6)：第一项对应之前的负对数似然最小化，第二项\\\(\-\\alpha\\lambda w\\\)将权重推向原点\\\(0\\\)。 \\\[ % https://tex\.stackexchange\.com/a/9477 \\def\\mathunderline\#1\#2\{\\color\{\#1\}\\underline\{\{\\color\{black\}\#2\}\}\\color\{black\}\} \\begin\{align\*\} &\{\\color\{YellowOrange\}\{w\}\} \\leftarrow \{\\color\{YellowOrange\}\{w\}\} \- \\alpha \\left\(\\mathunderline\{BrickRed\}\{\\frac\{\\partial \\color\{BrickRed\}\{L\}\}\{\\partial \\color\{YellowOrange\}\{w\}\}\} \+ \\mathunderline\{LimeGreen\}\{\\frac\{\\lambda\}\{2\} \\frac\{\\partial R\}\{\\partial \\color\{YellowOrange\}\{w\}\}\} \\right\)\\\\ \\Rightarrow &\{\\color\{YellowOrange\}\{w\}\} \\leftarrow \{\\color\{YellowOrange\}\{w\}\} \- \\alpha \\left\(\\mathunderline\{BrickRed\}\{\\frac\{\\partial \\color\{BrickRed\}\{L\}\}\{\\partial \\color\{YellowOrange\}\{w\}\}\} \+ \\mathunderline\{LimeGreen\}\{\\lambda \{\\color\{YellowOrange\}\{w\}\}\} \\right\)\\\\ \\Rightarrow &\{\\color\{YellowOrange\}\{w\}\} \\leftarrow \{\\color\{YellowOrange\}\{w\}\} \- \\alpha \\mathunderline\{BrickRed\}\{\\frac\{\\partial \\color\{BrickRed\}\{L\}\}\{\\partial \\color\{YellowOrange\}\{w\}\}\} \- \\alpha \\mathunderline\{LimeGreen\}\{\\lambda \{\\color\{YellowOrange\}\{w\}\}\} \\end\{align\*\}\\\] 这意味着正则化后的问题变成了： 总结：在损失中添加权重的平方和，等价于在每个梯度下降步中减去每个权重的缩放版本。这会将最小值移向权重更接近\\\(0\\\)的区域7 (https://irhum.github.io/blog/lorawd/#fn7)；即没有一个权重可以对模型预测产生极大影响。 全量微调高度灵活，但也*极其*占用内存：通常至少需要模型本身所需内存的 3 倍8 (https://irhum.github.io/blog/lorawd/#fn8)来存储梯度和优化器状态。这在模型参数量为\\\(O\(100M\)\\\)时不是问题，但在今天经常达到\\\(O\(10B\)\\\)到\\\(O\(100B\)\\\)参数时肯定如此。此外，如果你的应用中有 10 个子任务（需要为每个任务微调模型），全量微调要求你托管 10 个模型版本（好像托管一个还不够贵似的！）。 ### LoRA 微调 LoRA（低秩适配器）微调采用不同方法：不是直接调整 LLM 的巨大权重矩阵，而是对每个要调整的权重矩阵使用一对小型适配器矩阵，形式如下： 也就是说，对于每个初始冻结权重\\\(W\_\{\\text\{init\}\}\\\), 我们有适配器矩阵\\\(A\\\)和\\\(B\\\)。这两个矩阵相乘得到\\\(\\Delta W\\\)，它是\\\(W\_\{\\text\{init\}\}\\\)的一个低秩“调整”矩阵，形成调整后的矩阵\\\(W\\\)。这显著减少了自由参数的数量：假设原始矩阵\\\(W\_\{\\text\{init\}\}\\\)是\\\(4,096 \\times 16,384\\\)。在原始方法中，仅这一个权重矩阵就需要调整 6700 万个参数，如下所示： \\\[4,096 \\times 16,384 = 67,108,864 \\approx 67 \\text\{ million\}\\\] 使用秩为\\\(r=4\\\)的 LoRA，我们只有： \\\[4,096 \\times 4 \+ 4 \\times 16,384 = 81,920\\\] 这不到原始参数数量的 0.1%；存储这些值的 3 个变体（权重、梯度和优化器状态）的额外开销与模型本身使用的内存相比微乎其微。 此外，由于初始权重在所有微调运行中是“共享”的，在推理时我们只需加载一份初始模型，供多个微调版本共享，每个任务使用自己的任务特定适配器矩阵进行推理。这使得在应用中拥有“每个任务”微调的 LLM 不仅可行，而且容易。 ## 相互作用 既然我们已经介绍了 LoRA 是什么，就可以开始讨论它与权重衰减的相互作用如何产生一个特性/缺陷。由于\\\(A\\\)和\\\(B\\\)是我们实际进行梯度下降的矩阵，目标中的权重衰减项如下所示，它将最小值移向适配器矩阵更接近 0 的区域： \\\[R\(\\theta\)=\\sum\_i \\sum\_j\[A\_\{\{\\color\{RoyalBlue\}q\}\}^\{\\color\{PineGreen\}\{1\}\}\]\_\{ij\}^2\+ \\sum\_i \\sum\_j\[B\_\{\{\\color\{RoyalBlue\}q\}\}^\{\\color\{PineGreen\}\{1\}\}\]\_\{ij\}^2\+ \\cdots\\\] 让我们将其与全量微调中的公式进行对比： - 在全量微调中，我们有\\\(W \\rightarrow 0\\\)，即权重直接衰减到 0。 - 然而，在 LoRA 中，由于\\\(A\\\)和\\\(B\\\)衰减到 0，实际上我们有\\\(W \\rightarrow W\_\{\\text\{init\}\}\\\)。 这意味着 LoRA 的解偏向于原始冻结的权重矩阵，而全量微调中则偏向于零。并且这种行为不会随着 LoRA 秩\\\(r\\\)的增加而消失 —— 你可以将其增加到无穷大（！），优化过程仍然偏向于原始冻结权重而不是零。也就是说，即使在极限情况下，LoRA 也不会逼近全量微调，而是逼近一个不同的目标。 ### 一种修复方法 如果我们希望整个调整后的矩阵趋向于零（如同全量微调中那样），我们需要一个正则化项，其中*整个调整后的*权重矩阵趋向于零，如下所示： \\\[\\begin\{align\*\} R\(\\theta\)&=\\sum\_i \\sum\_j\[W\_\{\{\\color\{RoyalBlue\}q\}\}^\{\\color\{PineGreen\}\{1\}\}\]\_\{ij\}^2\+\\cdots\\\\ &=\\sum\_i \\sum\_j\[W\_\{\{\\color\{RoyalBlue\}q\\color\{Black\}\\text\{,init\}\}\}^\{\\color\{PineGreen\}\{1\}\} \+ A\_\{\{\\color\{RoyalBlue\}q\}\}^\{\\color\{PineGreen\}\{1\}\}B\_\{\{\\color\{RoyalBlue\}q\}\}^\{\\color\{PineGreen\}\{1\}\}\]\_\{ij\}^2\+\\cdots \\end\{align\*\}\\\] 这实际上很容易推导，并产生一对更新方程，可以实现为类似标准权重衰减的形式。首先，从权重衰减的核心定义开始，即计算权重关于正则化项的梯度： \\\[\{\\color\{YellowOrange\}\{w\}\} \\leftarrow \{\\color\{YellowOrange\}\{w\}\} \- \\alpha \\left\(\\frac\{\\partial \\color\{BrickRed\}\{L\}\}\{\\partial \\color\{YellowOrange\}\{w\}\} \+ \\frac\{\\lambda\}\{2\} \\frac\{\\partial R\}\{\\partial \\color\{YellowOrange\}\{w\}\} \\right\)\\\] 第二，计算\\\(A\\\)和\\\(B\\\)关于上述“修正”\\\(R\(\\theta\)\\\)的梯度9 (https://irhum.github.io/blog/lorawd/#fn9)。得到： \\\[\\begin\{align\*\} \\frac\{\\partial R\}\{\\partial \\color\{YellowOrange\}\{A\}\}&=2 \(W\_\{\\text\{init\}\} \+ \{\\color\{YellowOrange\}\{A\}\}\{\\color\{PineGreen\}\{B\}\}\) \{\\color\{PineGreen\}\{B^T\}\}\\\\ \\frac\{\\partial R\}\{\\partial \\color\{YellowOrange\}\{B\}\}&=2 \{\\color\{PineGreen\}\{A^T\}\}\(W\_\{\\text\{init\}\} \+ \{\\color\{PineGreen\}\{A\}\}\{\\color\{YellowOrange\}\{B\}\}\) \\end\{align\*\}\\\] 将其插入权重衰减的定义中，得到 \\\(A\\\)和\\\(B\\\)的具体更新方程： \\\[\\begin\{align\*\} \{\\color\{YellowOrange\}\{A\}\} &\\leftarrow \{\\color\{YellowOrange\}\{A\}\} \- \\alpha \\frac\{\\partial \\color\{BrickRed\}\{L\}\}\{\\partial \\color\{YellowOrange\}\{A\}\} \- \\alpha \\lambda \(W\_\{\\text\{init\}\} \+ \{\\color\{YellowOrange\}\{A\}\}\{\\color\{PineGreen\}\{B\}\}\) \{\\color\{PineGreen\}\{B^T\}\}\\\\ \{\\color\{YellowOrange\}\{B\}\} &\\leftarrow \{\\color\{YellowOrange\}\{B\}\} \- \\alpha \\frac\{\\partial \\color\{BrickRed\}\{L\}\}\{\\partial \\color\{YellowOrange\}\{B\}\} \- \\alpha \\lambda \{\\color\{PineGreen\}\{A^T\}\}\(W\_\{\\text\{init\}\} \+ \{\\color\{PineGreen\}\{A\}\}\{\\color\{YellowOrange\}\{B\}\}\) \\end\{align\*\}\\\] #### 代码实现 以下是 Optax（Babuschkin 等人，2020 (https://irhum.github.io/blog/lorawd/#ref-deepmind2020jax)）库中标准权重衰减公式的代码。它非常简洁：将参数`p`的`weight_decay`（\\\(\\lambda\\\)）缩放版本添加到其当前更新`g`10 (https://irhum.github.io/blog/lorawd/#fn10)。 `` # from https://github.com/google-deepmind/optax/blob/master/optax/_src/transform.py#L766 def update_fn(updates, state, params): if params is None: raise ValueError(base.NO_PARAMS_MSG) updates = jax.tree_util.tree_map( lambda g, p: g + weight_decay * p, updates, params) return updates, state `` 为了将其修改为实现我们刚才描述的数学公式，需要一些额外的代码，主要是提取`W_init`、`A`和`B`矩阵11 (https://irhum.github.io/blog/lorawd/#fn11)。核心逻辑只有第 18 行和第 20 行。 `` def update_fn(updates, state, params): def per_param_update_fn(path, update, param): # 获取整个层的参数字典。 param_name = path[-1].key # 如果当前参数是适配器矩阵。 if param_name in ['kernelA', 'kernelB']: layer_params = params for dict_key in path[:-1]: layer_params = layer_params[dict_key.key] # 提取初始权重矩阵和适配器矩阵。 W_init = layer_params['kernel'] A = layer_params['kernelA'] B = layer_params['kernelB'] # 计算修正后的衰减项。 if param_name == 'kernelA': decay_term = (W_init + A@B)@B.T else: decay_term = A.T@(W_init + A@B) # 如果当前参数*不是*适配器矩阵，使用默认的权重衰减版本。 else: decay_term = param return update + weight_decay * decay_term if params is None: raise ValueError(base.NO_PARAMS_MSG) updates = jax.tree_util.tree_map_with_path( per_param_update_fn, updates, params) return updates, state `` ## 结论 总结一下，LoRA 有一个不同于全量微调的隐含目标，但如果需要，也很容易纠正。就是这样，真的！ 据我所知，目前没有文献深入记载 LoRA 与权重衰减的相互作用。纯粹基于第一性原理的推测12 (https://irhum.github.io/blog/lorawd/#fn12)，我认为默认行为既是特性也*是*缺陷，取决于数据量 —— 当训练点非常少时，它是特性，*因为*它正则化更新后的模型，使其保持接近初始的“通用能力”模型。然而，当数据量很大时，它是*缺陷*，因为优化过程不易偏离基权重太远，*即使*这有助于最终任务性能。 话虽如此，尽管数学上很简洁，但经验结果才是唯一的真理。考虑到如此多的自由参数，在实践中可能确实存在与全量微调（正则化接近\\\(0\\\))同样好的解（当正则化接近\\\(W\_\{\\text\{init\}\}\\\)并赋予足够容量时）。 ## 附录 A：动量和权重衰减 你可能注意到的一个奇怪之处是，我花了很多时间显式推导正则化项\\\(R\(\\theta\)\\\)的梯度，而不是直接将其吸收到\\\(L\\\)中并让自动微分代劳。这是因为等价性（梯度权重衰减 = 添加\\\(L\_2\\\)正则化项