C ode F usion: A Pre-trained Diffusion Model for Code Generation (2023)

Abstract: Imagine a developer who can only change their last line of code—how often would they have to start writing a function from scratch before it is correct? Auto-regressive models for code generation from natural language have a similar limitation: they do not easily allow reconsidering earlier tokens generated. We introduce CodeFusion, a pre-trained diffusion code generation model that addresses this limitation by iteratively denoising a complete program conditioned on the encoded natural language. We evaluate CodeFusion on the task of natural language to code generation for Bash, Python, and Microsoft Excel conditional formatting (CF) rules. Experiments show that CodeFusion (75M parameters) performs on par with state-of-the-art auto-regressive systems (350M-175B parameters) in top-1 accuracy and outperforms them in top-3 and top-5 accuracy due to its better balance in diversity versus quality.

## introduction auto-regressive code generation models @xcite @xcite @xcite cannot easily reconsider tokens generated earlier in the decoding process. this limitation can lead to lower diversity generations @xcite in the related domain of text. to balance diversity and quality of candidates generated, prior work has explored decoding strategies such as grouped beam search @xcite or nucleus sampling @xcite . diffusion models, which have shown remarkable performance in image generation @xcite , have recently been extended to generate diverse text @xcite . these approaches use an embedding layer to convert discrete tokens to continuous embeddings, where gaussian noise can be added and predicted, to imitate the diffusion process. to map denoised embeddings back to discrete text, these approaches then select the vocabulary token with the closest embedding. in the code domain, where there are many syntactic and semantic constraints between tokens, independently projecting embeddings back to tokens can yield invalid programs. we propose codefusion, a natural language to code (nl-to-code) model that combines an encoder-decoder architecture @xcite with a diffusion process. the encoder maps the nl into a continuous representation, which is used by the diffusion model as an additional condition for denoising random gaussian noise input. to generate syntactically correct code, we then feed the denoised embeddings to a transformer decoder, with full self-attention and cross attention with the embedded utterance, to obtain probability distributions over code tokens. finally, we select the token with the highest probability at each index. to pre-train codefusion for code generation, we extend the continuous paragraph denoising (cpd) task introduced in @xcite to the code domain. specifically, we only apply noise to tokens that correspond to identifiers in code or to built-in keywords in the target language. this denoising task allows the model to learn relations between critical code tokens (like variable names, function names and control flow built-ins). we find that codefusion yields more diverse code (higher n-gram fraction, lower embedding similarity, and higher edit distance) than autoregressive models (see we evaluate codefusion on nl-to-code for three different languages: python @xcite , bash @xcite , and conditional formatting rules in microsoft excel @xcite . our results show that codefusion's (75m parameters) top-1 results are comparable or better than much larger state-of-the-art systems (350m-175b parameters). in top-3 and top-5, codefusion performs better than all baselines. this work makes the following contributions: nl-to-code is a popular and challenging sequence to sequence problem. previous techniques for nlto-code have used rnns (brunner and stockinger, 2021) and semantic parsers (lee et al., 2021). attention-based transformer models (vaswani et al., 2017) have been adapted for the code domain, and include codebert (feng et al., 2020) (encoder only), t5 (raffel et al., 2020) (encoderdecoder) and gpt-3 (brown et al., 2020) (decoder only). more recently, large scale transformer-based models like codegen @xcite and transcoder @xcite have shown promising results. instruction-tuning has also been used to further improve performance with models like starcoder @xcite and wizardcoder @xcite . these models can decode in an auto-regressive and non-auto-regressive fashion @xcite . recent surveys have discussed these models in more detail @xcite @xcite . diffusion models have been popularly applied to unsupervised @xcite ) and text conditioned image generation tasks @xcite . more recently, similar approaches have been extended to work in discrete domains such as text @xcite @xcite . ## methodology figure 1 shows codefusion's architecture. this section describes each component and our training and inference procedures. ## evaluation setup we briefly describe training, baselines, benchmarks and metrics. we provide further details for training and baselines in the appendix. ## evaluation we investigate the following questions. q1. does codefusion generate correct and diverse code? q2. how do different design decision impact performance? q3. how does the latent representation evolve during the diffusion steps? ## conclusion we propose codefusion, the first diffusion natural language to code (nl-to-code) generation model. with a decoder and code-focused pretraining, codefusion generates more syntactically correct programs than existing text diffusion models and more diverse programs than existing auto-regressive code models. our evaluation shows that codefusion competes with state-of-the-art transformer code generation models on python, bash and excel conditional formatting rules.