J ump C oder: Go Beyond Autoregressive Coder via Online Modification (2024)

Abstract: While existing code large language models (code LLMs) exhibit impressive capabilities in code generation, their autoregressive sequential generation inherently lacks reversibility. This limitation hinders them from timely correcting previous missing statements during coding as humans do, often leading to error propagation and suboptimal performance. We introduce JumpCoder, a novel model-agnostic framework that enables human-like online modification and non-sequential generation to augment code LLMs. The key idea behind JumpCoder is to insert new code into the currently generated code when necessary during generation, which is achieved through an auxiliary infilling model that works in tandem with the code LLM. Since identifying the best infill position beforehand is intractable, we adopt an infill-first, judge-later strategy, which experiments with filling at the k most critical positions following the generation of each line, and uses an Abstract Syntax Tree (AST) parser alongside the Generation Model Scoring to effectively judge the validity of each potential infill. Extensive experiments using six state-of-the-art code LLMs across multiple and multilingual benchmarks consistently indicate significant improvements over all baselines. Our code is available in the uploaded attachment.

## introduction recently, large language models (llms) @xcite have attracted considerable interest and achieved notable successes. as a prominent application of llms, numerous code llms @xcite @xcite @xcite @xcite are frequently released and exhibit exceptional efficacy in code generation tasks. despite their remarkable success, current llms work in an autoregressive style, generating text segments sequentially from left to right. this leads to an inherent limitation of irreversibility -they are incapable of revising previously generated text. notably, even skilled human programmers often struggle to write code in a linear style, since realworld code commonly undergoes iterative editing and refinement @xcite . for example, to accommodate the developing code, programmers need to retrospectively introduce key declarations, such as new variables, references, and functions in the previously written code. during the early stage of generation, code llms often miss key declarations due to the high generation uncertainty @xcite . because of the irreversibility limitation, code llms, unlike humans, cannot add missing declarations when they need these declarations later. this often results in producing code with undefined identifiers, known as nameerror in python. notably, for the widely-used codellama @xcite , nameerrors accounted for up to 6.0% to 18.6% of the total error cases on humaneval in our experiments. although these nameerrors are relatively easy to address by syntax checking or model constraints @xcite , the lack of necessary declarations or statements in the generated code segments often hinders llms from continuing to generate semantically correct code, causes the error to accumulate continuously along the generation, and limits the performance of llms. fig. 1 presents an intuitive example, demonstrating that in the absence of necessary variables, llms are prone to generating syntactically correct but semantically incorrect code, which is more difficult to detect and rectify. to address this limitation, the crux is to enable the model to jump back to previously generated code for inserting new code like a human whenever a missing declaration or statement is required. your goal is to separate those group into separate strings and return the list of those. your goal is to separate those group into separate strings and return the list of those. a major challenge to instantiate this idea is deciding whether (and where) to infill, or continue generation from the current code. to overcome this, we propose an infill-first, judge-later paradigm: after generating a line, let the infilling model experiment with filling at the start of the k most critical lines that have been generated, and subsequently judge their contributions to the current generation. this paradigm benefits from parallel generation and speculative infilling optimizations, which avoids a significant reduction in efficiency. the next challenge is how to implement the judgment method. we combine an abstract syntax tree (ast) parser and the generation model scoring to achieve this. ast parser is employed for the code using undefined identifiers, which accepts the infill that correctly adds the missing declaration. for other scenarios, the generation model scores the code following each infill position by comparing their mean token logit improvements. if an infill enhances the score of subsequent lines, we assume it indicates an improvement in overall generation quality. in cases outside these two scenarios, we infer that infilling is unnecessary, and opt to generate the next line. extensive experiments on six variants of codellama @xcite and @xcite on five code generation benchmarks indicated a consistent and significant enhancement across all baselines. notably, jump-coder is adaptable to various programming languages, and helps the code llms achieve a pass rate increase of 4.8% -8.2% in four multilingual humaneval benchmarks at most, while substan-tially reducing undefined identifier errors. to the best of our knowledge, our approach enables online modification for code generation for the first time. we hope our work can inspire further research into the irreversibility limitation of llms in the future. our contributions are threefold: code generation code generation, an important task of software engineering, focuses on automatically producing source code from software requirements (liu et al., 2022). the emergence of llms has propelled this field, demonstrating remarkable capabilities in code generation (rozière et al., 2024; wei et al., 2023; luo et al., 2023), especially when fine-tuned on domain-specific code datasets. code infilling conventional language models typically train on left-to-right next token prediction. although llms excel in code generation, they face limitations in code editing tasks like code infilling, where handling bidirectional context is crucial. to overcome this, the fill-in-the-middle (fim) training method (bavarian et al., 2022) has been proposed. fim divides input code into three parts (prefix, infix, suffix), then reorders them into a @xmath0. this trains the model on the conditional distribution p (infix | <@xmath1, enhancing its infilling capabilities. however, this approach doesn't directly translate to code generation, where subsequent contexts are absent. ## jumpcoder this section elaborates on jumpcoder, which is designed to address the irreversibility limitation of autoregressive generation. such limitation often leads to the generation of syntactically or semantically incorrect code. to address it, the key is to facilitate the model's ability to achieve online modification, e.g., jump back to the previously generated code to insert a new line like a human, which enables code llms to add missing statements. based on it, we introduce an innovative hybrid generation approach that enables the retrospective insertion of new lines by assisting the generation model with an infilling model, as shown in fig. 2 . however, identifying the optimal locations for infilling within the current code is a significant challenge. we discovered that while pinpointing the best infill location beforehand is intractable, assessing the suitability of an infill post-insertion is comparatively simpler. therefore, we propose an infill-first, judge-later strategy: after generating a line, the infilling model experiments with inserting a line before each of the k most critical lines that have been generated, followed by judging their impact on the current generation. this judging process is crucial as we found that unnecessary infills tend to degrade code quality (see appendix c for concise examples). specifically, jumpcoder iteratively updates code line by line. each iteration adds a new line to the current code, which encompasses three steps: hybrid generation ( § 3.1), involving simultaneous line generation by both models; judging ( § 3.2), assessing the infilling model's output to determine line selection; and combination ( § 3.3), merging this selection into the current code before continuing proceeding. the complete framework is outlined in fig. 3 , and the algorithm is summarized in appendix a. ## analysis ablation studies we evaluated jump-coder (vanilla) by excluding the ast parser (labeled as w/o ast) and generation model scoring (labeled as w/o scoring). fig. 6 (a) illustrates the performance improvements of these variants compared to the complete jumpcoder on various generation models. the results indicate both variants were less effective than the full jumpcoder. this is attributed to the complementary roles of the ast parser and generation model scoring: the ast parser deterministically accepts the infills that resolve undefined identifier errors, which can be probabilistically rejected by generation model scoring. conversely, generation model scoring is able to accept more good infills. this synergy allows more effective infills, highlighting the enhanced performance achieved by integrating both judging components in jumpcoder. ic -1 b c l-7b c l-i-7b c l-13 b c l-i-13 b infilling model 35 36 37 38 39 40 humaneval pass@1 (%) v f o (a) infilling model 0 1 2 3 36 37 38 39 humaneval pass@1 (%) w/o jumpcoder vanilla filtered oracle (b) improvement threshold τ 2 4 6 8 10 k 36 37 38 39 humaneval pass@1 (%) w/o jumpcoder vanilla filtered oracle (c) infill at top-k positions figure 7: impact of infilling models and hyperparameters on three variants of jumpcoder on humaneval. the generation model is codellama-instruct-7b, whose performance is marked in the green dashed line. @xmath2. @xmath3. cl-@xmath4. ## related work in this section, we review two lines of related work: code generation and code infilling. code generation training large language models with an extensive corpus of code is a common approach for code generation tasks @xcite . there are several large language models trained on massive code data, such as codex @xcite , codegen @xcite , @xcite , codegeex @xcite and code llama @xcite . these code generation models work in an autoregressive manner, which results in their inability to do online modification. some studies focus on improving the work of autoregressive code generation models. in self-debugging @xcite , the model generates multiple iterations of code based on the explanation and execution results of the initially generated code. in contrast, our approach enables the model for online modification during a single round. another concurrent study, self-infilling @xcite , explores enhancing autoregressive code generation through infilling capabilities. our approach differs from self-infilling in three key aspects: firstly, self-infilling requires the generation model to possess infilling abilities and mandates specific suffix formats for different languages, which limits its applicability. secondly, our method allows for online modifications, aligning more closely with human coding practices, whereas self-infilling merely alters the order of generation. lastly, the two methods are somewhat complementary: certain selfinfilling features, such as loops, can be integrated into our method, and our capabilities for online modification could also enhance their infilling or generation processes. several research targets improving llm-based code generation @xcite @xcite @xcite . @xcite utilizes one language model to generate both code snippets and test cases. subsequently, it ranks the code snippets by assessing the consistency between the code snippets and the associated test cases. codep @xcite enhances the existing code generation framework by incorporating an automatic pushdown automaton (pda) module to ensure the syntactic correctness of generated code. recently, many code generation approaches finetuned code llms with high-quality instruction finetuning datasets @xcite @xcite . our approach also harmonizes with these models and can be used to enhance their performance. code infilling various code generation models possess the capability of code infilling. while excelling in code generation, decoder-only architecture code generation models face limitations in understanding context due to their unidirectional attention mechanism. in response to this architectural limitation, these models employ methodologies akin to the fim @xcite concept. the model learns to fill these regions in a standard left-to-right autoregressive manner, acquiring infilling capabilities. this approach has been used to pre-train on various code generation models, including incoder @xcite , codegen 2 @xcite and @xcite , and has boosted downstream tasks such as code completion and document generation. in this work, we leverage these code-infilling capabilities to enable online modification and enhance code generation. ## conclusion this research analyzes the inherent irreversibility limitation of the autoregressive sequential generation, a critical yet underexplored issue. we introduce jumpcoder, a framework that enables human-like online code modification for existing llms through an infilling model without retraining. this framework is readily adaptable to diverse programming languages. through extensive experimentation, we demonstrate that our proposed jumpcoder significantly enhances code generation quality across various code llms and benchmarks. we hope our work can motivate more future investigations into this limitation of llms.