无需正交假设,可从黑盒中准确恢复多头注意力参数。
Provably Learning Multi-Head Attention with Queries
- 通过调整令牌数量构造有理函数样本,分离出等价的注意力头。
- 在精确输出下仅需 $4Hd^2-2H+1$ 次查询即可恢复所有头参数。
- 适用于已知头数或上界的情况,适合逆向工程Transformer模型。
我们研究从黑盒输入输出接口学习多头softmax注意力的问题。学习者可查询任意实值标记序列,并仅观察最终标记的标量输出。已有工作提出一种使用 $O(d^2)$ 值查询恢复单头参数 $(W,v)$ 的算法。对于多头情形,该工作在各头占据两两正交子空间的假设下建立了可识别性;若要分别应用单头恢复算法,则需已知这些子空间的基。本文通过合并相同 $W_h$ 的头、累加对应 $v_h$ 并在总和为零时丢弃,获得无需子空间假设的规范表示。通过改变令牌副本数,算法生成有理函数样本,其插值可分离出规范头。进一步通过添加特定标记向量的查询,实现跨查询的同头匹配。当预言机输出与后续计算均精确时,随机选择查询向量,以概率一恢复所有规范对 $iglackslash{(W_h,v_h):h\in[H]\bigr\}$(至排列)。若 $H$ 已知,仅需 $4Hd^2-2H+1$ 次最大长度为 $2H+1$ 的值查询;若仅有上界 $H_0$,则使用 $4H_0d^2-2H_0+1$ 次查询,最大长度 $2H_0+1$。对于近似预言机输出,给出条件使参数误差不超过模型与查询相关常数乘以输出误差。最后,将结果扩展至含多头注意力与无偏置ReLU前馈网络的一层Transformer,附加条件下可恢复功能等价的Transformer,无需单独学习前馈网络。
原文摘要 · Abstract (English)
We study the problem of learning multi-head softmax attention from black-box input-output access. The learner may query arbitrary real-valued token sequences and observe only the scalar output at the final token. Recent work gives an algorithm using $O(d^2)$ value queries to recover the single-head parameters $(W,v)$. For multiple heads, the same work establishes identifiability under the assumption that the heads occupy pairwise orthogonal subspaces. Applying the single-head recovery algorithm separately to the heads additionally requires bases for these subspaces to be known. We recover a canonical representation by merging heads with the same $W_h$, summing their corresponding $v_h$, and discarding a merged head when this sum is zero, without these subspace assumptions. By varying the number of copies of a token, our algorithm obtains samples of a rational function whose interpolation separates the canonical heads. Additional queries formed by adding selected token vectors then match the same head across different queries. When the oracle outputs and all subsequent computations are exact, the learner chooses its query vectors at random and recovers the canonical pairs $\{(W_h,v_h):h\in[H]\}$ up to permutation with probability one. When $H$ is known, it uses exactly $4Hd^2-2H+1$ value queries of maximum length $2H+1$. If only a known upper bound $H_0$ is available, the algorithm uses $4H_0d^2-2H_0+1$ value queries of maximum length $2H_0+1$. For approximate oracle outputs, we give conditions under which the parameter error is at most a model- and query-dependent constant multiple of the output error. Finally, we extend our result to a one-layer Transformer with multi-head attention followed by a bias-free ReLU feed-forward network. Under additional conditions, we recover a functionally equivalent Transformer without relying on a separate algorithm for learning the feed-forward network.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。