PaperlayerView exact revision on arXiv

Source-linked AI summary

UniXcoder: Unified Cross-Modal Pre-training for Code Representation

Daya Guo, Shuai Lu, Nan Duan, Yanlin Wang, Ming Zhou, Jian Yin

arXiv:2203.03850v1cs.CLcs.PLcs.SE

TL;DR

Existing unified encoder-decoder models are sub-optimal for auto-regressive code completion, which requires decoder-only inference. UniXcoder combines configurable attention with code, comments, and structure-preserving AST sequences, plus multimodal pre-training. It achieves state-of-the-art performance on most evaluated tasks, while its input representation assumes a one-to-one AST-to-sequence mapping.

  • Problem

    Unified encoder-decoder models support code understanding and generation but are sub-optimal for auto-regressive tasks such as code completion, which require decoder-only inference.

  • Method

    UniXcoder uses mask attention with prefix adapters, integrates code comments and ASTs through a one-to-one structure-preserving sequence mapping, and applies multimodal pre-training.

  • Results

    UniXcoder achieves state-of-the-art performance on most tasks evaluated across five code-related tasks and nine datasets.

  • Takeaways & Limitations

    AST and code comments both enhance UniXcoder's ability to capture code semantics, and zero-shot code-to-code search provides an additional evaluation of code-fragment representations.

  • Takeaways & Limitations

    The input representation assumes that the AST transformation is a one-to-one mapping so that different trees are not confused with one another.

Abstract

from arXiv · show

Pre-trained models for programming languages have recently demonstrated great success on code intelligence. To support both code-related understanding and generation tasks, recent works attempt to pre-train unified encoder-decoder models. However, such encoder-decoder framework is sub-optimal for auto-regressive tasks, especially code completion that requires a decoder-only manner for efficient inference. In this paper, we present UniXcoder, a unified cross-modal pre-trained model for programming language. The model utilizes mask attention matrices with prefix adapters to control the behavior of the model and leverages cross-modal contents like AST and code comment to enhance code representation. To encode AST that is represented as a tree in parallel, we propose a one-to-one mapping method to transform AST in a sequence structure that retains all structural information from the tree. Furthermore, we propose to utilize multi-modal contents to learn representation of code fragment with contrastive learning, and then align representations among programming languages using a cross-modal generation task. We evaluate UniXcoder on five code-related tasks over nine datasets. To further evaluate the performance of code fragment representation, we also construct a dataset for a new task, called zero-shot code-to-code search. Results show that our model achieves state-of-the-art performance on most tasks and analysis reveals that comment and AST can both enhance UniXcoder.

1 Introduction

UniXcoder addresses the tension between code understanding and generation by combining configurable attention with code, comments, and AST representations. It introduces structure-preserving AST encoding and multimodal pre-training, achieving strong results across diverse code tasks.

  • Contributions: UniXcoder unifies code understanding, generation, and auto-regressive tasks in a cross-modal pre-trained model.It uses code, comments, and AST information rather than code alone.
  • Model design: Mask attention matrices with prefix adapters control context access, supporting decoder-only inference for code completion.This design targets efficient auto-regressive inference while retaining support for other task types.
  • AST representation: A one-to-one mapping converts ASTs from trees into sequences while retaining their structural information for parallel encoding.The resulting sequence can be combined with source code and comments to enhance code representation.
  • Pre-training: Three language-modeling objectives and two multimodal pre-training tasks learn general task capabilities and code-fragment representations.Contrastive learning uses AST information, while cross-modal generation aligns representations across programming languages using comments.
  • Evaluation: UniXcoder is evaluated on five tasks across nine datasets and achieves state-of-the-art performance on most tasks.The study also introduces zero-shot code-to-code search and finds that AST and comments enhance code semantics.

2 Related Works

Prior code models divide into encoder-only, decoder-only, and encoder-decoder architectures, each emphasizing different capabilities. UniXcoder combines configurable attention with structure-preserving AST integration to support both understanding and generation, including auto-regressive completion.

  • Existing architectures: Pre-trained programming-language models are commonly categorized as encoder-only, decoder-only, or encoder-decoder systems.These categories differ in how tokens access context and which code-intelligence tasks they naturally support.
  • Encoder-only models: Encoder-only models use bidirectional attention and improve code understanding, but require an additional decoder for generation.Examples include CuBERT, CodeBERT, GraphCodeBERT, and SYNCOBERT.
  • Decoder-only models: Decoder-only models suit auto-regressive code completion, whereas their unidirectional attention is sub-optimal for understanding tasks.GPT-C and CodeGPT are identified as examples of this architecture.
  • Encoder-decoder models: Encoder-decoder models support understanding and generation through architectures such as PLBART, CodeT5, and TreeBERT.These systems use denoising or multitask-oriented pre-training strategies.
  • UniXcoder: UniXcoder uses mask attention with prefix adapters to control model behavior across understanding and generation tasks.Its AST-to-sequence mapping retains all AST structure and avoids requiring additional tasks to implicitly learn that structure.

3 UniXcoder

UniXcoder unifies code understanding and generation through shared Transformer parameters, mode-specific attention masks, and comment–AST inputs. Its pre-training combines language modeling with multimodal representation learning for code fragments.

  • Input Representation: A one-to-one AST mapping recursively surrounds non-leaf node contents with left and right markers, preserving tree structure in a sequence.The resulting flattened sequence can be encoded in parallel with source code and comments.
  • Model Architecture: UniXcoder combines comment and flattened AST inputs with shared Transformer parameters and prefix-controlled attention masks for multiple operating modes.The model uses [Enc], [Dec], and [E2D] prefixes to control context access for encoder-only, decoder-only, and encoder-decoder behavior.
  • Pre-training Tasks: The model pre-trains different modes with masked language modeling, unidirectional language modeling, and denoising objectives.Masked language modeling uses bidirectional context, unidirectional modeling predicts subsequent tokens, and denoising reconstructs randomly masked spans.
  • Pre-training Tasks: Masked language modeling uses comment semantics and AST syntax to infer masked code tokens from bidirectional context.The task samples 15% of input tokens and replaces most selected tokens with [MASK] or random tokens.
  • Code Fragment Representation Learning: UniXcoder learns code-fragment embeddings by mean-pooling mapped-AST hidden states and applying multimodal contrastive learning and cross-modal generation.Contrastive learning trains semantic embeddings, while comment generation supports semantic fusion and alignment across programming languages.

4 Experiments

UniXcoder is evaluated across understanding, generation, code completion, and zero-shot code-to-code search, with comparisons against multiple pre-trained model families. Results show strong performance across tasks, while ablations identify contributions from contrastive learning, multi-modal data, AST, comments, and the AST mapping function.

  • Understanding Tasks: UniXcoder outperforms encoder-only and encoder-decoder models on understanding tasks across all five datasets.
  • Understanding Tasks: Contrastive learning and multi-modal data provide the main improvements on understanding tasks.
  • Generation Tasks: UniXcoder achieves comparable generation performance to CodeT5-base and improves code generation accuracy by 0.3%.Its BLEU-4 scores are slightly worse on code summarization and generation; UniXcoder outperforms the smaller CodeT5-small model.
  • Code Completion: UniXcoder outperforms PLBART and CodeT5 for code completion and gains 2.3% absolute accuracy on the Java corpus.The comparison targets a decoder-only inference setting required for efficient code completion.
  • Zero-shot Code-to-Code Search: UniXcoder achieves state-of-the-art zero-shot code-to-code search performance with about 11 points of overall improvement over GraphCodeBERT.Removing contrastive learning reduces performance from 20.45% to 13.73%.
  • Model Analysis: Removing comments hurts understanding and generation, while AST improves understanding but not generation tasks.The proposed one-to-one AST mapping also outperforms BFS and DFS replacements, which can confuse distinct tree structures.
  • Model Analysis: In a CosQA case study, UniXcoder retrieves the correct dictionary-sorting function despite the query expressing value-based sorting indirectly.The example contrasts semantic retrieval with a prediction having higher lexical overlap but an incorrect dictionary input.

5 Conclusion

The conclusion presents UniXcoder as a unified cross-modal model that combines code, comments, and AST information for code understanding, generation, and auto-regressive tasks. It also introduces an AST sequence mapping, new code-representation pre-training tasks, and zero-shot code-to-code search, with experiments showing broad improvements.

  • UniXcoder incorporates code comments and ASTs to support code understanding, generation, and auto-regressive tasks.
  • Two new pre-training tasks learn code-fragment representations using multimodal information.
  • Zero-shot code-to-code search is introduced as a new downstream task with a dataset constructed for its evaluation.
  • Experiments show that UniXcoder significantly outperforms previous works on most tasks, while ablations support the contributions of AST, comments, mapping, and pre-training tasks.

A Proof for Mapping Function

The paper proves that the AST-to-sequence mapping function is one-to-one by showing that equal mapped sequences imply equal trees. The proof uses structural lemmas and induction on tree depth.

  • Structural lemmas: The mapped sequence's first element identifies the root, allowing equal sequences to establish equal roots.
  • Structural lemmas: Internal nodes occur with left and right suffixes, preserving matching structure and balancing the two suffix types.
  • The proof establishes the contrapositive: if two mapped sequences are equal, their source trees are equal.
  • Base case: The base case shows that a depth-one tree maps to a single root node, uniquely determining the tree.
  • Inductive proof: The induction assumes one-to-one mapping through depth n and proves it for depth n + 1.
  • Conclusion: By induction, the mapping function is one-to-one for all tree depths h ≥2.

B Pre-training Setting

UniXcoder uses a Transformer with a programming-language subword vocabulary and AST-specific tokens, trained on unimodal and multimodal data. Its training alternates objectives and language sampling to reduce high-resource-language bias.

  • UniXcoder uses 12 Transformer layers, 768-dimensional hidden states, and 12 attention heads.
  • Its vocabulary contains 50K programming-language subword units plus 1,416 special tokens for AST non-terminals.
  • Multimodal pre-training uses 2.3M comment-paired functions from CodeSearchNet across six programming languages.
  • The model is first trained with masked, unidirectional, and denoising objectives on C4 and additional unimodal data.
  • Multimodal pre-training continues for 100k steps, alternating objectives and sampling languages according to a distribution controlled by α=0.7.
  • The language-sampling distribution is intended to alleviate bias toward high-resource programming languages.

C.1 Clone Detection

Clone detection evaluates whether code fragments are similar on POJ-104 and BigCloneBench. The experiments use cosine similarity for POJ-104 retrieval and binary classification for BigCloneBench.

  • Clone detection measures similarity between two code fragments and is evaluated on POJ-104 and BigCloneBench.
  • POJ-104: POJ-104 contains 104 problems with 500 C/C++ programs per problem, split across 64/16/24 training, validation, and test problems.
  • POJ-104: For POJ-104, the task retrieves programs solving the same problem, using cosine similarity between mean vectors of final hidden states.
  • BigCloneBench: BigCloneBench provides 901,724/416,328/416,328 training, validation, and test examples spanning 10 functionalities.
  • BigCloneBench: BigCloneBench is treated as binary classification, with clone probability computed from cosine similarity between mean final-hidden-state vectors.

C.2 Code Search

Code search retrieves the most relevant code from candidate collections for natural-language queries. UniXcoder is evaluated on CSN, AdvTest, and CosQA using cosine similarity between code-query representations.

  • Code search retrieves the most relevant code from a candidate collection given a natural-language query.
  • The evaluation uses the CSN, AdvTest, and CosQA datasets.
  • AdvTest normalizes Python function and variable names to test models' understanding and generalization abilities.
  • Relevant scores are cosine similarities between mean vectors of the last hidden states, with other same-batch vectors used as negatives.

C.4 Code Generation

Code generation generates code snippets from natural-language descriptions using the CONCODE dataset, which includes code environments and Java code snippets.

  • Code generation produces a code snippet from a natural-language description.
  • CONCODE contains examples from about 33k Java projects on GitHub, split into 100k/2k/2k training, validation, and testing examples.
  • Each CONCODE example includes a natural-language description, code environment, and code snippet.
  • Fine-tuning uses a learning rate of 5e-5, batch size 32, maximum source and target lengths of 350 and 150, and beam size 3.

C.5 Code Completion

The paper focuses on line-level code completion evaluated on Python and Java corpora, using randomly selected lines as prediction targets.

  • Line-level code completion is evaluated on PY150 and the Github Java Corpus.
  • PY150 contains 150,000 Python source files, with 10,000 test examples created from different files.
  • For PY150, prediction lines are selected at random, and the average input and output lengths are reported in the dataset description.
  • The Github Java Corpus was collected from over 14 thousand Java projects, with 3,000 test examples created from different files.

C.6 Zero-shot Code-to-Code Search

Zero-shot code-to-code search evaluates whether code fragment embeddings retrieve code with the same semantics across programming languages without task-specific training.

  • Zero-shot code-to-code search retrieves code with the same semantics as a source-code query from candidate collections.
  • The dataset contains 11,744/15,594/23,530 Ruby/Python/Java functions drawn from CodeNet.
  • Each function solves one of 4,053 problems, and retrieval quality is measured with MAP after cosine-similarity ranking.
Loading 2203.03850v1…