TransformCode: A Contrastive Learning Framework for Code Embedding via Subtree Transformation

TL;DR

TransformCode presents an encoder- and language-agnostic framework for learning code embeddings through contrastive learning on abstract syntax trees. It introduces normalization and a suite of semantics-preserving subtree transformations to generate anchor samples, coupled with a MoCo-style momentum encoder and a queue to support large negative sets. A Transformer-based AST encoder with relative-position encoding and a projection head is trained via the InfoNCE objective, using a code-specific WordPiece tokenizer to improve efficiency. Empirically, TransformCode achieves competitive or superior results on code-clone detection, method-name prediction, and code classification while using smaller models and unlabeled data, suggesting practical benefits for scalable SE tasks and laying groundwork for future integration with large language models.

Abstract

Artificial intelligence (AI) has revolutionized software engineering (SE) by enhancing software development efficiency. The advent of pre-trained models (PTMs) leveraging transfer learning has significantly advanced AI for SE. However, existing PTMs that operate on individual code tokens suffer from several limitations: They are costly to train and fine-tune; and they rely heavily on labeled data for fine-tuning on task-specific datasets. In this paper, we present TransformCode, a novel framework that learns code embeddings in a contrastive learning manner. Our framework is encoder-agnostic and language-agnostic, which means that it can leverage any encoder model and handle any programming language. We also propose a novel data-augmentation technique called abstract syntax tree (AST) transformation, which applies syntactic and semantic transformations to the original code snippets, to generate more diverse and robust samples for contrastive learning. Our framework has several advantages over existing methods: (1) It is flexible and adaptable, because it can easily be extended to other downstream tasks that require code representation (such as code-clone detection and classification); (2) it is efficient and scalable, because it does not require a large model or a large amount of training data, and it can support any programming language; (3) it is not limited to unsupervised learning, but can also be applied to some supervised learning tasks by incorporating task-specific labels or objectives; and (4) it can also adjust the number of encoder parameters based on computing resources. We evaluate our framework on several code-related tasks, and demonstrate its effectiveness and superiority over the state-of-the-art methods such as SourcererCC, Code2vec, and InferCode.

Figures (15)

• Figure 1: An example of code normalization.
• Figure 2: An example of PermuteDeclaration.
• Figure 3: An example of SwapCondition.
• Figure 4: An example of ArithmeticTransform.
• Figure 5: An example of WhileForExchange.