AST-T5: Structure-Aware Pretraining for Code Generation and Understanding

W5. Any specific reason why Code Summarization tasks were not included?

Thank you for your insights! Yes, this is indeed a limitation of our work. AST-T5, as it currently stands, excels in code generation by using a unique masking strategy within AST subtrees during pretraining. This specificity, while beneficial for code-centric tasks, may not yield the best results in natural language generation. Acknowledging this limitation, future versions of AST-T5 could investigate strategies such as masking docstrings and comments to broaden its applicability. This would potentially improve performance across various tasks, including code summarization. We have also included the discussion of limitations of AST-T5 in the revised manuscript, specifically in Appendix A.1.

W6. If the authors are fine-tuning the model, why is a prompt required? Why can't the model be fine-tuned on input-output pairs directly?

Thank you for your insightful question regarding the use of prompts in our fine-tuning process. Our approach, diverging from CodeT5's single-task fine-tuning, uses a multi-task framework similar to those in renowned NLP models like T5 and T0. This necessitates prompts to differentiate between various tasks, ensuring consistency across all models, including AST-T5 and baselines, for fair experimental comparison.

Moreover, prompts in our fine-tuning cater to challenges in datasets like HumanEval, which lack training data. By applying prompt-based fine-tuning on CodeSearchNet, AST-T5's performance on HumanEval is notably enhanced. While CodeT5+ also adapts to NL-to-Code tasks using CodeSearchNet, it differs by using an alternate loss function, not prompts. However, prompt usage is a secondary aspect of our evaluation protocol; the primary focus and contribution of our paper is the AST-Awareness, as clearly demonstrated by the results in Table 1.

W1. Misused terms: tree-sitter library actually creates a Concrete Syntax Tree and not AST

Thank you for highlighting the distinction between Concrete Syntax Trees (CSTs) and Abstract Syntax Trees (ASTs). We acknowledge that Tree-Sitter produces CSTs, which, while related, are indeed more verbose and concrete compared to ASTs. Our model, AST-T5, indeed relies on CST parsing rather than AST-level abstractions.

To address this, we will amend our manuscript to replace all instances of "AST" with "Parse Tree" and "AST-aware" with "PT-aware" in the next version. We apologize for not making these changes in the current draft, as the title cannot be modified during the rebuttal phase, and we aim to maintain consistency in our discussion. Thank you for your valuable suggestions in improving the accuracy and clarity of our paper!

W2. What is the rationale behind including explicit NL data in the pre-training corpus?

We thank the reviewer for their insightful attention to the inclusion of explicit natural language (NL) data in our pre-training corpus. It's important to note that NL data has been a key component in training code-related models, as evidenced by various models in the field. For example, Codex leverages a rich NL foundation by fine-tuning from GPT-3, and models like CodeT5 and GraphCodeBERT use NL-PL bimodal corpora, including CodeSearchNet, to enhance their training datasets. Moreover, the practice of using the entire GitHub dump, which encompass NL data from sources such as Markdown, JavaDoc, Plain Text, or HTML, is a standard approach in this domain (refer to https://huggingface.co/datasets/bigcode/the-stack-dedup/tree/main/data/markdown). Additionally, CodeLLaMa [2] explicitly incorporate web-based NL data, such as StackOverflow discussions, into their training processes. Therefore, the use of NL data is a common thread across existing models.

Our inclusion of explicit NL data from external datasets compensates for the limitations in our data crawling pipeline, which focuses on programming languages such as Python, Java, C, C++, and C# through file extensions like ".py" or ".cpp". This approach, while effective for sourcing code, often misses rich NL content in Markdown and plain text files, essential in training processes and used by comparable models in the field. By integrating additional NL datasets, we ensure our pretraining corpus is comprehensive, addressing this critical gap. This decision is independent of our AST-aware masking strategy.

W3. Is it possible to do an ablation study with base T5 + AST Subtree Corruption to assess the performance?

Thank you for your suggestion regarding the additional experiment for "Base T5 + AST-Aware Subtree Corruption". However, AST-Aware Subtree Corruption, as outlined in Algorithm 2, is dependent on AST-Aware Segmentation. It operates recursively on each subtree within a segmented chunk. Without the AST-Aware Segmentation, the model would encounter highly fragmented spans, because standard Greedy Segmentation could result in a code segment comprising numerous disjoint subtrees. This fragmentation would likely degrade the performance of such "Base T5 + AST-Aware Subtree Corruption" model.

W4. Other fuzzy evaluation metrics such as CodeBLEU [4] or CodeBERTScore [5] scores might be a better fit for some tasks such as Concode.

Thank you for your valuable suggestion to incorporate alternative evaluation metrics. In response, we have included CodeBLEU scores in Appendix A.7 for Concode, because this metric has also been used in prior works like CodeT5. The result is also shown here, the reported numbers of other models are from [1]:

Our analysis reveals a strong correlation between CodeBLEU and Exact Match (EM) scores, and AST-T5 outperforms baselines in both metrics, underscoring the robustness of our model's performance. Regarding CodeBERTScore, we noted its dependency on a specific vocabulary, limiting its applicability compared to the baselines we have used, such as CodeT5.

[1] CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation, https://arxiv.org/abs/2109.00859

[2] Code Llama: Open Foundation Models for Code, https://arxiv.org/abs/2308.12950

W1. About other previously published ways of utilizing AST structures

We acknowledge the potential benefits of various AST-aware techniques in Transformer architecture. Indeed, methods like tree-position embeddings and AST-aware attention mechanisms have shown promise. However, these approaches often come with architectural changes and assumptions that may not align with all use cases, particularly those involving incomplete or erroneous code, which is common in real-world applications like bug fixing or code completion.

Our decision to focus on AST-T5's simple yet effective approach was strategic. By avoiding complex architectural changes and maintaining compatibility with existing T5 variants, AST-T5 is designed for seamless integration and broader applicability. This simplicity also negates the need for parsing or static analysis in downstream tasks, which is a crucial factor for practicality in varied coding scenarios.

While integrating additional AST-aware techniques like tree-position embeddings could potentially enhance the model, such modifications would compromise its general applicability and ease of use. AST-T5's focus on pretraining with AST-awareness, rather than relying on complex runtime processes, represents a balance between capturing code's structural nuances and maintaining a flexible, adaptable framework suitable for a wide range of applications.

We have expanded our discussion on related work in the manuscript's revised "Related Work" section.

W2. This paper uses a context length of 1024 tokens. The benefit of AST-aware segmentation will diminish with increasing context length

We appreciate your attention to the context length in AST-T5. Notably, AST-T5 uses a context length of 1024 tokens for both source and target. This exceeds the context lengths of our main baselines, CodeT5 and CodeT5+, which are 512/256 and 768/600 tokens respectively.

The AST-Aware segmentation at the core of AST-T5 is pivotal in maintaining the structural integrity of code during pretraining, offering two benefits:

1. Enhanced Pretraining Efficiency: The structured nature of inputs, thanks to AST-aware segmentation, boosts pretraining efficiency. While this advantage may diminish with longer context lengths (e.g., 2k, 8k, 32k tokens), its impact is significant in our current setup.

2. Reduced Mismatch with Downstream Tasks: AST-aware segmentation ensures a closer match between pretraining inputs and the complete code structures encountered in downstream tasks, like function definitions. This alignment minimizes the pretraining-evaluation mismatch, a benefit that remains crucial regardless of context length.

Although the first benefit might reduce as context lengths increase, the second benefit remains highly relevant. We acknowledge the value of experimenting with longer context lengths, and we plan to explore this in future work, once we have sufficient computational resources.

W3. The related work section also needs to be fleshed out, in particular w.r.t. "leveraging code structure in pretraining."

Thank you for highlighting the need for a more comprehensive discussion of leveraging code structure in pretraining in the related work section. In response to your insightful suggestion, we have extensively expanded this paragraph, doubling its length to provide more discussions and comparisons about relevant studies. Here is the updated version of this part:

Leveraging Code Structure in Pretraining. Code differs from natural language in two key aspects: its executability and strict structural syntax. Previous research leveraged execution traces for improving model performance (Chen et al., 2018; 2021b; Shojaee et al., 2023) but this approach faces scalability challenges when applied to large, web-crawled code datasets used in pretraining. Regarding code’s structured nature, various studies have integrated syntactic elements into neural network models. Li et al. (2018), Kim et al. (2021) and Zugner et al. (2021) add AST-Aware attention mechanisms in their models, while Alon et al. (2020) and Rabinovich et al. (2017) focus on modeling AST node expansion operations rather than traditional code tokens. In parallel, Guo et al. (2021) and Allamanis et al. (2017) explore DFG-Aware attention mechanisms and Graph Neural Networks (GNNs), to interpret code based on its Data Flow Graph (DFG). StructCoder (Tipirneni et al., 2023) enriches the code input by appending AST and DFG as additional features. These methods, however, necessitate parsing or static analysis for downstream tasks, which is less feasible for incomplete or incorrect code scenarios like bug fixing.

Our work, AST-T5, aligns with methods that utilize code structure only in pretraining, like DOBF (Roziere et al., 2021) and CodeT5 (Wang et al., 2021), which obfuscate inputs to force the model to grasp abstract structures. Our approach uniquely diverges by using AST-driven segmentation and masking in T5 span corruption during pretraining. This novel approach offers a more refined pretraining signal compared to structure-agnostic T5, equipping our model to proficiently encode and generate semantically coherent code structures.

We believe these enhancements not only address your concerns but also enrich the clarity of the manuscript. We are grateful for your feedback and hope that our revisions meet your expectations.

Q1. What are some potential limitations of this method?

We appreciate the opportunity to discuss the limitations of our work. AST-T5, as it currently stands, excels in code generation by using a unique masking strategy within AST subtrees during pretraining. This specificity, while beneficial for code-centric tasks, may not yield the best results in natural language generation. Acknowledging this limitation, future versions of AST-T5 could investigate strategies such as masking docstrings and comments to broaden its applicability. This would potentially improve performance across various tasks, including code summarization. We have also included the discussion of limitations of AST-T5 in the revised manuscript, specifically in Appendix A.1.

Q2. Were masking ratios tested between 25% and 100%? Why not go higher than 25%?

Thank you for your insightful question regarding the masking ratio in AST-T5. Masking ratio is an important hyperparameter in T5-style pretraining. Indeed, a 100% mask ratio would essentially make the model equivalent to a GPT-like, decoder-only architecture, underutilizing the encoder's parameters. Conversely, a 0% mask ratio is infeasible as it would negate any loss function application, rendering the model untrainable.

We rigorously tested various masking ratios, including 15%, 25%, and, as newly added per your suggestion, 50%. The results of the 50% mask ratio experiment are now added to Table 1 of our paper and also included below. We observed that increasing the masking ratio from 25% to 50% led to a marginal improvement in HumanEval performance (improving Pass@1 from 12.8 to 13.0). However, this increase adversely affected both transpilation and understanding tasks.

Our analysis suggests that the transpilation and understanding tasks particularly benefit from a strong bidirectional Transformer encoder. A higher mask ratio, while potentially strengthening the decoder, correspondingly weakens the encoder. Therefore, we concluded that a 25% mask ratio offers the most balanced performance across various code-related tasks, without compromising the model's overall efficacy in transpilation and understanding.

Q3. Was there work done on evaluating this model for code completion tasks? Why or why not?

We acknowledge the relevance of the evaluation of AST-T5 on code completion tasks. However, the existing benchmarks for code completion, notably CodeXGLUE and HumanEval Fill-In-the-Middle (FIM) [1], present inherent biases and practical limitations when applied to our model. Specifically, CodeXGLUE's evaluation depends heavily on its predefined token vocabulary, disadvantageous to models like ours that uses a different Byte Pair Encoding (BPE) vocabulary. More critically, the FIM benchmark's approach of masking individual lines can disrupt the integrity of Abstract Syntax Tree (AST) structures. For instance, in the following code:

for x in l:
    if x % 2 == 0:
        print(x)

masking a line such as if x % 2 == 0: in a Python loop breaks the AST, as the line is only part of an AST node (the if node). Doing so puts AST-aware models at a disadvantage as the key idea behind such models is to treat code as structured units based on its node type (if, while, for, etc) instead of collections of tokens as in typical NLP setting. Masking individual lines is less representative of practical scenarios compared to masking AST nodes like x % 2 == 0 or x % 2, where our model is likely to do better.

To address these limitations, we adapted the single-line infilling task from [1] to only mask each line only when it contains a complete AST node (e.g., a full assignment statement). This modification yielded a pass@1 rate of 54.2% for AST-T5. In contrast, CodeT5+-220M achieved a pass@1 rate of 18.5% under the same conditions. The result shows the effectiveness of AST-T5 in a more relevant setting, which we will explore more later on.

[1] Efficient Training of Language Models to Fill in the Middle, https://arxiv.org/abs/2207.14255

W1.1. Lack of models for comparison at this scale

Regarding the scale of model comparison, the limited computational resources indeed restricted our capacity for training larger models. Despite this, our study ensures a fair and rigorous comparison as detailed in Table 1, where all models, including AST-T5, were trained using the same dataset and environment. This approach enables a meaningful evaluation of AST-T5's performance against similar-sized models, showing its effectiveness in code-related tasks. Our focus on methodological rigor over scale underscores the value of AST-T5’s contributions.

W1.2. Potential data contamination of HumanEval and the need for results on more benchmarks

Thank you for your insightful comments on dataset overlap and benchmarks. Our pretraining data, sourced from Google BigQuery's GitHub dump as of Summer 2021, predates the release of Codex and HumanEval, thus mitigating concerns about dataset overlap.

Incorporating your valuable suggestions, we've added the results on the MBPP benchmark in Figure 4, and results on the Java split of the Multi-lingual HumanEval and MBXP benchmarks in Appendix A.6. Currently, AST-T5 supports Python, Java, C/C++, and C#, and we have already shown good results while running on these languages. Future work will extend our pretraining dataset to include all languages in these benchmarks. Due to the absence of results from models at a comparable scale (226M parameters), we compare AST-T5 with both the 350M-parameter CodeGen-multi (multi-lingual) and CodeGen-mono (Python-specialized) models, as well as larger models for a broader context.

Our results are summarized below:

Result on MBPP:

Result on multi-lingual HumanEval and MBXP, compared with reported numbers in [1]:

Despite its smaller size of 226M parameters, AST-T5 outperforms larger models like CodeGen-multi-350M, PaLM-8B, and InCoder-1.3B on MBPP, showing its efficiency. While CodeGen-mono-350M leads in MBPP because it is CodeGen-multi-350M further finetuned on Python, AST-T5, as a multi-lingual model, outperforms models like CodeGen-mono/multi-350M, BLOOM, and OPT in the Multilingual HumanEval and MBXP-Java benchmarks, underscoring its effectiveness in multiple programming languages.

[1] Multi-lingual Evaluation of Code Generation Models, http://arxiv.org/abs/2210.14868

W2. Writing of the abstract

Thank you for pointing out the need for clearer contextualization in our abstract! To address this, we have revised the sentence to: “Structure-awareness makes AST-T5 particularly powerful in code-to-code tasks, surpassing CodeT5 by 2 points in exact match score for the Bugs2Fix task and by 3 points in exact match score for Java-C# Transpilation in CodeXGLUE.” This revision aims to provide a more context of our results, clearly indicating the metrics used and the specific tasks where improvements were observed. Thank you for helping us improve the clarity of our abstract!

W3. Typo

We apologize for the typo and we have fixed it in the updated manuscript.

Q4. The abstract mentioned anonymized model + code, where is this available?

The codebase has now been uploaded to the supplementary material section of our submission. We will make it public as we publish the paper. As for the model, we are currently in the process of merging our code into HuggingFace to facilitate its integration. We will release our model on the HuggingFace Model Hub. We appreciate your patience and interest in our work!