PaperlayerView exact revision on arXiv

Source-linked AI summary

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

José Antonio Hernández López, Martin Weyssow, Jesús Sánchez Cuadrado, Houari Sahraoui

arXiv:2206.11719v2cs.CLcs.AIcs.LGcs.PLcs.SE

TL;DR

Earlier probes had not established whether pre-trained language models encode the complete grammatical structure of code. AST-Probe recovers whole ASTs through a syntactic subspace in model representations; the paper reports that this information appears in five models, mostly in middle layers, and occupies 64–128 dimensions.

  • Problem

    Earlier code probes addressed selected linguistic properties or simplified trees, leaving open whether models encode complete programming-language ASTs.

    Source: §1 · p. 1, §2.3 · p. 3

  • Method

    AST-Probe assumes a syntactic subspace in model representations and uses its geometry to reconstruct whole ASTs from code snippets.

    Source: §1 · p. 2, §7 · p. 10

  • Results

    The syntactic subspace appears in five pre-trained language models, is mostly encoded in middle layers, and has a dimension ranging from 64 to 128.

    Source: §6 · p. 9, §7 · p. 10

  • Takeaways & Limitations

    The models use few dimensions of their hidden representation spaces to encode AST-related syntactic information.

    Source: §7 · p. 10

    The study considers transformer-based models and three common programming languages; probe coverage depends on extracting vectors d, c, and u from input code.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Pre-trained language models support many code-related tasks, but whether their hidden representations encode complete programming-language grammar remained an open question. The paper introduces AST-Probe to recover whole ASTs and investigates the subspace and model layers that encode this information.

1 INTRODUCTION

  • Prior probing work had not tested whether pre-trained language models encode whole ASTs, including the full grammatical structure of input code.The paper frames this as an open question despite models’ demonstrated potential for code-related tasks.Source: §1 · pp. 1–2
  • AST-Probe projects code representations into an assumed syntactic subspace and uses its geometry to predict the input code’s whole AST.The method is presented as a way to assess whether models implicitly learn programming-language syntax.Source: §1 · p. 2
  • The evaluation covers five pre-trained language models and the programming languages Python, JavaScript, and Go.The listed models include CodeBERT, GraphCodeBERT, CodeT5, CodeBERTa, and RoBERTa.Source: §1 · p. 2
  • The paper also estimates the optimal syntactic-subspace size to assess how many dimensions hold AST-related syntax information.This is presented as an estimate of the subspace’s compactness.Source: §1 · p. 2
  • The contributions include representing ASTs as a compact tuple of vectors adapted from [34] and a language-agnostic method for recovering whole ASTs.The method is intended to assess models’ understanding of programming-language syntax.Source: §1 · p. 2

2 BACKGROUND AND RELATED WORK

The background reviews pre-trained language models and probing methods, then identifies that earlier code probes did not recover complete ASTs. AST-Probe addresses that gap by targeting whole trees through a syntactic subspace.

2.1 Pre-trained language models

  • Pre-trained language models learn representations from large text corpora using general objectives, then can be fine-tuned for task-specific data.The overview distinguishes autoregressive, masked-language, and encoder-decoder model objectives.Source: §2.1 · p. 2

2.2 Probing in NLP

  • Probing tests whether hidden representations encode a linguistic property; structural probes map word vectors so distances approximate dependency-tree distances.AST-Probe instead uses an orthogonal projection and targets labeled ASTs, which include non-terminal nodes absent from the input token sequence.Source: §2.2 · pp. 2–3

2.3 Probing in pre-trained language models of code

  • Earlier code probes predicted selected AST properties or unlabeled binary trees, but did not recover whole ASTs with terminal and non-terminal nodes.AST-Probe targets whole-tree recovery by assuming a syntactic subspace that encodes programming-language syntax.Source: §2.3 · p. 3

3 THE AST-PROBE APPROACH

AST-Probe assumes AST syntax is encoded in a subspace of a model’s hidden representations. It converts ASTs into vectors, learns an orthogonal projection, and uses projected representations to reconstruct the trees.

3 THE AST-PROBE APPROACH

  • AST-Probe learns a syntactic subspace from hidden representations and predicts a tuple of vectors whose information can be used to reconstruct the whole AST.The approach assumes the subspace encodes AST-related syntax and estimates its dimension.Source: §3 · p. 3

3.1 AST to binary tree

  • Unary nodes are merged with their child, n-ary nodes receive a ∅ node, and certain unary chains are removed with their label transferred to the terminal.The transformations produce a binary tree that can be converted back to the AST by removing ∅ nodes, reconnecting the tree, and expanding unary chains.Source: §3.1 · p. 4

3.2 Binary tree to tuple

  • The binary tree is encoded as (d, c, u): d carries structure, c labels lowest common ancestors of consecutive tokens, and u records unary labels.The tuple transformations are bidirectional, allowing recovery of the binary tree and AST.Source: §3 · p. 3, §3.2 · p. 4

3.3 Syntactic subspace and probing classifier

  • An orthogonal projection defines the syntactic subspace, and projected token representations predict d, c, and u using distances and learned label vectors.The probe learns an orthonormal basis and vectors C and U; its parameter complexity is governed by the subspace dimension.Source: §3.3 · p. 5
  • The training objective combines losses for structure and labels with orthogonal regularization; the structural loss uses rankings because d’s ranking suffices to recover the AST.The regularizer encourages the basis vectors to be orthonormal.Source: §3.3 · pp. 5–6

4 EXPERIMENTS

The experiments assess AST-Probe across code datasets, models, layers, and subspace dimensions. They define comparison questions and evaluation procedures, but these passages do not report experimental outcomes.

4 EXPERIMENTS

  • The experiments ask whether AST-Probe parses informative representations, which model and layers encode ASTs best, and how compact the syntactic subspace is.For the first question, baseline comparisons test whether performance differs from non-contextualized embeddings and a randomly initialized model.Source: §4 · p. 6

4.1 Data

  • Experiments use CodeSearchNet samples in Python, Go, and JavaScript, with 20,000 training, 4,000 test, and 2,000 validation samples per language.ASTs are extracted with tree-sitter and converted into vector tuples.Source: §4.1 · p. 6

4.2 Language models

  • The comparison includes CodeBERT, GraphCodeBERT, CodeT5, CodeBERTa, and RoBERTa, alongside CodeBERT-0 and CodeBERTrand baselines.Only CodeT5’s encoder layers are considered; RoBERTa tests a model trained on natural-language texts rather than code.Source: §4.2 · p. 6

4.3 Evaluation metrics

  • Effectiveness is assessed by comparing the recovered AST with the ground truth using precision, recall, and F1 over tree constituents.An AST constituent is a non-terminal and its token-position scope.Source: §4.3 · p. 6

4.4 Addressing the research questions

  • For RQ1–RQ3, the probe is trained for each language, model, and layer with syntactic-subspace dimension m2 = 128; RQ4 varies that dimension from 8 to 512.For RQ4, each model’s layer with its best F1 score is selected first.Source: §4.4 · p. 7

4.5 Training details

  • Training averages subword representations for whole-word tokens and uses 768-dimensional embeddings, orthogonal regularization λ = 5, and Adam optimization.The initial learning rate is 0.001; training runs for at most 20 epochs with early stopping after five epochs without validation-loss improvement.Source: §4.5 · p. 7

5 RESULTS

AST-Probe distinguishes informative code representations from baselines, and the strongest AST results come from CodeBERT and GraphCodeBERT. AST information peaks in middle layers and occupies a compact syntactic subspace.

5.1 RQ1 – Validity of the AST-Probe

  • Across all considered languages and metrics, baselines trail the other models, supporting AST-Probe validity.Table 2 reports each model’s best layer and includes the baselines.Source: §5.1 · p. 7, Table 2 · p. 8

5.2 RQ2 – Best model

  • CodeBERT and GraphCodeBERT lead F1 across languages, but GraphCodeBERT’s margins are +0.0032 F1 in Python, +0.0010 F1 in JavaScript, and +0.0084 F1 in Go.The differences are marginal, so the authors do not firmly conclude which model is better.Source: §5.2 · p. 8
  • RoBERTa, though trained on text, competes with code-trained models, matches CodeT5 across languages, and outperforms CodeBERTa in Python and JavaScript.The authors report similar results with CodeT5 for all languages.Source: §5.2 · p. 8

5.3 RQ3 – Best layers

  • AST-Probe F1 peaks in the middle layers across models; best layers are 5th and 8th for RoBERTa, 6th-7th-8th for CodeT5, and 4th-5th for CodeBERT, GraphCodeBERT, and CodeBERTa.Figure 6 plots F1 by layer.Source: §5.3 · p. 8, Figure 6 · p. 8

5.4 RQ4 – Estimating the dimension of S

  • Across languages and models, the syntactic subspace spans 64–128 dimensions of 768, and F1 falls when it is too small or approaches the full space.The authors attribute the decline near the full dimension to the projection retaining information beyond syntax.Source: §5.4 · pp. 8–9, Figure 8 · p. 9

6 DISCUSSION

The discussion concludes that AST-related syntax is encoded in middle layers and a compact subspace. It contrasts full-AST recovery with prior tree probes and identifies scope limits and future questions.

6 DISCUSSION

  • The experiments identify AST-related syntax in the hidden representations of five models, concentrated in middle layers and a 64–128-dimensional subspace.The discussion aligns the middle-layer finding with prior syntax-probe results and BERT findings for natural languages.Source: §6 · pp. 9–10
  • Unlike Wan et al.’s [38] probe, which recovers unlabeled binary trees, AST-Probe recovers full ASTs, a structure the authors describe as closer to programming-language syntax.The authors also note that their comparison evaluates ground-truth ASTs against predicted ASTs using a metric used for constituency parsers.Source: §6 · pp. 9–10
  • The authors estimate that 8.3%–16.7% of the 768 representation dimensions encode AST-related information, leaving the content of the rest as a future question.They describe this as the first analysis of this type.Source: §6 · p. 10
  • The authors propose studying correlations between syntax recovery and downstream-task performance, and using constituent-level metrics to locate tree-prediction errors.They present both as future research directions.Source: §6 · p. 10
  • The study covers transformer-based models and three common programming languages; applying the probe depends on extracting vectors d, c, and u and having a language grammar.The authors say the probe is architecture- and language-agnostic in principle, given representation spaces and a grammar definition.Source: §6 · p. 10

7 CONCLUSION AND FUTURE WORK

AST-Probe recovers whole ASTs by learning a syntactic subspace and using its geometry. Experiments find this subspace in five pre-trained language models, concentrated mostly in middle layers and using few dimensions; future work will broaden the study.

7 CONCLUSION AND FUTURE WORK

  • AST-Probe recovers whole ASTs from hidden representations by learning a syntactic subspace and using its geometry.The method predicts the AST of an input code snippet.Source: §7 · p. 10
  • The syntactic subspace exists in five pre-trained language models, with most AST information in middle layers and few dimensions used for syntax.The paper estimates the subspace's compactness.Source: §7 · p. 10
  • Future work will test more programming languages, models, and architectures, and examine syntax-task correlations and syntactic subspaces before and after fine-tuning.Source: §7 · p. 10

Related papers

Loading 2206.11719v2…