[Paper Review] CodePAD: Sequence-based Code Generation with Pushdown Automaton

Yihong Dong, Xue Jiang+3 more|arXiv (Cornell University)|Nov 2, 2022

Natural Language Processing Techniques4 citations

TL;DR

This paper proposes CodePAD, a sequence-based code generation framework that integrates a pushdown automaton (PDA) module to enforce grammatical correctness during code generation. By leveraging PDA state representations and joint prediction with state, CodePAD ensures 100% grammatical correctness on benchmark datasets, significantly improving CodeBLEU and EM scores over base models and pre-trained models like CodeGen-350M.

ABSTRACT

In the process of code generation, it is essential to guarantee the generated code satisfies grammar constraints of programming language (PL). However, neglecting grammar constraints is a fatal drawback of commonly used sequence-based code generation. In this paper, we devise a pushdown automaton (PDA)-based methodology to address this problem, exploiting the principle that PL is a subset of PDA recognizable language and code accepted by PDA is grammatical. Specifically, we construct a PDA module and design an algorithm to constrain the generation of sequence-based models to ensure grammatical correctness. Guided by this methodology, we further propose CodePAD, a sequence-based code generation framework equipped with a PDA module, to integrate the deduction of PDA into deep learning. Additionally, this framework can leverage states of PDA deduction (including state representation, state prediction task, and joint prediction with state) to assist models in learning PDA deduction. To comprehensively evaluate CodePAD, we construct a PDA for Python and conduct extensive experiments on four public benchmark datasets. CodePAD can leverage existing sequence-based models, and we show that it can achieve 100\% grammatical correctness percentage on these benchmark datasets. Thus, it relatively improve 17\% CodeBLEU on CONALA, 8\% EM on DJANGO, and 15\% CodeBLEU on JUICE-10K compared to base models. In addition, our method significantly enhances pre-trained models, e.g., CodeBLEU of CodeGen-350M improvement from 3.21 to 21.54 on MBPP in zero-shot setting.

Motivation & Objective

• To address the critical limitation of sequence-based code generation models, which often produce syntactically incorrect code despite strong semantic performance.
• To develop a method that guarantees grammatical correctness during code generation without relying on expensive parsing or AST-based approaches.
• To integrate PDA-based grammatical constraints into sequence models while preserving their efficiency and flexibility for real-world applications like code completion.
• To evaluate the effectiveness of PDA state representation, state prediction, and joint prediction in improving code generation quality.

Proposed method

• Design a PDA module that models the grammar of a programming language (e.g., Python) as a pushdown automaton, enabling recognition of valid code sequences.
• Develop an algorithm to simulate PDA deduction during generation, ensuring that each next token is selected only from the valid terminal symbol set recognized by the current PDA state.
• Introduce state representation to embed PDA states into the model’s hidden space, enabling end-to-end training with attention mechanisms.
• Train the model with a multi-task objective: joint prediction of the next token and the next PDA state, using a learnable coefficient α to balance the two losses.
• Use beam search with PDA state constraints during inference to ensure that only grammatically valid sequences are generated.
• Leverage the PDA’s accept state set to terminate generation only when a valid, complete code sequence is produced.

Experimental results

Research questions

• RQ1 RQ1: How well does CodePAD perform compared to state-of-the-art tree-based methods that guarantee syntactic correctness without pre-training?
• RQ2 RQ2: What is the contribution of each component of the PDA-based method—state representation, state prediction, and joint prediction—to the overall performance?
• RQ3 RQ3: How does the PDA module enhance the effectiveness of pre-trained sequence-based models like CodeGen and CodeT5?
• RQ4 RQ4: How does the PDA module enable CodePAD to generate grammatically correct code during inference?

Key findings

• CodePAD achieves 100% grammatical correctness on all four benchmark datasets, eliminating syntax errors in generated code.
• On the CONALA dataset, CodePAD improves CodeBLEU by 17% compared to the base model, and on DJANGO, it improves EM by 8%.
• For the MBPP benchmark, CodeGen-350M’s CodeBLEU score increases from 3.21 to 21.54 in a zero-shot setting when fine-tuned with CodePAD.
• The optimal hyperparameter α for the loss weighting is found to be 1.0 for LSTM-based CodePAD and varies per dataset for Transformer-based models, indicating strong effectiveness of the state prediction task.
• Compilation success rate on CONALA reaches 99.2%, with most compilation errors stemming from semantic or type mismatches rather than syntax, showing that grammatical correctness is a solid foundation for compilation.
• The PDA module enables efficient, real-time grammatical validation during generation without requiring full parsing of incomplete code, avoiding the high cost of token-by-token trial-and-error.