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

1 Introduction . The advent of sequence-to-sequence [ Sutskever et al. , 2014 ] and , more recently , text-to-text [ Raffel et al. , 2019 , Radford et al. , 2019 ] abstractions has provided a simplifying and unifying view of much work in natural-language processing . By casting problems in this framework , one can apply a single model architecture to many different problems . This benefits machine learning ( ML ) researchers by focusing attention on the core modeling problem . It also amplifies ML advances , because a general formulation allows practitioners to frame new problems in terms of a standardized abstraction and then to apply state-of-the-art architectures easily , which also facilitates multi-task and transfer learning . In this paper , we ask : What is the equivalent unifying abstraction for machine learning for source code ( ML4Code ) ? While many ML4Code tasks have been cast in terms of the sequence-to-sequence abstraction , source code is inherently different from natural language [ Hindle et al. , 2016 ] . It is ∗Work done during internship at Google . †Work done during visiting-faculty appointment at Google . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . syntactically more structured as it is designed to be machine-readable and unambiguously parsable . As such , it lends itself to more sophisticated automated analysis than natural language ( e.g. , static analysis ) . Furthermore , code tokens follow a power-law distribution due to many rare identifiers , and exhibit high rate of repetition across files and projects [ Allamanis and Sutton , 2013 , Casalnuovo et al. , 2019 ] . In addition , the software engineering community recognizes that significant effort goes into software maintenance involving small changes , rather than writing code from scratch [ Koskinen , 2003 ] . These considerations have led to a wealth of literature going beyond text-to-text models that make use of graphs , copy mechanisms , pointers , and other custom architectures [ Maddison and Tarlow , 2014 , Bielik et al. , 2016 , Allamanis et al. , 2016 , Mukherjee et al. , 2017 , Yin and Neubig , 2017 , Allamanis et al. , 2018 , Alon et al. , 2018 , Dinella et al. , 2020 , Yasunaga and Liang , 2020 , Nye et al. , 2020 ] . In particular , the GRAPH2TOCOPO formulation from Tarlow et al . [ 2020 ] pairs a graph encoder with a Transformer-style decoder augmented with pointers and a copy mechanism and has been shown to be effective for certain ML4Code tasks . Our proposal is thus to unify around GRAPH2TOCOPO as a common vernacular for ML4Code . This formulation strictly generalizes text-to-text , by additionally supporting graph structures over the input when available , allowing the generation of both pointers and tokens in the output—pointers to specify the input location to edit and tokens to specify the content of the edit—and it has a built-in copy mechanism to more easily enable the re-use of rare tokens from the input . This flexibility enables us to take 16 tasks and models from the ML4Code literature and convert them to a unified form . Using the same hyperparameter sweep and GRAPH2TOCOPO models , we achieve at or above state-of-the-art results on nearly all tasks , including improving over the Hoppity model [ Dinella et al. , 2020 ] for program repair , the GREAT model [ Hellendoorn et al. , 2020 ] for variable misuse , and the code2seq model [ Alon et al. , 2018 ] for variable naming . Having unified a number of previous tasks and achieved strong results , we can then answer a number of empirical questions : Can we remove the graph structure as input and just use a Transformer encoder ? How important are pointers and copy mechanisms ? We can also inspect the differences between our model and previous ones and observe the importance of different design choices . The framework also makes it easy to experiment with alternative graph representations of code , and we demonstrate the effect of five different choices for converting a set of abstract syntax tree ( AST ) paths into graphs in the code2seq domain . Finally , we demonstrate that due to the unified representation , the framework makes it easy to do multi-task learning across classification and repair tasks , yielding improved performance compared to training the same model separately on each task . In summary , the resulting framework , called PLUR for Program Learning , Understanding , and Repair , provides a powerful and simplifying perspective from which to interpret and build upon many recent ML4Code advances . Open-source code for the full PLUR framework will be available under an Apache 2 license at https : //github.com/google-research/plur . 2 Background and Related Work . Encoders for Source Code Much work has shown the benefit of structured representations of code based on abstract syntax trees ( ASTs ) or richer semantic graphs containing control flow and data flow [ Allamanis et al. , 2018 , Brockschmidt et al. , 2019 , Hellendoorn et al. , 2020 , Alon et al. , 2018 , Cvitkovic et al. , 2019 ] . A number of recent models appear to achieve the best of both worlds between graph neural networks ( GNNs ) and Transformers , by using the graph structure in conjunction with relative attention in Transformers [ Hellendoorn et al. , 2020 , Shaw et al. , 2018 , Wang et al. , 2020 ] . Decoders for ML4Code Tasks The literature on ML4Code proposes a wealth of decoder architectures , targeting a diverse array of tasks . Typical classification tasks ( e.g. , bug detection ) tend to require a simple softmax operation over an output vocabulary . Sequence prediction tasks ( e.g. , function-name prediction ) extend this simple mechanism with an autoregressive decoder [ Alon et al. , 2018 ] , or a more complex autoregressive structured decoder [ Chen et al. , 2018 ] . Other tasks require the ability to point to the input , e.g. , to localize a bug [ Vasic et al. , 2019 , Hellendoorn et al. , 2020 ] , or point to where an edit must take place [ Dinella et al. , 2020 ] . Decoders for such tasks use a pointer network based on attention [ Vinyals et al. , 2015 ] . To handle the very large vocabularies of program identifiers and string literals , pointers are also often used to copy content from the input , rather than extract it ( or reconstruct it ) from a limited vocabulary [ Allamanis et al. , 2016 ] —this has also been used in some natural-language contexts [ Gu et al. , 2016 ] . Unlike using pointer decoders to produce pointers required by the task output , such mechanisms use pointer decoders to identify content in the input from which to copy into the output , e.g. , for summarization [ See et al. , 2017 , Fernandes et al. , 2018 ] . Tarlow et al . [ 2020 ] proposed combining such decoders into a single decoder that produces TOCOPO , an autoregressive sequence where each element may be a TOken , a COpy , or a POinter . Source-Code Benchmarks and Frameworks Several recent efforts have created benchmarks of tasks and datasets for ML4Code . CuBERT [ Kanade et al. , 2020a ] released a benchmark of six tasks , for the purpose of evaluating BERT-style pre-trained code embeddings . CodeSearchNet [ Husain et al. , 2019 ] created a benchmark and a framework for competing implementations targeting code search and retrieval . CodeXGLUE [ Lu et al. , 2021 ] further expanded the benchmark to include tasks from code summarization , repair , code generation , and more . It also provided several baseline model implementations of different characteristics ( e.g. , CodeBERT [ Feng et al. , 2020 ] similar to BERT [ Devlin et al. , 2019 ] , GraphCodeBERT [ Guo et al. , 2021 ] incorporating data-flow edges , etc . ) , each adapted to a particular subset of task types ( e.g. , a GPT-based model for sequence generation , a BERT-based model for classification , etc. ) . In contrast , PLUR seeks a single encoder interface , accepting graphs , and a single decoder interface , producing a mixture of tokens , copies , and pointers , with the thesis that this single interface , regardless of model architecture underneath the encoder and decoder , is sufficient to serve a diverse set of ML4Code tasks . 3 A Unifying Framework for Program Understanding and Repair . Our goal from the modeling perspective is specifically not to innovate on machine-learning architectures . Our challenge is to package strong existing modeling components in a way that achieves ( near ) state-of-the-art performance but with maximum simplicity and generality . The novelty in our work comes from making these choices and demonstrating that the result performs strongly across the 16 tasks . Fig . 1 illustrates the approach . See details of the software architecture in the Appendix . 3.1 Models . Based on their strong performance across many tasks in machine learning , we start with Transformers as the base modeling architecture . The encoder is applied to the source-code input , and the decoder generates task-specific sequential outputs . However , we augment the encoder and decoder to incorporate components that have proven useful in recent work on modeling code , most notably relational information that can be used to represent syntax , data flow , and control flow . We adopt the GREAT model as an encoder [ Hellendoorn et al. , 2020 ] , an instance of a Transformer with relational attention bias . Similarly , we adopt the TOCOPO output formulation from Tarlow et al . [ 2020 ] to provide a flexible language that enables outputting arbitrary sequences of tokens , copies , and pointers , along with its modified Transformer decoder architecture . For comparison , we also instantiate this GRAPH2TOCOPO architecture with a vanilla Transformer encoder , and with a graph neural network encoder ( a Gated-Graph Neural Network [ Li et al. , 2016 ] , or GGNN ) . 3.2 Input-Output Representations . The next challenge is to adapt a wide variety of recent tasks and datasets into the proposed framework . In some cases , the translation is straightforward . For example , tasks that use GNN encoders and have already released graph representations can be directly used . In other cases , there are custom architectures that do not obviously fit into our proposed unification . However , we show via experiments in Sec . 5 that a light adaptation is able to preserve the favorable inductive biases of the custom approach without the custom modeling . While we use standard representations of inputs and outputs , we describe them explicitly here to make clear what the requirements are for mapping an existing task into the PLUR framework . The input is represented as a graph composed of nodes and edges . Each node has a discrete label , a discrete type , and an integer position . The label is used to represent token values like names of ( possibly subtokenized ) variables . The type is used to represent more abstract categories , e.g. , to distinguish that some nodes represent source code text and others represent internal AST nodes . The positions are used to impose a linear ordering on the nodes . This enables a vanilla Transformer encoder to process the graph like a token sequence . Edges ( s , t , e ) are directed , linking a source node index s to a target node index t. Additionally , edges are allowed to have a discrete type e that specifies the kind of relationship that the nodes have . Our graph edges are directed , but in preprocessing we always add a corresponding reverse edge ( t , s , ereverse ) with a fresh type ereverse . The TOCOPO output can be intuitively viewed as a script that describes the task output in terms of tokens , drawn from the output vocabulary , and pointers pointing to some input node , concluding with a DONE token marking the end of the output . Every task can make its own use of these facilities to express a grammar for its output . For example , a classification task can just use token outputs , one per expected class ; a sequence-prediction task can produce sequence of tokens ; a repair task can use a pointer to point at a particular input node , and a token output to replace that input node . The copy mechanism in the model enables the production of output by copying the label from an input node .

Summary: 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.