A polar coordinate system represents syntax in large language models

TL;DR

The paper investigates how syntactic trees are encoded in large language models and identifies limitations of the Structural Probe, which only captures the existence of relations via distances. It introduces the Polar Probe, a linear readout that uses distance and angles in a polar subspace to encode both the presence and the type/direction of dependencies, outperforming prior methods. Using the English Web Treebank and a controlled long-nested dataset, trained on Mistral-7B-v0.1 and Llama-2-7b-hf, the Polar Probe achieves high dependency-type accuracy (e.g., AUC around $95\%$) and improved LAS, with the optimal probe dimensionality around $k'' \approx 128$, revealing a compact geometry for syntax. The results suggest a geometry of neural activations that aligns with linguistic theory and has implications for understanding brain syntax representations and broader compositional cognition.

Abstract

Originally formalized with symbolic representations, syntactic trees may also be effectively represented in the activations of large language models (LLMs). Indeed, a 'Structural Probe' can find a subspace of neural activations, where syntactically related words are relatively close to one-another. However, this syntactic code remains incomplete: the distance between the Structural Probe word embeddings can represent the existence but not the type and direction of syntactic relations. Here, we hypothesize that syntactic relations are, in fact, coded by the relative direction between nearby embeddings. To test this hypothesis, we introduce a 'Polar Probe' trained to read syntactic relations from both the distance and the direction between word embeddings. Our approach reveals three main findings. First, our Polar Probe successfully recovers the type and direction of syntactic relations, and substantially outperforms the Structural Probe by nearly two folds. Second, we confirm that this polar coordinate system exists in a low-dimensional subspace of the intermediate layers of many LLMs and becomes increasingly precise in the latest frontier models. Third, we demonstrate with a new benchmark that similar syntactic relations are coded similarly across the nested levels of syntactic trees. Overall, this work shows that LLMs spontaneously learn a geometry of neural activations that explicitly represents the main symbolic structures of linguistic theory.

Figures (8)

• Figure 1: Dependency trees hypothesized in linguistics and in neural networks.A. According to the dependency grammar framework, the sentences can be described as linear sequences of words connected by an acyclic graph. B. More precisely, such acyclic graph is both labeled and directed, where each edge has a direction, representing the hierarchy of the syntactic relation, and a label, representing the type of syntactic relation. C. The Structural Probe Hewitt2019 finds a a linear transform (gray plane) of the language model's activations (here simplified as a 3D space), such that the distance between word embeddings is predicted by their dependency tree. In the Structural Probe subspace, however, it is not possible to distinguish whether "The cat chases the mouse" or "The mouse chases the cat." D. Our Polar Probe finds a linear transformation where the angle between syntactically-related word additionally represents the type and direction of these syntactic relations, and the distance codes its presence. The colored arrows indicate orthogonal directions in the Polar-Probe subspace.
• Figure 2: The Polar Probe reliably identifies dependency types.A. PCA visualization of edges linearly read by the Polar Probe. The color of each edge corresponds to one of three different dependency types ('nsubj', 'obj', 'det'): the linear readouts point in systematic directions. B. AUC and Balanced Accuracy metrics obtained for dependency type classification. C. Pairwise cosine similarity (0=orthogonal, 1=collinear) matrices obtained without a probe (left) the Structural Probe (middle) and the Polar Probe (right).
• Figure 3: The Polar Probe outperforms the Structural Probe at identifying labeled and directed dependencies.A. For dependency existence, the Polar Probe matches the UUAS performance of the Structural probe, peaking at layer 16. B. For dependency type, the Polar Probe outperforms in Label Accuracy the Structural (LAS) Probe by around 80% accross the different layers of Llama-2-7b-hf. C. For both dependency existence and type, the Polar Probe outperforms in LAS the Structural Probe by around 90% accross the different layers of Llama-2-7b-hf.
• Figure 4: The optimal dimensionality for the Polar Probe is an order of magnitude small than model's layer size. Polar Probe performance as a function of dimensionality, measured by A. UUAS, B. Dependency Type Accuracy and C. LAS for Llama-2-7b-hf as a function of $k"$, the dimensionality of the probe's space. The optimal dimensionality for the Polar Probe is 128, achieving the highest LAS.
• Figure 5: Visualization of the dependency tree uncovered by the Polar Probe on a set of sentences with increasingly complex hierarchical structures.A. We display a PCA visualization of the distributions of word embeddings (once linearly read out by the Polar Probe), for the different syntactic levels in the "Controlled Dataset". Each individual distribution corresponds to a specific role of the word in the sentence. The centroids are linked with colored lines, displaying the true syntactic tree of the corresponding sentence. B. Most frequent syntactic tree prediction by the Polar Probe for the different syntactic levels. The relations between words are color coded according to the type of syntactic dependency. The incorrectly predicted relations are represented with dashed arrows. That is, either a dependency relation existence (no arrow), or a dependency type (with arrow) was erroneously identified.