Directional Diffusion-Style Code Editing Pre-training

TL;DR

DivoT5 introduces a directional diffusion-style pre-training framework for code editing to mimic real-world incremental evolution. It employs token-level diffs, intermediate evolution states, and four evolution-focused pre-training tasks within an encoder-decoder Transformer initialized from CodeT5. The approach achieves state-of-the-art results on multiple NL-guided and code-only editing benchmarks and demonstrates generalization to non-editing tasks like code translation, while maintaining efficiency comparable to standard autoregressive models. Ablation studies confirm that each pre-training task and the diffusion-inspired data generation contribute meaningfully to performance, underscoring the value of modeling gradual code evolution in pre-training.

Abstract

Code pre-trained models have shown promising effectiveness in various software engineering tasks. Among these tasks, many tasks are related to software evolution and/or code editing. However, existing code pre-trained models often overlook the real-world code editing data and the evolutionary nature of the editing process. In this paper, to simulate the step-by-step code editing process of human developers, we propose DivoT5, a pre-trained model based on directional diffusion at the data level. In DivoT5, we adopt two categories of pre-training tasks. The first category is mask and denoising tasks augmented with a diffusion direction representing code evolution. That is, we first apply a noising process to the code snippets before evolution, and then ask the pre-training process to restore the snippets with noise into the code snippets after evolution. The second category is tasks aiming to reinforce the evolutionary direction. That is, we first generate various intermediate versions for each pair of snippets before and after evolution, and then ask the pre-training process to transform the intermediate versions into the snippet after evolution for each pair. We evaluate DivoT5 for two code-editing scenarios and one non-editing scenario using five downstream tasks. Given each downstream task, we fine-tune the pre-trained DivoT5 to evaluate its effectiveness. Our experimental results show that DivoT5 achieves state-of-the-art (SOTA) performance on most tasks in comparison to models of the same scale (220M), large scale (770M) models in fine-tuning, and billion-scale (6.7B, 8B, ChatGPT) models in few-shot settings. For one code-editing task (i.e., automated code review), DivoT5 pre-trained on top of CodeT5-small (60M) can even outperform CodeT5-base (220M) and other pre-trained models with 220M parameters except for DivoT5 pre-trained on top of CodeT5-base (220M).

Figures (4)

• Figure 1: Illustration of directional diffusion. The upper part depicts the transition of code editing states, where noise in the old code is gradually reduced from left to right until the new code is obtained. Starting with the old code ($X_T$), artificial noise ($X_{T+1}$) is introduced via masking strategies to guide the evolutionary direction. Intermediate states ($X_t$) represent evolutionary noise, reinforcing the editing direction. The lower part simulates the changes in code content during human editing, with green indicating additions and blue representing modifications. From the old code to the new code, the programmer introduces the package $p_2$, adds the variable $v_3$ and the function $f_{12}$, and updates the functions $f_{11}$ and $f_{21}$.
• Figure 2: Illustration of traditional diffusion and directional diffusion in text generation. The upper section showcases traditional diffusion, where 'ABC' is generated during pre-training, and 'FGH' during inference. The lower section highlights directional diffusion, showing 'CbDbEb' generated from 'CaDaEa' during pre-training, and 'FbGbHb' from 'FaGaHa' during inference.
• Figure 3: The process of modifying the old code (center) to fulfill the requirements outlined in the comments (top left) often necessitates changes across multiple dispersed locations. Achieving the desired evolution (right) requires incremental editing at all relevant points. The blue italics aim to highlight the parts that persist throughout the entire code editing process.
• Figure 4: Comparison of generation case using different models on the CodeReview dataset