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

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.

## introduction transformer-based models have demonstrated significant advances across numerous source code representation tasks, such as code summarization, clone detection, defect prediction, and semantic search. these models effectively leverage selfattention mechanisms, supplemented by positional embeddings, to encode the sequential order of tokens, thus achieving robust semantic understanding. however, a notable limitation arises from the fundamentally linear positional encodings typically employed by these models, which do not adequately capture the inherently hierarchical nature of programming languages @xcite @xcite . source code differs markedly from natural language in its explicit, structured hierarchy, most commonly represented as abstract syntax trees (asts) @xcite ). an ast encodes critical syntactic and semantic relationships, including nested scopes, parent-child dependencies, and sibling ordering among code constructs. although sequential positional embeddings, such as sinusoidal or learned encodings @xcite , effectively capture linear token sequences, they disregard hierarchical structure entirely. consequently, current transformer architectures may overlook important syntactic relationships, potentially limiting performance and generalization capabilities in source code understanding tasks. addressing this gap, we propose integrating tree-based positional embeddings into transformer models, explicitly encoding the hierarchical structure of source code. our approach introduces embeddings based on token depth within the ast hierarchy and their sibling positions, effectively guiding the self-attention mechanism to recognize and utilize structural context alongside semantic content. in this paper, we make the following contributions. our findings underscore the importance of integrating hierarchical structural information into transformer architectures, not only enhancing source code representation but also potentially improving models for broader structured data. ## related work transformer-based models, including gpt @xcite , bert @xcite , and roberta @xcite , have significantly advanced natural language processing (nlp) and have subsequently been adapted for source code understanding tasks. this section reviews the relevant literature, classified into transformer models for source code, positional embedding methods, and tree-based code representations. transformers for source code. pre-trained transformer models such as codebert and code-berta utilize broad datasets such as codesearch-net to effectively learn token-level semantic representations @xcite @xcite . although these models have achieved strong results, their effectiveness stems primarily from capturing semantic relationships between tokens without explicitly modeling the hierarchical structural relationships encoded in abstract syntax trees (asts). consequently, structural nuances, crucial for tasks that require deeper syntactic and semantic understanding, remain largely unaddressed. positional embeddings in transformers. the original transformer model @xcite employs sinusoidal positional embeddings to encode token positions within sequences, establishing order-aware representations. subsequent developments introduced learned positional embeddings, enabling dynamic adaptation during training @xcite . recently, rotary positional embeddings (rope) @xcite have emerged as an effective technique for capturing relative positional information in transformers, achieving superior generalization in sequence modeling tasks. despite these advancements, positional embeddings typically remain confined to linear positional encodings, which do not inherently capture hierarchical or structural relationships essential in structured data like source code @xcite . ## theoretical foundations traditional positional embeddings employed in transformer models typically assume a linear sequence of tokens, effectively capturing the sequential order but failing to represent the complex hierarchical structures inherent in many forms of structured data, notably source code. source code is commonly represented by abstract syntax trees (asts), which explicitly encode syntactic and semantic relationships such as nesting, parent-child dependencies, and sibling ordering. a linear positional encoding is inadequate for capturing such hierarchical nuances, motivating the need for hierarchical positional encodings. hierarchical positional embeddings provide a principled approach to representing each node (or token) by encoding its structural position within a tree. formally, each node's hierarchical position can be represented by a path from the root to that node. let f (x) denote the hierarchical position of node x, defined recursively as follows: figure 1 provides a visual representation of hierarchical positional encoding, illustrating how depth and sibling indices are assigned to each node in an ast. this structure enables the model to capture both global (depth) and local (sibling order) relationships, which are then transformed into learned embedding vectors. each hierarchical position uniquely encodes both local (sibling order) and global (depth in the hierarchy) structural context. these positional values are transformed into learned embedding vectors at each level and then aggregated to form a single positional embedding vector p (x): this formulation allows transformer models to inherently interpret structural relationships between tokens or nodes. tokens with similar structural contexts (e.g., siblings or nodes within the same subtree) naturally receive similar positional embeddings, guiding the self-attention mechanism to focus appropriately on structurally relevant elements. moreover, hierarchical positional embeddings enhance the model's capability to capture longrange dependencies inherent in structured data. by explicitly encoding tree positions rather than linear indices, hierarchical positional embeddings facilitate the model's understanding of relationships between tokens that may be distant in a linear sequence yet closely related structurally. ## methodology overview. we propose tree-enhanced code-berta, an extension of codeberta that integrates tree-based positional embeddings explicitly derived from abstract syntax trees (asts). by incorporating depth embeddings and sibling index embeddings, our model captures hierarchical relationships inherent in source code, enhancing both syntactic and semantic understanding. ## results we evaluate tree-enhanced codeberta on two tasks: masked language modeling (mlm) and clone detection. across both, our model consistently outperforms the baseline in accuracy, f1 score, precision, and recall. additionally, we report final training loss values to reinforce these improvements. on mlm, the weighted sum configuration achieves a lower loss of 0.41417 compared to 0.44388 for the original model. similarly, in clone detection, our model attains a loss of 0.21799 versus 0.25836 for the baseline. these reductions confirm that incorporating hierarchical positional embeddings not only improves task-specific performance but also facilitates more effective representation learning. ## limitations while our proposed tree-enhanced codeberta shows significant improvements in capturing hierarchical source code structures, our approach has several inherent limitations that must be acknowledged: ## conclusion we introduced tree-enhanced codeberta, a transformer-based model incorporating hierarchical positional embeddings from abstract syntax trees (asts). by integrating depth and sibling index embeddings, our approach captures structural nuances overlooked by traditional positional encodings. evaluations on masked language modeling (mlm) and clone detection confirm that these embeddings enhance representation learning, improving accuracy, f1 score, precision, and recall. the weighted sum integration strategy proves most effective, balancing structural and semantic information while maintaining efficiency.