[论文解读] What can a Single Attention Layer Learn? A Study Through the Random Features Lens

Hengyu Fu, Tianyu Guo+2 more|arXiv (Cornell University)|Jul 21, 2023

Domain Adaptation and Few-Shot LearningComputer Science被引用 3

一句话总结

本文研究了在随机特征设定下，单个多头注意力层的表征能力和泛化性能，其中查询和键矩阵随机初始化并固定，而值矩阵可训练。结果表明，此类模型能够高效学习键向量的置换不变函数，且在样本复杂度上优于标准两层网络，尤其当查询-键矩阵具有非零均值时，能更好地近似基于相关性的函数。

ABSTRACT

Attention layers -- which map a sequence of inputs to a sequence of outputs -- are core building blocks of the Transformer architecture which has achieved significant breakthroughs in modern artificial intelligence. This paper presents a rigorous theoretical study on the learning and generalization of a single multi-head attention layer, with a sequence of key vectors and a separate query vector as input. We consider the random feature setting where the attention layer has a large number of heads, with randomly sampled frozen query and key matrices, and trainable value matrices. We show that such a random-feature attention layer can express a broad class of target functions that are permutation invariant to the key vectors. We further provide quantitative excess risk bounds for learning these target functions from finite samples, using random feature attention with finitely many heads. Our results feature several implications unique to the attention structure compared with existing random features theory for neural networks, such as (1) Advantages in the sample complexity over standard two-layer random-feature networks; (2) Concrete and natural classes of functions that can be learned efficiently by a random-feature attention layer; and (3) The effect of the sampling distribution of the query-key weight matrix (the product of the query and key matrix), where Gaussian random weights with a non-zero mean result in better sample complexities over the zero-mean counterpart for learning certain natural target functions. Experiments on simulated data corroborate our theoretical findings and further illustrate the interplay between the sample size and the complexity of the target function.

研究动机与目标

• 理解单个多头注意力层在良性样本复杂度下可学习的函数类。
• 使用有限样本和有限但较大的头数，分析随机特征注意力（RFA）模型的泛化性能。
• 研究查询-键矩阵初始化方式（特别是非零均值分布）对相关性函数学习效率的影响。
• 将RFA的样本复杂度与标准两层随机特征网络（RFMLP）进行比较，识别注意力机制在哪些场景下具有优势。
• 通过模拟数据实验验证理论发现，考察样本量与目标函数复杂度之间的相互作用。

提出的方法

• 该研究建立了一个单注意力层机制模型，包含单个查询向量和N个键向量，采用随机冻结的查询和键矩阵，以及可训练的值矩阵。
• 提出一种随机特征注意力（RFA）模型，其中注意力机制基于查询-键点积计算的注意力分数，对值向量进行加权求和。
• 推导出形式为 $\widetilde{\mathcal{O}}(\sqrt{B(f_\star)/n})$ 的过剩风险界，其中 $B(f_\star)$ 衡量目标函数 $f_\star$ 的复杂度，$n$ 为样本数量。
• 提出一种有偏RFA变体（BRFA），其查询-键矩阵乘积具有非零均值（例如单位矩阵），以模拟实际注意力机制的初始化模式。
• 理论分析表明，对于依赖于查询与键向量之间相关性的函数，BRFA在样本复杂度上优于标准RFA，当 $d \gg p$ 时可节省 $\text{Poly}(d)$ 因子。
• 在模拟数据上的实验比较了RFA、BRFA与RFMLP在不同序列长度 $N$、输入维度 $d$ 和目标函数复杂度下的表现，验证了理论样本复杂度趋势。

实验结果

研究问题

• RQ1 单个多头注意力层在查询和键矩阵随机冻结的前提下，能否表达一类广泛的置换不变目标函数？
• RQ2 使用RFA模型学习此类目标函数的样本复杂度是多少？与标准两层随机特征网络相比如何？
• RQ3 查询-键矩阵乘积的初始化分布（特别是非零均值与零均值）如何影响相关性函数的学习效率？
• RQ4 理论样本复杂度界是否能在具有实际模型规模和目标函数的有限样本设置中得到验证？
• RQ5 在RFA和BRFA中，序列长度、输入维度与模型性能之间的经验相互作用关系如何？

主要发现

• RFA模型能够表达一类广泛的置换不变目标函数，这些函数是成对标记函数的平均值，包括多项式函数和基于相关性的函数。
• RFA的过剩风险界为 $\widetilde{\mathcal{O}}(\sqrt{B(f_\star)/n})$，其复杂度仅依赖于输入维度，而不依赖于键向量数量 $N$，因此在样本复杂度上优于RFMLP。
• 对于依赖于查询与键向量之间相关性的函数，具有非零均值查询-键矩阵的有偏RFA（BRFA）模型在样本复杂度上优于标准RFA，当 $d \gg p$ 时可节省 $\text{Poly}(d)$ 因子。
• 实验结果表明，RFA在学习对键排列不变的函数时优于RFMLP，尤其当序列长度 $N$ 增大时，验证了理论上的样本复杂度优势。
• BRFA在学习基于相关性的函数时显著优于RFA，且性能差距随相关性强度（由 $\gamma$ 控制）的增加而扩大，证实了非零均值初始化的理论优势。
• 分析发现，BERT的预训练查询-键矩阵呈现对角主导特征，表明现实世界中的注意力机制可能早已利用非零均值初始化的归纳偏差，与本文的理论发现一致。