Vector Arithmetic in Concept and Token Subspaces

TL;DR

This paper demonstrates that by transforming Llama-2-7b hidden states using attention weights from concept and token heads, one can perform semantic and surface-level arithmetic operations with significantly improved accuracy, revealing structured subspaces.

Contribution

The study introduces a method to identify and utilize semantic and surface-level subspaces in LLMs through attention-based transformations, enabling effective vector arithmetic.

Findings

Semantic transformations achieve 80% accuracy in nearest-neighbor tasks.

Raw hidden states achieve only 47% accuracy in similar tasks.

Attention-based transformations reveal coherent semantic and surface-level structures.

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".