Paper: AST-T5: Structure-Aware Pretraining for Code Generation and Understanding (ICLR-2024)

Reviewer comment (Experiments):
What is the rationale behind including explicit NL data in the pre-training corpus? Prior works in this space such as CodeT5, PLBART, and GraphCodeBERT all have only code data in their pre-training mix. Is this a side-effect of the AST-aware masking strategy, where comments in code are not masked, and thus rather than including a new task - such as the Bimodal dual generation task in CodeT5 - this work requires NL data?

Author response:
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.

Reviewer comment (Experiments):
I have the impression that the AST Subtree Corruption objective is more important than AST-aware segmentation because, for the defined objective, the model is learning to unmask spans masked in an AST-aware manner through the AST Subtree Corruption. The remaining context (code), even if incomplete from an AST point-of-view might not be affecting the learning significantly. Is it possible to do an ablation study with base T5 + AST Subtree Corruption to assess the performance? Since this work utilizes tree-sitter, which is an incremental parsing library, it is possible to parse even incomplete code.

Author response:
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.

Reviewer comment (Evaluation):
For Concode, Bugs2Fix, and Java-C# translation tasks, the work reports the Exact Match (EM) score alone. EM is not an ideal measure for code as it measures string equivalence. Other fuzzy evaluation metrics such as CodeBLEU [4] or CodeBERTScore [5] scores might be a better fit for these tasks.

Author response:
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]:

|                 | **Exact Match** | **CodeBLEU** |
|-----------------|----------------:|-------------:|
| GPT-2           |            17.4 |         29.7 |
| CodeGPT-2       |            18.3 |         32.7 |
| CodeGPT-adapted |            20.1 |         36.0 |
| PLBART          |            18.8 |         38.5 |
| CodeT5-Small    |            21.6 |         41.4 |
| CodeT5-Base     |            22.3 |         43.2 |
| AST-T5 (Ours)   |        **22.9** |     **45.0** |

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

Reviewer comment (Experiments):
While the work does a good job of covering a decent number of downstream tasks, any specific reason why Code Summarization tasks were not included? The pre-training dataset for the model does include explicit NL data, so I expected to see code-to-text tasks included as well.

Author response:
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.

Reviewer comment (Experiments):
In section 4, the authors mention utilizing PromptSource to create prompt templates for fine-tuning. I might be wrong, but it is my understanding that prompts are more useful in zero/few-shot settings. 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?

Author response:
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.

Reviewer comment (Experiments):
HumanEval has been known to have been in public datasets (such as Github Archive) commonly used for training code LLMs (https://arxiv.org/abs/2306.11644) and being that it is an incredibly small dataset of questions and only Python, it would be a stronger argument for performance to show the Pass@1 vs Model Size graph for other common reported downstream evaluations to support the claim of "competitiveness against larger models". The reviewer suggests at least adding MBPP (present in Llama, PaLM, GPTs, etc) and considering showing performance across more languages through benchmarks such as MBXP and Multilingual HumanEval (https://arxiv.org/abs/2210.14868).

Author response:
**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:

|               | **#Params** | **MBPP** |
|---------------|-------------|----------|
| CodeGen-multi | 350M        | 7.5      |
| AST-T5 (Ours) | 226M        | **11.3** |
|               |             |          |
| PaLM          | 8B          | 5.0      |
| PaLM          | 62B         | 21.4     |
| PaLM          | 540B        | 36.8     |
| CodeGen-mono  | 350M        | 14.6     |
| InCoder       | 1.3B        | 10.9     |
| InCoder       | 6.7B        | 19.4     |
| LaMDA         | 137B        | 14.8     |
| LLaMA         | 7B          | 17.7     |
| LLaMA         | 13B         | 22.0     |
| LLaMA         | 33B         | 30.2     |
| LLaMA         | 65B         | 37.7     |

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

|               | **#Params** | **HumanEval** | **HumanEval** | **MBXP** | **MBXP** |
|---------------|-------------|---------------|---------------|----------|----------|
|               |             | Python        | Java          | Python   | Java     |
| CodeGen-multi | 350M        | 7.3           | 5.0           | 7.5      | 8.2      |
| CodeGen-mono  | 350M        | 10.3          | 3.1           | **14.6**     | 1.9      |
| AST-T5 (Ours) | 226M        | **12.8**      | **10.6**      | 11.3     | **9.8**  |
|   |  |  | |  |   |
| BLOOM         | 7.1B        | 7.9           | 8.1           | 7.0      | 7.8      |
| OPT           | 13B         | 0.6           | 0.6           | 1.4      | 1.4      |
| CodeGen-multi | 2B          | 11.0          | 11.2          | 18.8     | 19.5     |
| CodeGen-mono  | 2B          | 20.7          | 5.0           | 31.7     | 16.7     |
| CodeGen-multi | 6B          | 15.2          | 10.6          | 22.5     | 21.7     |
| CodeGen-mono  | 6B          | 19.5          | 8.7           | 37.2     | 19.8     |
| CodeGen-multi | 16B         | 17.1          | 16.2          | 24.2     | 28.0     |
| CodeGen-mono  | 16B         | 22.6          | 22.4          | 40.6     | 26.8     |

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

Reviewer comment (Evaluation):
What are some potential limitations of using this method?

Author response:
**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.

Reviewer comment (Experiments):
Were masking ratios tested between 25% and 100%? If 25% showed no degradation in understanding tasks, why not go higher? From paper:  
```
we observed that raising the mask ratio from 15% to 25% significantly improved generation capabilities without noticeably compromising performance in understanding tasks. Thus, we settled on a 25% mask ratio for our AST-T5 model.
```

Author response:
**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.

|                           | **HumanEval** | **Concode** |  **Bugs2Fix** |       **Java-C#** | **Clone** | **Defect** |  **Avg** |
|---------------------------|--------------:|------------:|--------------:|------------------:|----------:|-----------:|---------:|
| T5                        |           5.2 |        18.3 |     21.2/13.8 |         65.5/68.4 |      96.9 |       64.1 |     44.2 |
| + AST. Segmentation       |           7.2 |        20.2 |     22.5/15.1 |         66.3/69.3 |      98.3 |       65.9 |     45.7 |
| + AST. Subtree Corrupt    |           9.6 |        22.1 | 23.3/**16.5** |         67.3/72.2 |  **98.6** |   **66.0** |     47.0 |
| + Mask 25\% (AST-T5)      |          12.8 |    **22.9** | **23.8**/16.1 | **68.9**/**72.3** |  **98.6** |       65.8 | **47.7** |
| + Mask 50\% (Newly Added) |      **13.0** |        22.0 |     21.9/15.0 |         66.5/70.1 |      97.1 |       64.2 |     46.2 |

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.

Reviewer comment (Experiments):
Code completion (generation with both prefix and suffix context) has become a very common capability in code models and the span corruption objective is effectively code completion. Was there work done on evaluating this model for code completion tasks? Why or why not?

Author response:
**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:

```python
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

Reviewer comment (Reproducibility):
The abstract mentioned anonymized model + code, where is this available?

Author response:
**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!