Source-linked AI summary
Multi-Modal Attention Network Learning for Semantic Source Code Retrieval
Yao Wan, Jingdong Shu, Yulei Sui, Guandong Xu, Zhou Zhao, Jian Wu, Philip S. Yu
TL;DR
Existing code retrieval methods may overlook structured code semantics and make retrieval decisions difficult to interpret. MMAN combines token, AST, and CFG representations with attention-based fusion, and experiments report improved retrieval performance over DeepCS and positive effects from attention. The approach is constrained by its need for large-scale training data and CFG extraction from whole programs.
Problem
Existing code retrieval methods may ignore structured AST and CFG features and lack explainability for determining which code parts contribute to retrieval results.
Method
MMAN combines LSTM-based token representation, Tree-LSTM-based AST representation, and GGNN-based CFG representation with multi-modal attention fusion.
Results
MMAN (Tok+AST+CFG) consistently obtains higher R@k and MRR performance than DeepCS, while attention improves performance relative to its ablated variants.
Takeaways & Limitations
Attention visualization provides an interpretable view of modality-specific contributions, including greater CFG attention on invoked-function nodes.
Takeaways & Limitations
The approach requires large-scale corpus training and CFG extraction from whole programs, limiting extension to snippets where CFGs cannot be extracted.
Abstract
from arXiv · showhide
Code retrieval techniques and tools have been playing a key role in facilitating software developers to retrieve existing code fragments from available open-source repositories given a user query. Despite the existing efforts in improving the effectiveness of code retrieval, there are still two main issues hindering them from being used to accurately retrieve satisfiable code fragments from large-scale repositories when answering complicated queries. First, the existing approaches only consider shallow features of source code such as method names and code tokens, but ignoring structured features such as abstract syntax trees (ASTs) and control-flow graphs (CFGs) of source code, which contains rich and well-defined semantics of source code. Second, although the deep learning-based approach performs well on the representation of source code, it lacks the explainability, making it hard to interpret the retrieval results and almost impossible to understand which features of source code contribute more to the final results. To tackle the two aforementioned issues, this paper proposes MMAN, a novel Multi-Modal Attention Network for semantic source code retrieval. A comprehensive multi-modal representation is developed for representing unstructured and structured features of source code, with one LSTM for the sequential tokens of code, a Tree-LSTM for the AST of code and a GGNN (Gated Graph Neural Network) for the CFG of code. Furthermore, a multi-modal attention fusion layer is applied to assign weights to different parts of each modality of source code and then integrate them into a single hybrid representation. Comprehensive experiments and analysis on a large-scale real-world dataset show that our proposed model can accurately retrieve code snippets and outperforms the state-of-the-art methods.
I. INTRODUCTION
Code retrieval seeks code fragments matching natural-language specifications, but existing methods often miss structured code semantics and provide limited interpretability. MMAN addresses these gaps by combining token, AST, and CFG representations with attention-based fusion.
- Motivation: Code retrieval matches natural-language specifications to reusable code fragments from open-source repositories.The paper frames retrieval as searching large code repositories for code with specified functionality.
- Challenges: Existing approaches may capture shallow information while ignoring structured features such as ASTs and CFGs that encode code semantics.The cited limitation concerns method names, tokens, and API sequences rather than richer structural information.
- Challenges: Different code views capture complementary semantics because AST nodes represent tokens while CFG nodes represent statements and control flow.The motivating example contrasts plain text, a type-augmented AST, and a CFG for interpreting code functionality.
- MMAN: MMAN represents code sequentially with an LSTM, structurally with a Tree-LSTM for ASTs and a GGNN for CFGs, then fuses the modalities.The model is designed to capture these modalities simultaneously in a comprehensive representation.
- MMAN: Attention assigns different weights to parts of each code modality, providing an interpretable component for the deep multi-modal representation.The paper identifies attention as the mechanism intended to improve explainability.
- Evaluation: The evaluation uses a real-world GitHub dataset containing 28,527 C code snippets and compares MMAN with state-of-the-art methods.The experiments are presented as validation of the proposed model’s effectiveness.
A. Deep Code Representation
Deep code representation research uses structured and multi-modal learning, but prior approaches may omit CFG information or lack unified fusion. This section frames code retrieval as learning representations of heterogeneous code and descriptions in a shared semantic space.
- Deep Code Representation: Tree-structured neural networks have been used to represent ASTs for code-related tasks, including program classification and code summarization.The cited examples use Tree-CNN and Tree-LSTM architectures for structured code representation.
- Deep Code Representation: Prior approaches may ignore CFGs, despite their semantic information, and lack a unified network for fusing multiple code modalities.The paper presents multimodal fusion as the response to this limitation.
- Multi-Modal Learning: Multi-modal learning commonly develops joint representations for applications such as image captioning, summarization, visual question answering, and dialogue systems.These examples motivate applying multi-modal learning ideas beyond their existing application domains.
- Multi-Modal Learning: MMAN combines multi-modal code representation with attention to learn which code parts contribute more to the final semantic representation.The paper explicitly extends multi-modal learning with an attention mechanism for code.
- Attention Mechanism: Attention mechanisms let models focus on necessary input parts and have been applied across translation, captioning, classification, and visual question answering.Related work includes visual and textual attention frameworks as well as co-attention variants.
- Problem Formulation: The paper formalizes retrieval as learning representations of heterogeneous code snippets and descriptions in an intermediate semantic space.At test time, the model returns a similarity vector for candidate code snippets given a query.
B. Multi-Modal Learning
The paper distinguishes joint representations, which combine unimodal signals, from coordinated representations, which separately process modalities while aligning them in an intermediate semantic space. Attention provides weighting over inputs or internal features, and MMAN applies inter-attention to code modalities.
- Joint and Coordinated Representations: Joint representations combine unimodal signals into one representation space, whereas coordinated representations process them separately while enforcing similarity constraints.The paper applies joint representation to complementary code modalities and coordinated representation to code snippets and descriptions from different sources.
- Joint and Coordinated Representations: For a code snippet, MMAN extracts token, AST, and CFG modalities as complementary representations of the same code.These modalities are combined through a function relying on their unimodal representations.
- Joint and Coordinated Representations: Coordinated representation projects code and descriptions into an intermediate semantic space using separate projection functions and a similarity constraint.MMAN adopts cosine similarity for coordination.
- Attention Mechanism: Attention assigns weights over inputs or internal features to control which information becomes visible to later network components.The paper describes attention as either inter-attention or intra-attention depending on whether the query differs from or equals the key.
- Attention Mechanism: MMAN uses inter-attention with token or AST/CFG hidden states as keys and corresponding context vectors as values.The attended representation is formed from the weighted values in the memory.
IV. MULTI-MODAL ATTENTION NETWORK
MMAN trains a retrieval network from annotated code-description pairs and retrieves related snippets for natural-language queries. Its hybrid code representation combines sequential tokens with AST and CFG structure through modality-specific neural encoders.
- A. An Overview: MMAN’s workflow has offline training and online retrieval stages using annotated <code, description> pairs.The trained network receives a natural-language query and retrieves related source-code snippets.
- A. An Overview: The architecture contains multi-modal code representation, multi-modal attention fusion, and model learning modules.The modules represent code, weight and fuse modality parts, and learn code-description representations in a common space with a ranking loss.
- B. Multi-Modal Code Representation: MMAN represents code tokens with an LSTM, its AST with a Tree-LSTM, and its CFG with a GGNN.This hybrid embedding approach incorporates both sequential and structural source-code information.
- 1) Lexical Level - Tokens: The lexical-level representation uses an LSTM over sequential code tokens, whose final hidden state represents the token modality.Each token is embedded before being processed by the LSTM.
2) Syntactic Level - AST:
MMAN represents code syntax with Tree-LSTM over ASTs and GGNN over CFGs, alongside token-based representations. These encoders preserve hierarchical and control-flow information for later code retrieval.
- AST and CFG encoding: Tree-LSTM represents the syntactic level of source code from AST embeddings.It uses input, memory, output, and multiple forget gates for child nodes.
- AST and CFG encoding: The hidden state of the AST root node serves as the AST modality representation.Missing child states are set to zero, and Tree-LSTM reduces to vanilla LSTM for chain-shaped trees.
- AST and CFG encoding: GGNN represents CFGs whose vertices are code statements and whose typed edges encode control flow.It initializes vertex states, aggregates neighbor messages, and updates states recurrently.
- AST and CFG encoding: After T message-passing rounds, MMAN sums the hidden states of all CFG vertices to obtain the CFG representation.Figure 6 illustrates the GGNN structure.
- AST and CFG encoding: The architecture learns token, AST, and CFG representations with LSTM, Tree-LSTM, and GGNN before attention fusion.These modalities are then integrated into a single representation for model learning.
C. Multi-Modal Attention Fusion
MMAN applies attention within each code modality so that different tokens, AST nodes, and CFG nodes receive different weights before fusion. The attended modality vectors are combined into one semantic code representation.
- Within-modality attention: The fusion layer assigns different weights to elements within each modality because they contribute unequally to code semantics.The final representation combines attention-weighted information across modalities.
- Token attention: Token attention identifies code tokens that are more important to representing the token sequence.Token scores use hidden states, a linear layer, and a learned context vector.
- AST attention: AST attention weights nodes differently so distinct code constructs, such as if-condition-then structures, are represented distinctly.AST scores use node hidden states and a learned AST context vector.
- CFG attention: CFG attention assigns weights to CFG nodes, and sigmoid weighting performs better than softmax weighting in the experiments.CFG scores use node hidden states and a learned CFG context vector.
- Multi-modal fusion: MMAN concatenates the attended modality representations and feeds them into a one-layer linear network for multimodal fusion.The resulting vector is the final semantic representation of the code snippet.
D. Description Representation
MMAN encodes code descriptions with a vanilla LSTM and trains code and description representations in a shared semantic space. A ranking objective increases similarity for correct descriptions relative to incorrect ones.
- Description encoding: During training, descriptions come from code comments, while testing treats descriptions as input queries.A vanilla LSTM represents each description.
- Description encoding: The final hidden state of the description LSTM provides the vector representation of the description.Words are embedded before sequential encoding.
- Joint training: MMAN trains on triples containing a code snippet, its positive description, and a randomly selected negative description.The negative description is sampled from the pool of positive descriptions.
- Joint training: The model predicts cosine similarities for code-description pairs and minimizes a hinge ranking loss with margin β = 0.05.The loss encourages correct-pair similarity to rise and incorrect-pair similarity to fall.
F. Code Retrieval
At retrieval time, MMAN embeds each code snippet with its multimodal code module and each query with the description LSTM, then ranks snippets by similarity. The experiments examine overall retrieval performance, modality contributions, sensitivity to input characteristics, and attention explainability.
- Retrieval procedure: Given a code base and query, MMAN ranks all code snippets by the similarity between their code and query representations.Higher similarity indicates greater relatedness to the query.
- Evaluation questions: The retrieval experiments compare MMAN with state-of-the-art approaches to evaluate whether it improves code-retrieval performance.This is the stated focus of research question RQ1.
- Evaluation questions: The study evaluates the effectiveness and contribution of token, AST, and CFG modalities individually and in combination.This is the stated focus of RQ2.
- Evaluation questions: Additional analyses vary code length, AST node count, CFG node count, and comment length to assess model sensitivity.This is the stated focus of RQ3.
- Evaluation questions: RQ4 evaluates the attention mechanism and the explainability of its visualizations.The paper positions attention analysis as an evaluation of interpretability.
A. Dataset Collection
The study builds a C corpus of commented methods from GitHub, extracts sequential and structured code features, and evaluates retrieval automatically using ranking metrics.
- Dataset construction: The authors crawl GitHub C repositories and retain methods with documentation comments, producing 28,527 commented C methods.Repositories with fewer than two stars are excluded.
- Feature extraction: The dataset pairs each method with its name, tokens, AST, CFG, and description for multimodal retrieval.Method names and code tokens are extracted alongside AST and CFG representations.
- Dataset construction: The corpus was crawled in October 2016 and contains repositories created from August 2008 through October 2016.
- Evaluation setup: The data are shuffled into 27,527 training samples and 1,000 evaluation samples, with automatic evaluation replacing manually labeled search results.The authors describe manual labeling as potentially introducing human bias.
- Evaluation setup: Retrieval effectiveness is measured with SuccessRate@k and Mean Reciprocal Rank, where higher values indicate better search performance.SuccessRate@k measures queries with a correct result in the top k; MRR uses the reciprocal rank of the first hit.
C. Implementation Details
The implementation compares MMAN variants and established baselines on automatic retrieval evaluation, emphasizing multimodal attention and overall ranking performance.
- Compared methods: The comparison includes CodeHow, DeepCS, and MMAN variants with or without attention across modality combinations.CodeHow is information-retrieval based, while DeepCS learns unified code and query representations.
- Evaluation protocol: The automatic evaluation uses 1,000 descriptions, treating each description as a query and its corresponding code snippet as ground truth.Results are reported with R@1, R@5, R@10, and MRR.
- Overall results: MMAN (Tok+AST+CFG) consistently outperforms DeepCS on SuccessRate@k and MRR under all experimental settings.
- Overall results: 26.18%, 16.06%, and 12.21% are MMAN's improvements over DeepCS for R@1, R@5, and R@10, respectively.The reported MRR improvement over DeepCS is 19.89%.
- Attention ablation: The attention mechanism has a positive effect when comparing attention-equipped and attention-free MMAN variants.
E. Q2: Effect of Each Modality
Experiments examine modality combinations, robustness, retrieved examples, and attention visualizations to assess MMAN's multimodal representation and interpretability.
- Modality combinations: Adding more source-code modalities improves performance, indicating complementary rather than conflicting modality information.Attention also positively affects fusion across modality combinations.
- Sensitivity analysis: MMAN remains stable in most cases as code length or structural node number varies, supporting robustness of its multimodal representation.
- Sensitivity analysis: Increasing comment length decreases performance across four metrics, suggesting greater difficulty in understanding longer comments.
- Qualitative retrieval: For the query “Print any message in the axel structure,” MMAN accurately retrieves the ground-truth snippet compared with DeepCS and modality variants.The qualitative comparison highlights the multimodal approach's more accurate code representation.
- Attention visualization: Token attention highlights the function name print_message, while AST attention emphasizes axel and BinaryOperator nodes.
- Attention visualization: CFG attention emphasizes invoked function nodes such as prinf and free, reflecting CFG's description of source-code control flow.
VI. DISCUSSION
MMAN’s effectiveness is attributed to comprehensive multi-modal code representations, attention-based interpretability, and a unified code–query learning framework. The paper reports that MMAN is effective and outperforms state-of-the-art approaches, while identifying metric and extensibility limitations and proposing broader future evaluation.
- Strength of MMAN: MMAN combines tokens, ASTs, and CFGs as complementary modalities, with attention weighting parts of each modality to expose their contributions.The model uses attention to support interpretability of its deep learning-based representation.
- Strength of MMAN: A unified end-to-end architecture learns heterogeneous source-code and natural-language representations in an intermediate semantic space.This framework jointly learns representations of source code and descriptions.
- Threats to Validity and Limitations: The evaluation uses only SuccessRate@R and MRR because labeling results for Precision@k would require manual relevance judgments that may introduce human bias.The authors also state that human evaluation is needed for fair comparison with DeepCS.
- Threats to Validity and Limitations: Extensibility is constrained by the need for large-scale training data, extensive sample filtering, and whole-program CFG extraction.The authors specifically identify code snippets from StackOverflow as contexts where CFG extraction may be unavailable.
- Discussion: MMAN is reported to be effective and to outperform state-of-the-art approaches in the experimental study.The paper describes this conclusion without reporting a numerical result in the supplied passage.
- Future Work: Future work includes experiments on Java and Python datasets, human evaluation, and applying multimodal code representation to code summarization and clone detection.These directions are presented as plans to further verify effectiveness and explore other software engineering tasks.