Vector Arithmetic in Concept and Token Subspaces

AI summary

Overview

Research area: Mechanistic interpretability / natural language processing — specifically, how large language models represent semantic versus surface-level (spelling, form) information about words inside their hidden states.

Technical level: Intermediate. The paper builds on linear-algebra machinery from prior work (attention head value/output matrices) but its central claims — better word-vector arithmetic in targeted subspaces — are explained with concrete analogy examples.

Scope in one sentence: This paper shows that re-using the weights of two previously identified types of attention heads (concept induction heads and token induction heads) as linear "lenses" lets Llama-2-7b perform word2vec-style parallelogram arithmetic far more accurately than using its raw hidden states.

What This Paper Is About

Mikolov et al.'s word2vec work claimed that word embeddings are structured so that man is to woman as king is to queen — i.e., that vector differences encode relations. The authors find this only somewhat works on raw Llama-2-7b hidden states, and hypothesize that irrelevant information in those activations creates "interference."

Their goal is to isolate subspaces of model activations that hold either semantic meaning or surface-level word form, and to test whether arithmetic becomes more reliable inside each of those subspaces.

Key Contributions

1. Reusing concept and token induction heads as transformations. The authors take the OV matrices (output-weight times value-weight products) of the top-k concept induction heads and top-k token induction heads identified by Feucht et al. (2025), and sum them into two linear maps they call the concept lens and the token lens.

2. Demonstrating that parallelogram arithmetic improves in these subspaces. In the concept subspace, Athens – Greece + China yields Beijing; in the token subspace, coding – code + dance yields dancing. The concept transformation reaches 80% nearest-neighbor accuracy versus 47% for raw hidden states.

3. Task-level characterization of which lens helps when. Across 14 word2vec tasks, the concept lens wins on semantic tasks (e.g., capital cities, family relations) while the token lens wins on surface-level transformations (e.g., adjective-to-adverb). A full-word2vec-derived benchmark set and an appendix of additional function-vector tasks from Todd et al. (2024) are reported.

4. A rank analysis showing the transformations act like low-dimensional projections. Although the assembled lenses are full-rank at k = 80, truncating them to their top-r singular components down to r = 256 does not damage accuracy — showing that effectively only a lower-dimensional subspace is needed.

Main Findings

• Raw hidden states underperform concept subspaces for semantics: The paper reports 80% nearest-neighbor accuracy when transforming hidden states with concept head weights, compared with 47% using raw hidden states directly (the Athens – Greece + China = Beijing example).

• Capital cities match the model's own in-context ability: For the capital-cities task, concept-space arithmetic is "just as good as the model's accuracy when asked to complete the task in an ICL setting with 5 shots."

• Token lens wins on form-based tasks: The top-k token head subspace is most effective for grammatical tasks involving changing a word's spelling, such as code → coding.

• Some tasks are unexpectedly harder: The approach is "less effective for tasks that seem simpler, like present participles of verbs." Errors are hard to interpret because the incorrect nearest neighbor is often one of the operands in the original expression.

• Pluralization splits by type: Pluralizing nouns ("gram8-plural") can be done in either concept space or token space (by adding 's'), but pluralizing verbs ("gram9-plural-verbs") can only be done in token space — possibly because it depends on verb agreement rather than word meaning.

• "Opposite" is expected to fail by construction: The authors note the task is cyclic — representing "opposite" would require possible – impossible = impossible – possible, which is incompatible with parallelogram arithmetic.

• Targeted subspaces usually beat using all heads: The concept and token lenses outperform the baseline L_all (the sum of all attention head OV matrices) for most tasks, with the exceptions of gram1, gram3, and gram4.

• Lenses are formally full-rank but effectively low-dimensional: At k = 80 the concept and token lenses are full-rank, yet setting all singular values below the top-r to zero maintains performance for ranks as low as r = 256; r = 4096 reproduces the Figure 1 results.

• Context prefixes matter unevenly: Without prefixes, certain tasks such as "currency" are much less accurate (e.g., "real" is not immediately recognizable out of context), while "capital-common-countries" and "gram6-nationality-adjective" are slightly better without prefixes.

• Many function-vector tasks fail: The paper attributes failures to relations that are many-to-one (e.g., "capitalize-first-letter") or too complex to be directly encoded in embedding space (e.g., "national-parks"). It also notes that "country-currency" covers 197 countries, versus 30 in the word2vec "currency" task.

Methodology in Plain English

1. Extract a word vector. Feed a sentence with a constant task-specific prefix (e.g., "She travelled to") followed by a target word, and save the last-token hidden state at a chosen layer. The prefix is identical across all words in a task so the comparison is fair.

2. Build the lenses. For a given attention head, the product of its output matrix and value matrix is a low-rank d × d matrix describing everything that head writes into the residual stream. The authors sum these products over the top-k concept induction heads to get the concept lens, and over the top-k token induction heads to get the token lens. They use k = 80, following Feucht et al. (2025).

3. Apply the lens. Multiply the word's hidden state by the lens. This is compared against three alternatives: doing nothing (raw hidden states), using the token lens, and using the sum of all attention head OV matrices as a baseline.

4. Test arithmetic by nearest neighbor. For a tuple pair such as (Athens, Greece) and (Beijing, China), compute L·Athens – L·Greece + L·China and check whether L·Beijing is the nearest neighbor among all candidate words in the task. Accuracy is reported per task, alongside random-chance and 5-shot in-context-learning baselines.

5. Probe dimensionality. To check whether the lenses genuinely act as subspace projections, the authors take rank-r approximations by zeroing out all but the top-r singular values and sweep r to see where accuracy degrades.

Why This Matters

Impact on research: The work reframes the question of whether "word2vec arithmetic works in LLMs" as a question of which subspace you do the arithmetic in. It connects attention-head-level circuit analysis (concept and token induction heads) with classical distributional-semantics claims, suggesting that the geometry of LLM activations is meaningful only relative to a specified type of information.

Real-world applications:

• Interpretability tooling: Concept and token lenses could serve as diagnostic read-out tools for inspecting what a model encodes about a given word, similar to how logit lenses and probes are used today.
• Controllable text generation: Separating meaning from surface form could support editing a model's output to change wording or spelling without disturbing its semantics.
• Model editing and debugging: Attribute-editing methods that modify facts could use the concept subspace to target semantics rather than accidentally perturbing surface-level representations.
• Evaluation and benchmarking: The task suite (word2vec tasks plus function-vector tasks from Todd et al., 2024) provides a concrete benchmark for measuring how much semantic structure a given model exposes.

Industry relevance: Anyone building retrieval, analogy search, or embedding-based recommendation on top of a specific transformer could, in principle, apply a lens derived from that model's own attention heads instead of relying on the raw residual stream — potentially improving relational queries such as "country → capital."

Future Directions

• Why do simple tasks fail? The paper flags unexpectedly poor results for tasks like present participles, where errors often land on an operand of the expression. Understanding these failure modes is an open question.
• How low can rank go? Accuracy holds at r = 256; the paper does not report how far the rank can be reduced below that before performance breaks down.
• Cross-model generalization: All experiments use Llama-2-7b. Whether concept and token lenses behave the same in other architectures or scales is not reported.
• Improving hard tasks: The authors suggest some function-vector tasks fail because relations are many-to-one or too complex for parallelogram structure; finding alternative geometric or non-linear formulations for those relations is a natural next step.

Target Audience

Mechanistic interpretability researchers and NLP practitioners with some familiarity with transformer internals will get the most from this paper. It is also useful for anyone working on word embeddings or analogy-based retrieval who wants to know whether classical vector arithmetic still holds inside modern LLMs — and under what transformation it does. Readers without background in attention heads and residual streams will need some supplementary reading, since the method depends on OV matrix machinery from Elhage et al. (2021) and Feucht et al. (2025).

Authors’ abstract

In order to predict the next token, LLMs must represent semantic and surface-level information about the current word. Previous work identified two types of attention heads that disentangle this information: (i) Concept induction heads, which copy word meanings, and (ii) Token induction heads, which copy literal token representations (Feucht et al., 2025). We show that these heads can be used to identify subspaces of model activations that exhibit coherent semantic structure in Llama-2-7b. Specifically, when we transform hidden states using the attention weights of concept heads, we are able to more accurately perform parallelogram arithmetic (Mikolov et al., 2013) on the resulting hidden states, e.g., showing that "Athens" - "Greece" + "China" = "Beijing". This transformation allows for much higher nearest-neighbor accuracy (80%) than direct use of raw hidden states (47%). Analogously, we show that token heads allow for transformations that reveal surface-level word information in hidden states, allowing for operations like "coding" - "code" + "dance" = "dancing".