[Paper Review] 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 Techniques4 citations

TL;DR

This paper proposes a mutual learning framework to improve tree-structured code generation by jointly training two Seq2Tree models—one using pre-order traversal and the other breadth-first traversal for decoding. By alternately distilling knowledge between the models during training, both benefit from complementary context modeling, leading to significant performance gains on code generation benchmarks.

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.

Motivation & Objective

• Address the limitation of pre-order traversal-based decoders in code generation, which rely only on vertical dependencies and may miss critical horizontal context.
• Explore the potential of breadth-first traversal-based decoding as a complementary approach to improve context modeling in neural code generation.
• Develop a mutual learning framework that enables two models with different traversal strategies to enhance each other through synchronous knowledge distillation during training.
• Demonstrate that the joint training framework improves model performance without increasing inference complexity or model size.
• Validate the effectiveness and generalization of the framework across multiple benchmark datasets for code generation.

Proposed method

• Train two Seq2Tree models: one using pre-order traversal (standard approach) and another using breadth-first traversal for action sequence generation.
• Implement a mutual distillation mechanism where, at each training step, one model acts as the student and the other as the teacher.
• During each step, the student model is trained to fit both the ground-truth actions and the predicted action distributions from the teacher model.
• Alternate the roles of student and teacher across training steps to ensure symmetric knowledge transfer and mutual improvement.
• Use the Tranx model as the base architecture, with modifications to support breadth-first traversal decoding while maintaining the same decoder structure.
• Apply knowledge distillation loss to align the student’s action prediction distribution with the teacher’s, enhancing generalization and robustness.

Experimental results

Research questions

• RQ1 Can breadth-first traversal-based decoding provide complementary context modeling to pre-order traversal in code generation?
• RQ2 How does mutual knowledge distillation between models with different traversal strategies affect code generation performance?
• RQ3 Does the proposed mutual learning framework improve model generalization and robustness on diverse code generation benchmarks?
• RQ4 Can the mutual learning framework enhance model performance without increasing model size or inference complexity?
• RQ5 How do the two models with different traversal strategies compare in handling specific code generation errors (e.g., incorrect function names or argument generation)?

Key findings

• The proposed mutual learning framework significantly improves code generation performance across multiple benchmark datasets, including Django, Python-300, and Java-100.
• Models enhanced via mutual learning outperform their respective baselines on almost all test groups, with consistent improvements in exact match and BLEU scores.
• Case studies show that the mutual learning model successfully overcomes common errors—such as incorrect function names and argument generation—by combining strengths from both traversal strategies.
• The breadth-first traversal model alone achieves competitive performance, indicating that non-preorder traversals are viable and effective for code generation.
• The mutual learning framework enables both models to surpass the performance of the state-of-the-art doubly-recurrent decoder (Alvarez-Melis and Jaakkola, 2017), despite using a single LSTM decoder.
• The framework demonstrates strong generalization, with consistent improvements across diverse programming languages and code complexity levels.