C ode F ort: Robust Training for Code Generation Models (2024)

Abstract: Code generation models are not robust to small perturbations, which often lead to incorrect generations and significantly degrade the performance of these models. Although improving the robustness of code generation models is crucial to enhancing user experience in real-world applications, existing research efforts do not address this issue. To fill this gap, we propose CodeFort, a framework to improve the robustness of code generation models, generalizing a large variety of code perturbations to enrich the training data and enabling various robust training strategies, mixing data augmentation, batch augmentation, adversarial logits pairing, and contrastive learning, all carefully designed to support high-throughput training. Extensive evaluations show that we increase the average robust pass rates of baseline CodeGen models from 14.79 to 21.74. We notably decrease the robustness drop rate from 95.02% to 54.95% against code-syntax perturbations.

## introduction code generation models @xcite @xcite @xcite have demonstrated impressive performance in generating code from natural language descriptions, completing sections of code, and even tackling complex coding contest challenges. these models can potentially assist software engineers and increase their productivity. however, code generation models are not robust to minor perturbations in the input prompts (e.g., inserting whitespaces/typos in docstrings or substituting variable names in code), i.e., they often generate incorrect outputs, thus significantly degrading their impressive performance on nominal prompts and hurting user experience when deployed in realworld applications @xcite . figure 1 shows that the performance of the state-of-the-art * work done when the author was an intern at amazon. 4 in the appendix for details.) public code models @xcite @xcite significantly declines under semantic-preserving program transformations, particularly under code-syntax perturbations. thus, it is crucial to improve the robustness of models before they are universally deployed. despite extensive research efforts to improve the robustness of code-related tasks, beyond code generation, such as vulnerability prediction, clone detection, and code summarization, existing work has not tackled two unique challenges of improving the robustness of code generation models, primarily trained using casual language modeling (clm). challenge 1: distinct robustness definition unlike traditional classification tasks like vulnerability detection, where models produce a single classification, code generation models generate sequences, leading to a shift in the definition of robustness for certain perturbations. in code generation, model robustness is defined by generating a coherent output given an input perturbation. in contrast, in classification tasks, models are expected to maintain the same classification before and after perturbation. for instance, if a variable i is renamed to b in the input prompt (as illustrated in figure 2d ), a robust code generation model should generate completions with the variable b instead of the original variable i. this shift necessitates a new category of perturbations for code generation models and corresponding robust training approaches to tackle this new category. challenge 2: designing robust training approaches as code perturbations can insert dead code or typos into training data, directly training using data augmentation could adversely affect the model performance, leading to issues like generating dead code or typos. furthermore, applying more deliberate robust training approaches such as adversarial logits pairing (alp) @xcite and contrastive learning (cl) to clm presents unique challenges. alp, designed for single-class classification, requires careful alignment between original and perturbed sequences, complicated by potential differences in sequence lengths. although cl has demonstrated efficacy in improving the robustness of code representations in masked language modeling (mlm) @xcite , its applicability to improving the robustness of code generation models remains unexplored. to tackle the above two challenges and improve the robustness of code generation models, we introduce a structured definition of code perturbations (section 3) and design a novel robust training framework named codefort (section 4). to address challenge 1, we classify existing code perturbations into two categories: context-free and context-sensitive, based on the formal definition of code perturbations provided in section 3. context-free perturbations, such as the docstring perturbation in figure 2b , follow the traditional notions of robustness, whereas context-sensitive perturbations, such as the code-syntax perturbation in figure 2d , specific the distinct robustness definition highlighted in challenge 1. the distinction of these two categories allows codefort to employ different robust training methods according to each category. to address challenge 2, codefort employs example-level and sequence-level pairing to enrich the training set. these two pairing levels allow 1) a masking mechanism to mask unnatural perturbed tokens and 2) a careful alignment between the original and perturbed token sequences, addressing the crucial challenge of applying alp and cl to clm. we utilize codefort to extensively evaluate various strategies, mixing data augmentation, batch augmentation, alp, and cl. our approach, combining batch augmentation with the masking mechanism, alp, and alpd, significantly enhances the model robustness and surpasses the sub-optimal results of data augmentation. notably, our approach significantly improves the robustness under codesyntax perturbations, the type of perturbation that hurts the model robustness the most, as shown in figure 1 . our ablation studies show that contraseq, the cl objective used in previous work for mlm, has negligible robustness improvements on clm. we summarize our contributions: 1) a framework, codefort, for improving the robustness of code generation models trained by clm, 2) designs of robust training approaches, data augmentation, batch augmentation, alp, and cl tailored to clm, 3) an extensive evaluation of different robust training approaches, 4) a perturbed training set for future studies on the robustness of code generation models, and 5) a surprising finding that the con-traseq cl objective, which is known to be beneficial for improving robustness of other code related tasks, has negligible robustness improvements on clm. robust training on code-related tasks existing work typically enhances model robustness through data augmentation and adversarial training @xcite . bielik and vechev (2020) refine model representations by feeding only pertinent program parts to the model; @xcite def largest_divisor (n: int ) -> int : ## problem definition we address the robustness challenge in a code generation model f trained using causal language modeling (clm). clm predicts the next token in a sequence, and the model can only attend to tokens on the left. formally, given a sequence of tokens @xmath0, . . . , x n , the generation model f captures p f (• | x :i ), representing the conditional probabilities of the i-th token given the preceding tokens utilizing a given decoding strategy, such as greedy or temperature sampling @xcite , the generation model f produces a sequence of tokens by iteratively predicting the next tokens until a specified stop criterion is reached. we denote f (x :i ) = xi: as the generated token sequence by f . the terms prompt, completion, and ground truth refer to the input x :i , the output xi: , and the original completion x i: = x i , . . . , x n , respectively. in code generation, the token sequence x represents a code snippet. figure 2a shows a problem in humaneval @xcite . each prompt x :i in a problem contains a function signature and a corresponding docstring description, and each groundtruth completion x i: is the correct function implementation. given a completion xi: generated by a model f , if the completed function x :i + xi: passes all the hidden test cases, the completion is deemed correct, denoted as cor(x :i + xi: ) = true. otherwise, if any of the tests fail, cor(x :i + xi: ) = false. ## conclusion we propose codefort, a framework to improve the robustness of code generation models, generalizing a large variety of code perturbations to enrich the training data and enabling various robust training strategies. our approach significantly enhances the model robustness and surpasses the sub-optimal results of data augmentation. notably, our approach significantly improves the robustness under codesyntax perturbations, the type of perturbation that hurts the model robustness the most. our ablation studies show that contraseq, the cl objective used in previous work for mlm, has negligible robustness improvements on clm.