Functional Subspace, where language models can use vector algebra to solve problems

Jung H. Lee; Sujith Vijayan

cs.CLFeb 2, 2026

Functional Subspace, where language models can use vector algebra to solve problems

Authors: Jung H. Lee, Sujith Vijayan

Organizations: Pacific Northwest National Laboratory Seattle, WA · School of Neuroscience Virginia Tech Blacksburg, VA

Abstract

Large language models (LLMs) were invented for natural language tasks such as translation, but they have proved that they can perform highly complex functions across domains. Additionally, they have been thought to develop new skills without being trained on them. These learning capabilities lead to LLMs adoption in a wide range of domains. Thus, it is imperative that we understand their operating mechanisms and limitations for proper diagnostics and repair. The earlier studies proposed that high level concepts are encoded as linear directions in LLMs activation space and that the geometry of embeddings have semantic meanings. Inspired by these studies, we hypothesize that LLMs may use subspaces and vector algebra in subspaces to perform tasks. To address this hypothesis, we analyze LLMs' functional modules and residual streams collected from LLMs engaging in in-context learning (ICL), one of the emergent abilities. Our analyses suggest that 1) LLMs can create subspaces, where evidence can be accumulated and 2) ICL tasks can be solved via simple algebraic operations in subspaces.

Figures & tables

Figure 1: ICL prompts. (A), The structure of the used ICL prompts. (B), The structure of a single example ICL with uniform attention score

A_{ij}

. This panel illustrates how symbols

q_{i}

s_{i}

a_{i}

are mapped onto ICL tokens.

Figure 2: Explained variance accounted for by

30

principal components.

x

-axis denotes principal components, and

y

-axis denotes the cumulative explained variance. (A), Explained variance evaluated from LLMs engaging the antonym task. In the panel, 6 models are clarified using different colors (see the inset). (B)-(F), the same as (A), but the task is synonym, country-capital and English-French, country-capital and person-sport, respectively.

Model	ANT	SYN	CC	EF	PC	PS
Neox	(0.97, 0.02)	(0.97, 0.01)	(0.98, 0.01)	(0.97, 0.01)	(0.96, 0.01)	(0.99, 0.02)
Olmo	(0.86, 0.01)	(0.8, 0.01)	(0.89, 0.01)	(0.9, 0.01)	(0.85, 0.02)	(0.95, 0.02)
Pyt	(0.96, 0.01)	(0.97, 0.01)	(0.99, 0.02)	(0.97, 0.01)	(0.98, 0.01)	(0.99, 0.01)
GPT	(0.97, 0.01)	(0.95, 0.01)	(0.98, 0.01)	(0.96, 0.01)	(0.96, 0.01)	(0.97, 0.01)
Llma	(0.94, 0.01)	(0.81, 0.01)	(0.96, 0.01)	(0.88, 0.01)	(0.85, 0.01)	(0.97, 0.02)
Gemma	(0.91, 0.01)	(0.91, 0.01)	(0.92, 0.01)	(0.65, 0.01)	(0.89, 0.01)	(0.94, 0.02)

Table 1: Comparison of

R^{2}

with the baselines. We compares the mean values (

V_{1}

) of top-3 from our original experiments and those (

V_{2}

) from the control experiments. For each model and task, we display (

V_{1}

V_{2}

). That is, the first number in the parenthesis denotes the mean value of top 3

R^{2}

from the original experiment, whereas the second number denotes the mean value of top 3

R^{2}

from the control experiment. ANT, SYN, CC, EF, PC, PS stand for antonym, synonym, country-capital, English-French, Product-Campany and Person-Sport, respectively. Neox, Olmo, Py, GPT, Llma and Gamma denote GPT-neox-20B, OLMo-2-

0325

-32B, pythia-12B Llma-3.1-8B and gemma-3-27b-it, respectively.

Figure 3: Quality of linear regression evaluated. $R^{2}$ is evaluated using LLMs engaging in ‘antonym’ task. $x$ -axis denotes the transformer layer, and $y$ -axis denotes the principal components. (A)-(F) shows $R^{2}$ estimated from GPT-j-6B, Meta-Llama-3.1-8B, OLMo-2- $0325$ -32B, Phythia-12B, gemma-3-27b-it and GPT-NEOX- $20$ B.

Figure 4: The same as Fig. 3 , but the task is country-capital.

Figure 5: Query (que), separator (sep) and answer (ans) tokens in subspace spanned by 3 principal components with the highest $R^{2}$ s. Red, green and blue dots represent separator, answer and query tokens, respectively (see the inset). All tokens are collected from $6$ LLMs engaging in the antonym task. The model is specified by the name above each plot.

Model	PS	CC	PC
pythia-12B	(1.0, 0.1)	(1.0, 0.25)	(1.0, 0.16)
gpt-neox-20b	(1.0, 0.09)	(0.96, 0.08)	(1.0, 0.18)
gpt-j-6B	(0.98, 0.12)	(0.99, 0.01)	(1.0, 0.08)
gemma-3-27b-it	(0.7, 0.02)	(0.72, 0.0)	(0.09, 0.06)

Table 2: This table shows how often LLMs change their answers when the last token is perturbed via top-3 and bottom-3 components. For each case, the first number in the parenthesis shows the ratio of the trials, in which LLMs change their answers when the perturbation is delivered via top-3 components. The second number shows the ratio of the trials, in which LLMs change their answers when perturbation is delivered via bottom-3 components. PS, CC and PC stand for Person-Sport, Country-Capital and Product-Company, respectively.

Figure 6: The same as Fig. 5 , but the task is country-capital.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Supplementary Figure 1: Comparison of query

q^{l}

in Panel (A), separator

s^{l}

in Panel (B), answer

a^{l}

in Panel (C), and

\Delta

in Panel (D) when LLMs make the correct predictions and when they make incorrect ones. In this experiment, we use

t

-test and report

p

-values. The blue, orange and green lines denote country-capital task, product-company task and person-sport task, respectively.

Supplementary Figure 2: The same as Fig. 3 , but the task is the English-French.

Supplementary Figure 3: The same as Fig. 3 , but the task is synonym.

Supplementary Figure 4: The same as Fig. 3 , but the task is Product-Company.

Supplementary Figure 5: The same as Fig. 3 , but the task is Person-Sport.

Supplementary Figure 6: The same as Fig. 5 , but the task is English-French.

Supplementary Figure 7: The same as Fig. 5 , but the task is synonym.

Supplementary Figure 8: The same as Fig. Fig. 5 , but the task is product-company.

Supplementary Figure 9: The same as Fig. Fig. 5 , but the task is person-sport.

Model	ANT	SYN	CC	EF	PC	PS
Neox	(0.97, 0.04)	(0.97, 0.04)	(0.98, 0.05)	(0.97, 0.04)	(0.96, 0.08)	(0.99, 0.06)
Olmo	(0.86, 0.05)	(0.8, 0.03)	(0.89, 0.02)	(0.9, 0.08)	(0.85, 0.02)	(0.95, 0.02)
Pyt	(0.96, 0.03)	(0.97, 0.04)	(0.99, 0.03)	(0.97, 0.02)	(0.98, 0.09)	(0.99, 0.02)
GPT	(0.97, 0.06)	(0.95, 0.04)	(0.98, 0.03)	(0.96, 0.03)	(0.96, 0.09)	(0.97, 0.05)
Llama	(0.94, 0.06)	(0.81, 0.04)	(0.96, 0.02)	(0.88, 0.03)	(0.85, 0.05)	(0.97, 0.03)
Genna	(0.91, 0.03)	(0.91, 0.03)	(0.92, 0.02)	(0.65, 0.02)	(0.89, 0.04)	(0.94, 0.02)

Appendix

Supplementary Table 1: We compare the mean values (

V_{1}

) of top-3

R^{2}

from our original experiments and those (

V_{2}

) from the control experiments. For each model and task, we display (

V_{1}

V_{2}

). That is, the first number in the parenthesis denotes the mean value of top-3 from the original experiments, whereas the second number denotes those from the control experiments. ANT, SYN, CC, EF, PC, PS stand for antonym, synonym, country-capital, English-French, Product-Campany and Person-Sport, respectively.Neox, Olmo, Pyt, GPT, Llama and Gamma denote GPT-neox-20B, OLMo-2-

0325

-32B, pythia-12B Llma-3.1-8B and gemma-3-27b-it, respectively.

Explore similar work

Large Language Models Develop Belief State Geometry In-Context
Sep 15, 2026Daniel Balcells, Andrew Jun Lee, Chirag Rastogi +3In-Context Learning In-Context
Independent · UCLA · UIUC +2
Large language models reorganize representational geometry during in-context learning
May 16, 2026Hua-Dong Xiong, Li Ji-An, Robert C. Wilson +2In-Context Learning Representational Capacity
School of Psychological and Brain Sciences, Georgia Tech · Department of Psychology, New York University · School of Psychological and Brain Sciences, Georgia Tech 3 Center of Excellence for Computational Cognition, Georgia Tech +2
Fast Multi-dimensional Refusal Subspaces via RFM-AGOP
Jul 2, 2026Thomas WinningerLLM Reasoning Strategies Subspace
Télécom SudParis, Évry-Courcouronnes, France · ENS Paris-Saclay, Gif-sur-Yvette, France

Sep 15, 2026, cs.LGJ/K move · Enter open · S save

Large Language Models Develop Belief State Geometry In-Context

Daniel Balcells, Andrew Jun Lee, Chirag Rastogi, Paul M. Riechers+2

Independent · UCLA · UIUC+2

In-Context Learning In-Context

Open paper AskPDF arXiv

Large language models (LLMs) trained on next-token prediction exhibit remarkable in-context learning (ICL) abilities, yet the representations that support ICL remain poorly understood. We consider such representations in a controlled setting: prompting LLMs with data emitted from hidden Markov models (HMMs) and probing for the corresponding belief state -- the posterior distribution over the HMM's hidden states given the observed token history. Across six open-source LLMs prompted with data from 40 HMMs selected for non-trivial belief structure, we find that belief states are linearly decodable from residual stream activations, with peak probe

R^2

-values from 0.83-0.99 across HMM and LLM combinations, ranging from early to late layers. To establish functional relevance, we intervene directly on the probe-identified subspace via patching and steering, resulting in downstream prediction quality on the order of the untampered model, while controls degrade performance substantially. Together, these results provide representation-level evidence that ICL in open-source LLMs approximates optimal Bayesian prediction over a context-inferred generative model. More broadly, our findings extend prior results linking input-distribution structure to activation geometry: from toy networks trained explicitly on HMM data to production-scale LLMs.

Abstract

Explore similar work

Large Language Models Develop Belief State Geometry In-Context

Large language models reorganize representational geometry during in-context learning

Fast Multi-dimensional Refusal Subspaces via RFM-AGOP