Seamlessly Integrating Tree-Based Positional Embeddings into Transformer Models for Source Code Representation

2025-07-05

AI Summary

Standard Transformer positional encodings inadequately model the hierarchical syntactic structure of source code—particularly abstract syntax trees (ASTs)—limiting performance in code representation tasks. To address this, we propose an AST-aware hierarchical positional encoding that explicitly incorporates node depth and sibling order information, seamlessly integrated into the CodeBERTa architecture. Our structure-aware embedding design bridges the gap between conventional positional encodings and fine-grained syntactic hierarchy modeling. In masked language modeling (MLM) pretraining, our method achieves lower MLM loss; in downstream code clone detection, it consistently outperforms baselines across accuracy, F1-score, precision, and recall. This work is the first to systematically encode fine-grained AST hierarchical relations—such as depth and sibling ordering—directly into positional embeddings, establishing a novel paradigm for syntax-aware code pretraining.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Machine Learning: Hardware-aware ML

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Transformer-based models have demonstrated significant success in various source code representation tasks. Nonetheless, traditional positional embeddings employed by these models inadequately capture the hierarchical structure intrinsic to source code, typically represented as Abstract Syntax Trees (ASTs). To address this, we propose a novel tree-based positional embedding approach that explicitly encodes hierarchical relationships derived from ASTs, including node depth and sibling indices. These hierarchical embeddings are integrated into the transformer architecture, specifically enhancing the CodeBERTa model. We thoroughly evaluate our proposed model through masked language modeling (MLM) pretraining and clone detection fine-tuning tasks. Experimental results indicate that our Tree-Enhanced CodeBERTa consistently surpasses the baseline model in terms of loss, accuracy, F1 score, precision, and recall, emphasizing the importance of incorporating explicit structural information into transformer-based representations of source code.

Problem

• Traditional positional embeddings fail to capture source code hierarchy
• Propose tree-based embeddings encoding AST depth and sibling indices
• Enhance CodeBERTa with structural info for better code representation

Innovation

• Tree-based positional embeddings for hierarchical encoding
• Integration into CodeBERTa transformer architecture
• Enhanced performance in MLM and clone detection