The Structure of Relation Decoding Linear Operators in Large Language Models

AI summary

Overview

Research area: Mechanistic interpretability and knowledge representation in large language models — specifically the linear relational embeddings (LREs) introduced by Hernandez et al. (2023).

Technical level: Advanced. The paper assumes familiarity with transformer internals, affine/linear algebra, low-rank approximation, and tensor network decompositions (order-3 tensors, contractions, bond/inner dimensions).

Scope (one sentence): The paper studies how a whole collection of relation-decoding linear operators in transformer language models is organized, showing that these operators can be compressed into small order-3 tensor networks because they encode recurring coarse-grained semantic properties rather than distinct relations.

What This Paper Is About

Prior work showed that for a single relation (e.g., "capital city of"), the map from a subject's embedding to the object token can be approximated by one affine operator — a Linear Relational Embedding (LRE) matrix. This paper asks what happens when you look at many such operators together: do they encode separate, fine-grained relations, and can an entire collection be stored in a far more compact form? The authors build compression tools and a cross-evaluation protocol to answer both questions, concluding that the operators behave as shared property extractors.

Key Contributions

1. A cross-evaluation protocol as a semantic-closeness measure. Each relation's decoder is applied to every other relation's subjects, producing a k-by-k faithfulness matrix that acts as an empirical similarity kernel over relations. This reveals that decoders encode coarse-grained properties (e.g., country, gender, antonym) rather than relation-specific mappings.

2. Order-3 tensor network compression of an entire relation collection. Two architectures are proposed (SimpleOrder3Network and TriangleTensorNetwork), which take relation embeddings as input and output full linear operators, with optional additional relation embedders. Collections of 47 relation matrices are compressed from roughly 788 million parameters to under one million without significant loss in decoding accuracy.

3. A property-based reinterpretation of linear relational decoding. Across the original 47-relation dataset, a new 79-relation extended dataset, and both GPT-J, Llama-3.1-8B, and GPT-NeoX-20B, the cross-evaluation matrices show block structure tied to shared properties.

4. An analysis of generalization to held-out relations. Tensor networks generalize weakly on general language relations (only to relations semantically close enough to be interchangeable with training relations), but on a purpose-built mathematical dataset of arithmetic relations they reach an average test faithfulness of 0.96 (± 0.031).

Main Findings

• Relation decoders are highly compressible. In GPT-J, embedding dimension d = 4096, so a single LRE matrix has more than 16 million parameters; a collection of 100 relations amounts to approximately 1.6 billion parameters. Stacking the 47 relation matrices of Hernandez et al. (2023) gives roughly 788 million parameters, and a Jacobian-based baseline using 8 examples reaches a mean faithfulness of 0.41 at that size. Tensor networks with fewer than one million parameters significantly outperform these baselines in both parameter count and faithfulness.

• Linear structure is sufficient. Tensor network models without a separate relation embedder consistently outperformed those with one, and adding a feed-forward relation embedder (three dense layers with ReLU) yielded no substantial improvement in generalization.

• Cross-evaluation exposes off-diagonal structure. In the 47-relation matrix, many off-diagonal entries are larger than zero, several exceeding a faithfulness of 0.7. Blocks appear among relations sharing a property — for example, characteristic gender, university degree gender, and occupation gender show competitive faithfulness in any permutation (around 0.65), with the occupation gender decoder even outperforming the characteristic gender decoder when evaluated on characteristic gender.

• Syntactic as well as semantic overlap exists. Some cross-compatibility is syntactic rather than semantic — for instance, the first letter of a word relation proves effective for adjective superlative, likely because many superlatives begin with the same token as their base adjective.

• Block structure in the extended 79-relation dataset. Landmark in country and primary language spoken in a country cluster under a country block; adjective antonym and noun antonym cluster under an antonym block; zero inter-block scores appear where no common property exists. A separate identity-like block appears for relations without common semantic correspondence, such as semiconductor chip manufactured by company and mathematical theorem named after mathematician. The fruit inside color and fruit outside color operators reach a faithfulness of only around 0.3 but maintain similar performance during cross-evaluation.

• Sample-wise generalization is retained. On the original dataset, the tensor network outperforms or equals the majority baseline on 34 of 47 relations. On the extended dataset it outperforms the majority baseline on 49 relations and equals it on 8 of 79, with an overall mean test faithfulness of 0.42 versus a majority-guess baseline of 0.30.

• Compression goes beyond low-rank structure. The original dataset has a near-diagonal cross-evaluation matrix (little semantic overlap), yet substantial compression is still possible relative to the individual low-rank baseline. This indicates tensor networks capture regularities beyond the low-rank property of individual relation matrices.

• Held-out relation generalization is limited on language data. Models produce high-faithfulness decoders for some test relations but fail on others, generalizing only to relations whose decoder matrices are effectively interchangeable with those in the training set — i.e., the model identifies the correct coarse-grained property extractor but does not generalize more broadly.

• Arithmetic relations generalize well. On the mathematical dataset, split 75%–25% into training and test relations and run with three random seeds, training-set faithfulness averages 0.992 (± 0.012). Test-set decoders reach an average faithfulness of 0.96 (± 0.031) with a maximum of 0.991.

• Ablations confirm the role of embeddings. Randomizing relation embeddings causes a drastic drop on held-out (subject, object) pairs while training examples are still memorized perfectly. Randomizing subject and object embeddings causes failure to memorize the training set or generalize, with faithfulness close to 0 on small tensor networks.

Methodology in Plain English

The authors begin from the affine decoder formulation: a relation R is represented by a matrix W_R and bias b_R that transform a subject embedding into an object prediction. Instead of deriving these matrices from gradients (the Jacobian approach of Hernandez et al., 2023), they fit them by training a parameterized model against the cross-entropy loss for the correct object token, keeping the language model frozen.

To compress a whole set of relations at once, they stack the individual matrices into a single order-3 tensor and factor it with tensor networks — small collections of interconnected tensors whose shared "inner" (bond) dimensions act as bottlenecks. Two layouts are tested: a simple network with one central order-3 tensor and three projection matrices, and a triangular network with three interconnected order-3 tensors. Bias is folded in by enlarging a bond dimension by one and appending a constant 1 to subject vectors. Parameter counts are set by the embedding dimension and the chosen inner dimensions, which were grid-searched.

Two baselines anchor the results: individually trained low-rank LRE matrices per relation, and the Jacobian-based procedure of Hernandez et al.

To explain why compression works, the authors define cross-evaluation: apply each relation's decoder to another relation's subjects and record faithfulness, assembling a square matrix. High off-diagonal values mean two relations are functionally interchangeable; near-zero values mean they are unrelated. They then run this on an extended 79-relation dataset deliberately built to contain shared properties, and on a mathematical dataset of arithmetic operators, and finally test relation-level generalization by training on a subset of relations and evaluating decoders generated for unseen relation embeddings.

Why This Matters

Impact on research. The work shifts the interpretation of linear relational decoding from relation-specific mappings to property-based ones. That reframing connects mechanistic interpretability to prototype and family-resemblance theories of cognition, and it offers a concrete explanation of why LREs transfer only to semantically adjacent relations. It also demonstrates that approximate relation decoders — not just model weight matrices — are a viable target for collective tensor decomposition.

Real-world applications (drawn from the paper's stated implications):

• Memory- and compute-efficient storage of factual knowledge, reducing parameter counts by orders of magnitude and broadening access in resource-constrained settings.
• Debugging and refining factual knowledge in deployed systems by exposing how facts are organized into coarse-grained properties.
• Probing and potentially editing broad classes of attributes — the paper explicitly names gender, nationality, and religion as isolatable properties.
• Building compact, interpretable modules that a practitioner can inspect rather than treating factual recall as opaque weights.

Industry relevance. The compression angle is directly relevant to model deployment and fine-tuning infrastructure, including the paper's suggested extension of treating collections of LoRA matrices across layers as a single tensor network, and the connection it draws to mixture-of-experts architectures where structured, modular compression matters.

Future Directions

• Close the gap between the two generalization regimes. The authors explicitly state that reconciling strong generalization on the mathematical dataset with weak generalization on general language relations goes beyond this paper's scope, and flag it as promising future work with implications for performance, compression, and interpretability.

• Scale and diversify the evaluation. The study uses models of 6, 8, and 20 billion parameters; how the findings hold for larger, instruction-tuned, or mixture-of-experts architectures remains open, as does coverage of the far broader space of relations.

• Explore the larger design space of tensor networks. The authors note the number of possible internal structures of three-legged tensor networks is theoretically infinite and that they examine only two architectures here.

• Extend to LoRA and mixture-of-experts settings. The paper proposes representing collections of LoRA matrices — across layers or across independent adaptations — with a unified tensor network, and incorporating contextual information on a tensor leg.

Target Audience

Mechanistic interpretability researchers and NLP practitioners working on knowledge representation in transformers will get the most from this paper, particularly those already familiar with LREs and with tensor decomposition methods for neural network compression. Readers interested in tensor networks, low-rank adaptation, or model compression will find the order-3 architectures directly applicable, while cognitive-science-adjacent readers may be interested in the property-based framing of relational knowledge. Beginners without linear algebra and tensor-network background will find the methodology difficult to follow without the referenced appendices.

Authors’ abstract

This paper investigates the structure of linear operators introduced in Hernandez et al. [2023] that decode specific relational facts in transformer language models. We extend their single-relation findings to a collection of relations and systematically chart their organization. We show that such collections of relation decoders can be highly compressed by simple order-3 tensor networks without significant loss in decoding accuracy. To explain this surprising redundancy, we develop a cross-evaluation protocol, in which we apply each linear decoder operator to the subjects of every other relation. Our results reveal that these linear maps do not encode distinct relations, but extract recurring, coarse-grained semantic properties (e.g., country of capital city and country of food are both in the country-of-X property). This property-centric structure clarifies both the operators' compressibility and highlights why they generalize only to new relations that are semantically close. Our findings thus interpret linear relational decoding in transformer language models as primarily property-based, rather than relation-specific.