RETROcode: Leveraging a Code Database for Improved Natural Language to Code Generation

TL;DR

RETROcode introduces a retrieval-augmented memory mechanism for NL-to-code generation by integrating a large code database into a seq2seq transformer. The model retrieves neighbouring code chunks via a frozen codeBERT encoder and merges them into the decoder through either sequential or parallel cross-attention, with a hybrid database further guiding initial decoding. Across CoNaLa and CodeXGlue, RETROcode with a classic database improves over a strong baseline, and the hybrid database configuration achieves state-of-the-art-like BLEU scores on CoNaLa, approaching Codex with far fewer parameters and training data. The work demonstrates that memory-augmented decoding can significantly boost code generation efficiency and quality, while also highlighting considerations around initialization, chunk size, and database quality for practical deployment.

Abstract

As text and code resources have expanded, large-scale pre-trained models have shown promising capabilities in code generation tasks, typically employing supervised fine-tuning with problem statement-program pairs. However, increasing model size and data volume for performance gains also raises computational demands and risks of overfitting. Addressing these challenges, we present RETROcode, a novel adaptation of the RETRO architecture \cite{RETRO} for sequence-to-sequence models, utilizing a large code database as an auxiliary scaling method. This approach, diverging from simply enlarging model and dataset sizes, allows RETROcode to leverage a vast code database for prediction, enhancing the model's efficiency by integrating extensive memory. Our findings indicate that RETROcode not only outperforms similar-sized traditional architectures on test sets but also approaches the effectiveness of the much larger Codex model, despite being trained from scratch on a substantially smaller dataset.

Figures (3)

• Figure 1: Process of $Query_{k}(C_q)$ to obtain k-nearest neighbours and their continuation. Here, the chunk length $m$ to construct the database is equal to 8.
• Figure 2: Illustration of the RETROcode architecture, which includes two variations for integrating neighbour encoding into the baseline model . (a) Sequential aggregation: we incorporate the information from the neighbours into the code generation process using a two-step process. First, we use the classic cross-attention mechanism to combine the information from the natural language. Then, we perform a second cross-attention between the output of the first cross-attention and the neighbours. This process is described in equation \ref{['eqn:sequentialaggregation']}. (b) Parallel aggregation: we separately compute the information from the neighbours and the natural language with the decoder using cross-attention, and then merge the results with a linear layer as described in equation \ref{['eqn:parallelaggregation']}.
• Figure 3: Illustration of chunk-cross attention mechanism with chunk length $m=4$. This illustration introduces a variation of the database, discussed in \ref{['databasecreation']}, featuring a hybrid database.