Exploiting Code Symmetries for Learning Program Semantics

We really appreciate your constructive comments. We address each of your points below.

Weakness 1.1: The definition of the Aut(PDG) distance matrix to be a bit hard to reason about

We will clarify the definition of the distance matrix used in our paper. In general, there can be various distance matrices to represent the graph in self-attention layers and preserve automorphisms as long as it satisfies the two properties of Lemma 4 & 5 (see proofs in Lemma 6 in Appendix). We chose to represent the distance between two nodes using their distance to their lowest common ancestor, based on the intuition that this can potentially represent how they are related based on program dependencies, e.g., branch condition, same source of dataflow, etc.

Q1 & Weakness 1.2: Relationship between this distance matrix and the standard use of positional encodings

SymC does not use absolute positional embeddings fed to the input sequences. Instead, it follows a similar setting as relative positional embedding, which biases the attention computation, i.e., the product of query and key, with the program structure. Unlike regular relative positional embedding, where the distance between the pair of nodes is still computed based on their relative distance in the linear sequence, e.g., the distance between token 2 and token 5 is (5-2=3), we compute the relative distance based on the customized distance matrix using PDG, and we prove the self-attention biased by this distance matrix preserves the automorphism of PDG, thus the semantics-preserving code permutations.

Q4: Illustrate cases where the relaxation of the interpretation graph to the program dependence graph is too conservative

As described in our response to Reviewer 6HuJ, consider the case in binary analysis:

…
1       mov rcx, [rbx]
2       mov [rax], rdx
…

We have 2 nodes in PDG, i.e., {1,2} denoting 2 instructions, and the edge {1->2} because we can not easily eliminate the possibility that [rbx] and [rax] would point to the same memory location. This can lead to the spurious edge, rendering the permutation of 1 and 2 an invalid automorphism. In memory-bound applications, memory-accessing instructions can be very frequent. For these programs, the size of the automorphism group can indeed be very small (not many valid semantics-preserving permutations), restricting the number of possible permutations.

We appreciate your time and effort in reviewing the paper. We address each of your points below.

Q1: As stated in the paper, current code LLMs struggle with generalization to new code. Have you tried to evaluate the robustness of SYMC using the adversarial attack methods based on code transformations? The evaluation may make your method more convincing.

Thanks for bringing up the robustness angle against adversarial attacks. The key approach we started off from this paper is the development of a code model with provable robustness by construction against semantics-preserving code permutations. Guided by this framework, we proved that our approach is verifiably robust against any possible adversarial permutations, and thus not bound by minimal perturbations assumed in many adversarial attacks [1]. We leave exploring other code symmetry groups and their robustness guarantees as future work.

[1] Pierazzi, Fabio, Feargus Pendlebury, Jacopo Cortellazzi, and Lorenzo Cavallaro. "Intriguing properties of adversarial ML attacks in the problem space." In 2020 IEEE Symposium on Security and Privacy (S&P), pp. 1332-1349. IEEE, 2020.

Indeed, when evaluating other code transformations beyond permutations, we agree that adversarial attacks can potentially lead to a greater number of violations when they are guided properly, e.g., by the input gradient. Therefore, we evaluate SymC and the newly suggested code models against the adversarial attack for function name prediction, i.e., Averloc [2]. The adversarial transformations implemented in Averloc are comprehensive, and include the adversarial code transformations proposed in prior work by Yefet et al. [3], e.g., variable renaming, dead code insertion, etc., with additional transformations such as loop unrolling. We follow the setting in Averloc by computing the adversarial attacks against a seq2seq model trained by the Averloc authors, and evaluate SymC and the baselines (GrahCodeBERT as suggested by the reviewer, and DOBF and CodeT5 as suggested by reviewer Srm2) on the generated adversarial examples. This ensures a fair comparison by evaluating all the models on the same set of adversarial examples. The baselines are trained on the same exact training set as SymC.

The table shows that SymC outperforms the second-best baseline, GraphCodeBERT, by 10.7% and 49%, in F1 and violation rate on the adversarial examples, respectively, against adversarial transformations (beyond semantics-preserving permutations). This indicates the strong robustness of SymC against adversarial code transformations, even though the attacks are not statement permutations.

[2] Henkel, Jordan, Goutham Ramakrishnan, Zi Wang, Aws Albarghouthi, Somesh Jha, and Thomas Reps. "Semantic robustness of models of source code." In 2022 IEEE International Conference on Software Analysis, Evolution and Reengineering (SANER), pp. 526-537. IEEE, 2022.

[3] Yefet, Noam, Uri Alon, and Eran Yahav. "Adversarial examples for models of code." Proceedings of the ACM on Programming Languages 4, no. OOPSLA (2020): 1-30.

Dear reviewer nM8R,

Thanks again for making the effort to review our paper. We would be grateful if you could let us know whether our responses, i.e., the experiments on adversarial attacks and GraphCodeBERT, satisfactorily answer your questions and if so, would you like to increase the score. Thanks!

• New results

However, we agree with the reviewer that there are many ML4code tasks, like defect prediction, that can benefit from code symmetry. To demonstrate that effect, we follow the reviewer’s suggestion by incorporating the defect prediction and include our preliminary results in the following.

Particularly, we consider defects4j dataset [7] and follow CodeXGLUE to formulate the task as a binary classification (buggy/non-buggy) given the function. We ensure the number of buggy and non-buggy samples are the same, and measure the prediction accuracy on the original and permuted samples. Our baselines include GraphCodeBERT [8], DOBF [9], CodeT5 [10], UnixCoder [11], and CodeBert [12].

The above table shows that SymC outperforms all the baselines, even when they are pre-trained on large-scale code datasets, by 11.4% on average. For example, SymC outperforms the second-best model, UnixCoder, by 2.5% in accuracy on permuted samples. Similar to function name prediction, all the baselines are susceptible to simple statement permutations. We will add more discussions to the related works and experiments in the paper.

[7] Just, René, Darioush Jalali, and Michael D. Ernst. "Defects4J: A database of existing faults to enable controlled testing studies for Java programs." In Proceedings of the 2014 International Symposium on Software Testing and Analysis, pp. 437-440. 2014.

[8] Guo, Daya, Shuo Ren, Shuai Lu, Zhangyin Feng, Duyu Tang, Shujie Liu, Long Zhou et al. "GraphCodeBERT: Pre-training code representations with data flow." arXiv preprint arXiv:2009.08366 (2020).

[9] Roziere, Baptiste, Marie-Anne Lachaux, Marc Szafraniec, and Guillaume Lample. "Dobf: A deobfuscation pre-training objective for programming languages." arXiv preprint arXiv:2102.07492 (2021).

[10] Wang, Yue, Weishi Wang, Shafiq Joty, and Steven CH Hoi. "Codet5: Identifier-aware unified pre-trained encoder-decoder models for code understanding and generation." arXiv preprint arXiv:2109.00859 (2021).

[11] Guo, Daya, Shuai Lu, Nan Duan, Yanlin Wang, Ming Zhou, and Jian Yin. "Unixcoder: Unified cross-modal pre-training for code representation." arXiv preprint arXiv:2203.03850 (2022).

[12] Feng, Zhangyin, Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng, Ming Gong, Linjun Shou et al. "Codebert: A pre-trained model for programming and natural languages." arXiv preprint arXiv:2002.08155 (2020).

We really appreciate your constructive comments. We address each of your points below, with additional experiments. We think that this clarification might significantly affect your view of the paper. We hope that you will take this into consideration.

Q1 & Weakness 1: The paper could benefit from a more comprehensive comparison with related works, such as DOBF (https://arxiv.org/abs/2102.07492), which exploits code symmetry in pretraining through a deobfuscation objective, and CodeT5 (https://arxiv.org/abs/2109.00859), which leverages code symmetry in pretraining with identifier-aware data augmentation. These related works were not discussed or compared to the proposed method in the paper.

Thanks for introducing the related works to us. As shown in our response to Reviewer nM8R, Q2, we include the new results comparing SymC to DOBF, CodeT5, and GraphCodeBERT. These baselines are pre-trained, while SymC is not. Therefore, to enable a more fair comparison, we consider both pre-trained and non-pre-trained baselines and also a pre-trained version of SymC (similar to the experiment setup in Section 6.3, Figure 4(b)). In particular, the pre-trained version of SymC was pre-trained using masked language modeling on the functions provided by code2seq (java-small in the paper).

Not pre-trained:

Pre-trained:

The results in the above tables show that, without pre-training, SymC outperforms the second-best model, CodeT5, by 43.2% and 42.8%, in original and permuted test samples, respectively. The baselines have 24.3% violation rate on average when the samples are permuted, while SymC enjoys provable invariance by construction. In addition, we observe that pre-trained CodeT5 has 10.7x higher violation rate than that when it is not pre-trained. This is likely due to that the model learns spurious positional biases not robust to statement permutations during pre-training.

We will add more discussion to the related work, and complete Table 2.

Weakness 3: The paper does not discuss potential limitations of the SymC model, such as its requirement for input code to be processed by a parser and a static analysis tool. This requirement may limit the model's applicability when dealing with incomplete or syntactically incorrect code, such as in code completion tasks or when faced with an empty Python block. While it is acceptable to establish certain assumptions for input code, these assumptions must be explicitly discussed rather than overlooked.

As shown in our response to Reviewer 6HuJ, statically analyzing the input code incurs additional overhead, e.g., 250.2ms per sample, but still much less than the existing model. That being said, we agree that efficiently analyzing incomplete or syntactically incorrect code on the fly can be an interesting future direction. We will add the discussion of limitations in the paper.

Dear reviewer Srm2,

Thanks so much again for your constructive comments. We would be grateful if you could let us know whether our responses, i.e., the experiments on comparing to code baselines, e.g., DOBF, CodeT5, GraphCoderBERT, CodeBERT, and UnixCoder, and the evaluation on a new task, i.e., defect prediction, satisfactorily answer your questions and if so, would you like to increase the score. Thanks!

Q2 & Weakness 3: How costly is obtaining the code graphs? Computation of graphs associated to code may be costly and restrictive.

To incorporate the inductive bias of code, many code models involve extracting and encoding code structures, including our baselines, as computing PDGs statically is not overly expensive. Below, we add the runtime performance (in milliseconds) per code sample of SymC and PalmTree on binary analysis, training, and inference, using the same exact hardware (an AMD EPYC 7502 processor, 128 virtual cores, 256GB RAM, Nvidia RTX 3090 GPUs).

SymC’s cheap computation of PDG incurs 17.8x less runtime overhead. It is also an interesting research problem to incorporate optimization, e.g., caching, to improve the efficiency of the PDG computation during inference. We will update the paper to include both source and binary analysis tasks.

Q3: Why does SymC require "40,162 lines of code"? I'm curious as to what makes it require so much.

Besides implementing the model and algorithm (~4.5k lines), we include the lines of code for doing program analysis (both source and binary code) for obtaining PDGs (2.4k lines) and generating ground truths for the analysis tasks (3k lines). We also include the code for setting up the baselines (4.3k lines), implementing program transformations (2.8k lines), data processing (1k), and all the scripts for experiments and result analysis (22k lines). We will clarify in the paper by providing a breakdown for each of these components.

We really appreciate your constructive comments. We address each of your points below.

Q1 & Weakness 1: Could you illustrate an example program interpretation graph and program dependence graphs? This would be quite helpful for understanding.

Consider the following Java example:

……
1       int a = 1;
2       int b = 2;
3       if (a > b) {
4           return 1;
5       } else {
6           return 0;
7       }
……

We will have 5 nodes in IG, i.e., {1,2,3,4,6} denoting 5 statements, and the edges {1->3, 2->3, 3->4, 3->6}. In such a case, IG is equivalent to PDG.

Consider the case in binary analysis:

…
1       mov rcx, [rbx]
2       mov [rax], rdx
…

We have 2 nodes in IG, i.e., {1,2} denoting 2 instructions, and there is no edge in actual IG, as there is no dataflow from 1 to 2. However, in PDG, we conservatively consider the edge {1->2} by marking them as a write-after-read dependence. This is because we can not easily eliminate the possibility that [rbx] and [rax] would point to the same memory location.

Weakness 2: Experimental details are lacking. What is the training procedure for SymC, is it just supervised training on the downstream task? How about the other models? For Function Name prediction, do the LLMs take in just the text as input, and what exactly does SymC take as input there?

SymC adopts only supervised training and trains the model from scratch without any pre-training. While Figure 4(b) demonstrates that adding pre-training brings extra improvement, we aim to demonstrate having symmetry-preserving architecture alone leads to better performance than many pre-trained models.

For other baselines, we respect their own training paradigm. In particular, PalmTree, GPT4, and WizardCoder are pre-trained on programs in the wild. code2vec, code2seq, and GGNN use supervised training. The new models that we have experimented with during the rebuttal, i.e., CodeBert, GraphCodeBERT, UnixCoder, DOBF, and CodeT5 (see our responses to reviewer Srm2), are pre-trained with large-scale code datasets.

For function name prediction, both LLMs and SymC take as input the function (function body and function signature) with the function name stripped.

Q3: Relatedly, models like GGNN may be a better comparison. Notably that model has lower invariance violation than LLMs. Given that GGNN also works on graphs associated to code, could you explain why it is not invariant to the code symmetry transformations you consider?

We will answer this question from both conceptual and experimental perspectives.

Conceptually, while GGNN incorporates the data and control flow graph, its particular dataflow construction does not preserve invariance to semantics-preserving permutations. For example, GGNN maintains edges for all reads and writes to the same variable token, e.g., using LastRead and LastWrite edge types [1]. This will prevent valid permutation between two statements that only read the same variable (Read-after-Read). Moreover, the permuted nodes can potentially lead to changed edge types. Such designs can result in the changed model’s output while its input transformation is semantics-preserving, breaking the robustness guarantees that SymC offers.

[1] Allamanis, Miltiadis, Marc Brockschmidt, and Mahmoud Khademi. "Learning to represent programs with graphs." arXiv preprint arXiv:1711.00740 (2017).

However, we agree that GNNs based strictly on our PDG will naturally preserve the Aut(PDG)-invariance. We aim to study how such a GNN architecture compares to SymC in our future work. The line of works on the graph Transformer suggests that Transformer with graph structures outperforms both the vanilla Transformer and GNN [2,3], but it remains interesting to study their robustness guarantees to code [3].

[2] Ying, Chengxuan, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie-Yan Liu. "Do transformers really perform badly for graph representation?." Advances in Neural Information Processing Systems 34 (2021): 28877-28888.

[3] Hellendoorn, Vincent J., Charles Sutton, Rishabh Singh, Petros Maniatis, and David Bieber. "Global relational models of source code." In International conference on learning representations. 2019.

Experimentally, we observe the particular low-invariance-violation on GGNN in our experiments is because it keeps predicting the same meaningless label, i.e., its F1 score remains at 0.016 across all the permutation percentages (Table 2). For example, in more than 60% of the samples, GGNN implementation provided by [4] (the original GGNN paper [1] did not perform function name prediction) keeps predicting a meaningless name “SPR9486” on disparate samples.

[4] Rabin, Md Rafiqul Islam, Nghi DQ Bui, Ke Wang, Yijun Yu, Lingxiao Jiang, and Mohammad Amin Alipour. "On the generalizability of Neural Program Models with respect to semantic-preserving program transformations." Information and Software Technology 135 (2021): 106552.

1  private void applyNormalPattern(String modeStr, Matcher matcher) {
2      boolean commaSeperated = false;
3      for (int i = 0; i < 1 || matcher.end() < modeStr.length(); i++) {
4          if (i > 0 && (!commaSeperated || !matcher.find())) {
5              throw new IllegalArgumentException(modeStr);
6          }
7          String str = matcher.group(2);
8          char type = str.charAt(str.length() - 1);
9          boolean user, group, others, stickyBit;
10        user = group = others = stickyBit = false;
11  }

1  private void applyNormalPattern(String modeStr, Matcher matcher) {
2      boolean commaSeperated = false;
3      for (int i = 0; i < 1 || matcher.end() < modeStr.length(); i++) {
4          if (i > 0 && (!commaSeperated || !matcher.find())) {
5              throw new IllegalArgumentException(modeStr);
6          }
7          boolean user, group, others, stickyBit;
8          String str = matcher.group(2);
9          char type = str.charAt(str.length() - 1);
10        user = group = others = stickyBit = false;
11  }

Q1: The example and runtimes are helpful, I definitely recommend putting them in the new versions of the paper.

We are grateful for your prompt response. We will definitely include them in the paper - we are working on drafting the figure/table and adding the discussion. We will update you as soon as we finish the draft.

Q2: One issue I have with the paper is the framing of SymC as an LLM (this is why I asked the question about the training procedure). The model is a Transformer with graph information included via relative positional encodings. But it cannot be trained or used for next-token prediction, but it can be used for supervised tasks on parts of code.

We appreciate you pointing out this issue. We agree that our focus in this work is primarily on the Transformer encoder with bidirectional self-attention, without explicitly demonstrating its application in the context of a Transformer decoder, as commonly used in most LLMs today.

However, we would like to highlight Section 3.5, “Token-level predictor” and Lemma 3, and the proof in Appendix B in the paper. We demonstrate how SymC preserves $Aut(\mathcal{IG})$-invariance at the individual token level. This ensures end-to-end $Aut(\mathcal{IG})$-invariant property in a Transformer decoder. As the autoregressive decoder generates a single token one at a time, attending to all previous tokens is in the same manner as the token-level predictor, and it is $Aut(\mathcal{IG})$-invariant according to Lemma 3.

We will edit the framing of the paper, by focusing more on how SymC augments the Transformer architecture, and leave a detailed evaluation of LLM (autoregressive decoder) as future work. We will edit Section 3.5 to make the discussion of Transformer decoder more explicit. Thanks for helping us improve the clarity of our paper.