CSA-Trans: Code Structure Aware Transformer for AST

2024/04/07 by Saeyoon Oh, Shin Yoo, Oh, Saeyoon +1 · 2 citations
Computer Science · #Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Natural Language Processing Techniques #Software Engineering (cs.SE)

paper · pdf · doi:10.48550/arxiv.2404.05767

openalex publication_date 2024/04/07 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

When applying the Transformer architecture to source code, designing a good self-attention mechanism is critical as it affects how node relationship is extracted from the Abstract Syntax Trees (ASTs) of the source code. We present Code Structure Aware Transformer (CSA-Trans), which uses Code Structure Embedder (CSE) to generate specific PE for each node in AST. CSE generates node Positional Encoding (PE) using disentangled attention. To further extend the self-attention capability, we adopt Stochastic Block Model (SBM) attention. Our evaluation shows that our PE captures the relationships between AST nodes better than other graph-related PE techniques. We also show through quantitative and qualitative analysis that SBM attention is able to generate more node specific attention coefficients. We demonstrate that CSA-Trans outperforms 14 baselines in code summarization tasks for both Python and Java, while being 41.92% faster and 25.31% memory efficient in Java dataset compared to AST-Trans and SG-Trans respectively.

Cited by

• Clarifying Semantics of In-Context Examples for Unit Test Generation
• Seamlessly Integrating Tree-Based Positional Embeddings into Transformer Models for Source Code Representation

Related

• Source Code Summarization with Structural Relative Position Guided Transformer
• AST-Transformer: Encoding Abstract Syntax Trees Efficiently for Code Summarization
• What does Transformer learn about source code?
• Code Summarization with Structure-induced Transformer
• Code Structure–Guided Transformer for Source Code Summarization