Intervention-Based Alignment of Code Search with Execution Feedback (2023)

Abstract: One of the fundamental goals in code search is to retrieve a functionally correct code for a given natural language query. As annotating for correctness requires executing test cases (i.e. obtaining execution feedback), existing code search training datasets approximate text-code co-occurrences as positive execution feedback. However, this approximation may misalign models’ retrieval decisions from ground-truth correctness. To address such limitation, we propose Code Intervention-based Reinforcement Learning (CIRL) that perturbs training code to result in misalignment (i.e. code intervention), then tests models’ decisions and corrects them with the execution feedback by reinforcement learning. The first technical contribution of CIRL is to induce the execution feedback from perturbation, without actual execution. Secondly, CIRL introduces structural perturbations using abstract syntax trees, going beyond simple lexical changes. Experimental results on various datasets demonstrate the effectiveness of CIRL compared to conventional approaches.

## introduction code search aims to retrieve code snippets that are most relevant to a given query. however, the notion of relevance often falls short of capturing the ultimate goal of functional correctness, focusing on whether the code would execute as the programmer intended @xcite . for instance, programmers would prefer c (in figure 1a ) over c 1 (in figure 1b ), as the entire "else" clause is missing in c 1 . ideally, ground truth labels for functional correctness should be obtained from execution feedback, collected after executing each code c with a sufficient amount of test cases for each query q. we denote such execution feedback as efpq, cq. however, conventional training datasets, such as ˚corresponding author. from bisect import bisect_left def find_lis(array): @xmath0: @xmath1, num) if @xmath2: dp.append(num) else: dp[idx] = num return len(dp) codesearchnet (csn; @xcite , collect data from github by regarding commented text descriptions as queries then labeling code snippets by query-code co-occurrences. this collection process can introduce an observation bias of lexical dissimilarity between positive and negative training snippets. consequently, models can learn such observation bias, resulting in misaligned retrieval decisions with ef of returning a negative code with high lexical similarity as a false positive. for example, figure 1b shows that a code search model θ, trained from a conventional dataset, decides to retrieve c foot_0 , or r θ pq, c 1 q " 1, but this model decision is misaligned with the execution feedback, or efpq, c 1 q " 0. the objective of this paper is to resolve the misalignment, i.e., r θ ‰ ef. a straightforward solution is to expose models to misaligned code pairs and correct the model decisions by ef. in reinforcement learning (rl), this trial-and-error process is achieved by interacting with agents while perturbing input states, referred to as intervention @xcite . we employ this intervention technique by perturbing the positive training code c to generate its negative counterpart c 1 , where both code snippets yield the same model decision. formally, r θ pq, c 1 q " r θ pq, cq " 1 but efpq, cq ą efpq, c 1 q. however, obtaining ef to correct decisions for each c 1 is expensive, and to make matters worse, there are infinite possible perturbations satisfying the above misalignment condition. to identify a representative c 1 , existing solution known as counterfactual perturbation @xcite @xcite , suggests to find c 1 that is lexically closest to c, but with different ef. however, adopting this strategy would expose models solely to trivial syntax errors, such as omitting a colon from figure 1a , causing a syntax error (and thus different ef). a realistic intervention, such as figure 1a to figure 1b , induces a larger lexical change, such as removing the "else" clause. our key technical contribution is a sampleefficient code intervention for rl. our proposed approach, namely code intervention-based reinforcement learning (cirl), perturbs a positive code into negative with ϵ structural changes, which is designed to include the intuition of minimal lexical edit, but generalizes beyond to enfold structural similarity. specifically, we leverage the abstract syntax tree (ast) representation of the code, treating subtrees such as statements, clauses and expressions, as units of structural changes. a structural perturbation is done by masking out a subtree from a positive code then filling it out unless model prediction changes. we stress that cirl subsumes counterfactual perturbation (when subtree is a leaf), and also adapts to the test distribution, by selecting subtrees with a realistic granularity (e.g. "else" clause in figure 1a ). further, we ensure syntac-tic validity while replacing subtrees, to discourage exposures to trivial syntax errors. we summarize our contributions as follows: first, we employ reinforcement learning with intervention that performs perturbation to simulate misaligned code without feedback from real code execution. second, we propose cirl, which utilizes ast information to inject ϵ structural changes to code to alter ef while preserving model decision. third, our experimental results on various datasets show that cirl effectively contributes to aligning code search with ef. all our implementation and datasets are publicly available 1 for future research purposes. ## preliminary in this section, we will begin by defining the code search task and addressing the misalignment issue in existing supervised code search training.