[论文解读] Thinking Like Transformers

Gail Garfinkel Weiss, Yoav Goldberg+1 more|arXiv (Cornell University)|Jun 13, 2021

Advanced Neural Network Applications被引用 11

一句话总结

本文提出了 RASP（受限访问序列处理语言），一种符号编程语言，通过将注意力和前馈操作抽象为一等原语，建模 Transformer 编码器的计算行为。该语言支持对 Transformer 能力的正式推理，表明 RASP 程序可预测解决排序、直方图计算和 Dyck 语言识别等任务所需的最小层数与头数，且当模型匹配 RASP 衍生的注意力模式时，训练后的 Transformer 准确率超过 99%。

ABSTRACT

What is the computational model behind a Transformer? Where recurrent neural networks have direct parallels in finite state machines, allowing clear discussion and thought around architecture variants or trained models, Transformers have no such familiar parallel. In this paper we aim to change that, proposing a computational model for the transformer-encoder in the form of a programming language. We map the basic components of a transformer-encoder -- attention and feed-forward computation -- into simple primitives, around which we form a programming language: the Restricted Access Sequence Processing Language (RASP). We show how RASP can be used to program solutions to tasks that could conceivably be learned by a Transformer, and how a Transformer can be trained to mimic a RASP solution. In particular, we provide RASP programs for histograms, sorting, and Dyck-languages. We further use our model to relate their difficulty in terms of the number of required layers and attention heads: analyzing a RASP program implies a maximum number of heads and layers necessary to encode a task in a transformer. Finally, we see how insights gained from our abstraction might be used to explain phenomena seen in recent works.

研究动机与目标

• 为 Transformer 编码器提供一个形式化的计算模型，类似于有限自动机对 RNN 的作用。
• 通过在领域特定语言中将任务表达为符号程序，实现对 Transformer 能力的高层级推理。
• 预测在 Transformer 中实现特定序列处理任务所需的最小层数和注意力头数。
• 验证训练后的 Transformer 是否能学习与 RASP 衍生解完全匹配的注意力模式，从而建立符号推理与神经网络行为之间的联系。
• 利用 RASP 抽象解释 Transformer 变体中的经验现象，并揭示高效架构的局限性。

提出的方法

• 设计 RASP 为一种最小化编程语言，包含对序列进行选择、过滤和聚合的原语，以捕捉注意力和前馈的约束。
• 将 RASP 程序编译为等效的注意力和 MLP 操作，从而定义 Transformer 中所需的注意力模式和层结构。
• 使用双重监督训练 Transformer：在输出上使用标准交叉熵损失，同时在注意力模式上使用 MSE 损失以匹配 RASP 解。
• 利用 RASP 程序推导给定任务所需层数和头数的理论下界。
• 通过在 RASP 预测的最小尺寸下训练 Transformer 并测量减少头数或层数时的准确率下降，分析 RASP 程序的可行性。
• 将 RASP 应用于重新表达并改进先前针对 Dyck-k 语言识别的构造，证明 Transformer 可对所有 k 实现完全识别。

实验结果

研究问题

• RQ1 能否构建一种符号编程语言，以类似于有限自动机对 RNN 的方式，建模 Transformer 编码器的计算行为？
• RQ2 一个符号程序所预测的、在 Transformer 中实现给定序列处理任务所需的最小层数和注意力头数是多少？
• RQ3 训练后的 Transformer 是否能学习与给定任务的 RASP 程序所推导出的模式完全匹配的注意力模式？
• RQ4 当模型规模减小（层数或头数减少）时，性能如何变化？RASP 能否预测这些失效点？
• RQ5 RASP 是否可用于解释或预测 Transformer 变体中的经验行为，例如高效架构中的性能下降？

主要发现

• 针对双直方图、排序和最常见标记等任务的 RASP 程序，成功编译为注意力模式，训练后的 Transformer 在测试中实现了超过 99% 的准确率。
• 在输出和注意力监督双重训练下，Transformer 在双直方图任务上达到 99.9% 的准确率，在排序任务上达到 99.8%，与 RASP 指定的注意力模式高度一致。
• 当模型规模低于 RASP 预测的最小值时，准确率显著下降——例如，减少一层后反向任务准确率降至 41.2%，表明 RASP 对架构限制具有强大的预测能力。
• RASP 抽象使 Dyck-k 语言识别的构造更加简洁且正确，优于先前工作，证明了 Transformer 可对所有 k 实现完全识别。
• 对于排序任务，单层单头的 Transformer 由于均匀注意力模式，实现了接近完美的准确率，表明其有效实现了某种形式的桶排序。
• 本研究发现，具有受限注意力机制的‘高效’Transformer 可能在需要复杂信息流的任务上失败，这与 RASP 分析的预测一致。