# CSA-Trans: Code Structure Aware Transformer for AST

arXiv: https://arxiv.org/abs/2404.05767 · HTML: https://mlsanity.com/paper/2404.05767/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 47 references that are in the map.
- First posted on arXiv 7 Apr 2024.

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://mlsanity.com/paper/2006.03654/) (2020)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation](https://mlsanity.com/paper/2109.00859/) (2021)
- [Benchmarking Graph Neural Networks](https://mlsanity.com/paper/2003.00982/) (2020)
- [Invariant and Equivariant Graph Networks](https://mlsanity.com/paper/1812.09902/) (2018)
- [Graph Neural Networks with Learnable Structural and Positional Representations](https://mlsanity.com/paper/2110.07875/) (2021)
- [Do Transformers Really Perform Bad for Graph Representation?](https://mlsanity.com/paper/2106.05234/) (2021)
- [Session-based Recommendation with Graph Neural Networks](https://mlsanity.com/paper/1811.00855/) (2018)
- [Pure Transformers are Powerful Graph Learners](https://mlsanity.com/paper/2207.02505/) (2022)
- [code2seq: Generating Sequences from Structured Representations of Code](https://mlsanity.com/paper/1808.01400/) (2018)
- [A Transformer-based Approach for Source Code Summarization](https://mlsanity.com/paper/2005.00653/) (2020)
- [DeepInf: Social Influence Prediction with Deep Learning](https://mlsanity.com/paper/1807.05560/) (2018)
- [Improving Automatic Source Code Summarization via Deep Reinforcement Learning](https://mlsanity.com/paper/1811.07234/) (2018)
- [Equivariant and Stable Positional Encoding for More Powerful Graph Neural Networks](https://mlsanity.com/paper/2203.00199/) (2022)
- [Retrieval-Augmented Generation for Code Summarization via Hybrid GNN](https://mlsanity.com/paper/2006.05405/) (2020)
- [Guiding Neural Machine Translation with Retrieved Translation Pieces](https://mlsanity.com/paper/1804.02559/) (2018)
- [Code Generation as a Dual Task of Code Summarization](https://mlsanity.com/paper/1910.05923/) (2019)
- [Automatically Generating Commit Messages from Diffs using Neural Machine Translation](https://mlsanity.com/paper/1708.09492/) (2017)
- [Transformers Generalize DeepSets and Can be Extended to Graphs and Hypergraphs](https://mlsanity.com/paper/2110.14416/) (2021)
- [Code Structure Guided Transformer for Source Code Summarization](https://mlsanity.com/paper/2104.09340/) (2021)
- [Tree-to-Sequence Attentional Neural Machine Translation](https://mlsanity.com/paper/1603.06075/) (2016)
- [Transformers meet Stochastic Block Models: Attention with Data-Adaptive Sparsity and Cost](https://mlsanity.com/paper/2210.15541/) (2022)

## Built on by (most cited first, 3 of 3)

- [CodeCrash: Exposing LLM Fragility to Misleading Natural Language in Code Reasoning](https://mlsanity.com/paper/2504.14119/) (2025)
- [Clarifying Semantics of In-Context Examples for Unit Test Generation](https://mlsanity.com/paper/2510.01994/) (2025)
- [Seamlessly Integrating Tree-Based Positional Embeddings into Transformer Models for Source Code Representation](https://mlsanity.com/paper/2507.04003/) (2025)

## Abstract (arXiv)

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.