Enhancing Tree-Based Language Model via Function Dependencies for Code Search

Recent advancements in pre-trained programming language models have significantly enhanced code search tasks. However, many models struggle to fully leverage the syntactic and structural information in source code, especially when parsing syntax trees and understanding function-level semantics. In this paper, we present CodeBERT-TD,a novel multi-modal code search model that combines documentation, code, syntax trees, and function dependency information for improved retrieval performance. Our approach employs a tree-based attention mechanism and hybrid masking strategy to capture hierarchical relationships within the code. By flattening syntax trees through a node-wrapped traversal and using tree-structured attention masks, we improve the model's focus on syntactic dependencies. Additionally, incorporating function dependency relationships enhances performance on downstream tasks. To tackle sequence length limitations from multiple modalities, we utilize rotary positional encoding to boost extrapolation capabilities. These components come together to create a cooperative design, with extensive evaluations demonstrating that our model significantly enhances code search performance, laying a solid foundation for future advancements in code representation learning.