(论文速读)IPA-GNN:用指令指针注意图神经网络学习执行程序
论文题目Learning to Execute Programs with Instruction Pointer Attention Graph Neural Networks用指令指针注意图神经网络学习执行程序会议NIPS2020摘要图形神经网络(GNN)已经成为学习软件工程任务的强大工具包括代码完成、错误查找和程序修复。它们受益于利用控制流图之类的程序结构但它们不太适合像程序执行这样的任务这些任务需要比GNN传播步骤数量多得多的顺序推理步骤。另一方面递归神经网络(RNN)非常适合于长顺序的推理链但它们并不自然地融入程序结构在上述任务中表现通常较差。我们的目标是两全其美我们通过引入一种新的GNN体系结构来实现这一点指令指针注意图神经网络(IPA-GNN)实现了对使用控制流图学习执行程序的任务的改进的系统泛化。该模型是将运行在带有分支决策的程序踪迹上的RNN视为潜在变量而产生的。IPA-GNN既可以被视为RNN模型的持续放松也可以被视为更适合执行的GNN变体。为了测试这些模型我们建议使用控制流图来评估学习执行的系统泛化以测试顺序推理和程序结构的使用。更实际的是我们在学习执行部分程序的任务中评估这些模型如果将该模型用作程序综合中的启发式函数可能会出现这种情况。结果表明IPA-GNN在这两个任务上的表现都优于各种RNN和GNN基线。一、论文要解决什么问题1.1 普通 GNN 能看到程序结构但难以完成长程执行推理程序可以表示成控制流图Control Flow GraphCFG每一条语句对应一个节点控制流边表示执行完当前语句后可能跳转到的下一条语句。GNN 很适合利用这种图结构但普通消息传递 GNN 存在一个直接限制经过 T 层传播后一个节点通常只能整合 T 跳范围内的信息。真实程序的执行轨迹可能包含循环执行步数可以远大于源代码行数也可能远大于模型的图传播层数。因此“拥有控制流图”并不等于“能够按照程序语义执行控制流图”。〖论文 Figure 1一个程序的逐行表示、四元组分词结果以及语句级控制流图。该图适合作为控制流图和程序表示的入门示意图〗1.2 RNN 能按顺序推理但不能自然处理分支和循环结构最简单的方法是把程序逐行输入 RNN。对于没有分支的直线程序这种做法与解释器的执行过程比较接近模型读一条语句更新一次隐藏状态再继续处理下一条语句。但一旦程序出现 if、else、while、break 或 continue源代码中的下一行不一定是实际执行的下一条语句。逐行 RNN 会处理没有被执行的分支也难以显式利用程序的控制流结构。另一种方法是让 RNN 沿真实执行轨迹处理程序也就是 Trace RNN。但这需要提前知道真实执行轨迹相当于使用了一个“轨迹预言机”。在静态分析场景中模型不能调用解释器因此真实轨迹并不可用。1.3 真正的问题不是拟合短程序而是能否系统性泛化到更复杂程序如果训练集和测试集中的程序复杂度相同模型可能只是在记忆局部模式。论文采用更严格的设置训练时只看长度不超过 10 的程序测试时直接处理长度为 20、30、……、100 的程序。这种设置考察的是系统性泛化模型是否学到了可以组合、重复使用的程序语义而不是只在训练分布内完成模式匹配。1.4 部分程序执行程序缺一条语句时还能否预测原程序输出论文还提出部分程序执行任务。作者从完整程序中随机选取一条非控制流表达式并将其替换为[MASK]但监督目标仍然是原始完整程序正确执行后的输出。这个任务对应程序合成中的一个实际需求在搜索完整程序之前需要对中间的部分程序进行评价从而判断它是否值得继续扩展。需要注意的是部分程序本身通常不能唯一决定原程序输出。因此该任务并不是严格意义上的“执行一个语义完整的程序”而是结合训练分布对被遮盖语句及最终输出进行预测。二、论文的核心创新是什么2.1 用解释器的因果结构设计神经网络经典解释器在每个执行步骤中维护两类信息程序状态例如变量当前取值指令指针即下一条需要执行的语句。执行当前语句后解释器更新程序状态如果当前语句包含条件判断还要决定沿控制流图的哪一条边前进。IPA-GNN 的主要设计原则不是简单地把注意力机制加到图上而是让神经网络的信息流尽量对应解释器的因果过程执行语句、作出分支决策、移动指令指针、聚合新的程序状态。2.2 从离散指令指针到软指令指针对于逐行 RNN隐藏状态更新为其中是第 t 步的指令指针。如果直接学习离散分支Hard IP-RNN 可以写成但 argmax 不可微无法直接使用标准反向传播进行端到端训练。IPA-GNN 将离散指令指针替换为软指令指针。表示在第 t 个计算步骤时模型位于语句 n 的概率或注意力质量。这样模型不必立即选择唯一分支而是可以把概率质量分配给多个后继节点。〖论文 Figure 2Line-by-Line RNN、Trace RNN、Hard IP-RNN、IPA-GNN 和 GGNN 的信息传播路径对比。该图重点展示离散分支与连续软分支的区别〗2.3 每个语句节点都维护一个条件化程序状态IPA-GNN 不只维护一个全局隐藏状态而是在每个时间步、每个语句节点上维护隐藏状态。可以理解为假设模型在第 t 步执行到语句 n此时程序状态应该是什么。对于每个节点模型先使用 RNN 单元执行当前语句产生状态提案这里的 RNN 在实验中采用两层 LSTM。2.4 学习软分支决策当语句 n 只有一个后继节点时分支概率固定为 1。当语句 n 有两个可能后继节点和时模型根据当前状态提案产生软分支概率表示模型在第 t 步位于语句 n 时下一步沿控制流边转移到语句 n 的概率。2.5 沿控制流图同时传播状态和指令指针某个前驱节点对当前节点的贡献由三部分共同决定前驱节点当前拥有多少指令指针概率质量前驱节点选择当前节点作为后继的分支概率前驱节点执行语句后产生的状态提案。新的节点状态为软指令指针按相同的控制流权重传播因此IPA-GNN 的一层计算可以拆成三个阶段Execute每个节点执行自己的语句并产生状态提案Branch根据状态提案预测分支概率Aggregate按照软指令指针和分支概率聚合到后继节点。〖论文 Figure 3单层 IPA-GNN 的 Execute、Branch、Aggregate 三阶段流程图。建议放在公式之后帮助理解状态和概率质量如何在图上流动〗2.6 IPA-GNN 同时连接了 RNN 和 GNN 两条模型路线论文给出了 IPA-GNN 与几类 RNN 的关系当软分支概率饱和为 one-hot 时IPA-GNN 等价于 Hard IP-RNN当 Hard IP-RNN 的分支判断全部正确时它等价于使用真实轨迹的 Trace RNN当程序不存在分支时Trace RNN 又退化为逐行 RNN。从另一个方向看IPA-GNN 也属于消息传递 GNN。它与 GGNN 的关键区别有两点IPA-GNN 使用节点级 RNN 模拟“执行一条语句”IPA-GNN 使用软指令指针和分支概率模拟“解释器的控制流”。论文通过替换这两个组件构造了 NoControl 和 NoExecute 两个消融模型NoControl 保留语句执行模块但移除 IPA-GNN 的指令指针控制机制NoExecute 保留软指令指针控制机制但移除节点级语句执行模块两者都替换后得到普通 GGNN。〖论文 Table 1IPA-GNN、NoControl、NoExecute 和 GGNN 的计算组件对照。该表适合用于说明两个消融模型究竟删除了什么〗三、任务和数据集如何构造3.1 程序语言范围论文使用概率文法生成一个受限的 Python 子集程序中包含变量赋值多位数初始值加法、减法和乘法if-else 条件分支while 循环break、continue 和 pass控制流结构的嵌套。变量名限制为 v0 到 v9常数和条件范围也受到限制。〖论文 Figure 6生成数据所使用的概率程序文法包括 Program、Block、Statement、Condition 和 Expression 等产生式〗3.2 训练集和测试集论文用程序长度作为复杂度 c(x)阈值 C10。训练集包含 500 万个程序所有程序长度均不超过 10。测试集包含 4500 个程序长度分别为 20、30、40、50、60、70、80、90 和 100每个长度有 500 个样本。这意味着所有测试程序都比训练期间见过的程序更长。3.3 完整程序执行任务输入是完整程序及其控制流图目标是预测程序执行结束后 v0 的值模 1000选择模 1000 的目标是为了把实验重点放在“控制流和长执行轨迹的泛化”上而不是同时测试任意精度数值计算能力。3.4 部分程序执行任务作者从每个完整程序中随机选择一条非控制流表达式将其替换为[MASK]。模型仍然需要预测原始完整程序执行后的 v0 mod 1000。该任务比完整程序执行更难因为模型既要推断控制流还要处理缺失语句带来的不确定性。3.5 有限计算预算IPA-GNN 并没有获得足够的传播层数去完整模拟大多数程序的真实执行轨迹。作者只保证每个循环体至少有机会传播两次。这个设置迫使模型学习“执行捷径”它可能不逐步复现每次循环而是通过学习程序模式在更少计算步数内预测最终结果。这也是论文结果中最值得注意、同时最需要谨慎解释的部分IPA-GNN 学到的是对任务有效的可微分执行近似不一定是逐步忠实的解释器模拟。四、实验设置4.1 对比模型论文使用以下模型进行比较Line-by-Line RNN按照源代码行顺序处理程序Trace RNN按照真实执行轨迹处理程序需要轨迹预言机仅用于完整程序任务R-GAT关系图注意力网络NoControl移除 IPA-GNN 的控制流注意力机制NoExecute移除 IPA-GNN 的节点级执行机制GGNN门控图神经网络IPA-GNN本文模型。R-GAT 经过额外超参数调节后仍未达到其他模型的竞争水平因此没有进入论文的主要结果表。4.2 训练细节优化器Adam损失函数交叉熵batch size32训练 3 个 epoch隐藏维度 H 从 200 和 300 中选择学习率从 0.003、0.001、0.0003 和 0.0001 中选择RNN、Trace RNN 和 IPA-GNN 的基础循环单元为两层 LSTM最优超参数根据长度恰好为 10 的留出验证集准确率选择。五、主要实验结果〖论文 Table 2各模型在完整程序执行和部分程序执行测试集上的总体准确率。建议将该表放在本节开头〗5.1 总体准确率完整程序执行任务Trace RNNOracle66.4%Line-by-Line RNN32.0%NoControl28.1%NoExecute50.7%GGNN16.0%IPA-GNN62.1%。部分程序执行任务Line-by-Line RNN11.5%NoControl8.1%NoExecute20.7%GGNN5.7%IPA-GNN29.1%。在完整程序任务上IPA-GNN 比最强的非预言机基线 NoExecute 高 11.4 个百分点在部分程序任务上高 8.4 个百分点。IPA-GNN 的完整程序准确率比使用真实轨迹的 Trace RNN 低 4.3 个百分点但 Trace RNN 在推理时需要真实执行轨迹不符合静态分析设置因此不能作为可部署的同条件基线。5.2 随程序长度增加IPA-GNN 的下降速度更慢Figure 4 按程序长度给出了准确率曲线。在训练长度附近Line-by-Line RNN 与 IPA-GNN 的差距并不大但程序长度超过训练范围后各基线准确率快速下降IPA-GNN 保持了明显更好的稳定性。在完整程序任务中IPA-GNN 在部分程序长度上甚至超过了拥有真实控制流轨迹的 Trace RNN。论文将其归因于有限计算预算下学到的捷径Trace RNN 必须逐步处理真实轨迹而 IPA-GNN 可以学习压缩某些重复计算。〖论文 Figure 4(a)完整程序执行准确率随程序长度变化的曲线〗〖论文 Figure 4(b)部分程序执行准确率随程序长度变化的曲线〗5.3 指令指针注意力比节点级执行模块更关键消融实验中NoExecute 在完整任务上达到 50.7%在部分任务上达到 20.7%NoControl 仅达到 28.1% 和 8.1%。由于 NoExecute 保留软指令指针而 NoControl 移除了软指令指针这一结果说明模型能否把信息限制在与当前执行路径相关的控制流上是 IPA-GNN 性能提升的主要来源。节点级 RNN 执行模块仍然有价值因为完整 IPA-GNN 又进一步超过 NoExecute但从消融幅度看控制机制的贡献更大。5.4 模型学出的软分支经常接近离散决策Figure 5 可视化了不同输入下的。亮色位置表示某个时间步的指令指针概率集中在哪条语句上。作者观察到尽管模型使用连续概率训练很多分支决策最终接近 one-hot也就是概率集中在单一路径上。这说明连续松弛并没有让模型始终在所有路径上平均传播而是能够形成近似离散的执行路径。5.5 模型没有机械复现真实轨迹而是学会短路执行Figure 5 的第一个例子中模型主要关注真正影响结果的控制流路径。第二个例子包含一个循环真实执行轨迹需要访问循环体 7 次但 IPA-GNN 的注意力只进入循环体 1 次仍然预测出了正确结果。这说明 IPA-GNN 学到了一种“短路执行”或“压缩执行”它不要求每一步都对应解释器的一次真实执行而是在有限传播深度内提取足以预测最终输出的语义信息。〖论文 Figure 5两个程序在不同 v0 初始值下的软指令指针热力图展示近似离散分支和循环短路现象〗六、如何理解这篇论文的贡献6.1 真正的创新不是简单增加注意力而是加入解释器式归纳偏置IPA-GNN 的技术组件并不复杂RNN、softmax、图上的消息聚合都是成熟方法。论文的关键贡献在于重新组织这些组件使信息流对应程序执行的三个因果步骤执行语句、选择分支、更新指令指针。这种结构归纳偏置直接针对系统性泛化问题。模型不再把所有邻居一视同仁而是根据当前程序状态动态决定哪些控制流边应该传递信息。6.2 论文建立了 RNN 与 GNN 之间清晰的模型联系IPA-GNN 既可以被理解为 Hard IP-RNN 的连续松弛也可以被理解为针对程序执行重新设计的消息传递 GNN。这种双重解释提高了模型的可理解性它不是一个完全孤立的新结构而是在“沿执行轨迹递归计算”和“在控制流图上传播消息”之间建立了连续过渡。6.3 部分程序执行为程序合成启发式函数提供了一个原型任务在 Programming by Example 中搜索算法需要评价尚未补全的候选程序。部分程序执行模型可以作为启发式函数的一部分对候选程序可能产生的输出进行估计。不过论文只验证了合成数据上的输出预测没有直接将 IPA-GNN 接入真实程序合成系统因此这里仍然属于潜在应用而不是已经完成的工程验证。七、论文的局限性与结果边界7.1 实验使用的是受限的合成 Python 子集程序变量只有 v0 到 v9运算主要是小常数范围内的加、减、乘输出也被压缩为 v0 mod 1000。实验没有包含函数调用、递归、复杂数据结构、异常处理、对象状态、库依赖或真实项目中的跨过程控制流。因此论文证明的是IPA-GNN 在受控程序文法上对更长程序具有更强的长度外推能力。它不能直接证明模型已经能够执行真实 Python 程序。7.2 训练数据规模很大训练集包含 500 万个合成程序。模型的结构归纳偏置确实提高了泛化但结果并不是小样本学习结果。实际使用时如何生成或获得覆盖真实语义的大规模训练数据仍然是问题。7.3 测试的是长度外推而不是所有类型的分布外泛化训练和测试来自同一个概率文法主要差异是程序长度。论文没有同时测试未见过的操作符未见过的控制流结构更大的数值范围新的数据类型不同代码风格或真实代码分布。因此“系统性泛化”在本文中主要指同一语言生成分布下的复杂度外推。7.4 正确输出不等于忠实执行IPA-GNN 在循环中可能只访问一次循环体却得到正确输出。这对有限预算预测是优势但也说明模型内部轨迹不能直接解释为真实执行轨迹。如果下游任务要求逐步可验证的执行状态、精确中间变量值或安全关键分析仅有最终输出准确率是不够的还需要对中间状态和控制流路径增加监督或一致性约束。7.5 部分程序任务存在不可辨识性一条表达式被遮盖后剩余程序通常不能唯一确定原语句及原输出。模型能够取得 29.1% 的准确率说明它学到了数据生成文法中的统计规律但不能把这一结果等同于从不完整源代码中逻辑推导出唯一答案。八、补充材料中的可视化论文补充材料还给出了更多随机程序的注意力图可用于观察模型是否沿合理路径传播。〖论文 Figure 7四个随机完整程序上的 IPA-GNN 软指令指针热力图四个示例均预测正确〗〖论文 Figure 8将 Figure 7 中每个程序随机遮盖一条语句后的部分程序注意力图其中前三个预测正确第四个预测错误〗这些图说明注意力可视化能够帮助分析模型行为但单个热力图并不能充分证明模型学习了完整语言语义。更严格的验证仍需要中间变量监督、执行轨迹一致性测试和对抗性程序测试。九、总结本文提出 IPA-GNN用软指令指针把程序状态沿控制流图传播并通过节点级 RNN 模拟单条语句执行。其最大价值在于把经典解释器的因果结构转化为可微分神经网络结构从而同时利用 RNN 的顺序推理能力和 GNN 的程序结构表示能力。实验中模型只在长度不超过 10 的程序上训练却在长度 20 到 100 的程序上取得 62.1% 的完整程序执行准确率和 29.1% 的部分程序执行准确率。与普通 GGNN、逐行 RNN 以及两个消融模型相比IPA-GNN 表现出更强的长度外推能力。论文最重要的经验是对于算法执行类任务模型结构是否与目标算法的因果过程对齐可能比单纯扩大隐藏维度或增加通用消息传递层更重要。但结果也应限定在论文的实验范围内这是受限合成语言上的输出预测和长度泛化不是对真实 Python 解释器的替代。IPA-GNN 学到的“短路执行”可以提高有限预算下的预测效率却不保证内部计算过程与真实执行轨迹一致。

相关新闻