[论文解读] Improving Tree-Structured Decoder Training for Code Generation via Mutual Learning

Binbin Xie, Jinsong Su+5 more|arXiv (Cornell University)|May 31, 2021

Natural Language Processing Techniques被引用 4

一句话总结

本文提出一种互知识蒸馏框架，通过联合训练两个Seq2Tree模型（一个使用前序遍历解码，另一个使用广度优先遍历解码），以提升树结构代码生成性能。在训练过程中，两个模型交替进行知识蒸馏，从而实现互补的上下文建模，显著提升代码生成基准测试的性能。

ABSTRACT

Code generation aims to automatically generate a piece of code given an input natural language utterance. Currently, among dominant models, it is treated as a sequence-to-tree task, where a decoder outputs a sequence of actions corresponding to the pre-order traversal of an Abstract Syntax Tree. However, such a decoder only exploits the preorder traversal based preceding actions, which are insufficient to ensure correct action predictions. In this paper, we first throughly analyze the context modeling difference between neural code generation models with different traversals based decodings (preorder traversal vs breadth-first traversal), and then propose to introduce a mutual learning framework to jointly train these models. Under this framework, we continuously enhance both two models via mutual distillation, which involves synchronous executions of two one-to-one knowledge transfers at each training step. More specifically, we alternately choose one model as the student and the other as its teacher, and require the student to fit the training data and the action prediction distributions of its teacher. By doing so, both models can fully absorb the knowledge from each other and thus could be improved simultaneously. Experimental results and in-depth analysis on several benchmark datasets demonstrate the effectiveness of our approach. We release our code at https://github.com/DeepLearnXMU/CGML.

研究动机与目标

• 解决基于前序遍历的解码器在代码生成中的局限性，即仅依赖垂直依赖关系，可能遗漏关键的水平上下文信息。
• 探索基于广度优先遍历的解码作为互补方法，以改善神经网络代码生成中的上下文建模能力。
• 构建一种互知识蒸馏框架，使采用不同遍历策略的两个模型在训练过程中通过同步知识蒸馏相互增强。
• 证明联合训练框架可在不增加推理复杂度或模型参数量的前提下提升模型性能。
• 在多个代码生成基准数据集上验证该框架的有效性与泛化能力。

提出的方法

• 训练两个Seq2Tree模型：一个使用前序遍历（标准方法）进行动作序列生成，另一个使用广度优先遍历进行动作序列生成。
• 实现一种互知识蒸馏机制，即在每个训练步骤中，一个模型作为学生，另一个作为教师。
• 在每个步骤中，学生模型同时学习真实动作和教师模型预测的动作分布。
• 在训练过程中交替变换学生与教师的角色，以确保对称的知识迁移与相互提升。
• 以Tranx模型为基础架构，对模型进行修改以支持广度优先遍历解码，同时保持相同的解码器结构。
• 应用知识蒸馏损失，使学生模型的动作预测分布与教师模型对齐，从而提升泛化能力与鲁棒性。

实验结果

研究问题

• RQ1 基于广度优先遍历的解码是否能为代码生成中的前序遍历提供互补的上下文建模？
• RQ2 采用不同遍历策略的模型之间进行互知识蒸馏，对代码生成性能有何影响？
• RQ3 所提出的互知识蒸馏框架是否能提升模型在多样化代码生成基准上的泛化能力与鲁棒性？
• RQ4 该互知识蒸馏框架是否能在不增加模型大小或推理复杂度的前提下提升模型性能？
• RQ5 采用不同遍历策略的两个模型在处理特定代码生成错误（如函数名错误或参数生成错误）时表现如何？

主要发现

• 所提出的互知识蒸馏框架在Django、Python-300和Java-100等多个基准数据集上显著提升了代码生成性能。
• 经互知识蒸馏增强的模型在几乎所有测试组中均优于各自的基线模型，且在精确匹配率与BLEU分数上均表现出一致的提升。
• 案例研究显示，互知识蒸馏模型通过融合两种遍历策略的优势，成功克服了常见错误（如函数名错误或参数生成错误）。
• 仅使用广度优先遍历的模型也达到了具有竞争力的性能，表明非前序遍历方法在代码生成中具有可行性与有效性。
• 尽管仅使用单个LSTM解码器，该互知识蒸馏框架使两个模型均超越了SOTA双循环解码器（Alvarez-Melis and Jaakkola, 2017）的性能表现。
• 该框架展现出强大的泛化能力，在多种编程语言及不同代码复杂度水平下均保持一致的性能提升。