[论文解读] On the Expressive Power of Self-Attention Matrices

Valerii Likhosherstov, Krzysztof Choromański+1 more|arXiv (Cornell University)|Jun 7, 2021

Gene expression and cancer classification参考文献 40被引用 5

一句话总结

本文证明，通过仅修改输入，固定自注意力机制可以近似任意稀疏注意力矩阵——其中每行和每列的非零元素数量有界——且隐藏维度 $ d $ 随序列长度 $ L $ 对数增长。关键结果是基于随机投影和 Johnson-Lindenstrauss 引理的构造性证明，表明 $ d = O(\log L) $ 足够实现对输入相关稀疏模式的高精度近似。

ABSTRACT

Transformer networks are able to capture patterns in data coming from many domains (text, images, videos, proteins, etc.) with little or no change to architecture components. We perform a theoretical analysis of the core component responsible for signal propagation between elements, i.e. the self-attention matrix. In practice, this matrix typically exhibits two properties: (1) it is sparse, meaning that each token only attends to a small subset of other tokens; and (2) it changes dynamically depending on the input to the module. With these considerations in mind, we ask the following question: Can a fixed self-attention module approximate arbitrary sparse patterns depending on the input? How small is the hidden size $d$ required for such approximation? We make progress in answering this question and show that the self-attention matrix can provably approximate sparse matrices, where sparsity is in terms of a bounded number of nonzero elements in each row and column. While the parameters of self-attention are fixed, various sparse matrices can be approximated by only modifying the inputs. Our proof is based on the random projection technique and uses the seminal Johnson-Lindenstrauss lemma. Our proof is constructive, enabling us to propose an algorithm for finding adaptive inputs and fixed self-attention parameters in order to approximate a given matrix. In particular, we show that, in order to approximate any sparse matrix up to a given precision defined in terms of preserving matrix element ratios, $d$ grows only logarithmically with the sequence length $L$ (i.e. $d = O(\log L)$).

研究动机与目标

• 理解变换器中自注意力矩阵的表达能力，特别是在稀疏性和动态输入依赖等实际约束下的表现。
• 确定固定自注意力模块为近似任意稀疏注意力模式所需最小隐藏维度 $ d $。
• 建立近似精度与序列长度依赖关系的理论边界，重点关注输入相关的稀疏模式。
• 提供一种构造性算法，用于寻找输入和固定注意力参数，以生成期望的稀疏注意力矩阵。

提出的方法

• 使用随机投影技术将输入标记嵌入到低维空间，以便近似注意力模式。
• 应用 Johnson-Lindenstrauss 引理以在投影下保持标记表示之间的成对距离，确保注意力权重比值的稳定性。
• 从期望的注意力模式构造目标稀疏矩阵 $ B $，其每行和每列的非零条目数量有界。
• 将注意力矩阵变换到对数空间以控制其大小和稀疏性，从而实现浓度界限。
• 对所有相关标记对使用并集界，以限制近似误差超过阈值 $ \epsilon_2 $ 的概率。
• 推导出 $ d $ 的充分条件，使得成功近似的概率超过 $ 1 - \delta $，从而得出 $ d = O(\log L) $。

实验结果

研究问题

• RQ1 通过输入变化，隐藏维度较小的固定自注意力模块（$ d \ll L $）能否近似任意稀疏注意力矩阵？
• RQ2 为在给定精度下近似任意稀疏注意力矩阵，所需最小隐藏维度 $ d $ 是多少？
• RQ3 在近似输入相关的稀疏模式时，所需 $ d $ 如何随序列长度 $ L $ 变化？
• RQ4 能否实现构造性近似，即提供输入与参数选择的显式算法？
• RQ5 Johnson-Lindenstrauss 引理是否能实现注意力矩阵在低维空间中的稳定近似，同时保持元素级比值？

主要发现

• 固定自注意力机制可利用 $ d = O(\log L) $ 近似任意稀疏注意力矩阵，其中 $ d $ 为隐藏维度，$ L $ 为序列长度。
• 近似结果在高概率 $ \geq 1 - \delta $ 下保持注意力权重比值，精度由用户定义的 $ \epsilon_2 $ 控制。
• 所需 $ d $ 随 $ L $ 对数增长，显著优于线性或多项式缩放。
• 该证明具有构造性，可生成算法以计算输入嵌入和固定注意力参数，从而生成期望的稀疏注意力矩阵。
• 该方法依赖随机投影与 Johnson-Lindenstrauss 类型的浓度界限，确保低维嵌入空间中的稳定性和准确性。
• 实验验证表明，当 $ k=2 $ 且 $ L $ 达到数百时，$ d $ 在 300–1000 范围内已足够，与理论预测一致。