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.

    Source: §1 · p. 1

  • 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.

    Source: §1 · pp. 1–2

  • Results

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

    Source: §1 · p. 2, §5 · p. 9

  • 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.

    Source: §1 · p. 2, §5 · p. 9

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

    Source: §3.1 · p. 3

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.Source: §1 · pp. 1–2

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.Source: §1 · p. 1

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.Source: §1 · pp. 1–2

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.Source: §1 · p. 2

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.Source: §1 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · pp. 2–3

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.Source: §3.1 · pp. 3–4

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.Source: §3.2 · p. 4

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.Source: §3.3 · pp. 4–5
  • 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.Source: §3.3 · p. 4

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.Source: §3.3 · p. 5

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.Source: §4.2 · p. 6
  • Contrastive learning and multi-modal data provide the main improvements on understanding tasks.Source: §4.2 · p. 6

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.Source: §4.3 · p. 7

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.Source: §4.4 · p. 8

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%.Source: §4.5 · p. 8, §4.6 · p. 8

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.Source: §4.6 · pp. 8–9
  • 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.Source: §4.6 · p. 9

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.Source: §5 · p. 9
  • Two new pre-training tasks learn code-fragment representations using multimodal information.Source: §5 · p. 9
  • Zero-shot code-to-code search is introduced as a new downstream task with a dataset constructed for its evaluation.Source: §5 · p. 9
  • Experiments show that UniXcoder significantly outperforms previous works on most tasks, while ablations support the contributions of AST, comments, mapping, and pre-training tasks.Source: §5 · p. 9

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

Base case

Inductive proof

Conclusion

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.

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.Source: §C.1 · p. 12

POJ-104

  • POJ-104 contains 104 problems with 500 C/C++ programs per problem, split across 64/16/24 training, validation, and test problems.Source: §C.1 · p. 12
  • For POJ-104, the task retrieves programs solving the same problem, using cosine similarity between mean vectors of final hidden states.Source: §C.1 · p. 12

BigCloneBench

  • BigCloneBench provides 901,724/416,328/416,328 training, validation, and test examples spanning 10 functionalities.Source: §C.1 · p. 12
  • BigCloneBench is treated as binary classification, with clone probability computed from cosine similarity between mean final-hidden-state vectors.Source: §C.1 · p. 12

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.Source: §C.2 · p. 12
  • The evaluation uses the CSN, AdvTest, and CosQA datasets.Source: §C.2 · p. 12
  • AdvTest normalizes Python function and variable names to test models' understanding and generalization abilities.Source: §C.2 · pp. 12–13
  • Relevant scores are cosine similarities between mean vectors of the last hidden states, with other same-batch vectors used as negatives.Source: §C.2 · p. 13

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.Source: §C.4 · p. 13
  • CONCODE contains examples from about 33k Java projects on GitHub, split into 100k/2k/2k training, validation, and testing examples.Source: §C.4 · p. 13
  • Each CONCODE example includes a natural-language description, code environment, and code snippet.Source: §C.4 · p. 13
  • 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.Source: §C.4 · p. 13

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.Source: §C.5 · p. 13
  • PY150 contains 150,000 Python source files, with 10,000 test examples created from different files.Source: §C.5 · p. 13
  • For PY150, prediction lines are selected at random, and the average input and output lengths are reported in the dataset description.Source: §C.5 · p. 13
  • The Github Java Corpus was collected from over 14 thousand Java projects, with 3,000 test examples created from different files.Source: §C.5 · p. 13

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.Source: §C.6 · p. 14
  • The dataset contains 11,744/15,594/23,530 Ruby/Python/Java functions drawn from CodeNet.Source: §C.6 · p. 14
  • Each function solves one of 4,053 problems, and retrieval quality is measured with MAP after cosine-similarity ranking.Source: §C.6 · p. 14

Related papers

Loading 2203.03850v1…