引言: 效率困境

在大语言模型 (LLMs) 的世界里，规模已成为智能的代名词。从 GPT-3 到 Llama 3，趋势显而易见: 扩大参数规模能带来更优的性能。然而，这种“越大越好”的理念代价高昂。训练这些庞然大物需要天文数字般的计算资源、巨大的能源消耗以及漫长的时间。

研究人员长期以来一直在寻求打破这种依赖的方法。我们知道像 LoRA (低秩适应) 这样的技术在 微调 方面效果显著——允许我们使用极少的训练参数来调整巨大的模型。这自然引发了一个紧迫的问题: 如果低秩方法在微调中如此有效，为什么我们不能用它从头开始预训练模型?

理想情况下，我们可以设计一个天生就是低秩的模型，从第一天起就减少参数数量 (从而减少计算负载) 。然而，以前尝试这样做的努力大多失败了。在预训练期间直接对整个网络应用低秩分解，通常会导致模型“变笨”，难以有效地进行学习。

最近的一篇论文《Scalable Efficient Training of Large Language Models with Low-dimensional Projected Attention》 (基于低维投影注意力的可扩展大语言模型高效训练) 提供了一个突破性的解决方案。研究人员发现，低秩预训练本身并没有缺陷；我们只是把它用错了地方。通过针对 Transformer 架构的特定组件，他们提出了一种称为 低维投影注意力 (Low-dimensional Projected Attention, LPA) 的方法。这种方法不仅将训练速度提高了 12% 以上，而且令人惊讶的是，其实际表现甚至可以超越标准的 Transformer。

在这篇文章中，我们将剖析 LPA 的工作原理，以前尝试失败的原因，以及这对 AI 未来的可扩展性意味着什么。

背景: 低秩直觉

要理解 LPA，我们首先需要理解 低秩分解 的概念。

在神经网络内部，层本质上是巨大的数字矩阵 (权重) 。当我们说一个矩阵具有“满秩”时，意味着每一行和每一列都创造了一个独特的信息维度。然而，研究人员发现深度学习模型通常具有“低内在秩”。这意味着即使矩阵很大，它处理的实际信息也存在于一个小得多的子空间中。

LoRA 的联系

你可能熟悉 LoRA，这是一种流行的微调方法。LoRA 冻结了巨大的预训练权重，并添加了两个微小的矩阵 \(A\) 和 \(B\) 来表示变化。如果原始权重矩阵是 \(1000 \times 1000\) (100 万参数) ，LoRA 可能使用一个 \(1000 \times 8\) 的矩阵和一个 \(8 \times 1000\) 的矩阵 (\(16,000\) 参数) 来近似更新。变量 \(r\) (在此例中为 8) 即为 秩。

逻辑很简单: \(A \times B\) 生成一个与原始矩阵大小相同的矩阵，但它是由一个大小为 \(r\) 的“瓶颈”形成的。

预训练的挑战

如果这种瓶颈在微调期间能节省显存，为什么不将这些瓶颈永久地安装在模型中呢？

历史上，在预训练期间尝试用低秩模块 (\(W = W_A \times W_B\)) 替换 整个 Transformer 中的标准密集层，结果都导致了性能下降。模型根本无法捕捉数据的复杂性。LPA 背后的研究人员意识到，Transformer 架构并不是铁板一块；它有两个截然不同的部分，承担着非常不同的工作: 注意力机制 (Attention Mechanism) 和 前馈神经网络 (FFN) 。

他们的关键洞察是什么? 位置至关重要。

核心方法: 低维投影注意力 (LPA)

这篇论文的核心创新在于发现: 注意力层对低秩分解有很高的容忍度，而 FFN 层则不然。

基于此，他们提出了 LPA 架构。在这种设置中，注意力机制中的标准权重矩阵被低维模块替换，而 FFN 层则基本保持不变 (或区别对待) 。

架构可视化

让我们看看 LPA 如何重构标准的 Transformer 块。

图 1: 低维投影注意力 (LPA) 的示意图。Softmax 函数中的计算衡量了输入 Token 之间的关系。

如图 1 所示，输入 \(x\) 进入注意力机制。在标准的 Transformer 中，查询 (\(Q\)) 、键 (\(K\)) 和值 (\(V\)) 是通过将 \(x\) 乘以大的方阵 (\(W_Q, W_K, W_V\)) 生成的。

在 LPA 中，这些大矩阵被成对的较小矩阵替换。注意图中高亮的“Low Dimensional Space” (低维空间) 部分。输入首先被 向下 投影到一个较小的维度 (瓶颈) ，然后再 向上 投影到目标维度。这一过程发生在 \(Q\)、\(K\)、\(V\) 以及输出投影 (\(O\)) 中。

LPA 的数学原理

让我们通过数学公式看看具体在哪里节省了参数。

标准注意力: 在标准 Transformer 中，注意力输出 \(z\) 是使用满秩矩阵计算的。

标准注意力的方程。

这里，\(W_Q, W_K, W_V, W_O\) 是巨大的、参数密集的矩阵。

LPA 注意力: 在 LPA 中，这些矩阵中的每一个都被分解了。例如，\(W_Q\) 变成了 \(W_{Q1} \times W_{Q2}\)。

LPA 注意力的方程。

因为连接 \(W_{Q1}\) 和 \(W_{Q2}\) 的内部维度很小 (秩 \(r\)) ，所以总参数数量显著下降。如果输入维度是 \(d_{in}\) 且秩是 \(r\)，复杂性从 \(d_{in}^2\) 下降到 \(2 \cdot d_{in} \cdot r\)。由于 \(r\) 通常远小于 \(d_{in}\)，节省是巨大的。

为什么不压缩 FFN？

你可能会问，“如果这对注意力层有效，为什么不对前馈神经网络 (FFN) 也这样做呢？”FFN 通常包含模型的大部分参数，因此压缩它将产生更大的节省。

研究人员严格地测试了这个假设。他们比较了将低秩模块应用于:

  1. Low Attn: 仅注意力层 (LPA 方法) 。
  2. Low FFN: 仅 FFN 层。
  3. Low All: 注意力和 FFN 层。

结果很有说服力。将低秩模块应用于 FFN 层 (或所有层) 始终会损害模型的困惑度 (Perplexity，衡量模型困惑程度的指标——越低越好) 。然而，仅 将其应用于注意力层，通常会导致性能 优于 原始的标准 Transformer。

为什么会有这种差异? 作者基于这些层处理数据的方式提供了一个令人信服的理论解释。

1. FFN 独立处理 Token。 FFN 层对每个 Token 独立应用非线性变换。它将 Token 投影到高维空间以解开特征，然后再将它们压缩回来。

标准 FFN 的方程。

LPA FFN 的方程。

如果我们强迫 FFN 通过一个低维瓶颈 (如上面第二个方程所示) ，我们就削弱了它将单个 Token 映射到那个必要的高维特征空间的能力。

2. 注意力处理关系。 注意力层则根本不同。它计算 Token 之间 的关系。

显示注意力依赖于所有 Token 的方程。

如上式所示 (具体为论文中的引理 1) ，输出 \(z_i\) 取决于 \(x_i\) 与 所有其他 Token 之间的交互。Softmax 函数将这些关系归一化为概率分数。研究人员认为，这种关系映射对向量空间的绝对维度不太敏感。“低维投影注意力”作为一个两步投影，捕捉了必要的关系，而无需密集矩阵的完整参数量。

实验与结果

理论上的优雅固然很好，但在实践中真的有效吗？研究人员进行了广泛的实验，将他们的方法扩展到了 30 亿参数规模。

1. 预训练性能 (困惑度)

LLM 预训练的主要指标是留出测试数据 (held-out test data) 上的困惑度。作者将 LPA 与两个基线进行了比较:

  • Transformer (Same-Dim): 具有相同隐藏层大小 (但参数更多) 的标准模型。
  • Transformer (Same-Param): 缩小规模以具有与 LPA 模型相同参数数量的标准模型。

表 3: 所有 130M 和 370M 参数规模模型的测试困惑度。

表 3 显示了 130M 和 370M 参数模型的结果。

  • 观察: LPA 模型始终优于“同参数 (Same-Param) ”Transformer。
  • 关键发现: 在 370M 规模下，LPA 甚至优于“同维度 (Same-Dim) ”Transformer (12.89 vs 13.65) ，尽管它的参数更少。这表明注意力中的低秩结构可能起到了一种有益的正则化作用，防止过拟合或使模型专注于最重要的特征。

2. 可扩展性: 30 亿参数测试

任何效率方法的真正考验是可扩展性。在小模型上有效的方法在扩大规模时往往会失效。作者训练了一个 2.43B 参数的 LPA 模型，并将其与一个 3.23B 的标准 Transformer 进行了比较。

图 2: 2.43B LPA 模型、3.23B 同维度 Transformer 和与 LPA 模型参数量几乎相同的 2.49B Transformer 的训练损失。

图 2 中的训练损失曲线令人震惊。蓝绿色线 (LPA) 始终保持比蓝色线 (同参数 Transformer) 更低的损失，并且有效地匹配或击败了红色线 (同维度 Transformer) 。

这证实了 LPA 是可扩展的 。 它允许我们训练一个性能像 32 亿参数模型的模型，而实际上只有 24 亿参数模型的计算占用。

3. 效率收益

这对实际训练时间意味着什么？因为 LPA 减少了浮点运算 (FLOPs) 的数量，它直接转化为挂钟时间 (wall-clock time) 的加速。

论文报告称，对于 3B 规模的模型，LPA 节省了大约 12.4% 的训练时间 并减少了 GPU 显存消耗。虽然这看起来并不多，但当训练运行成本高达数百万美元时，12% 的减少是巨大的。此外，该模型创建的检查点更小，节省了存储空间。

4. 下游任务

低困惑度并不总是意味着更聪明的模型。为了确保模型真正理解语言，研究人员在 GLUE 基准测试上对其进行了测试，其中包括情感分析和问答等任务。

表 5: 预训练 LPA 和 Transformer 模型在 GLUE 基准上的测试结果。

如表 5 所示，LPA 模型获得了比标准 Transformer (67.47) 更高的平均分 (70.72) 。这验证了注意力层中的参数减少并没有损害模型的推理或理解上下文的能力；事实上，它似乎增强了这种能力。

5. 对秩 (\(r\)) 的敏感度

这种方法的一个关键超参数是 \(r\)——即瓶颈的大小。如果 \(r\) 太小，信息就会丢失。如果 \(r\) 太大，我们就无法节省任何参数。

图 3: 不同 r 值的 Transformer 和 LPA 模型的训练损失。

图 3 展示了不同 \(r\) 值下的训练损失。

  • 鲁棒性: 模型具有惊人的鲁棒性。即使将 \(r\) 降低到 64 或 128 (深色线) ，其性能也优于基线 Transformer。
  • 极限: 只有当 \(r\) 降至 32 时，性能才会显著下降，这表明我们可以压缩注意力机制的程度存在一个下限。

6. 分配盈余参数

由于 LPA 节省了参数，研究人员问道: “如果我们把这些节省下来的参数重新投入到其他地方会怎样？”

他们尝试通过增加 FFN 大小或添加更多层来加回参数。然而，最有效的策略是 增加注意力维度 (输出大小) 。通过保持低秩结构但使投影空间稍微宽一点，他们实现了效率和性能的最佳平衡。

表 8: 通过参数重新分配获得的变体模型的测试困惑度。

表 8 显示，“Attn Dim.”策略 (将参数重新分配给注意力宽度) 产生了最低的困惑度 (12.85) ，进一步证明了注意力层是最有效的优化位置。

结论与启示

这篇题为《基于低维投影注意力的可扩展大语言模型高效训练》的研究挑战了预训练需要满秩密集矩阵的假设。

主要收获:

  1. 精准是关键: 盲目应用低秩预训练会失败。当专门针对 注意力层 时，它才有效。
  2. LPA 架构: 通过用低维 (\(W_A, W_B\)) 模块替换标准注意力矩阵，LPA 减少了参数数量和计算成本。
  3. 卓越的性能: LPA 不仅能匹敌标准 Transformer；它在困惑度和下游任务上通常表现更好。
  4. 可扩展: 该方法在数十亿参数规模上仍然有效，为更廉价地训练大型模型提供了一条可行的路径。

这项工作为 LLM 的民主化打开了令人兴奋的大门。通过降低预训练有效模型的硬件门槛，LPA 表明我们尚未触及架构效率的天花板。展望未来，将 LPA 与其他效率技术 (如量化或稀疏注意力) 相结合，可能会催生新一代功能强大、轻量级的 AI 模型。