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