Tree Metric Embedding

structures / Linear Structures / Tree Metric Embedding

TOOLmetric-modellearned-metricintermediatehow it's classified →

A learned inner product (equivalently, a linear transform B) under which squared distance between two representations approximates graph distance in a discrete tree, and squared norm approximates depth from the root — embedding an entire tree's edge structure into the metric of a linear subspace, not just a single decodable relation.

Replicationcomputed from the corpus — never hand-assigned

14 papers · no shared authors4 architecture classes · across papers2 domains · across papers15 model families · across papers

Filled = two or more values reported by papers that share no author — replication. Outlined = two or more values, but all from a single study — breadth, not replication. Grey = a single value. Derived from paper authorship and each model's architecture class, domain and family; it updates itself when a paper is added.

Definition

Given a tree $T$ over $n$ nodes (e.g. a sentence’s dependency parse) with graph distance $d_T(u,v)$ (number of edges on the path between $u,v$ ) and depth $\mathrm{depth}_T(u)$ (edges from $u$ to the root), a tree metric embedding is a linear map $B \in \mathbb{R}^{k\times m}$ of node representations $h_u \in \mathbb{R}^m$ such that

$d_B(h_u,h_v)^2 = \big(B(h_u-h_v)\big)^\top\big(B(h_u-h_v)\big) \approx d_T(u,v), \qquad \|h_u\|_B^2 = (Bh_u)^\top(Bh_u) \approx \mathrm{depth}_T(u)$

$B$ is fit (by gradient descent on the discrepancy between predicted and true tree distances/depths, summed over all node pairs in a training corpus of trees) to recover the entire tree’s edge and depth structure simultaneously from one shared linear transform — not a single pairwise relation, and not per-node classification.

Terminology (corrected). This is a learned linear metric model / structural probe: $B$ induces a positive semi-definite form $M=B^\top B$ and a squared readout $d_B(h_u,h_v)^2=(h_u-h_v)^\top M (h_u-h_v)$ . Two precise points:

The regression target is the squared Euclidean distance $\|B(h_u-h_v)\|^2$ . This squared quantity is not itself a metric — it violates the triangle inequality in general. The genuine metric is the unsquared $\delta_B(h_u,h_v)=\|B(h_u-h_v)\|$ ; fitting the square is a modeling choice (and it is the square that Pythagorean tree embeddings target, see below), so one should not call $d_B^2$ a distance.
The depth readout $\|Bh_u\|^2$ pins the root to the origin of the transformed space and is therefore not translation-invariant: replacing $h_u \mapsto h_u+c$ preserves all pairwise distances but changes every depth prediction. Depth recoverability is thus a property of the probe’s chosen origin, not an intrinsic, coordinate-free property of the representation geometry. (Distances, by contrast, are translation-invariant.)

Because the points $h_u$ are fixed and only $B$ is trained, an exact solution need not exist; the model is an approximation, and its success is the empirical claim.

Intuition

A classifier probe answers “is word $u$ the head of word $v$ ?” one pair at a time. A tree metric embedding instead asks for a single coordinate system in which ordinary squared Euclidean distance — a single global quantity, defined the same way for every pair of points — already encodes how far apart every pair of words is in the parse tree, and vector length already encodes how deep each word sits. If such a $B$ exists, the tree isn’t just “decodable” from the representation; the representation’s own metric geometry already is (an approximation of) the tree.

Properties

A global, not pairwise, claim. Because $d_B$ is a genuine metric (bilinear form composed with a norm), fitting $B$ to reproduce $d_T(u,v)$ for some pairs constrains its predictions for every pair simultaneously (triangle-inequality-like consistency) — this is a stronger structural claim than a set of independent pairwise classifiers could make, even if each pairwise classifier were individually accurate.
Distance and depth are two separate probes sharing one representational picture. The distance probe (edges between any two nodes) and the depth probe (edges to the root) are fit independently but are proposed as complementary readouts of the same underlying embedded tree — together they recover both a tree’s undirected shape and its rooted orientation.
Effective rank is typically far below the ambient dimension. The transform $B$ need not be full-rank; a “how many dimensions does the tree actually need” sweep (constraining $B\in\mathbb{R}^{k\times m}$ for small $k$ ) is the operational way to ask whether the tree occupies a genuine low-dimensional subspace of representation space, distinct from just asking whether some decodable signal exists.
A specifically Euclidean (linear-transform) instance of tree-distance embedding — not the only possible metric space for the job. Trees embed with much lower distortion in curved, negatively-curved geometry (see Hyperbolic Manifold‘s exponential-volume-growth property) than in flat Euclidean space; a tree metric embedding under a linear map $B$ is the flat-space baseline against which curved-space alternatives (a Poincaré/hyperbolic probe) are compared, not evidence that the representation space itself is Euclidean rather than curved.
Recoverability shifts differently under fine-tuning depending on the downstream task. Perez-Mayos, Carlini, Ballesteros & Wanner (2021) track the same structural probe across BERT fine-tuning checkpoints for six tasks: parsing tasks (constituency, dependency) reinforce tree-metric-probe recoverability early and preserve it; PoS tagging steadily erodes it; semantics-oriented tasks (SRL, QA, paraphrase identification) lose some absolute recoverability (UUAS, Root %) while keeping relative distance/depth ordering comparatively stable (DSpr, NSpr). See tree-metric-evolution-during-finetuning.
Sensitivity to constituency depth, not just dependency structure — but only partially confirmed. Kennedy (2025) applies the original Hewitt & Manning probe, unmodified and out-of-domain, to sentences with identical dependency parses but differing constituency-tree complement size (vP/TP/CP); two of four tested conditions show statistically significant predicted-distance differences, two do not, and an unresolved finiteness confound limits how cleanly this can be read as evidence for constituency-depth sensitivity specifically. See constituency-depth-partially-recoverable-from-tree-metric.
Bottom-up construction across layers, but not universally. Someya, Yoshida, Yanaka & Oseki (2025) fit the structural probe separately per layer and compute an expected-construction-layer for sub-tree structures: in BERT, local (micro-syntactic) structure is consistently built before global (macro-syntactic) structure is integrated, and mistimed integration (too early in BERT-base, too late in BERT-large) predicts subject-verb agreement failures — but GPT-2 shows a more parallel construction pattern instead, so bottom-up derivation is architecture-dependent, not a universal property of the tree-metric-probe structure itself. See bottom-up-syntactic-derivation.
Does not by itself establish causal use. Recovering tree distances via a fitted probe shows the geometry is consistent with an embedded tree; as with any probe (see Linear probing‘s “what it does and doesn’t establish”), it does not on its own show the model’s own computation routes through this metric structure, absent a separate causal intervention.

Extension from trees to general in-context relational graphs

Diego-Simón, Orhan, Chemla, Lakretz & King (2026, “Polar Probe”) generalize the distance/depth pair from parse trees to arbitrary in-context-constructed relational graphs across five domains (arithmetic ordinality, spatial layouts, thematic roles, family trees, metro maps): a jointly-trained linear probe subspace decomposes an entity-pair’s difference vector into an existence readout (Euclidean distance, directly analogous to this node’s distance term) and a type readout (cosine similarity to learned per-relation-type prototype vectors) — a genuinely new element beyond the tree-metric- probe’s single depth-from-origin readout, since “type” indexes which of several relation categories holds, not a scalar. On Llama-3.1-8B, OLMo-2-7B, and the Pythia family (70M-6.9B), performance saturates by probe rank ~32 — directly replicating this node’s own “effective rank far below ambient dimension” property in a non-tree setting — and scales monotonically with pretrained model size while randomly- initialized controls stay near chance. Causally, the learned prototype vectors are mapped back into residual-stream space via the probe’s pseudoinverse and used to steer QA-token probabilities up or down. See diego-simon-etal-2026-polar-probe-decomposes-in-context-relational-graphs-into-a-distance-readout-for-existence-and-a-per-relation-cosine-prototype-readout-for-type-saturating-by-rank-32.

Extension to programming-language syntax (whole ASTs, not just distance/depth)

Hernandez Lopez, Weyssow, Sanchez Cuadrado & Sahraoui (2022) extend the structural-probe idea from natural-language dependency trees to programming-language Abstract Syntax Trees, fitting a learned subspace that predicts a (distance, label, marker) tuple bidirectionally convertible to/from the complete labeled AST - a richer target than Hewitt & Manning’s distance/depth pair alone. Across five pre-trained code/text models (CodeBERT, GraphCodeBERT, CodeT5, RoBERTa, CodeBERTa), they estimate the syntactic subspace’s dimensionality directly at 64-128 out of each model’s 768 ambient dimensions, with AST information concentrated in middle layers - both findings paralleling Hewitt & Manning’s own rank-sweep and layer-wise results, now confirmed in an entirely different (non-natural-language) tree domain. See ast-probe-recovers-the-full-abstract-syntax-tree-of-code-from-a-64-to-128-dimensional-syntactic-subspace-in-five-code-language-models.

Mathematical justification for the squared-distance form

Coenen, Reif, Yuan, Kim, Pearce, Viégas & Wattenberg (2019) give a theoretical answer to a question Hewitt & Manning (2019) leave open — why squared Euclidean distance specifically, rather than distance itself or some other power? They prove a general isometric tree embedding is essentially impossible: even a 4-node star tree (one root, three children) cannot be isometrically embedded in any $\mathbb{R}^n$ without collapsing two children onto the same point. But a power-2 (“Pythagorean”) embedding — one satisfying $\|f(x)-f(y)\|^2=d(x,y)$ rather than $\|f(x)-f(y)\|=d(x,y)$ — always exists for any tree with $n$ nodes, into $\mathbb{R}^{n-1}$ : assign each edge to its own orthogonal unit basis vector, so that the embedded position of a node is the sum of unit steps from the root along its ancestor path; any two nodes are then connected by $d(x,y)$ mutually perpendicular unit segments, giving $\|f(x)-f(y)\|^2 = d(x,y)$ by the Pythagorean theorem exactly. They further prove no power- $p$ embedding exists for any $p<2$ for a sufficiently large star-shaped tree — so $p=2$ is not an arbitrary choice among many that would have worked equally well; it is the smallest power for which a tree embedding is possible in general. A randomized version of the same construction (i.i.d. Gaussian “branch” vectors instead of exactly orthogonal ones) gives an approximate Pythagorean embedding using only local, per-edge information, whose distance distribution has the right mean and shrinking-with-dimension variance — offering a plausible mechanism for how a trained network might approximate this construction without knowing the tree’s global shape in advance.

The structural probe extended to a broader synthesis, plus attention heads that separately track coreference clusters. Manning, Clark, Hewitt, Khandelwal & Levy (2020) apply the same structural-probe metric embedding to real, pretrained BERT (base/large), confirming parse-tree distance recovery layer-by-layer, and additionally show specific attention heads track coreference clusters — a companion synthesis paper rather than an independent new method. See manning-etal-2020-the-structural-probe-metric-embedding-recovers-syntax-tree-distances-from-real-bert-representations-and-attention-heads-separately-track-coreference-clusters.
A kernelized (RBF) structural probe significantly outperforms the linear baseline, recovering syntax-tree-distance geometry in real mBERT across six languages. White, Pimentel, Saphra & Cotterell (2021) fit a nonlinear-kernel variant of the Hewitt & Manning structural probe to real multilingual BERT representations across six typologically diverse languages, finding the RBF-kernel probe consistently and significantly outperforms the original linear transform $B$ at recovering tree distance — evidence that the tree-metric-probe’s underlying geometry is only partially linear, with a genuinely nonlinear component still capturable at the metric level. Purely a geometric-shape measurement; no causal intervention. See white-etal-2021-a-kernelized-rbf-structural-probe-significantly-outperforms-the-linear-baseline-recovering-syntax-tree-distance-geometry-in-real-mbert-across-6-languages.
A shared cross-lingual tree-metric probe works nearly as well as a per-language probe for typologically close languages, but needs a language-specific orthogonal transform for typologically distant ones. Limisiewicz & Mareček (2021) apply an orthogonal structural probe to real mBERT across nine typologically diverse languages, finding a single shared (ALLLANGS) probe loses only 0.027-0.048 correlation relative to a per-language probe for Indo-European languages but drops sharply (e.g. -0.305 for Arabic lexical depth) for non-Indo-European ones — direct evidence that mBERT’s tree-metric geometry is nearly shared across related languages but requires a language-specific orthogonal rotation to recover the same geometry for typologically distant ones. See limisiewicz-marecek-2021-a-shared-cross-lingual-syntax-tree-metric-in-real-mbert-needs-a-language-specific-orthogonal-transform-only-for-typologically-distant-languages.
The unmodified structural probe transfers to a Minimalist syntax abstraction (clausal finiteness) beyond its dependency-distance training target, with causal confirmation via activation patching. Chen & Chin (2026) apply the same frozen Hewitt & Manning structural probe (trained only on UD-EWT dependency distances) to 13 base LLMs across four families, finding a phase-count gradient (bare < infinitival < finite embedded clauses) in probe distance at each model’s canonical layer in 12/13 models, plus a 13/13 sign asymmetry on a token pair whose raw UD-dependency distance is fixed — evidence the probe’s geometry tracks a finer syntactic abstraction than dependency-tree edges alone specify. Activation patching (replacing the embedded-subject’s residual-stream representation with the infinitival-condition version) causally shifts probe distance in the predicted direction in 12/13 models, with a negative-control patching site staying near zero. See an-unmodified-structural-probes-tree-metric-distance-encodes-a-phase-count-gradient-for-clausal-finiteness-in-12-13-base-llms-and-activation-patching-causally-confirms-the-representation-in-12-13.

Exercises

Base

Why does the structural probe use squared Euclidean distance $d_B(h_u,h_v)^2$ to approximate tree distance, rather than fitting $d_B(h_u,h_v)$ (unsquared) directly? What is lost by using the squared quantity?

Solution

The squared quantity $(B(h_u-h_v))^\top(B(h_u-h_v))$ is a smooth, easy-to-optimize quadratic function of $B$ (a genuine inner-product-induced quantity), whereas the true metric would require a square root, complicating the gradient. What is lost: the squared distance no longer obeys the triangle inequality (a true metric’s defining property), so $d_B^2$ is technically only a pseudo-distance-squared. In terms of which tree structure is recovered, nothing is lost: since squaring is a monotonic transform of a non-negative quantity, the same relative ordering of distances (and hence the same minimum spanning tree) is recovered whether one fits the distance or its square, and the square root can be applied to the fitted quantity afterward.

A structural probe achieves near-perfect distance recovery (UUAS close to 100%) using a transform $B$ with $k=768$ (full ambient dimension, no rank constraint). Is this alone strong evidence for a genuine low-dimensional tree-metric-probe structure? Why or why not?

Solution

Not by itself — an unconstrained, full-rank $B$ has enough free parameters that it might be fitting a complex, possibly overfit transform that only incidentally reproduces tree distances on the training/test data, without the representation genuinely organizing itself around a compact tree-metric-probe structure. The stronger, more falsifiable claim (“the tree occupies a genuinely low-rank subspace”) requires also showing that performance is preserved (or barely degrades) as $k$ is swept down to a much smaller value — only then is the “few dimensions suffice” claim about the representation’s own structure, rather than about the probe’s flexibility, supported.

Middle

Two tree-metric-probe embeddings are fit on the same model: one recovers tree distance with Spearman correlation 0.85 using a Euclidean transform $B$ , another recovers it with correlation 0.90 using a curved (e.g. hyperbolic) metric space instead of a linear $B$ . Does the second result establish that the representation space is intrinsically curved? What additional test would be needed to distinguish “a curved probe merely fits slightly better” from “the representation is intrinsically negatively curved”?

Solution

No — outperforming a Euclidean probe on a distance-recovery metric only shows the curved metric space is a better-fitting operational tool for this specific task, since curved spaces have strictly more geometric degrees of freedom (curvature itself, plus the usual linear/rotational ones) than a flat linear transform, so some improvement is expected even if the underlying representation has no genuine negative curvature. To distinguish the two, one would need a test targeting the defining signature of negative curvature directly — e.g. measuring whether the number of representable points within a geodesic ball of radius $r$ grows exponentially in $r$ (the hyperbolic volume-growth law) rather than polynomially, as described in Hyperbolic Manifold — a correlational distance-recovery improvement alone cannot establish this.

Suppose a depth probe achieves high Spearman correlation between predicted and true tree depth, but a separate distance probe (fit independently) achieves poor UUAS on the same model and layer. Are these two results necessarily inconsistent with each other, given that both are described as readouts of “the same” embedded tree?

Solution

Not necessarily inconsistent — the two probes are fit independently (potentially with different transforms $B$ and different loss functions) and test different structural properties of the tree (depth-from-root is a coarser, 1-dimensional ordering; pairwise distance requires the finer relational structure between arbitrary node pairs). A representation could encode “how deep is this word” as an easily linearly-recoverable global property (e.g. correlated with something like sentence position or embedding norm generically) while still failing to encode the finer pairwise adjacency structure needed for accurate tree reconstruction — the two findings would then reveal a partial embedding of the tree (depth recovered, full shape not), which is itself an informative, falsifiable outcome rather than a contradiction.

Pro

The originating paper finds that performance plateaus once the transform’s rank $k$ exceeds roughly 64–128, for both BERT and ELMo, despite these models having different total hidden dimensions (768 vs. 1024). Propose one hypothesis this convergence would support, and one alternative explanation that would not implicate anything about the representation’s structure specifically.

Solution

Supporting hypothesis: if two different architectures independently converge on a similar effective rank for encoding the same external structure (English dependency syntax), this is suggestive evidence that the complexity of the tree-metric-probe task itself (not an architecture-specific quirk) sets the effective dimensionality — i.e. syntax trees over natural sentences may simply require on the order of 64-128 dimensions to embed with low distortion regardless of which model is being probed, similar to how a fixed graph family has a fixed minimum embedding dimension for a given distortion tolerance. Alternative, non-representational explanation: the plateau could instead reflect a limitation of the training procedure shared across both probes (e.g. both are optimized with similar amounts of data, similar optimizers, and similar stopping criteria), so that the apparent rank convergence reflects “how much rank this specific fitting procedure can make use of before overfitting or optimization difficulty sets in,” rather than a genuine fact about either model’s representational geometry — distinguishing the two would require varying the training procedure independently of the models being probed and checking whether the same effective-rank plateau persists.

Found in (12 observations · 14 families)

Gemma

A Minimalist Structural Probe for Phase-Count Abstraction in Transformer Language Models (2026)measured

A structural probe's tree-metric encodes a phase-count gradient for finiteness

Details

Chen & Chin apply the unmodified Hewitt-Manning structural probe (trained per layer on UD-EWT) to 13 base LLMs across four families, regressing probe distance on a Minimalist-syntax phase-count condition (bare/infinitival/finite embedded clauses) [chen-chin-2026-minimalist-phase-structure-probe] A phase-count gradient (beta_fin > beta_inf > 0) holds at each model's canonical layer in 12/13 models, the sole exception being a sign-reversed Qwen3-4B [chen-chin-2026-minimalist-phase-structure-probe] A 13/13 sign asymmetry appears on a token pair whose raw UD-dependency distance is identical across conditions, tracking a phase-internal cohesion the treebank formalism cannot represent [chen-chin-2026-minimalist-phase-structure-probe] Activation patching the embedded-subject representation from the infinitival source into the bare target raises probe distance in 12/13 models, while a wh-position negative control stays near zero [chen-chin-2026-minimalist-phase-structure-probe]

models: Gemma 3 27B · method: Structural probing (metric/distance probe), Activation patching

Llama

A Minimalist Structural Probe for Phase-Count Abstraction in Transformer Language Models (2026)measured

A structural probe's tree-metric encodes a phase-count gradient for finiteness

Details

models: Llama-3.1-8B · method: Structural probing (metric/distance probe), Activation patching

Polar probe linearly decodes semantic structures from LLMs (2026)measured

A polar probe splits in-context relational graphs into existence and type readouts

Details

Diego-Simon et al. extend the structural probe's distance/depth decomposition from parse trees to arbitrary in-context relational graphs across five domains (arithmetic ordering, spatial layouts, thematic roles, family trees, metro maps) [diego-simon-etal-2026-polar-probe-semantic-structures] Within a learned probe subspace, a pair's difference vector is decomposed into an existence readout (Euclidean distance) and a type readout (cosine similarity to learned per-relation-type prototype vectors) [diego-simon-etal-2026-polar-probe-semantic-structures] Tested on Llama3-8B, OLMo-7B (9 pretraining checkpoints), and the Pythia sweep (10M-6.8B), existence decoding peaks around 0.80 and type around 0.50-0.70 in middle layers, near chance in a random-init control and scaling monotonically with model size [diego-simon-etal-2026-polar-probe-semantic-structures] Performance saturates by probe rank ~32, and the five domains' probe subspaces are largely disjoint except a pronounced spatial-layout/variable-ordinality overlap [diego-simon-etal-2026-polar-probe-semantic-structures] Mapping each prototype back into the residual stream via the probe's pseudoinverse and adding or subtracting it causally raises or lowers the correct-token probability, though the method underperforms on non-commutative or many-to-many graphs [diego-simon-etal-2026-polar-probe-semantic-structures]

models: Llama-3-8B · method: Structural probing (metric/distance probe)

A polar coordinate system represents syntax in large language models (2024)measured

Syntactic relations are encoded as directions, not just distances

Details

- A supervised "Polar Probe" on frozen LLM activations encodes each dependency edge by both distance and relative direction, making same-type relation embeddings collinear and different-type ones orthogonal. [diego-simon-etal-2024] - This extends Hewitt-Manning's distance-only tree metric with an angular component, so a dependency tree's relations become readable as directions in activation space. [diego-simon-etal-2024] - The Polar Probe reaches AUC 95%, versus 74% for a distance-only Structural Probe, 80% for no probe, and 50% on a randomly initialized model (optimal probe dimensionality 128). [diego-simon-etal-2024] - Tested on Llama-2-7B, Mistral-7B-v0.1 and BERT-large (casing unspecified); the collinearity is imposed by the probe's supervised objective rather than discovered, and there is no causal intervention. [diego-simon-etal-2024]

models: Llama-2-7B · method: Polar probing (distance + angle), PCA

Mistral

A Minimalist Structural Probe for Phase-Count Abstraction in Transformer Language Models (2026)measured

A structural probe's tree-metric encodes a phase-count gradient for finiteness

Details

models: Mistral-7B-v0.3 · method: Structural probing (metric/distance probe), Activation patching

A polar coordinate system represents syntax in large language models (2024)measured

Syntactic relations are encoded as directions, not just distances

Details

models: Mistral-7B-v0.1 · method: Polar probing (distance + angle), PCA

Qwen

A Minimalist Structural Probe for Phase-Count Abstraction in Transformer Language Models (2026)measured

A structural probe's tree-metric encodes a phase-count gradient for finiteness

Details

models: Qwen3-14B, Qwen3-4B · method: Structural probing (metric/distance probe), Activation patching