PLUR: A Unifying, Graph-Based View of Program Learning, Understanding, and Repair

Abstract: Machine learning for understanding and editing source code has recently attracted significant interest, with many developments in new models, new code representations, and new tasks.This proliferation can appear disparate and disconnected, making each approach seemingly unique and incompatible, thus obscuring the core machine learning challenges and contributions.In this work, we demonstrate that the landscape can be significantly simplified by taking a general approach of mapping a graph to a sequence of tokens and pointers.Our main result is to show that 16 recently published tasks of different shapes can be cast in this form, based on which a single model architecture achieves near or above state-of-the-art results on nearly all tasks, outperforming custom models like code2seq and alternative generic models like Transformers.This unification further enables multi-task learning and a series of cross-cutting experiments about the importance of different modeling choices for code understanding and repair tasks.The full framework, called PLUR, is easily extensible to more tasks, and will be open-sourced (https://github.com/google-research/plur).


Reviews and author responses:

 Thanks for your detailed response! This makes sense and is very insightful. I'll keep my review of Accept.

 Thank you, this is very insightful! I increased my score to a clear accept.

The paper presents a general architecture which unifies fifteen previously
separate ML4Code tasks and achieves state-of-the-art or near state-of-the-art
results on them.

The inputs of the architecture are graphs where the nodes have labels and
positions, and both the nodes and the edges have types. The outputs follow the
ToCoPo output formulation where each output element can be a token, a copy, or a
pointer. This makes the architecture flexible enough to accomodate most ML4Code
tasks.
  I believe that the paper is an excellent and important contribution to ML4Code
as it unifies previously separate tasks utilizing different algorithms into a
single consistent state-of-the-art framework.

The framework will be open-sourced and can be the base of future research: new
models can be tested and ablated across the diverse 15 tasks, and new tasks can
be incorporated into the framework.

This is demonstrated in the paper with various ablation studies which give
interesting, novel insights. The authors compare different encoders across all
the tasks, the usefulness of pointers and copies across two tasks, multi-task
learning, and different graph representations for code2seq.

Overall I really like the paper, and only have a few concerns/questions:
1. The reported values for grid search in Appendix C
   suggest |learning_rate|\*|dropout_rate|\*|batch_size_per_device| = 3\*2\*2 = 12
   hyperparameter combinations for Table 4, and 2*2 = 4 combinations for Table
   5, yet the paper reports 24 combinations.
2. The hyperparameters were set separately for each task with grid search. The
   values tried by grid search are reported in Appendix C, but the concrete
   hyperparameters are not.
   - It would be good to see some summary about how different the
     hyperparameters are across the tasks. For example, it would be instructive
     to know whether dropout is useful for all the tasks, or there are
     exceptions.
   - Are there hyperparameters for which the architecture performs well across
     all tasks?
3. I think that the cells should be bolded differently in Table 2 (d). The
   purpose of that table is to compare between using pointers and copy and not
   using them, so the best result in each column should be bolded in my opinion.
   For example, the bolding of the cell CuBERT-VMR/GREAT/No pointer is
   particularly misleading, as that is the only bolded cell in the column and it
   suggests to me that No pointer is better than Full at a first glance.
4. Maybe the Discussion would be better as the Conclusion?
5. There is a small typo in line 102: "See and details of the software
   architecture"

# Update

I would like to thank the authors for their insightful and transparent answers
and their corrections to the paper. I'm increasing my score to a clear accept.
 The authors address the limitations and the societal impact appropriately .

This paper proposes a framework, called PLUR to unify a number of code-related tasks into a general form. Specifically, it pairs a graph encoder with a Transformer-style decoder augmented with pointers and a copy mechanism as a unified form and uses GREAT as the graph encoder.  On 15 different tasks e.g., function name prediction, program repair, variable misuse, PLUR has achieved at or above state-of-the-art results. Furthermore, PLUR also investigates multi-task learning by the proposed framework and yields promising results.  This paper attempts to unify many ML4Code tasks into a general formulation, which is interesting and meaningful. However, it is also very complicated since different code task has its characteristic and needs customized network architecture to achieve promising performance. This paper utilizes the transformer as the base modelling architecture combing with GREAT for learning and it has performed a large number of experiments over 15 tasks to illustrate the effectiveness of the proposed framework compared with the original paper. The improvements on some tasks e.g., Hoppity, CuBERT-VMR are significant.

 - The architecture of the network is not novel, GREAT and TOCOPO have already proved the effectiveness in the related papers.

- Missing the important baselines CodeBERT and GraphCodeBERT for comparison since CodeBERT has already achieved state-of-the-art performance on many ML4Code tasks. In the chosen tasks, PLUR has already compared with CuBERT, to make the approach solid, it is better to compare with the latest techniques i.e., CodeBERT and GraphCodeBERT.

- The selected tasks can roughly be categorized into four groups i.e.,  function classification (ManySStuBs4J, CuBERT), function name prediction (code2seq, ogb-code, convattn), code summary (funcom) and bug repair (VarMisuseH, Hoppity),  whether the selected tasks are diverse and fair enough to cover all tasks? 

- It seems pointer and copy mechanisms just severs for VarMisuse task,  is it necessary to equip in the decoder?



 We would like to post a small correction to the hparams table in the comment above to the CuBERT-VMR, CuBERT-FD and CuBERT-VM rows. 

 |Dataset        | Model       | Batch size | Learning rate | Dropout |
|----------------|-------------|------------|---------------|---------|
| CuBERT-VMR     | GREAT |          4 |      0.0001 |     0.1 |
| CuBERT-FD      | GREAT |         16 |      1.00E-05 |     0.1 |
| CuBERT-VM      | GGNN |         16 |      1.00E-05 |     0.1 |

 We thank the reviewer for their time and comments. 

> The architecture of the network is not novel

PLUR explicitly does not seek to innovate in model architectures (See Sec 3, Paragraph 1: "Our goal from the modeling perspective is specifically not to innovate on machine-learning architectures."). We don't see this as a weakness, because our goal is to unify and simplify the landscape of ML4code approaches while demonstrating that significant new insights can be derived from a unified formulation.

> PLUR has already compared with CuBERT

There appears to be a misunderstanding here. We'd like to clarify that we are evaluating with the CuBERT benchmarks (the 6 tasks and datasets used in that paper), but we are not comparing to methods that use pre-training (like CuBERT’s fine-tuned models on those benchmarks). Our baselines on these tasks are the base Transformer models without pre-training.

The point of the current paper is the unifying approach, which we view as an orthogonal choice to whether pre-training is used. In future work, we are very excited about a version of PLUR that incorporates pre-training. However, there are some design decisions and research questions that need to be figured out in order to do this properly. For example, there is significant variability in the representation of code across tasks, ranging from graphs of AST paths in Code2Seq to ASTs in Hoppity to token sequences for CuBERT’s tasks. We'd like to study whether there is an existing or new pre-training approach that leads to good performance across the variability in code representations in PLUR, and then that would be a natural comparison point for CuBERT, CodeBERT, and GraphCodeBERT. Thank you for raising this. We will expand discussion around this point in the paper. 

> It is better to compare with the latest techniques i.e., CodeBERT and GraphCodeBERT

For the above reasons, we do not compare to the CodeBERT and GraphCodeBERT fine-tuned models, but leave a comparison on pre-training/fine-tuning to future work. We note that both CodeBERT’s and GraphCodeBERT’s pre-training objectives can be directly mapped to PLUR’s output DSL, thanks to the expressive use of input pointers (for predicting edges and variable mappings in GraphCodeBERT) and, of course, tokens/copies for MLM-style objectives.

> Representativeness of the selected tasks

PLUR is an on-going project and will be an open-sourced framework that can grow with more tasks. We would be happy to add auto-completion as a task to PLUR in the next iteration (as gAPh suggested). 

> Pointer and copy for the VarMisuse task

The pointer network and copy mechanism is common to all tasks, but some tasks don’t require pointers. Among the tasks in this paper, VarMisuseH, CuBERT-VMR, and hoppity all use pointers. Section 3.3 and Appendix B provide details.


 We thank the reviewer for their valuable comments and their strong support for this paper.

> Number of hyperparameter combinations

Thank you for pointing this out. On review, we realized that there was an error in our reporting of the number of hyperparameters per task. As pointed out, the correct number is  12 (Tab 4) + 4 (Tab 5) = 16 hyperparameters per model class. We will fix this along with the compute budget in the updated version of the paper.


> Performance of different hyperparameters on different tasks

We have collected the best hyperparameters from all tasks and compiled them into a table. Here are some insights:
* Dropout should almost always be used, but only at 10% level.
* Larger learning rates are preferable, with exception on some CuBERT datasets.
* No strong signal for the batch size on different datasets.

| Dataset        | Model       | Batch size | Learning rate | Dropout |
|----------------|-------------|------------|---------------|---------|
| code2seq-small | GREAT       |          4 |        0.0001 |     0.1 |
| code2seq-med   | GREAT       |         16 |        0.0001 |     0.1 |
| code2seq-large | GREAT       |         16 |      3.00E-05 |     0.1 |
| convattn       | GREAT       |          4 |        0.0001 |     0.1 |
| funcom         | GREAT       |         16 |      3.00E-05 |     0.1 |
| great          | GREAT       |          4 |      3.00E-05 |     0.1 |
| hoppity        | GREAT       |          4 |      3.00E-05 |     0.1 |
| ogb-code       | GREAT       |          4 |        0.0001 |     0.1 |
| manysstubs4j   | GREAT       |          4 |        0.0001 |     0.1 |
| CuBERT-VMR     | Transformer |          4 |      3.00E-05 |     0.1 |
| CuBERT-VM      | GREAT       |          4 |      3.00E-05 |     0.1 |
| CuBERT-SO      | GREAT       |         16 |      3.00E-05 |     0.1 |
| CuBERT-EC      | Transformer |          4 |      1.00E-05 |     0.1 |
| CuBERT-FD      | Transformer |         16 |      3.00E-05 |     0.1 |
| CuBERT-WB      | GGNN        |         16 |      3.00E-05 |     0.1 |

> Text format, section title and typo:

We have fixed these issues and would like to thank the reviewer for the careful review.


 We thank the reviewer for their valuable comments and their strong support for this paper.

> Include DeepFix?

DeepFix uses a Seq2Seq-like architecture that predicts a line number and then a replacement value for the line. This would be straightforward to represent within PLUR by using a pointer to predict the line number and tokens and the copy mechanism to predict the replacement line. It could be added in the future, but we prioritized other program repair tasks (Hoppity and the VarMisuse tasks) in the initial release for the following reasons:
* Hoppity and VarMisuseH are similar tasks that also come with graph structures, which exercises more of PLUR's capabilities.
* CuBERT-VMR already provides an example program repair task without graph structure.
* It's a minor technical detail, but DeepFix creates training data on the fly by mutating correct programs to introduce errors. For reproducibility in PLUR, we have opted for datasets that are fixed and archived, and don’t depend on executing code to retrieve or on randomness. Although creating an archival fixed dataset for DeepFix is not technically challenging, we focused on datasets that are already fixed.

It would certainly be possible to include it, but for the above reasons and due to the fact that we already have 15 tasks to bring in and run experiments on, we chose not to prioritize it.

> Include code completion?

Thanks for the suggestion. A code completion task would be a nice addition, as the PLUR formulation generalizes text-to-text Transformers, and this is an effective current approach to code completion. We are look into adding a code completion task for the next revision of the paper.

> What kinds of tasks are not included or do not fit in the current framework

More broadly, one of the main limitations is the encoder-decoder formulation of PLUR. While it's a very general framework, it doesn't encapsulate all tasks. For example, it's not clear how retrieval tasks, knowledge graph-based tasks, and retrieval-augmented models could be beneficially implemented within PLUR. We also haven't added support for tasks with real-valued outputs. Thanks for this question. We'll add further discussion of these limitations to the text.

>  Technical challenges

Developing a unified framework such as PLUR definitely has its challenges. We touched on some of the key ones in the paper related to graph design and what API and modeling framework to unify around. Some additional considerations are as follows:
* Unification: Trying to unify the input pipeline, ML model and the evaluation process as much as possible, so that each task only requires minimum additional implementation, while keeping the framework extensible.
* Graph representation: Certain datasets come with custom graph structures that make sense in their custom ML model but require post-processing to fit a general GNN/transformer setup (see, e.g., our Code2Seq experiment in Sec 5.4).
* Large graph sizes and filtering: Very large graph sizes may lead to out-of-memory errors.  We either needed to filter these and treat them as failed examples (while implementing proper accounting) or prune these by heuristically selecting the most important part of the input.
* Larger batch sizes increase the total training epochs, however large batch sizes also sometimes suffer from out of memory issues.
* Evaluation: While not strictly a technical challenge, evaluating the models in the same way as reported in the original papers was not always straightforward, and it required care to get right. For example many function name prediction approaches use f1-score as the evaluation metric, but they can differ on how they tokenize, if they lower-case all characters, etc.

> Typo

Thanks, we have fixed this typo.

 We thank the reviewer for their valuable comments and their strong support for this paper. We sincerely hope that this paper will be a valuable contribution to the field and boost our common understanding of deep learning for code.

 We thank the reviewers for their time and thoughtfulness in the reviews.

The overall response to the paper is positive, with reviewers appreciating the originality and significance of our unifying formulation for solving ML for Code tasks: "it is original because it is the first unified framework for solving multiple types of ML4Code tasks ... answers several important research questions and demonstrated that PLUR can be used for multi-task learning, which is also novel" (ZqMx); "opening up a new research avenue ... notable insight for the ML4Code field" (gAPh); "an excellent and important contribution to ML4Code ... interesting, novel insights" (ZF27).

There were some questions raised by the reviewers, which we address in the responses below. Notably:
* Reviewer gAPh asks about extensibility to new tasks and suggests including code completion and C code repair like DeepFix. Reviewer 9ZHz asks whether the tasks we've chosen are representative enough. These are good suggestions and it would be straightforward to add more tasks to PLUR. We explain our rationale for preferring Hoppity and Variable Misuse to DeepFix in the individual reviewer response. We'll aim to add a code completion task for the next revision of the paper. More generally, PLUR is designed to be extensible to new tasks, and we are open to creating an expanded "PLUR v2" in the future, perhaps with help from the broader community if there is interest. However, for the current work we had to draw the line somewhere, and 15 tasks felt like a sufficiently large number to us.
* Reviewer gAPh and ZF27 have some writing and presentation suggestions and questions about hyperparameters. These are good questions and suggestions. We respond to the detailed points in the individual responses, and we'll incorporate the responses into the next revision.
* Reviewer 9ZHz thinks there are some missing BERT-based baselines. We clarify in our response that we are not considering any pretrained models in this work, and we explain in the individual response why it's appropriate to consider pretraining in future work. This is certainly an exciting future direction.
* Reviewer 9ZHz asks about how applicable the pointer and copy mechanisms are. We clarify in the individual response that the copy mechanism is enabled for all tasks and that pointers enable a number of tasks in the literature beyond just those that we've chosen for the initial paper (e.g., the DeepFix task discussed above).

We thank the reviewers again for their overall support for the work, and for the helpful feedback we will use to improve the paper.

The paper proposed a general method called PLUR for solving ML4code problems. PLUR is based on mapping a graph to a sequence of tokens and pointers. The author showed that 15 different ML4Code tasks can be casted in the form of PLUR. Experiments show that PLUR achieves near or above state-of-the-art results for nearly all tasks. In addition, the unified framework enables the author to study a few important research questions, including 1) performance of different encoders (GREAT, Transformer, GGNN), 2) importance of the pointer mechanism, 3) effect of different graph representations, 4) effect of multi-task learning. To be more specific, the author showed that 1) GREAT and Transformer encoders outperform GGNN, 2) multi-task learning with PLUR improves the performance in most datasets, 3) copy mechanism and pointer network improves the performance.
  
- Originality

Although the PLUR framework is based on the Graph2ToCoPo formulation, it is original because it is the first unified framework for solving multiple types of ML4Code tasks. In addition, the author answers several important research questions and demonstrated that PLUR can be used for multi-task learning, which is also novel.

- Technical Quality

The paper is technically strong. The author evaluated PLUR on 15 different datasets and conducted ablation study to answer the research questions. In addition, the multi-tasking results proved the advantage of the PLUR formulation.


- Clarity

The paper is well-written.

- Significance

The study conducted in the paper is significant and will boost the community's understanding of how to build a deep learning model for code data.

 The authors have adequately addressed the limitations and potential negative societal impact.

This work proposes a unified encoder-decoder framework, PLUR, that can cover 15 recently published source code tasks (e.g. code classification, code repair, function name prediction). The encoder generalizes sequence-based modeling (e.g. Transformer) and relation/graph-aware modeling. The decoder includes copying and pointer mechanisms. Their unified model, PLUR, outperforms task-specific models in previous literature, and also enables multi-task learning. The authors also study important research questions such as Transformer vs GNN. I think this paper deserves acceptance to NeurIPS.  Reasons to Accept:
- Various tasks and models have been proposed in the field of ML4Code, but the format is not standardized and hence it has been difficult to make systematic comparison of models (e.g. what model choices really matter in code understanding?). This work provides a nice step to unify source code tasks and modeling framework. Under this modeling framework, one can study research questions in a more controlled and reproducible way
- The unified model also facilitates multi-task learning for source code, opening up a new research avenue.
- The author also shows that their general model can outperform existing models designed for specific tasks, which is a notable insight for the ML4Code field.
- The paper is clear and well-written. The experiments are extensive.

Weakness/questions:
- PLUR covers diverse tasks, but there are some common tasks not included (e.g. code completion, C code repair like DeepFix). Is it because they do not fit the PLUG framework? It'd be great to discuss a bit more 1) what kinds of tasks are not included or do not fit in the current framework, and 2) if there is a prospect to cover more tasks by extending the PLUG framework.
- Were there any technical challenges that needed to be addressed when developing the unified framework? If there were, it might be worthwhile to mention these insights in the paper, as they may help future work to build on the proposed framework.

Typos/grammar:
- L373: "one risk is that" The authors discuss limitations and broader impact of their work in L373-379.

Meta-review: This paper presents a unifying, graph-based ML architecture for understanding and editing source code. The method, called PLUR, takes the general approach of mapping a graph to a sequence of tokens and pointers, a form into which many ML4Code tasks can be cast. The authors show that PLUR admits 15 recently published, diverse tasks, enabling a single model architecture to achieve near or above state-of-the-art results on them. PLUR is extensible to new problems beyond these initial 15. Significantly, its unifying framework enables approaches like multi-task learning and experiments on the importance of different modelling choices for ML4Code. These experiments, several of which are conducted in the present work, should guide the community toward robust improvements in ML4Code model development.

Reviewers are in consensus that the paper is above the NeurIPS acceptance threshold, with one championing it as a top-50% submission. They agree that it is technically strong, clear, and significant, as do I. The variety of tasks covered by the experiments is persuasive. Two reviewers mentioned that the framework does not exactly innovate on the modelling side, since it is based on the existing Graph2ToCoPo formulation; however, its unification of previously disconnected ML4Code approaches simplifies this important problem space in an insightful, clarifying manner that will almost certainly impact the field going forward. I recommend the paper for acceptance.

Decision: accept