​耶鲁团队揭示多头自注意力结构的上下文学习机制，证明梯度流算法的收敛性

来源：DeepTech深科技

众所周知，目前的大模型大多基于 Transformer 架构。Transformer 的核心结构是多头自注意力模型（multi-head self-attention model）。

大模型的一个重要能力是所谓的“上下文学习”。具体来说，当大模型的参数训练好之后，用户和大模型的交互方式，是通过提供上文来获得大模型的下文，这时大模型的参数是固定的。

当所提供的上文包含一些关于同一主题的输入输出例子时，大模型可以根据给的这些例子，学到这些例子背后的主题，从而可以在给到一个新输入时，回答正确的输出。

比如，上文可以是：

（来源：资料图）

这时，Claude 3 大模型的回答是：

（来源：资料图）

由图可知，Claude 根据这些例子意识到“+”其实代表着减法，故能针对“10+5=？”这一新问题给出正确回答。

上下文学习，是大模型的一个基础能力。使用大模型时的其他更复杂方式比如 Chain-of-thought reasoning，都是以此为基础。

但是，从原理来看上下文学习的机制并不是很清楚。很大原因在于大模型作为一个系统，它不仅非常复杂，而且模型参数非常多，训练数据也非常大。

为了更好地理解上下文学习，美国斯坦福大学团队曾在 GPT-2 架构之下，针对大模型如何使用上下文，学习解决简单的回归预测进行了研究。

其发现当使用简单函数的数据来训练大模型时，训练好的大模型，可以通过上下文学到这些简单函数。

一个特别的例子便是线性函数。这时的训练数据是一些线性数据 x_1，w x_1，…，x_n，w x_n，其中 w 是高斯随机向量。

换句话说，每个“句子”里都有 n 个线性函数的例子，而这个线性函数是随机的。

（来源：https://arxiv.org/pdf/2208.01066.pdf）

以此为启发，美国耶鲁大学助理教授杨卓然和团队，希望可以从理论上研究这种训练过程是否收敛、以及收敛到哪里，也希望厘清多头自注意力结构到底是如何实现上下文学习的。

随后，他和所在团队考虑了一个最简单的模型：一层多头自注意力模型。

（来源：arXiv）

具体来说在本次课题之中，他们研究了训练多头自注意力模型（multi-head self-attention model）的优化问题。

尤其是，他们回答了这样一个问题：在使用一层多头自注意力模型（one-layer multi-head self-attention model）进行上下文学习时：

首先，梯度优化算法是否能够收敛？

其次，梯度优化算法收敛到的解统计性质如何？

再次，从网络结构的角度看，多头自注意力模型是如何进行上下文学习的？

期间，他们所使用的训练数据是多任务线性模型（multi-task linear regression）。

特别地，每个线性模型的参数 G 在一个固定的正交基下，有一个分块对角的分解。

也就是说如果能找到这组基，这个线性模型就可以分解成 H 个独立的线性模型。

对于每个参数 G，能够生成 L 个（x，y）对，并且可以让 Transformer 推测一个随机的 q 所对应的 y 是什么。

（来源：arXiv）

在这种多任务线性数据上，课题组使用梯度流来训练 Transformer，进而研究这一算法的收敛问题。

通过此，他们发现：梯度流算法的确是收敛的。并且收敛有三个阶段——（a）预热阶段、（b）任务分配阶段 、以及（c）最终收敛阶段。

在（a）这一预热阶段，损失函数缓慢下降。

在（b）这一“任务分配阶段”，损失函数迅速下降。并且，softmax 函数使得每一个自注意力头只关注多任务线性模型的一个任务，该团队把这一现象称为“任务分配”。

在（c）最终收敛阶段，每个自注意力头继续对它被分配的任务求解，最终达到收敛。

（来源：arXiv）

此外，他们还描述了梯度流学习的极限模型的上下文学习预测误差。

当 （d/L） 趋于零时，误差衰减到零，其中 d 是线性模型的维数，L 是上下文学习中（x，y）例子的数量。

并且，该团队还证明多头自注意力模型，显著好于单头自注意力模型。所预测的误差相差 H 倍之多，其中 H 是注意力头的个数。

换句话说，注意力头的个数越多，性能差距越大。

据介绍，该团队的分析主要基于对自注意力权重的分解。

自注意力机制里主要有两类权重：QK 权重（query-key）和 OV 权重（output-value）。其中，QK 权重反应着 query 和 key 的关系。

简单来说，就是给定了 query q（新的输入）和过去的例子（x，y）的关系。

而 QK 权重反应着 attention 对每一个过去的例子的重视程度。

OV 权重反应着输出和每一个输入例子（x，y）的关系，即 attention 如何通过组合上下文学习中的例子从而得到输出。

需要注意的是在回归问题里面，q 是一个输入，和 x 有一样的维度，输出和 y 有一样的维度。

课题组发现，QK 权重和 OV 权重都是分块的，并且 QK 权重的 X-X 分块和 OV 权重的 Y 分块最为重要。

也就是说，在回归问题里只需使用 q 和例子里的 x 比较得到注意力值（attention score）。

在输出时，只需要根据注意力值（attention score）来合并例子里的那些 y。

而通过利用数据的线性结构，他们发现 QK 权重和 OV 权重的分块结构，可以被梯度流算法保持。

更特别的是，因为多任务线性模型的参数 G 可以在某个基下分解，凭借此他们证明 QK 权重和 OV 权重也是可以被分解的。

这样一来，就可以把参数的梯度流化，简为奇异值的梯度流，这时就只需要分析奇异值的变化。

其中，总共有 H*（d_x + d_y）个奇异值，每个头的 QK 权重有 d_x 个奇异值，OV 权重有 d_y 个奇异值。

而 d_x 是 x 的维度，d_y 是 y 的维度，也就是多任务线性模型的任务数量。

并且，每个自注意力头的 OV 权重的奇异值，反应着自注意力头对于对应任务的重视程度。

随后，该团队开始分析这些奇异值的变化。他们发现自注意力头的任务分配基于“OV 权重–每个任务”的原则，来将最大的自注意力头分给对应的奇异值。

比如，第一个任务被分配给了第一个奇异值最大的自注意力头。

最终，在梯度流达到收敛之后，根据任务分配机制，每个自注意力头的 OV 权重只有唯一一个非零的奇异值。

（来源：arXiv）

陈思宇和王天浩分别是第一作者和第三作者，杨卓然担任通讯作者。其中，王天浩将于 2025 年秋入职美国加州大学圣地亚哥分校。

图 | 杨卓然（来源：杨卓然）

不过，课题组仍然觉得自己对于 transformer 和上下文学习的理解还非常粗浅。

目前，他们只研究了一层自注意力模型。后续，他们希望能够研究多层的自注意力模型。

与此同时，目前他们只研究了线性模型。因此，他们也非常希望研究非线性的上下文学习问题。

此外，目前课题组给到 transformer 的输入，是独立同分布的（x，y）输入输出数据对，这里输入并没有任何复杂的前后依赖结构。

但是，实际用来训练 transformer 的数据都是文本数据，里面有复杂的依赖结构，针对此他们也将继续加以探索。

参考资料：

1.https://arxiv.org/pdf/2403.00993

排版：罗以