# AST-Probe: Recovering abstract syntax trees from hidden representations of pre-trained language models

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

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 47 references that are in the map.
- First posted on arXiv 23 Jun 2022.

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [CodeBERT: A Pre-Trained Model for Programming and Natural Languages](https://mlsanity.com/paper/2002.08155/) (2020)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [CodeT5: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation](https://mlsanity.com/paper/2109.00859/) (2021)
- [Probing Classifiers: Promises, Shortcomings, and Advances](https://mlsanity.com/paper/2102.12452/) (2021)
- [A Primer in BERTology: What we know about how BERT works](https://mlsanity.com/paper/2002.12327/) (2020)
- [CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation](https://mlsanity.com/paper/2102.04664/) (2021)
- [CodeSearchNet Challenge: Evaluating the State of Semantic Code Search](https://mlsanity.com/paper/1909.09436/) (2019)
- [GraphCodeBERT: Pre-training Code Representations with Data Flow](https://mlsanity.com/paper/2009.08366/) (2020)
- [Recent Advances in Natural Language Processing via Large Pre-Trained Language Models: A Survey](https://mlsanity.com/paper/2111.01243/) (2021)
- [What you can cram into a single vector: Probing sentence embeddings for linguistic properties](https://mlsanity.com/paper/1805.01070/) (2018)
- [Unified Pre-training for Program Understanding and Generation](https://mlsanity.com/paper/2103.06333/) (2021)
- [A Survey of Machine Learning for Big Code and Naturalness](https://mlsanity.com/paper/1709.06182/) (2017)
- [Fine-grained Analysis of Sentence Embeddings Using Auxiliary Prediction Tasks](https://mlsanity.com/paper/1608.04207/) (2016)
- [What do Neural Machine Translation Models Learn about Morphology?](https://mlsanity.com/paper/1704.03471/) (2017)
- [ReACC: A Retrieval-Augmented Code Completion Framework](https://mlsanity.com/paper/2203.07722/) (2022)
- [Finding Universal Grammatical Relations in Multilingual BERT](https://mlsanity.com/paper/2005.04511/) (2020)
- [What Do They Capture? -- A Structural Analysis of Pre-Trained Language Models for Source Code](https://mlsanity.com/paper/2202.06840/) (2022)
- [Self-Supervised Contrastive Learning for Code Retrieval and Summarization via Semantic-Preserving Transformations](https://mlsanity.com/paper/2009.02731/) (2020)
- [When Deep Learning Met Code Search](https://mlsanity.com/paper/1905.03813/) (2019)
- [A Tale of a Probe and a Parser](https://mlsanity.com/paper/2005.01641/) (2020)
- [Fast and Memory-Efficient Neural Code Completion](https://mlsanity.com/paper/2004.13651/) (2020)

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

- [Deep Learning for Code Intelligence: Survey, Benchmark and Toolkit](https://mlsanity.com/paper/2401.00288/) (2023)
- [Towards Efficient Fine-tuning of Pre-trained Code Models: An Experimental Study and Beyond](https://mlsanity.com/paper/2304.05216/) (2023)
- [Exploring Code Analysis: Zero-Shot Insights on Syntax and Semantics with LLMs](https://mlsanity.com/paper/2305.12138/) (2023)
- [Better Context Makes Better Code Language Models: A Case Study on Function Call Argument Completion](https://mlsanity.com/paper/2306.00381/) (2023)
- [Evaluating and Explaining Large Language Models for Code Using Syntactic Structures](https://mlsanity.com/paper/2308.03873/) (2023)
- [The EarlyBIRD Catches the Bug: On Exploiting Early Layers of Encoder Models for More Efficient Code Classification](https://mlsanity.com/paper/2305.04940/) (2023)
- [Substance or Style: What Does Your Image Embedding Know?](https://mlsanity.com/paper/2307.05610/) (2023)
- [Meta-Models: An Architecture for Decoding LLM Behaviors Through Interpreted Embeddings and Natural Language](https://mlsanity.com/paper/2410.02472/) (2024)
- [Pop Quiz! Do Pre-trained Code Models Possess Knowledge of Correct API Names?](https://mlsanity.com/paper/2309.07804/) (2023)
- [Where's the Bug? Attention Probing for Scalable Fault Localization](https://mlsanity.com/paper/2502.13966/) (2025)
- [Hyperdimensional Probe: Decoding LLM Representations via Vector Symbolic Architectures](https://mlsanity.com/paper/2509.25045/) (2025)
- [INSPECT: Intrinsic and Systematic Probing Evaluation for Code Transformers](https://mlsanity.com/paper/2312.05092/) (2023)
- [Towards Understanding What State Space Models Learn About Code](https://mlsanity.com/paper/2602.06774/) (2026)
- [Latent Programming Horizons in Coding Agents](https://mlsanity.com/paper/2607.05188/) (2026)
- [How to Select Pre-Trained Code Models for Reuse? A Learning Perspective](https://mlsanity.com/paper/2501.03783/) (2025)
- [Toward Neurosymbolic Program Comprehension](https://mlsanity.com/paper/2502.01806/) (2025)

## Abstract (arXiv)

The objective of pre-trained language models is to learn contextual representations of textual data. Pre-trained language models have become mainstream in natural language processing and code modeling. Using probes, a technique to study the linguistic properties of hidden vector spaces, previous works have shown that these pre-trained language models encode simple linguistic properties in their hidden representations. However, none of the previous work assessed whether these models encode the whole grammatical structure of a programming language. In this paper, we prove the existence of a syntactic subspace, lying in the hidden representations of pre-trained language models, which contain the syntactic information of the programming language. We show that this subspace can be extracted from the models' representations and define a novel probing method, the AST-Probe, that enables recovering the whole abstract syntax tree (AST) of an input code snippet. In our experimentations, we show that this syntactic subspace exists in five state-of-the-art pre-trained language models. In addition, we highlight that the middle layers of the models are the ones that encode most of the AST information. Finally, we estimate the optimal size of this syntactic subspace and show that its dimension is substantially lower than those of the models' representation spaces. This suggests that pre-trained language models use a small part of their representation spaces to encode syntactic information of the programming languages.