[论文解读] Can Transformers Learn to Solve Problems Recursively?

Shizhuo Dylan Zhang, Curt Tigges+3 more|arXiv (Cornell University)|May 24, 2023

Adversarial Robustness in Machine Learning被引用 5

一句话总结

本文研究了小型变压器模型是否能通过在输入输出示例上进行训练，学习模仿结构递归函数——这类函数在程序合成和形式化验证中至关重要。利用机械可解释性方法，作者重构了模型所学习的“捷径”算法，成功预测了高达91%的失败案例，揭示出变压器模型通过有缺陷的、与深度相关的启发式方法近似递归，而非真正的语义泛化。

ABSTRACT

Neural networks have in recent years shown promise for helping software engineers write programs and even formally verify them. While semantic information plays a crucial part in these processes, it remains unclear to what degree popular neural architectures like transformers are capable of modeling that information. This paper examines the behavior of neural networks learning algorithms relevant to programs and formal verification proofs through the lens of mechanistic interpretability, focusing in particular on structural recursion. Structural recursion is at the heart of tasks on which symbolic tools currently outperform neural models, like inferring semantic relations between datatypes and emulating program behavior. We evaluate the ability of transformer models to learn to emulate the behavior of structurally recursive functions from input-output examples. Our evaluation includes empirical and conceptual analyses of the limitations and capabilities of transformer models in approximating these functions, as well as reconstructions of the ``shortcut" algorithms the model learns. By reconstructing these algorithms, we are able to correctly predict 91 percent of failure cases for one of the approximated functions. Our work provides a new foundation for understanding the behavior of neural networks that fail to solve the very tasks they are trained for.

研究动机与目标

• 研究变压器模型是否能够从输入输出示例中学习模仿结构递归函数，这类函数是程序合成和形式化验证中的核心程序类别。
• 分析并理解变压器模型为何在明确为其训练的任务上失败，特别是在建模递归语义方面。
• 利用机械可解释性技术逆向工程变压器模型所学习的“捷径”算法。
• 理解注意力机制和模型深度在塑造这些学习到的次优行为中的作用。
• 为改进在递归和语义推理任务上的神经模型训练、提示工程与评估提供基础。

提出的方法

• 从零开始训练小型变压器模型，任务为涉及结构递归的合成任务，包括二元后继和树遍历。
• 使用抽象状态机（ASMs）作为形式化框架，以建模和分析变压器的计算行为。
• 应用机械可解释性技术，如反事实补丁、电路追踪和注意力头分析，以逆向工程所学习的算法。
• 追踪各层和注意力头中的注意力模式，识别模型如何处理递归子结构、括号信号和节点复制。
• 使用 UNROLL 和 EMPTY 标记作为符号信号，以追踪模型在递归约简和子树复制过程中的行为。
• 通过分析注意力依赖关系和关键点（如括号闭合）处的行为变化，重构“捷径”算法。

(a) Decoder self-attention. Reverse Order.

实验结果

研究问题

• RQ1 尽管缺乏对递归结构的显式监督，小型变压器模型是否能从输入输出示例中学习模仿结构递归函数？
• RQ2 变压器模型学习了哪些类型的“捷径”算法，而非真正的递归泛化？这些算法在何处失败？
• RQ3 变压器中的注意力机制在推理过程中如何编码和操作递归子结构？
• RQ4 机械可解释性技术在多大程度上能够预测在递归任务上训练的变压器模型的失败案例？
• RQ5 模型深度和超参数选择在多大程度上影响递归推理中深度特定启发式的出现？

主要发现

• 模型能够成功近似结构递归函数，但其方式是通过学习到的捷径，而非真正的递归泛化。
• 对于某一函数，高达91%的失败案例可通过重构模型所学习的捷径算法被正确预测。
• 早期层中的注意力机制表现出任务特定的模式：一个注意力头追踪前向符号（如括号），另一个则线性关注编码后的标记。
• 对于更深的树结构，模型依赖解码器自注意力机制来追踪父节点，并使用 UNROLL 标记作为递归组合的信号。
• 模型学习到与深度相关的技巧——例如在子树展开后复制根节点——这些技巧在更复杂或更深的递归结构中会失效。
• 不同的超参数设置会导致不同的学习算法，表明即使在简单的递归任务上，模型行为也对训练设置高度敏感。

(b) Encoder self-attention. Reverse Order