WhyGen: Explaining ML-powered Code Generation by Referring to Training Examples

Understanding which training samples are more relevant to a certain generation is challenging, because neural networks are usually regarded as black boxes that are difficult to interpret. The training examples are used to compute gradients that accumulate into millions of model weights. It is hard to distinguish the contribution of each training example after the model parameters are learned.

Instead of analyzing which training examples contribute the most to the code generation, we analyze which training examples trigger similar decision logic as the user query. We assume the training examples with similar decision logic are the relevant examples for the generated code. This assumption, though not formally provable, is intuitive because human brains also process relevant concepts with similar decision pattern.

We introduce a data structure, named inference fingerprint, to represent the decision logic of the neural network and compare across different data examples. An inference fingerprint is a vector of activation values produced by a set of intermediate neurons in the network during the inference pass. The same set of intermediate neurons is used to produce the fingerprints, and thus the fingerprints are comparable across different data examples. Prior work has attempted to use intermediate neurons to represent the decision logic of DNN (Zhang et al., 2020; Liu et al., 2020), but they are mainly designed for other purposes (such as adversarial detection, data distribution estimation, etc.) and the computation of critical neurons is relatively slow.

In our work, the selection of the intermediate neurons for producing fingerprints must meet two objectives. First, the number of selected intermediate neurons must be small, since the total number of neurons in a code generator model is too huge to compute. Second, the selected intermediate neurons should be representative, so that the relevant code examples can be grouped together.

Modern code generators are mostly based on the Transformer architecture (Lu et al., 2021; Wang et al., 2021; Ahmad et al., 2021).A typical inference step of a Transformer-based code generator is illustrated in Figure 2, in which the input is a sequence of preceding code tokens, and the output is the predicted next token. Each piece of generated code is produced in a token-by-token manner, where each token is predicted by an inference step. The predicted token in a step is appended to the query sequence and used as the input to predict the subsequent token in the next step.

Taking CodeGPT (Lu et al., 2021) as an example, it takes a sequence of tokens as the input and predicts the next token step by step until the <end> identifier is predicted. In each step of next-token prediction, CodeGPT uses the Beam Search algorithm to retain the top-k candidate tokens with the highest scores. Then for each of these top-k candidates, it further runs the inference pass and finds the top-k highest-score candidate tokens, resulting in $K^{2}$ candidate combinations. Among them, only the top-k candidate combinations with the highest scores are kept in the next step, and the process repeats until the end of decoding. Finally, the candidate token combination with the highest score is returned as the final prediction.

We combine the heuristic understanding of the model and quantitative methods to locate the intermediate neurons. We first narrow down the selection of intermediate neurons to the activation layers after each encoder module, because they are designed as the result of each independent encoding stage. Moreover, we focus on the activation values corresponding to the first generated token since they have encoded all user-generated input tokens and are more explicitly related to the generated code.

To further locate the neurons that can better represent the decision process, we use a profiling phase to understand the behavior of the neurons in activation layers. The training samples are fed into the model and the neuron output values are recorded. We compute several statistics based on the profiling results and compare several criteria to select the critical neurons. We find that the most high-variance neurons are more representative, and their output values are concatenated together as the inference fingerprint.

Next, we compute the inference fingerprints for all training examples and save them to a database. The inference fingerprint generation process for the training examples is consistent with the process for user input (as described in Section 2.1), in order to ensure that the inference fingerprints of training examples can be compared and searched with the fingerprint generated by the user input at the test time. Each record in the database includes the inference fingerprint, the code snippet, and the original source (e.g., repository URL and/or file path) of the code. The fingerprint vectors are indexed to speed up the process of searching for the most relevant training examples.

When the code generator produces a prediction, we compute the inference fingerprint for the prediction, and find the most similar fingerprints in the database. The similarity is measured as the Euclidean distance between the two vectors. The training examples corresponding to the most similar inference fingerprints are returned to the user as the relevant training examples.

We implement the prototype of WhyGen with an open-source DNN-powered code generator CodeGPT (Lu et al., 2021), which is based on an advanced language model GPT-2 (Radford et al., [n. d.]) and fine-tuned on the PY150 dataset (Raychev et al., 2016). The state-of-the-art closed-source code generator, Codex or Copilot (Chen et al., 2021), is based on GPT-3 architecture. While larger in size, GPT-3 is conceptually and structurally similar to GPT-2. Thus, we believe our method can be applied to it as well.

To index and search for the fingerprints, we use the Faiss open-source library (Johnson et al., 2017). The size of the inference fingerprint is set to 100 in our implementation, and the number of returned relevant training examples is set to 10 by default.

Since the relevance of training examples is a subjective concept, directly evaluating it is difficult. Thus, we take an indirect approach instead - we first find some reciting behaviors of the code generator (i.e., the generator generates code exactly the same as in the training set). The recitations are regarded as the ground truth of relevant examples, so the effectiveness of WhyGen can be evaluated by examining whether the recited code snippets appear in the results produced by WhyGen.

To find the recitations, we randomly pick 10,000 code snippets from the test set and use the code generator to predict the next line for each snippet. For each predicted line of code, we search the training dataset to find the most similar line, i.e., the line with the shortest edit distance to the predicted line. If the edit distance is 0 and the code line is unique enough (number of occurrences is smaller than 10), we consider it as a recitation. In the end, we obtain 3,842 cases of recitations. We use the top-k accuracy metric to evaluate WhyGen, which means the probability that the recited training example is among the top k examples returned by WhyGen.

Based on the found recitations, we evaluate the effectiveness of WhyGen. Due to the lack of baselines in this area, we compare the default configuration of WhyGen with several variants. Each variant uses a different strategy to select the critical neurons to compute the inference fingerprint. For example, “Random” means to randomly select the intermediate neurons, “Maximum” and “Minimum” mean to select the neurons with maximum or minimum output values, and “FFN” means to select high-variance neurons from the feed-forward network layer rather than the self-attention layer.

The accuracy results are shown in Table 1. Clearly, our default configuration of WhyGen achieves the best results with a top-10 accuracy of 81.21% and top-1 accuracy of 73.84%, which is significantly better than using other criteria to select the fingerprint neurons. Selecting critical neurons from the FFN layer can achieve competitive results, but it is still slightly less effective than using the self-attention layers.

The accuracy results imply that the inference fingerprint computed by WhyGen does a good job in encoding important information about the decision-making process during the code generation, and it can effectively be used to find the training samples that share the similar decision logic with the query sample.

Figure 3 shows an example of relevant training examples returned by WhyGen when generating the next line for a given query code. We can see that the returned five training examples are almost all very relevant to the query code and generated code, and the first example is an exact recitation. In practice, the returned relevant training examples can serve as a reminder or guidance for the user. If the generated code recites or highly imitates the copyrighted code, the user can modify or abandon the generated code to avoid legal and ethical concerns. WhyGen will also provide the source path of the returned training examples, so that users can learn more about the code predicted by the code generator and decide whether to use it in their own software.

We further measure the overhead of WhyGen in training and serving scenarios using a Linux server with an AMD EPYC 7742 CPU.

In the training stage, WhyGen needs to compute the fingerprints for all training examples and build an index for the fingerprints. The whole process takes around 20 hours, which is shorter than the training time of code generator models (around 25 hours). We believe the training overhead is acceptable since it is a one-time offline cost.

In the serving stage, WhyGen needs to compute the inference fingerprint and retrieve relevant examples for each prediction made by the code generator. The overhead is around 6 ms, which is minimal as compared to the code generation process (360 ms). Thus, we believe our tool can be used in real-time to give meaningful prompts to the code generator users.