Text-to-Code Generation with Modality-relative Pre-training (2024)

Abstract: Large pre-trained language models have recently been expanded and applied to programming language tasks with great success, often through further pre-training of a strictly-natural language model–where training sequences typically contain both natural and (linearised) programming language. Such approaches effectively map both modalities of the sequence into the same embedding space. However, programming language keywords (e.g. “while”) often have very strictly defined semantics. As such, transfer learning from their natural language usage may not necessarily be beneficial to their code application and vise versa. Assuming an already pre-trained language model, in this work we investigate how sequence tokens can be adapted and represented differently, depending on which modality they belong to, and to the ultimate benefit of the downstream task. We experiment with separating embedding spaces between modalities during further model pre-training with modality-relative training objectives. We focus on text-to-code generation and observe consistent improvements across two backbone models and two test sets, measuring pass@ k and a novel incremental variation.

## introduction increasingly more large pre-trained language models @xcite @xcite based on the transformer @xcite architecture have been proposed and shown to achieve stateof-the-art results on a variety of natural language processing (nlp) tasks. lately, such models have been adapted to more specific language domains, e.g. biomedical @xcite @xcite , legal @xcite , cyber security @xcite and finance @xcite , while also expanding to include signals from modalities other than natural language, such as vision @xcite @xcite @xcite , proteins @xcite , time series @xcite and code @xcite @xcite @xcite . in this work, we experiment with pre-trained language models specifically created for text-to-code generation, i.e. the task of program synthesis given natural language (nl) descriptions (e.g. problem definitions or docstrings). while some of the proposed models for this task adopt the encoderdecoder architecture @xcite , most are trained as decoder-only transformer models @xcite @xcite @xcite . the latter approach is often preferred, since single-component architectures are well-suited for continuous training over vast amounts of raw non-annotated data, such as github. 2 when employing a single-component architecture, all tokens in the sequence are often vectorised descr. sign. def gcd(a: int, b: int) -> int: body while b: a, @xmath0, a % b return a tests gcd(3, 5) = 1 gcd(25, 15) = 5 as our test case, we focused on python and carried out zero-shot experiments demonstrating that our modality-relative embeddings and training objectives achieve consistent improvements over two baselines on the humaneval @xcite and mbpp @xcite datasets. we also introduce a novel variation of the pass@k @xcite evaluation metric for program synthesis, namely incremental pass@k, that combines both synthesis and completion tasks to better differentiate models capabilities. ## methodology in this work, we focus on the task of text-tocode generation, i.e. to synthesise running code that successfully solves the problem described in a nl description. ## zero-shot results we first perform an ablation analysis on the 100m model variants to determine the gains based on our training data, the training objective and the embedding separation approaches. our primary comparisons are against two baseline model variants (marked with a gray background), i.e. after modality-agnostic (mapt) training and after standard next-token-prediction (text-code) on the entire sequence. the latter facilitates comparison against continual training on the given data for the same number of epochs as the other objectives. in the same datasets and methods. firstly, we notice that we get an improvement of +2.4, +1.9 and +0.6 points for @xmath2,10,100 over mapt training on hu-maneval and correspondingly +0.2, +1.1 and +0.8 on mbpp. we attribute this gain to the formulation of the data in a format that exactly matches the task at hand. overall, embedding separation consistently offers additional improvements up to +1.2, +1.3, +2.9 on humaneval and +0.8, +0.9, +1.3 on mbpp. in terms of incremental pass@k, again we observe larger gains over text-code, mostly on mbpp with +3.7, +7.3 and +4.5 points. across training objectives, we notice that in the majority of settings, code-clm outperforms across the board. across datasets, separation offers the most notable gains in the prefix-code-clm objective, probably because it helps the model to better adapt to bidirectional attention on the docstring. we note that most of the gains from em-bedding separation seem to affect pass@10 and pass@100, which implies that separation helps increase the expressiveness of the model but not the map (maximum a posteriori) solution as much. ## visualising separated embeddings we attempt to qualitatively check the embedding space separation by visualising token embeddings on a 2d plane via t-sne plots. figure 3 illustrates the 20 closest neighbours of the tokens open and join in the two modality spaces, as measured using cosine similarity for pangu-coder 350m model with partial embedding separation (pes). for open, in the docstring space we observe that its representation is separate from close or other build-in operations, e.g. get. on the contrary, in the code space open and close are grouped together with other operations that are used in a similar way in code. analogously, join in the docstring space is close to words with a similar natural language meaning and surface form, while in the code space its representation is close to functionalike tokens in python such as remove, get, split, etc. we observe similar behaviours for other tokens; see additional plots in appendix e. ## generalisation analysis and discussion we performed an analysis over the performance of all the checkpoints of our trained models (checkpoints were cached every 10k steps), to determine whether performance on any test set was predictive on the performance of other sets; we included humaneval and both development and test sets of mbpp. pearson analysis between checkpoint performance suggests very weak correlation between model performance overall; consult this seems to indicate that the distribution of problems that are covered by any one test set is relatively unique, and raises the question whether comparisons over any particular dataset should be expected to generalise on unseen problems when models are deployed. potentially, this is partially explained by the small size of the test sets, but could be indicative of a larger problem; that due to the specificity of problems contained in these test sets, they do not evaluate generalisable model coding capabilities as much as they measure whether they have been exposed to very similar problems and definitions during pre-training. to offer an example, if the model has not been exposed to problems regarding fibonacci sequences during pre-training, it is unrealistic to expect it to generate a relevant code solution in a zero-shot setting. similarly, whether a model can generate a solution to a fibonacci problem offers no indication of the its capabilities to solve other problems. ## related work recently, there has been an increasing interest in extending proven natural language understanding and generation methods to code understanding and generation tasks. codebert @xcite , for instance, was trained using a combination of masked language modelling inspired by @xcite and replaced token detection from @xcite . codet5 @xcite and pymt5 @xcite were built on top of @xcite while unixcoder @xcite was based on unilm @xcite and combines three pre-training objectives using different attention masks to control access to surrounding context for a token to be predicted. focused on the task of text-to-code generation, @xcite introduced codex, a set of gpt-based language models trained on publicly available code from github, up to 12b parameters in size. @xcite introduced alphacode, a set of sequence-to-sequence models with up to 41b parameters, trained on data from programming competitions, e.g. codeforces 8 as well as github code in several programming languages. codegen @xcite was proposed as a conversational text-to-code approach using large language models with sizes of up to 16b parameters. the authors proposed three model variants, trained on the pile @xcite , continuously trained on bigquery (with 6 programming languages) and finally trained on pythononly data. incoder @xcite extends left-to-right code generation with an infilling training objective, similar to @xcite , and is able to predict spans of partial programs as well. santacoder @xcite is one of the latest code lms, sizing up to 1.1b parameters, supporting python, java, and javascript. it was trained with multi-query attention and the fill-inthe-middle (fim) objective @xcite on the stack @xcite , a 3 tb publicly available dataset supporting 358 programming languages collected from permissively-licensed source code files from github. starcoder (li et al., 8 https://codeforces.com/ 2023) is a 15.5b parameter model, similarly pretrained with multi-query attention and fim on 80+ programming languages from the stack. with the exception of codex @xcite calculating clm loss on code exclusively, this is the first work to consider modality-relative embedding separation and training objectives. methods that consider nl and code as different modalities mostly focus on taking into account different views of code. for instance, graphcode-bert @xcite , noted that previous pretrained models treat code snippets as sequences of tokens while ignoring the inherent structure of code. they presented graphcodebert, which showed that incorporating the data flow, i.e. a semanticlevel structure of code extracted from the abstract syntax tree (ast), leads to richer code representations. @xcite instead, used the actual ast together with code snippets. to the best of our knowledge we are not aware of any work assigning different embeddings to tokens of the same sequence depending on which type of text (natural language/code) they appear. ## conclusion existing codelms consider both code and natural language as a single modality, mapping them into a shared embedding space. however, in this work, we posit that the semantics and usage of tokens can differ between code and nl, requesting for a possible space separation. as such, we proposed to consider code and natural language as different modalities for the task of text-to-code generation by introducing modality-aware embedding separation strategies and training objectives. in detail, assuming a general codelm trained on raw data-where modalities are shared-we continue training on taskspecific data with separated embedding spaces. we present partial separation, which targets languagespecific tokens, as well as full separation that duplicates the entire model's vocabulary. in addition, we proposed incremental pass@k, as a variant of the standard metric that evaluates the code completion capabilities of models. zero-shot evaluation on the humaneval and mbpp datasets, with two 100m and one 350m parameters models, indicate that embedding space separation improves code generation across different objectives. we also observe that further pretraining on data formatted to match the target task consistently boosts performance. although there is no clear winner among partial versus full space separation, each method comes with advantages and drawbacks. while partial separation minimally increases the model's vocabulary, it requires manual effort for each target programming language. on the other end, full separation is automatic but significantly increases vocabulary size. given these observations, future work could target efficient approaches for separating only certain tokens to strike a balance between the two.