可证明地学习多头注意力层
Provably learning a multi-head attention layer
AI总结：
本文首次从随机样本中可证明地学习多头注意力层，给出$(dk)^{O(m^3)}$时间算法及指数下界，方法基于凸体雕刻，适用于离散和连续输入。
摘要：
多头注意力层是Transformer架构中区别于传统前馈模型的关键组件之一。给定序列长度$k$，注意力矩阵$\boldsymbol{\Theta}_1,\ldots,\boldsymbol{\Theta}_m\in\mathbb{R}^{d\times d}$和投影矩阵$\mathbf{W}_1,\ldots,\mathbf{W}_m\in\mathbb{R}^{d\times d}$，对应的多头注意力层$F: \mathbb{R}^{k\times d}\to \mathbb{R}^{k\times d}$通过$F(\mathbf{X}) \triangleq \sum^m_{i=1} \mathrm{softmax}(\mathbf{X}\boldsymbol{\Theta}_i\mathbf{X}^\top)\mathbf{X}\mathbf{W}_i$将长度为$k$的$d$维词元序列$\mathbf{X}\in\mathbb{R}^{k\times d}$进行变换。在这项工作中，我们开创了从随机样本中可证明地学习多头注意力层的研究，并给出了该问题的首个非平凡上下界：\n- 在$\{\mathbf{W}_i, \boldsymbol{\Theta}_i\}$满足某些非退化条件的前提下，我们给出一个$(dk)^{O(m^3)}$时间复杂度的算法，该算法能从均匀取自$\{\pm 1\}^{k\times d}$的随机标记样本中学习$F$到小误差。\n- 我们证明了计算下界，表明在最坏情况下，对$m$的指数依赖是不可避免的。\n我们专注于布尔型$\mathbf{X}$以模拟大型语言模型中词元的离散性质，尽管我们的技术自然可以扩展到标准连续设置，例如高斯分布。我们的算法以利用样本雕刻一个包含未知参数的凸体为核心，这与现有的可证明学习前馈网络的算法有显著不同，后者主要利用高斯分布的代数和旋转不变性。相比之下，我们的分析更加灵活，因为它主要依赖于输入分布及其“切片”的各种上尾和下尾界。
英文摘要：
The multi-head attention layer is one of the key components of the transformer architecture that sets it apart from traditional feed-forward models. Given a sequence length $k$, attention matrices $\mathbfΘ_1,\ldots,\mathbfΘ_m\in\mathbb{R}^{d\times d}$, and projection matrices $\mathbf{W}_1,\ldots,\mathbf{W}_m\in\mathbb{R}^{d\times d}$, the corresponding multi-head attention layer $F: \mathbb{R}^{k\times d}\to \mathbb{R}^{k\times d}$ transforms length-$k$ sequences of $d$-dimensional tokens $\mathbf{X}\in\mathbb{R}^{k\times d}$ via $F(\mathbf{X}) \triangleq \sum^m_{i=1} \mathrm{softmax}(\mathbf{X}\mathbfΘ_i\mathbf{X}^\top)\mathbf{X}\mathbf{W}_i$. In this work, we initiate the study of provably learning a multi-head attention layer from random examples and give the first nontrivial upper and lower bounds for this problem: - Provided $\{\mathbf{W}_i, \mathbfΘ_i\}$ satisfy certain non-degeneracy conditions, we give a $(dk)^{O(m^3)}$-time algorithm that learns $F$ to small error given random labeled examples drawn uniformly from $\{\pm 1\}^{k\times d}$. - We prove computational lower bounds showing that in the worst case, exponential dependence on $m$ is unavoidable. We focus on Boolean $\mathbf{X}$ to mimic the discrete nature of tokens in large language models, though our techniques naturally extend to standard continuous settings, e.g. Gaussian. Our algorithm, which is centered around using examples to sculpt a convex body containing the unknown parameters, is a significant departure from existing provable algorithms for learning feedforward networks, which predominantly exploit algebraic and rotation invariance properties of the Gaussian distribution. In contrast, our analysis is more flexible as it primarily relies on various upper and lower tail bounds for the input distribution and "slices" thereof.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
