Latent Concept Disentanglement in Transformer-based Language Models

## Abstract

Abstract When large language models (LLMs) use in-context learning (ICL) to solve a new task, they must infer latent concepts from demonstration examples.
This raises the question of whether and how transformers represent latent structures as part of their computation.
Our work experiments with several controlled tasks, studying this question using mechanistic interpretability. First, we show that in transitive reasoning tasks with a latent, discrete concept, the model successfully identifies the latent concept and does step-by-step concept composition. This builds upon prior work that analyzes single-step reasoning. Then, we consider tasks parameterized by a latent numerical concept. We discover low-dimensional subspaces in the model’s representation space, where the geometry cleanly reflects the underlying parameterization.
Overall, we show that small and large models can indeed disentangle and utilize latent concepts that they learn in-context from a handful of abbreviated demonstrations.

## 1 Introduction

Transformer-based Large Language Models (LLMs) demonstrate remarkable in-context learning (ICL) abilities: with only a handful of source-target demonstrations, they can generalize to new queries without any parameter updates (Brown et al., 2020). These successes hint that the models might be inferring latent rules or concepts implicit in the prompt. Understanding this goes beyond studying a specific question about ICL; it is a key step to deciphering how the attention mechanism encapsulates the influence of prior tokens on posterior tokens, through latent intermediate representations.

Our work takes a systematic view of core questions on how transformers process and use latent concepts to perform ICL.
By latent, we refer to unstated variables or rules that are necessary for the computation.
We also want to understand how the model represents the implicit functions, going beyond just measuring the accuracy of different tasks. Given the breadth of domains where ICL appears to work, we are faced with a large canvas of experiment design: Is the model inferring an elementary function or latently performing chained reasoning? Are the hidden arguments instance-specific, or abstract and reusable?
How is ICL performed when “memorized world knowledge” is entailed, versus when the task relies on more elementary logic?

To explore these diverse dimensions, we design and implement two sets of experiments.
The first addresses the question of whether and how transformers resolve hidden, intermediate entities in ICL when “memorized world knowledge” is entailed. An example is mapping an arbitrary city in a country to the capital city of this country. The second studies the structure of representations when the model performs numerical or geometrical computations. For example, the model may need to output the next point in a traversal of a circle or a rectangle. While these tasks are fairly simple, we believe that studying them sheds new light on more fundamental questions, including whether the model is taking shortcuts or performing abstract reasoning in its hidden activations.

For tasks involving world knowledge, we use pre-trained Gemma-2 models. We apply causal and correlation techniques to study how the hidden activations map to certain key parts of the solution process.

- •
For the largest 27B model, we find that it relies on step-by-step composition of latent concept representations to obtain the result. In particular, we discover that a sparse set of attention heads is responsible for resolving the intermediate latent concept (e.g., the country when going from city to capital, or the company when going from product to headquarter city); the latent concepts exhibit orthogonality in the embedding space. Then, we show that there is a set of heads and MLPs deeper in the model responsible for realizing the concrete output (such as the capital) from the intermediate concept.
- •
In contrast to the large 27B model, we find that the smaller 2B variant contains a much weaker and noisier version of the 27B model’s circuit. This corroborates the general wisdom that model size significantly impacts latent-concept disentanglement and composition abilities in the LLMs.
- •
We also show that, as expected, adding more in-context examples leads to higher accuracy. We attribute this to a strengthening of the model circuit’s causal importance and an increase in concept representation disentanglement. In other words, the model more fully utilizes its relevant sub-circuits when we add more ICL examples.

Following these tasks involving world knowledge, we study quantitative self-contained tasks. This allows us to perform more fine-grained experiments on two-layer models trained from scratch.
The ICL tasks here are single-step "arithmetic": add-$k$,
Circular-Trajectory, and Rectangular-Trajectory.

- •
Recent work has identified linear task vectors for problems such as basic arithmetic, single-step reasoning, and linguistic mappings (Todd et al., 2024, Hendel et al., 2023). Our key finding is that models not only have task vectors, but these representations reflect the geometry of the latent variable. For example, the task vectors for add-$k$ almost entirely project onto a line. More surprisingly, we can intervene on which function the model computes: interpolating along the task vector line approximately interpolates on the latent parameter $k$ for the task.
- •
For the Circular-Trajectory and Rectangular-Trajectory tasks, we see a similar geometry in a 2-dimensional space. This aligns with the linear representation hypothesis, but provides more nuanced evidence for it by showing a more continuous parameterization along the representation direction. As with the two-hop task, we provide evidence through both correlational analysis and localization of task vectors in the model’s representation.

We designed these settings to be clean enough for controlled experimentation, yet hint at some more general phenomena that deserve further study. For example, models can encapsulate the latent structures of the tasks they learn. Moreover, these structures may be localized and interpretable in the model. We posit that even for much more complex tasks, we will be able to find that the latent concepts are captured by a sparse set of attention heads or a relatively low-dimensional encoding.

### 1.1 Preliminaries

To focus on analyzing whether and how transformers utilize latent concept representations for solving ICL, we work with prompts that contain *demonstration-only* specifications of latent functions. Intuitively, we have a source $x$ and target $y$, along with a hidden function $F$. As input, the model only sees a handful of pairs $(x_{i},y_{i})_{i=1}^{n}$ where $y_{i}=F(x_{i})$. Then, on a new $x^{\prime}$ the model has to compute $F(x^{\prime})$ to obtain a correct answer. The subtle part is that $F=R\circ C$, where $C$ maps input $x$ to a low-dimensional “concept space”, which is then refined by $R$ to produce the output.

The concept map is marked by its re-usability over different instances of sources and targets: for example, a function that maps cities and landmarks to Country representations, or one that maps points on a circle to the same radius representation. Such *abstract* representations enable lower-complexity inference. To test for the existence and utilization of such maps in transformers performing ICL, we consider several varieties of demonstration-only prompts, under two categories. The first category focuses on concept maps relying on world knowledge, with Countries and Companies serving as the hidden concepts connecting the source to the target (Fig. [1](https://arxiv.org/html/2506.16975v2#S1.F1) top half). The second category focuses purely on self-contained numerical puzzles, with Radius, Offset and other numerical values serving as the hidden intermediate concepts (Fig. [1](https://arxiv.org/html/2506.16975v2#S1.F1) bottom half).

More precisely, we consider the following ICL settings:

- •
Factuality-based 2-hop Reasoning. This setting involves demonstration examples with under-specified reasoning steps and discrete latent concepts. Specifically, we consider $2$-hop factual recall tasks where the model must map a “source” entity to a “target” entity (e.g., (non-capital) city$\to$capital) by first latently inferring the hidden “bridge” entity (e.g., country), allowing us to ask whether the LLM first resolves the “bridge”, then refines it to a “target” entity via (causal) concept compositions. An instance looks like “Toronto, Ottawa. Mumbai, New Delhi. Shanghai, ”, where the answer is Beijing. The motivation here is to understand what the model does after reading the input city: does it jump directly to the capital, or does it first invoke the country as an intermediate reasoning step? While both are plausible strategies, only the latter captures the latent causal structure that goes through the latent variable (the country).
- •
Numerical Tasks with Latent Parameters.
We generate demonstrations based on quantitative parameters. We will see that this determines task similarity and induces smooth geometric relationships in the task space, suggesting that the model encodes such parameters along a low-dimensional geometry. We first consider the numerical *add-$k$* task (also studied by (Hu et al., 2025)) where the demonstration examples are $(x_{i},y_{i})_{i=1}^{n}$ with $y_{i}=x_{i}+k$, where $k$ is the latent parameter.
An example of add-$4$ is “5, 9. 3, 7. 1, 5. 2, ”, where the answer is 6. We then study geometric tasks, such as Circular-Trajectory where points lie on a circle of varying radii and a related Rectangular-Trajectory task. Our goal is to localize, and more importantly, understand the geometry of the task vectors: does it reflect the geometry of the latent parameter?

Both sets of tasks are demanding enough to require intricate latent concept disentanglement and manipulation, yet sufficiently controlled to permit causal, feature and circuit-level analysis.

Figure: Figure 1: An illustration of our main findings. We primarily focus on how decoder-only transformer-based language models disentangle and manipulate latent concepts for solving in-context learning (ICL) problems. In the discrete multi-hop setting, we discover that transformers compose latent concept representations for predicting the answer. For example, as shown in the upper half of the figure, by intervening on certain “bridge-concept” attention heads, we can push the model’s “belief” (reflected in logit and rank) in the original answer Canberra (the capital of Australia, which the city Sydney belongs to) to the “type-corrected” alternative answer Ottawa (the capital of Canada, which the landmark Niagara Falls belongs to). In the continuous-parameterization ICL setting, we discover that transformers’ hidden embeddings capture the geometry of the latent concepts for our prediction tasks. For instance, for a transformer trained to predict circular trajectories whose radius is randomly chosen from a continuous interval, not only do we obtain causal evidence for task vectors (TVs) which control the trajectory’s radius, but they also fall on a smooth 2D manifold.
Refer to caption: x1.png

### 1.2 Related Work

In-context Learning Interpretation. ICL abilities of transformer-based models were first observed by Brown et al. (2020), which sparked work in analyzing this ability. This includes analyzing how pretrained LLMs solve ICL tasks requiring abilities such as copying, single-step reasoning, basic linguistics (Olsson et al., 2022, Min et al., 2022, Zhou et al., 2023, Hendel et al., 2023, Todd et al., 2024, Yin & Steinhardt, 2025a), and smaller models trained on synthetic tasks like regression
(Garg et al., 2022, Von Oswald et al., 2023, Akyürek et al., 2023, Bai et al., 2023, Guo et al., 2025), discrete tasks (Bhattamishra et al., 2024), and mixture of Markov chains (Edelman et al., 2024, Rajaraman et al., 2024, Park et al., 2025a). These setups enable discovery of relations between in-context and in-weight learning (Lin & Lee, 2024, Singh et al., 2025, Russin et al., 2025), and internal algorithms that models implement (Olsson et al., 2022, Edelman et al., 2024, Li et al., 2023, Park et al., 2025a, Yin & Steinhardt, 2025a). We contribute to this line of work, by shedding light on how transformers solve ICL problems which have more intricate latent structures.

Linear Representation Hypothesis (LRH). Our results are also connected to the LRH, which essentially speculates that LLMs represent high-level concepts in (almost) linear latent directions (Park et al., 2024, Merullo et al., 2024, Park et al., 2025b, Huh et al., 2024, Ilharco et al., 2023, Li et al., 2025a, Dumas et al., 2025, Beaglehole et al., 2025). Many papers motivated by the LRH then find “concept” vectors that can capture directions of truthfulness (Marks & Tegmark, 2024, Arditi et al., 2024), sentiment (Tigges et al., 2024), humor (von Rütte et al., 2024), toxicity (Turner et al., 2025), etc. We deepen this study, asking how LLMs’ representations capture/disentangle latent concepts, and compose them during inference. In addition, the LRH is rooted in the field of mechanistic interpretability, which aims to reverse engineer mechanisms in transformer-based LMs (Elhage et al., 2021, Olsson et al., 2022, Singh et al., 2025, Wu et al., 2023, Wang et al., 2023, Hong et al., 2024, Brinkmann et al., 2024, Bakalova et al., 2025, Heindrich et al., 2025, AlKhamissi et al., 2025b, Vig et al., 2020, Baek & Tegmark, 2025, Wang & Xu, 2025).

Task and Function Vectors. A specific line of work in analyzing ICL mechanisms focus on task or function vectors. They show that there exist certain causal patterns which capture the input-output relationship of the ICL task, on relatively simple problems such as “Country to Capital”, “Antonyms”, “Capitalize a Word” (Todd et al., 2024, Hendel et al., 2023, Davidson et al., 2025, Yin & Steinhardt, 2025b). Similarly, Liu et al. (2024), Merullo et al. (2024), Li et al. (2025b), AlKhamissi et al. (2025a), Venhoff et al. (2025), Soligo et al. (2025) observed that LLMs compress certain task or context information into sparse sets of vectors. We work with ICL problems with more complex latent structures, and our focus is not solely on (high-level) task vectors, but on how the model disentangle and manipulate latent concepts useful to answering the query. In addition, our work complements the function vector analysis from contemporaneous work of Hu et al. (2025), providing add-$k$ results for smaller models where we have full control over training and can hence conclude that the geometry of the task vector only arises from the latent task structure. We also compare results from add-$k$ with other ICL tasks, giving additional insights.

## 2 Disentanglement of Latent Concepts in 2-hop Reasoning

We start with a 2-hop task based on connecting two facts through a latent entity, where the model must infer the relationships from in-context demonstrations. This task builds on prior work that analyzes problem with a single step of reasoning (1-hop) over world knowledge, such as geography puzzles “Country$\to$Capital” or “National Park$\to$Country” (Minegishi et al., 2025, Todd et al., 2024, Yin & Steinhardt, 2025a).
Extending to two steps enables us to explore how LLMs solve ICL problems with under-specified reasoning steps while requiring latent multi-hop reasoning.

Our goal is to mechanistically analyze how a pre-trained LLM solves a “source$\to$target” problem when the “bridge” concept is latent. The model needs to understand the nature of the bridge from the demonstrations, which only contain the sources and targets. Interestingly, models achieve reasonably high accuracy with sufficient in-context demonstrations. This leads to two competing hypotheses:

- •
Hypothesis 1 (Shortcut): The LLM maps the queried input directly to the answer (going from source to target), without internally computing the bridge entity in a discernible manner.
- •
Hypothesis 2 (Latent two-hop): The LLM first resolves the latent bridge concept (e.g., “country”) in its hidden representations, then composes it with the output concept to obtain the answer.

We first present the problem definition and experimental set-up. Then, we present our main causal and correlation evidence favoring Hypothesis 2.

The “Source$\to$Target” Problem. We create two-hop ICL puzzles by composing two facts linked by a common “bridge” entity. That is, we sample fact tuples {($S_{i}$, $r_{1}$, $B_{i}$, $r_{2}$, $T_{i}$)}${}_{i=1}^{n}$, where the *source* entity $S_{i}$ is related to the bridge entity $B_{i}$ via relation $r_{1}$, and $B_{i}$ is related to the *target* entity $T_{i}$ via $r_{2}$. We then create the ICL puzzle in the form [$S_{1},T_{1}.\,S_{2},T_{2}\dots S_{n},$] [Answer: $T_{n}$](^1^11We think of this as a systematic ICL version of TwoHopFact (Yang et al., 2024). Here, the model must figure out relations between the source-bridge and bridge-target facts from the ICL examples.). Note that the bridge entities $B_{i}$’s are never specified in the prompt. An example City$\rightarrow$ Capital problem is

Sydney, Canberra. Nantes, Paris. Oshawa,

Here, $r_{1}$ is “belongs to the country of”, the (unspecified) bridge entities for this example are “Australia”, “France”, “Canada”, and $r_{2}$ is “has capital”. Therefore, the prompt’s answer is Ottawa, the capital of Canada, the country Oshawa is in.

Figure: Figure 2: Accuracy of Gemma-2-27B on the two-hop “Source$\to$Target” ICL problems.
Refer to caption: figs/fig2nn.png

In the main text, we work with geography puzzles, with source types {City, University, Landmark}, and output types {Capital, Calling code}. In addition to this Country dataset, we work with the Company dataset in App. [A.4](https://arxiv.org/html/2506.16975v2#A1.SS4) for generality(^2^22For our Country & Company data, we collect source and target entity data for 40 countries, 36 companies.). See App. [A.2](https://arxiv.org/html/2506.16975v2#A1.SS2) for details on the data generation process.

We focus on Gemma-2-27B (Gemma Team, 2024) in the main text.
To establish a baseline, we evaluate Gemma-2-27B on Source$\to$Target ICL puzzles. Fig. [2](https://arxiv.org/html/2506.16975v2#S2.F2) reports its accuracy: although the model finds these tasks more challenging than one-hop counterparts—reflected by the need for more in-context examples—it achieves high accuracy at 20 shots.

We provide new causal and correlational evidence suggesting that the model indeed performs latent multi-hop reasoning via sequential concept composition: it first infers an abstract bridge concept representation (e.g. a “Canada” concept), and then specializes to a specific output type (e.g. the “capital” of “Canada”) deeper in the model. Surprisingly, we find that the bridge-resolving mechanism is enabled by a sparse set of attention heads in our set of problem settings.

Methodology. We use causal mediation analysis (CMA), also known as activation patching (Pearl, 2022, Zhang & Nanda, 2024) to obtain causal evidence for our claims (see App. [A.1](https://arxiv.org/html/2506.16975v2#A1.SS1) for details). We discuss the bridge-resolving mechanism in the main text, and delay the analysis of the output-concept component to App. [A.2](https://arxiv.org/html/2506.16975v2#A1.SS2). For concreteness, consider two prompts with different source-target types:

- •
A normal prompt, with type [City$\rightarrow$Capital]: “Okinawa, Tokyo. Sydney, Canberra. Chicago, ” [Answer: Washington; Bridge: USA].
- •
An alternative prompt, with type [Landmark$\rightarrow$Calling Code]: “Chapel bridge, 41. The Grand Canyon, 1. The Great Wall, ” [Answer: 86; Bridge: China].

We perform activation patching on normal and alternative problem pairs with different bridge entities, across source-target types, at the final token position.
Our central experimental assumption is that if a component computes an abstract representation of the bridge, this should transfer across different source and target types.

This should hold when running the model on the normal prompt, and replacing a selected component’s activation by the corresponding activation obtained on the alternative prompt should make the model favor the alternative prompt’s bridge, without bending the output type. For our example, the alternative-to-normal “patching” pushes the model’s answer on the normal prompt from Washington to Beijing (the output type remains Capital, but the bridge switches from USA to China). We refer to this as the “type-corrected” version of the alternative answer. This evaluation is slightly unorthodox: rather than judging an intervention by whether it reproduces the literal alternative answer (86), we assess movement toward the type-corrected alternative (Beijing). By contrast, under the Shortcut hypothesis, patching activations would at worst break the model and yield nonsensical outputs, or at best, push it to emit the literal alternative answer (86).

We now formalize the intervention with CMA, which enables us to localize components inside the LLM that specialize in handling the different (latent) reasoning steps of the problem.

Causal Mediation Analysis. Formally, let the normal and alternative prompts be denoted as ${\bm{p}}_{\text{norm}}=[S^{(\text{norm})}_{1},T^{(\text{norm})}_{1}...\,S^{(\text{norm})}_{n},]$ and ${\bm{p}}_{\text{alt}}=[S^{(\text{alt})}_{1},T^{(\text{alt})}_{1}...\,S^{(\text{alt})}_{n},]$. Let $\hat{T}^{(\text{alt})}_{n}=\text{Type}_{\text{norm}}(T^{(\text{alt})}_{n})$ denote the type-corrected output. We run the LLM on both the prompts and cache the LLM’s hidden activations at the last token position, denoted as $({\bm{h}}_{\text{norm}},{\bm{h}}_{\text{alt}})$. To perform CMA on a model component of interest, say an attention head with activations indexed by $({\bm{a}}_{\ell,h}^{(\text{norm})},{\bm{a}}_{\ell,h}^{(\text{alt})})$, we run the LLM on the normal prompt again, but this time, intervene by replacing its normal activation with the alternative ${\bm{a}}_{\ell,h}^{(\text{alt})}$, and let the remainder of the forward pass execute. We then contrast the model’s output behavior on the normal and patched runs by, for example, examining the logit differences between the normal and (type-corrected) alternative answer, and their ranking pre- and post-intervention. Further details on CMA and its suitability to our problem setting are deferred to App. [A.1](https://arxiv.org/html/2506.16975v2#A1.SS1).

Figure: Figure 3: We present causal evidence for the bridge-resolving mechanism in (a) and (b), and correlational evidence in (c) for Gemma-2-27B. In (a) and (b), we run activation patching experiments across prompts with problems of different source and target types, at the final token position. (a) shows the percentage logit difference variation of attention heads with the strongest causal influence, on select intervention experiments. Head group (24,30;31)’s representation exhibits strong transferability across source and target types. (b) further examines the causal role of (24,30;31), in the patching experiment [University, Code]$\to$[City, Capital]. Note that the reciprocal rank of $10^{-1}$ for the type-corrected alternative answer after patching (24,30;31) is in the 27th percentile ($>73\%$ in top 10 after intervention). In (c), we show a cosine similarity plot of bridge-representation disentanglement for head group (24,30;31), computed on a collection of samples with several combinations of source and target types, and with “Italy” and “Spain” as the bridge values in the prompts.
Refer to caption: x2.png

### 2.1 Results and Analysis for the Bridge-resolving Mechanism

Causal evidence. Fig. [3](https://arxiv.org/html/2506.16975v2#S2.F3) provides evidence favoring the Latent Two-hop Hypothesis rather than the Shortcut Hypothesis: we observe causal transferability of the bridge representation. In Fig. [3](https://arxiv.org/html/2506.16975v2#S2.F3)(a), we report results on a select set of patching experiments: on both tasks with and without overlap in source and target types, we observe that a sparse set of attention heads consistently exhibits very strong causal effects in pushing the model’s “belief” from $T^{(\text{norm})}_{n}$ towards $\hat{T}^{(\text{alt})}_{n}$ (reflected in logit difference); the head group (24,30;31) is especially dominant.(^3^33In our CMA experiments, we account for grouped-query attention by patching heads in groups of 2 on Gemma-2-27B. We noticed that this tends to produce stronger causal effects than with individual heads.) To further understand whether intervening on (24,30;31) during the normal run really boosts the type-corrected alternative answer $\hat{T}^{(\text{alt})}_{n}$ (instead of only decreasing model’s confidence on the normal answer $T^{(\text{norm})}_{n}$, which logit difference might not tell), in Fig. [3](https://arxiv.org/html/2506.16975v2#S2.F3)(b), we show an example patching experiment result of [University, Code]$\to$[City, Capital]. Surprisingly, at least 73% of the time, patching this head group boosts the model’s rank of the alternative prompt’s answer into top 10 (and directly become the top-1 answer more than 40% of the time!), when its original, intervention-free rank is typically in the hundreds to thousands. In consideration of limited space, we visualize the full set of results in App. [A.2](https://arxiv.org/html/2506.16975v2#A1.SS2).

Correlational evidence. To understand the nature of (24,30;31)’s output embeddings better, in Fig. [3](https://arxiv.org/html/2506.16975v2#S2.F3)(c), we visualize an example cosine similarity matrix of this attention head, with either “Italy” or “Spain” as the bridge values for $S_{n}$ (the query source entity) in the prompts, across a total of 12 different combinations of bridge, source and target types. Specifically, for each combination of the bridge and source-target type shown in the grid, we sample 10 prompts which obey such requirement,(^4^44We only specify the bridge entity for $S_{n}$ in the prompt; the bridge for $S_{i}$ for all $i<n$ are randomly chosen.) giving us a total of 120 prompts. We then obtain head group (24,30;31)’s embedding of these prompts at the last token position, and compute the pairwise cosine similarities. Observe that the embedding consistently exhibits strong disentanglement with respect to the bridge value in the prompt, regardless of source and target types. We provide detailed statistics of disentanglement strength in App. [A.3](https://arxiv.org/html/2506.16975v2#A1.SS3) and Fig. [27](https://arxiv.org/html/2506.16975v2#A1.F27), where we discuss its relation with the number of ICL examples.

Additional Insights. We highlight some additional mechanistic insights below (see Appendix for details).
First, the Gemma-2-2B model has a weak and noisy version of the 27B’s bridge resolving mechanism. This suggests that increasing model size likely benefits latent concept disentanglement and utilization in the LLM. We provide details in App. [A.2](https://arxiv.org/html/2506.16975v2#A1.SS2), and Fig. [26](https://arxiv.org/html/2506.16975v2#A1.F26). Second, in App. [A.3](https://arxiv.org/html/2506.16975v2#A1.SS3), we show that disentanglement strengthens with more ICL examples, which is reflected in both the causal importance of the bridge-resolving components and the angular (dis-)similarities of the bridge embeddings. Lastly, to complement our study on the geography dataset, we experiment with a similar, albeit smaller-scaled Company dataset in App. [A.4](https://arxiv.org/html/2506.16975v2#A1.SS4), where company name is the bridge entity. We again find causal and correlational evidence for the presence of a bridge-resolving mechanism in the model, and observe overlap in the attention heads driving this mechanism on the two datasets.

## 3 Disentanglement for Numerical Latent Variables

In this section, we consider two problems with numerical or continuous parameterization. For these experiments we study a very small transformer, with a similar architecture to GPT-2 (Radford et al., 2019).
We use a $2$-layer $1$-head transformer, with embedding dimension $128$, trained with AdamW.

Figure: Figure 4: Illustration of an input sequence for the circle trajectory problem. Here, radius $r\!=\!3$, period $p\!=\!2$, sequence length $n\!=\!13$. Every $p$ consecutive steps on the trajectory are equal. We first sample $\lfloor\tfrac{n}{p}\rfloor\!+\!1$ unique step-sizes in $[0,1]$, and get the full sequence $\{{\color[rgb]{0.23046875,0.45703125,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.23046875,0.45703125,0.6875}a_{1}},{\color[rgb]{0.23046875,0.45703125,0.6875}\definecolor[named]{pgfstrokecolor}{rgb}{0.23046875,0.45703125,0.6875}a_{2}},{\color[rgb]{0.9375,0.52734375,0.2109375}\definecolor[named]{pgfstrokecolor}{rgb}{0.9375,0.52734375,0.2109375}a_{3}},{\color[rgb]{0.9375,0.52734375,0.2109375}\definecolor[named]{pgfstrokecolor}{rgb}{0.9375,0.52734375,0.2109375}a_{4}},\dots\}$, where same colors denote equal step-sizes. Then, we generate the trajectory by rotating point ${\mathbf{x}}_{i}$ clockwise by angle $a_{i}\!\cdot\!\tfrac{2\pi}{n}$ (see text for formal description).
Refer to caption: figs/offset/traj_new2.png

add-$k$ Problem. Each task is a sequence consisting of pairwise examples $\{(x_{i},y_{i})\}_{i=1}^{n+1}$, where $y_{i}=x_{i}+k$, for a given offset $k$. Here, we use integer inputs and offsets; all values are in $\{0,\dots,V-1\}$, each treated as a distinct token. We consider a collection of $K$ tasks parameterized by different offset values in $\{k_{i}\}_{i=1}^{K}$, where $k_{1}=1$ and we fix $k_{i+1}-k_{i}=3$. The model is trained autoregressively to predict the label for each example in the sequence. At test time, the model observes the first $n$ examples and should predict $y_{n+1}=x_{n+1}+k$ for the last example.

Circular-Trajectory Problem. Here a task consists of a sequence $\{{\mathbf{x}}_{i}\}_{i=1}^{n+1}$ of points on a circle centered at the origin. Each task is parameterized by the circle’s radius $r$; for $K$ tasks, the set of radii $\{r_{i}\}_{i=1}^{K}$ is sampled uniformly from $[1,4]$. A task sequence is generated as follows. We first sample $\theta_{0}$ uniformly at random in $[0,\tfrac{\pi}{2}]$, so ${\mathbf{x}}_{1}=r[\cos\theta_{0},\sin\theta_{0}]^{T}$. Then, we select the period $p$ randomly from $\{2,3,4\}$, which determines the number of equal consecutive step-sizes. Specifically, we first sample a sequence of $\lfloor\tfrac{n}{p}\rfloor+1$ unique step-sizes uniformly between $[0,1]$, and then get the full sequence of steps $\{a_{i}\}_{i=1}^{n}$, where $a_{j}=a_{j+1}=\dots=a_{j+p-1}$ for $j\in\{0,p,2p,\dots\}$. Here, context length $n=12m+1$ for integer $m$. We also sample $c\in\{\pm 1\}$, which denotes if the trajectory is clockwise or counterclockwise. Next, we generate a sequence of angles $\{\theta_{i}\}_{i=1}^{n}$, where $\theta_{i}=\theta_{0}+c\tfrac{2\pi}{n}\sum_{j\leq i}a_{j}$. Using the sequence of angles, we generate the sequence, ${\mathbf{x}}_{i+1}=rR(\theta_{i}){\mathbf{x}}_{i}$, where $R(\theta)$ is the 2D rotation matrix for $\theta$. [Fig.˜4](https://arxiv.org/html/2506.16975v2#S3.F4) shows an example. As in the previous problem, we train the transformer autoregressively on these types of sequences. Additionally, in App. [B](https://arxiv.org/html/2506.16975v2#A2), we consider another shape problem, namely the Rectangular-Trajectory problem, where the trajectories contain points on lying on axis-aligned rectangles centered at the origin.
The Circular-Trajectory problem is parameterized by one continuous parameter (radius), whereas the Rectangular-Trajectory problem has two parameters, namely the lengths of the two sides of the rectangle.

### 3.1 Results

Existence of Task Vectors. We first outline the process to identify the task vectors for the add-$k$ problem. We set $V=100$, $n=4$, and $K=2$. [Fig.˜6](https://arxiv.org/html/2506.16975v2#S3.F6) shows the cosine similarities between the layer-$2$ attention embeddings at the last position for $200$ input sequences from each of the two tasks. We observe clustering between intra-task embeddings. Hence, the model disentangles the concept of different offset values in its representation. To provide causal evidence for disentanglement, and to locate where the task vectors emerge in the model, we linear probe the embeddings of the model to predict the final output and the offset/task type. We probe embeddings from the output of the MLP at the first layer, the attention block at the second layer, and the hidden and output layers of the MLP at the second layer. The results are in [Fig.˜6](https://arxiv.org/html/2506.16975v2#S3.F6). We see that task type is disentangled at layer-2 attention, and the output is computed at layer-2 MLP. For each task, we treat the layer-$2$ attention embeddings averaged across $200$ input sequences from that task as the task vector.

Figure: Figure 5: Cosine similarities between the layer-$2$ attention embeddings for $200$ input sequences from two tasks/offsets for the *add-k* problem. Strong clustering between intra-task embeddings shows that the model disentangles the concept of different offsets.
Refer to caption: figs/offset/cos_sim.png

We visualize the task vectors by performing PCA and projecting them onto the first two principal components;
[Fig.˜7](https://arxiv.org/html/2506.16975v2#S3.F7) presents the task vectors for the add-$k$ problem, for $K=4,8,16$ tasks/offsets. In all three settings, the vectors lie on a 1D linear manifold. More than $99.9\%$ of the variance is explained by the first PC. Notably, the model compresses the concept of offsets into a line with the ordering of the offsets (lower to higher) preserved on the manifold (left to right).
To corroborate these results, we study the effect of steering using the task vectors. Specifically, let ${\mathbf{t}}_{1}$ and ${\mathbf{t}}_{K}$ denote the task vectors for offsets $k_{1}$ and $k_{K}$, respectively. Then, for offset $k_{1}$ ($k_{K}$), we steer with $(1-\beta){\mathbf{t}}_{1}+\beta{\mathbf{t}}_{K}$ ($(1-\beta){\mathbf{t}}_{K}+\beta{\mathbf{t}}_{1}$) and evaluate the accuracy for predicting the output based on the original offset $k_{1}$ ($k_{K}$), the ‘opposite’ offset $k_{K}$ ($k_{1}$), or the target offset $(1-\beta)k_{1}+\beta k_{K}$ ($(1-\beta)k_{K}+\beta k_{1}$), where $\beta\in[0,1]$. We use the first/last offsets to interpolate them and steer towards intermediate offsets. [Fig.˜8](https://arxiv.org/html/2506.16975v2#S3.F8) presents the top-$1$ and top-$3$ accuracies for each case. High top-$1$ accuracies and $\approx 100\%$ top-$3$ accuracies for the target
for all considered values of $\beta$ indicate that the model output is steered toward the target. Interpolating along the top principal direction is successful at interpolating values of $k$ in the task space, showing that the model captures the concept’s geometry.

Figure: Figure 7: 2D PCA projection of the task vectors for the *add-k* problem. The task vectors lie on a 1D linear manifold. Here the number of tasks refers to the number of values of $k$.
Refer to caption: figs/offset/tvs2.png

Figure: Figure 8: Steering with the task vectors for tasks $k_{1}$ and $k_{K}$ for the *add-k* problem (see text for details). We plot the top-$1$ and top-$3$ accuracies for predicting the output based on the original offset $k_{1}$ ($k_{K}$), the ‘opposite’ offset $k_{K}$ ($k_{1}$), or the target offset $(1-\beta)k_{1}+\beta k_{K}$ ($(1-\beta)k_{K}+\beta k_{1}$), where $\beta\in[0,1]$, The result shows that the model output can be steered toward the target.
Refer to caption: figs/offset/steer.png

Figure: Figure 9: 2D PCA projection of the task vectors for the Circular-Trajectory problem. The task vectors lie on a smooth low-dimensional manifold. Here the number of tasks refers to the number of radius values used for training.
Refer to caption: figs/offset/tvs_circles.png

Figure: Figure 10: Steering with the task vectors for tasks $r_{1}$ and $r_{K}$ for the Circular-Trajectory problem (see text for details). The MSE between the radius inferred from the model output and the original radius $r_{1}$ ($r_{K}$), the ‘opposite’ radius $r_{K}$ ($r_{1}$), or the target radius $(1-\beta)r_{1}+\beta r_{K}$ ($(1-\beta)r_{K}+\beta r_{1}$), where $\beta\in[0,1]$, indicates that the model output can be steered toward the target.
Refer to caption: figs/offset/circ_steer.png

[Fig.˜9](https://arxiv.org/html/2506.16975v2#S3.F9) presents the 2D PCA projection of the task vectors for the Circular-Trajectory problem, for $K=16,32,64$ (training) tasks/radii. In this setting, we consider $K=24$ radii, spaced evenly between $[1,4]$ to visualize the task vectors, since this task is continuous. We observe that in all three settings, the task vectors lie on a low-dimensional manifold. The variance explained by the first two PCs in the three cases is $97.05\%,96.44\%,93.68\%$, respectively. Similar to the previous setting, the order of the radii (lower to higher) is preserved in the compressed representation.

[Fig.˜10](https://arxiv.org/html/2506.16975v2#S3.F10) presents the results for steering the model output using the task vectors for radii $r_{1}$ and $r_{K}$. We follow the same procedure as in the add-$k$ problem, with a different evaluation metric. We compute the norm of the generated output after steering as the model’s radius (since the center of the circles is fixed at the origin), and consider the MSE between these radii and the original radius $r_{1}$ ($r_{K}$), the ‘opposite’ radius $r_{K}$ ($r_{1}$), or the target radius $(1-\beta)r_{1}+\beta r_{K}$ ($(1-\beta)r_{K}+\beta r_{1}$), where $\beta\in[0,1]$, averaged over $200$ sequences from each task. We observe that the MSE with the target radius is the lowest, which indicated the task vector can steer the model’s output toward the target.

In App. [B](https://arxiv.org/html/2506.16975v2#A2), we examine the task vectors for the Rectangular-Trajectory problem. Here the model has to reason over two latent continuous parameters, which are the real-valued side lengths of the unknown rectangle. We find that the first 2 PCs lie on a two-dimensional manifold in that case as well. This provides a second example of how the model captures the underlying geometry in terms of both separating the two orthogonal parameters and smoothly interpolating across hidden trajectory shapes.

## 4 Conclusion

Our work provides causal and correlational evidence that transformer-based models disentangle and manipulate latent concepts from in-context demonstrations in a structured, interpretable manner. For two-hop tasks, models contain sparse sets of attention heads responsible for first inferring the bridge entity and then resolving the output. For numerical tasks, the model uses task vectors which closely capture the underlying parameterization.
Future work could leverage these identified concept representations directly, for instance, by steering model behavior to improve ICL performance or by composing known concept vectors to enable zero-shot generalization to novel compositional tasks.

## Acknowledgments

This work was supported in part by NSF CAREER Award CCF-$2239265$. The authors acknowledge the use of USC CARC’s Discovery cluster.

## References

- Akyürek et al. (2023)
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou.
What learning algorithm is in-context learning? Investigations with linear models.
In *Int. Conference on Learning Representations*, 2023.
URL [https://openreview.net/forum?id=0g0X4H8yN4I](https://openreview.net/forum?id=0g0X4H8yN4I).
- AlKhamissi et al. (2025a)
Badr AlKhamissi, Greta Tuckute, Antoine Bosselut, and Martin Schrimpf.
The LLM language network: A neuroscientific approach for identifying causally task-relevant units.
In Luis Chiruzzo, Alan Ritter, and Lu Wang (eds.), *Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)*, pp. 10887–10911, Albuquerque, New Mexico, April 2025a. Association for Computational Linguistics.
ISBN 979-8-89176-189-6.
doi: 10.18653/v1/2025.naacl-long.544.
URL [https://aclanthology.org/2025.naacl-long.544/](https://aclanthology.org/2025.naacl-long.544/).
- AlKhamissi et al. (2025b)
Badr AlKhamissi, Greta Tuckute, Antoine Bosselut, and Martin Schrimpf.
The llm language network: A neuroscientific approach for identifying causally task-relevant units, 2025b.
URL [https://arxiv.org/abs/2411.02280](https://arxiv.org/abs/2411.02280).
- Arditi et al. (2024)
Andy Arditi, Oscar Obeso, Aaquib Syed, Daniel Paleka, Nina Panickssery, Wes Gurnee, and Neel Nanda.
Refusal in language models is mediated by a single direction, 2024.
URL [https://arxiv.org/abs/2406.11717](https://arxiv.org/abs/2406.11717).
- Baek & Tegmark (2025)
David D. Baek and Max Tegmark.
Towards understanding distilled reasoning models: A representational approach, 2025.
URL [https://arxiv.org/abs/2503.03730](https://arxiv.org/abs/2503.03730).
- Bai et al. (2023)
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei.
Transformers as statisticians: Provable in-context learning with in-context algorithm selection, 2023.
URL [https://arxiv.org/abs/2306.04637](https://arxiv.org/abs/2306.04637).
- Bakalova et al. (2025)
Aleksandra Bakalova, Yana Veitsman, Xinting Huang, and Michael Hahn.
Contextualize-then-aggregate: Circuits for in-context learning in gemma-2 2b.
*arXiv preprint arXiv:2504.00132*, 2025.
- Beaglehole et al. (2025)
Daniel Beaglehole, Adityanarayanan Radhakrishnan, Enric Boix-Adserà, and Mikhail Belkin.
Aggregate and conquer: detecting and steering llm concepts by combining nonlinear predictors over multiple layers, 2025.
URL [https://arxiv.org/abs/2502.03708](https://arxiv.org/abs/2502.03708).
- Bhattamishra et al. (2024)
Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade.
Understanding in-context learning in transformers and LLMs by learning to learn discrete functions.
In *The Twelfth International Conference on Learning Representations*, 2024.
URL [https://openreview.net/forum?id=ekeyCgeRfC](https://openreview.net/forum?id=ekeyCgeRfC).
- Brinkmann et al. (2024)
Jannik Brinkmann, Abhay Sheshadri, Victor Levoso, Paul Swoboda, and Christian Bartelt.
A mechanistic analysis of a transformer trained on a symbolic multi-step reasoning task.
*arXiv preprint arXiv:2402.11917*, 2024.
- Brown et al. (2020)
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.
Language models are few-shot learners.
In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), *Advances in Neural Information Processing Systems*, volume 33, pp. 1877–1901. Curran Associates, Inc., 2020.
URL [https://proceedings.neurips.cc/paper_files/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf](https://proceedings.neurips.cc/paper_files/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf).
- Davidson et al. (2025)
Guy Davidson, Todd M. Gureckis, Brenden M. Lake, and Adina Williams.
Do different prompting methods yield a common task representation in language models?, 2025.
URL [https://arxiv.org/abs/2505.12075](https://arxiv.org/abs/2505.12075).
- Dumas et al. (2025)
Clément Dumas, Chris Wendler, Veniamin Veselovsky, Giovanni Monea, and Robert West.
Separating tongue from thought: Activation patching reveals language-agnostic concept representations in transformers, 2025.
URL [https://arxiv.org/abs/2411.08745](https://arxiv.org/abs/2411.08745).
- Edelman et al. (2024)
Ezra Edelman, Nikolaos Tsilivis, Benjamin L. Edelman, Eran Malach, and Surbhi Goel.
The evolution of statistical induction heads: In-context learning markov chains.
In *The Thirty-eighth Annual Conference on Neural Information Processing Systems*, 2024.
URL [https://openreview.net/forum?id=qaRT6QTIqJ](https://openreview.net/forum?id=qaRT6QTIqJ).
- Elhage et al. (2021)
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah.
A mathematical framework for transformer circuits.
*Transformer Circuits Thread*, 2021.
https://transformer-circuits.pub/2021/framework/index.html.
- Garg et al. (2022)
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant.
What can transformers learn in-context? a case study of simple function classes.
In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), *Advances in Neural Information Processing Systems*, 2022.
URL [https://openreview.net/forum?id=flNZJ2eOet](https://openreview.net/forum?id=flNZJ2eOet).
- Gemma Team (2024)
Gemma Team.
Gemma 2: Improving open language models at a practical size, 2024.
URL [https://arxiv.org/abs/2408.00118](https://arxiv.org/abs/2408.00118).
- Guo et al. (2025)
Tianyu Guo, Hanlin Zhu, Ruiqi Zhang, Jiantao Jiao, Song Mei, Michael I Jordan, and Stuart Russell.
How do llms perform two-hop reasoning in context?
*arXiv preprint arXiv:2502.13913*, 2025.
- Heindrich et al. (2025)
Lovis Heindrich, Philip Torr, Fazl Barez, and Veronika Thost.
Do sparse autoencoders generalize? a case study of answerability, 2025.
URL [https://arxiv.org/abs/2502.19964](https://arxiv.org/abs/2502.19964).
- Hendel et al. (2023)
Roee Hendel, Mor Geva, and Amir Globerson.
In-context learning creates task vectors.
In *The 2023 Conference on Empirical Methods in Natural Language Processing*, 2023.
URL [https://openreview.net/forum?id=QYvFUlF19n](https://openreview.net/forum?id=QYvFUlF19n).
- Hong et al. (2024)
Guan Zhe Hong, Nishanth Dikkala, Enming Luo, Cyrus Rashtchian, Xin Wang, and Rina Panigrahy.
How transformers solve propositional logic problems: A mechanistic analysis.
*arXiv preprint arXiv:2411.04105*, 2024.
- Hu et al. (2025)
Xinyan Hu, Kayo Yin, Michael I Jordan, Jacob Steinhardt, and Lijie Chen.
Understanding in-context learning of addition via activation subspaces.
*arXiv preprint arXiv:2505.05145*, 2025.
- Huh et al. (2024)
Minyoung Huh, Brian Cheung, Tongzhou Wang, and Phillip Isola.
The platonic representation hypothesis.
*arXiv preprint arXiv:2405.07987*, 2024.
- Ilharco et al. (2023)
Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi.
Editing models with task arithmetic, 2023.
URL [https://arxiv.org/abs/2212.04089](https://arxiv.org/abs/2212.04089).
- Kantamneni & Tegmark (2025)
Subhash Kantamneni and Max Tegmark.
Language models use trigonometry to do addition.
*arXiv preprint arXiv:2502.00873*, 2025.
- Li et al. (2023)
Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak.
Dissecting chain-of-thought: Compositionality through in-context filtering and learning.
In *Thirty-seventh Conference on Neural Information Processing Systems*, 2023.
URL [https://openreview.net/forum?id=xEhKwsqxMa](https://openreview.net/forum?id=xEhKwsqxMa).
- Li et al. (2025a)
Yuxiao Li, Eric J. Michaud, David D. Baek, Joshua Engels, Xiaoqing Sun, and Max Tegmark.
The geometry of concepts: Sparse autoencoder feature structure.
*Entropy*, 27(4), 2025a.
ISSN 1099-4300.
doi: 10.3390/e27040344.
URL [https://www.mdpi.com/1099-4300/27/4/344](https://www.mdpi.com/1099-4300/27/4/344).
- Li et al. (2025b)
Zhuowei Li, Zihao Xu, Ligong Han, Yunhe Gao, Song Wen, Di Liu, Hao Wang, and Dimitris N. Metaxas.
Implicit in-context learning.
In *The Thirteenth International Conference on Learning Representations*, 2025b.
URL [https://openreview.net/forum?id=G7u4ue6ncT](https://openreview.net/forum?id=G7u4ue6ncT).
- Lin & Lee (2024)
Ziqian Lin and Kangwook Lee.
Dual operating modes of in-context learning.
In *Proceedings of the 41st International Conference on Machine Learning*, ICML’24. JMLR.org, 2024.
- Liu et al. (2024)
Sheng Liu, Haotian Ye, Lei Xing, and James Zou.
In-context vectors: making in context learning more effective and controllable through latent space steering.
In *Proceedings of the 41st International Conference on Machine Learning*, ICML’24. JMLR.org, 2024.
- Marks & Tegmark (2024)
Samuel Marks and Max Tegmark.
The geometry of truth: Emergent linear structure in large language model representations of true/false datasets, 2024.
URL [https://arxiv.org/abs/2310.06824](https://arxiv.org/abs/2310.06824).
- Meng et al. (2022)
Kevin Meng, David Bau, Alex J Andonian, and Yonatan Belinkov.
Locating and editing factual associations in GPT.
In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), *Advances in Neural Information Processing Systems*, 2022.
URL [https://openreview.net/forum?id=-h6WAS6eE4](https://openreview.net/forum?id=-h6WAS6eE4).
- Merullo et al. (2024)
Jack Merullo, Carsten Eickhoff, and Ellie Pavlick.
Language models implement simple Word2Vec-style vector arithmetic.
In Kevin Duh, Helena Gomez, and Steven Bethard (eds.), *Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)*, pp. 5030–5047, Mexico City, Mexico, June 2024. Association for Computational Linguistics.
doi: 10.18653/v1/2024.naacl-long.281.
URL [https://aclanthology.org/2024.naacl-long.281/](https://aclanthology.org/2024.naacl-long.281/).
- Min et al. (2022)
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer.
Rethinking the role of demonstrations: What makes in-context learning work?
In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), *Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing*, pp. 11048–11064, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
doi: 10.18653/v1/2022.emnlp-main.759.
URL [https://aclanthology.org/2022.emnlp-main.759/](https://aclanthology.org/2022.emnlp-main.759/).
- Minegishi et al. (2025)
Gouki Minegishi, Hiroki Furuta, Shohei Taniguchi, Yusuke Iwasawa, and Yutaka Matsuo.
Beyond induction heads: In-context meta learning induces multi-phase circuit emergence.
*arXiv preprint arXiv:2505.16694*, 2025.
- Olsson et al. (2022)
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah.
In-context learning and induction heads.
*Transformer Circuits Thread*, 2022.
https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html.
- Park et al. (2025a)
Core Francisco Park, Ekdeep Singh Lubana, and Hidenori Tanaka.
Competition dynamics shape algorithmic phases of in-context learning.
In *The Thirteenth International Conference on Learning Representations*, 2025a.
URL [https://openreview.net/forum?id=XgH1wfHSX8](https://openreview.net/forum?id=XgH1wfHSX8).
- Park et al. (2024)
Kiho Park, Yo Joong Choe, and Victor Veitch.
The linear representation hypothesis and the geometry of large language models.
In *International Conference on Machine Learning*, pp. 39643–39666. PMLR, 2024.
- Park et al. (2025b)
Kiho Park, Yo Joong Choe, Yibo Jiang, and Victor Veitch.
The geometry of categorical and hierarchical concepts in large language models, 2025b.
URL [https://arxiv.org/abs/2406.01506](https://arxiv.org/abs/2406.01506).
- Pearl (2022)
Judea Pearl.
*Direct and Indirect Effects*, pp. 373–392.
Association for Computing Machinery, New York, NY, USA, 1 edition, 2022.
ISBN 9781450395861.
URL [https://doi.org/10.1145/3501714.3501736](https://doi.org/10.1145/3501714.3501736).
- Radford et al. (2019)
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.
Language models are unsupervised multitask learners.
In *OpenAI blog*, 2019.
URL [https://api.semanticscholar.org/CorpusID:160025533](https://api.semanticscholar.org/CorpusID:160025533).
- Rajaraman et al. (2024)
Nived Rajaraman, Marco Bondaschi, Ashok Vardhan Makkuva, Kannan Ramchandran, and Michael Gastpar.
Transformers on markov data: Constant depth suffices.
In *The Thirty-eighth Annual Conference on Neural Information Processing Systems*, 2024.
URL [https://openreview.net/forum?id=5uG9tp3v2q](https://openreview.net/forum?id=5uG9tp3v2q).
- Russin et al. (2025)
Jacob Russin, Ellie Pavlick, and Michael J. Frank.
The dynamic interplay between in-context and in-weight learning in humans and neural networks, 2025.
URL [https://arxiv.org/abs/2402.08674](https://arxiv.org/abs/2402.08674).
- Singh et al. (2025)
Aaditya K. Singh, Ted Moskovitz, Sara Dragutinovic, Felix Hill, Stephanie C. Y. Chan, and Andrew M. Saxe.
Strategy coopetition explains the emergence and transience of in-context learning, 2025.
URL [https://arxiv.org/abs/2503.05631](https://arxiv.org/abs/2503.05631).
- Soligo et al. (2025)
Anna Soligo, Edward Turner, Senthooran Rajamanoharan, and Neel Nanda.
Convergent linear representations of emergent misalignment, 2025.
URL [https://arxiv.org/abs/2506.11618](https://arxiv.org/abs/2506.11618).
- Tigges et al. (2024)
Curt Tigges, Oskar John Hollinsworth, Neel Nanda, and Atticus Geiger.
Language models linearly represent sentiment, 2024.
URL [https://openreview.net/forum?id=iGDWZFc7Ya](https://openreview.net/forum?id=iGDWZFc7Ya).
- Todd et al. (2024)
Eric Todd, Millicent Li, Arnab Sen Sharma, Aaron Mueller, Byron C Wallace, and David Bau.
Function vectors in large language models.
In *The Twelfth International Conference on Learning Representations*, 2024.
URL [https://openreview.net/forum?id=AwyxtyMwaG](https://openreview.net/forum?id=AwyxtyMwaG).
- Turner et al. (2025)
Edward Turner, Anna Soligo, Mia Taylor, Senthooran Rajamanoharan, and Neel Nanda.
Model organisms for emergent misalignment, 2025.
URL [https://arxiv.org/abs/2506.11613](https://arxiv.org/abs/2506.11613).
- Venhoff et al. (2025)
Constantin Venhoff, Iván Arcuschin, Philip Torr, Arthur Conmy, and Neel Nanda.
Understanding reasoning in thinking language models via steering vectors, 2025.
URL [https://arxiv.org/abs/2506.18167](https://arxiv.org/abs/2506.18167).
- Vig et al. (2020)
Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber.
Investigating gender bias in language models using causal mediation analysis.
In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), *Advances in Neural Information Processing Systems*, volume 33, pp. 12388–12401. Curran Associates, Inc., 2020.
URL [https://proceedings.neurips.cc/paper_files/paper/2020/file/92650b2e92217715fe312e6fa7b90d82-Paper.pdf](https://proceedings.neurips.cc/paper_files/paper/2020/file/92650b2e92217715fe312e6fa7b90d82-Paper.pdf).
- Von Oswald et al. (2023)
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov.
Transformers learn in-context by gradient descent.
In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), *Proceedings of the 40th International Conference on Machine Learning*, volume 202 of *Proceedings of Machine Learning Research*, pp. 35151–35174. PMLR, 23–29 Jul 2023.
URL [https://proceedings.mlr.press/v202/von-oswald23a.html](https://proceedings.mlr.press/v202/von-oswald23a.html).
- von Rütte et al. (2024)
Dimitri von Rütte, Sotiris Anagnostidis, Gregor Bachmann, and Thomas Hofmann.
A language model’s guide through latent space, 2024.
URL [https://arxiv.org/abs/2402.14433](https://arxiv.org/abs/2402.14433).
- Wang et al. (2023)
Kevin Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt.
Interpretability in the wild: a circuit for indirect object identification in GPT-2 small.
In *The Eleventh International Conference on Learning Representations*, 2023.
URL [https://openreview.net/forum?id=NpsVSN6o4ul](https://openreview.net/forum?id=NpsVSN6o4ul).
- Wang & Xu (2025)
Zijian Wang and Chang Xu.
Functional abstraction of knowledge recall in large language models, 2025.
URL [https://arxiv.org/abs/2504.14496](https://arxiv.org/abs/2504.14496).
- Wu et al. (2023)
Zhengxuan Wu, Atticus Geiger, Thomas Icard, Christopher Potts, and Noah Goodman.
Interpretability at scale: Identifying causal mechanisms in alpaca.
In *Thirty-seventh Conference on Neural Information Processing Systems*, 2023.
URL [https://openreview.net/forum?id=nRfClnMhVX](https://openreview.net/forum?id=nRfClnMhVX).
- Yang et al. (2024)
Sohee Yang, Elena Gribovskaya, Nora Kassner, Mor Geva, and Sebastian Riedel.
Do large language models latently perform multi-hop reasoning?
In *Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)*, pp. 10210–10229, 2024.
- Yin & Steinhardt (2025a)
Kayo Yin and Jacob Steinhardt.
Which attention heads matter for in-context learning?
In *Forty-second International Conference on Machine Learning*, 2025a.
URL [https://openreview.net/forum?id=C7XmEByCFv](https://openreview.net/forum?id=C7XmEByCFv).
- Yin & Steinhardt (2025b)
Kayo Yin and Jacob Steinhardt.
Which attention heads matter for in-context learning?, 2025b.
URL [https://arxiv.org/abs/2502.14010](https://arxiv.org/abs/2502.14010).
- Zhang & Nanda (2024)
Fred Zhang and Neel Nanda.
Towards best practices of activation patching in language models: Metrics and methods.
In *The Twelfth International Conference on Learning Representations*, 2024.
URL [https://openreview.net/forum?id=Hf17y6u9BC](https://openreview.net/forum?id=Hf17y6u9BC).
- Zhou et al. (2023)
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuurmans, Claire Cui, Olivier Bousquet, Quoc Le, and Ed Chi.
Least-to-most prompting enables complex reasoning in large language models, 2023.
URL [https://arxiv.org/abs/2205.10625](https://arxiv.org/abs/2205.10625).
- Zhou et al. (2024)
Tianyi Zhou, Deqing Fu, Vatsal Sharan, and Robin Jia.
Pre-trained large language models use fourier features to compute addition.
In *NeurIPS*, 2024.

## Appendix

## Appendix A Experimental Details and Additional Experiments for Section ˜ 2

### A.1 Details on experimental setup and techniques

Figure: Figure 11: Basic illustration of CMA. $X$ = input (exposure), $M$ = mediator, $Y$ = output (outcome).
Refer to caption: x3.png

Causal mediation analysis. We primarily rely on causal mediation analysis (CMA), a.k.a. activation patching in the mechanistic interpretability literature, to obtain causal evidence for our claims in the LLM studies.

At a high level, CMA is about the study of indirect effects (IE) and direct effects (DE) in a system with causal relations (Pearl, 2022). Consider the following classical diagram of CMA, in Figure [11](https://arxiv.org/html/2506.16975v2#A1.F11).

Suppose we wish to understand whether a certain mediator $M$ plays an important role in the causal path from the input $X$ to the outcome $Y$. We decompose the “total effect” of $X$ on $Y$ into the sum of direct and indirect effects (DEs and IEs), as shown in the figure. The indirect effect measures how important a role the mediator $M$ plays in the causal path $X\to Y$. To measure it, we compute $Y$ given $X$, except that we artificially hold $M$’s output to its “corrupted” version, which is obtained by computing $M$ on a counterfactual (“corrupted”) version of the input. A significant change in $Y$ indicates a strong IE, which implies that $M$ is important in the causal path. On the other hand, a weak IE implies a strong DE, meaning that the mediator does not play a strong causal role in the system (for the distribution of inputs of interest).

There are two common classes of interventions in mechanistic interpretability for localizing model components with strong IE in the causal graph. The first class is simple ablation, such as mean ablation (replace activation of the mediator by its average output on a distribution of interest) (Wang et al., 2023) or “noising” (Meng et al., 2022). While this type of intervention is easy to perform, it typically leads to poor localization, surfacing low-level processing components irrelevant to the study (Zhang & Nanda, 2024).

The other class, which we employ, is “interchange” intervention: it requires construction of alterative prompts which differ from the normal prompt in subtle ways, requiring careful consideration of the problem’s nature, but allows “causal surgery” which surfaces model components with specific functional roles. Technically speaking, we are measuring the natural indirect effects of the mediator. In particular, it works as follows. We first run the system (the LLM) on both normal and alternative (or sometimes called counterfactual) inputs, and cache the output of the mediator $M$. We then hold $M$’s output to its alternative version, as we run the full system (the LLM) on the normal prompt. Everything downstream in the causal graph from $M$ are also influenced, up to the output $Y$. This helps us measure how the mediator $M$ causally implicate the answer. Or more intuitively, it measures how “flipping” the output of $M$ causally influences the LLM’s “belief” in the alternative answer over the normal answer.

What makes our intervention experiments somewhat novel lies in exactly how we measure the IE. In particular, as we briefly discussed Section 2.1 and 2.2, we do not directly use the alternative prompt’s ground truth answer to measure how well we are “bending” the model’s “belief” through intervention. We discuss our method in greater detail here.

First, to understand whether certain attention heads have functional roles in processing the query source entity $S_{n}$ which transcend source-target types of the two-hop problems, we work with normal-alternative prompts with distinct source and target types, such as sampling an alternative prompt “EPFL, 41. … University of Tokyo, ” ([University, Code] problem), and a normal prompt “Okinawa, Tokyo. … Chicago, ” ([City, Capital] problem). We hypothesize that there are certain model components which output the bridge concept, which is then composed with the target/output concept of the problem. For the normal example, this means “Chicago”$\to$“USA” is resolved first, then the model executes $\text{Capital}(\text{USA})=$ Washington D.C. as the output. This means that, patching a model component’s activation from the alternative prompt onto its activation on a normal prompt, would cause the model to favor the answer of the alternative prompt, but with the same target semantic type as the normal prompt. In our running example, this would be “Tokyo”, the capital of “Japan”, the country (bridge) of the university “University of Tokyo”.

It follows that, to evaluate the “causal effects” of such a bridge-resolving component, we should set $\hat{T}^{(\text{alt})}_{n}=\text{Type}_{\text{norm}}(T_{n}^{(\text{alt})})$. We then measure the (expected) intervened logit difference

$$ $\Delta_{\text{alt}\to\text{norm}}=\mathbb{E}\left[\text{logit}^{\text{alt}\to\text{norm}}({\bm{p}}_{\text{norm}})[T^{(\text{norm})}_{n}]-\text{logit}^{\text{alt}\to\text{norm}}({\bm{p}}_{\text{norm}})[\hat{T}^{(\text{alt})}_{n}]\right],$ (1) $$

where ${\bm{p}}_{\text{norm}}=[S^{(\text{norm})}_{1},T^{(\text{norm})}_{1}...\,S^{(\text{norm})}_{n},]$ is the normal prompt, $\text{logit}^{\text{alt}\to\text{norm}}({\bm{p}}_{\text{norm}})$ indicates the logits of the model obtained after intervention while running the model on the normal prompt, and $\text{logit}({\bm{p}}_{\text{norm}})$ indicates the logits of the model running naturally (un-intervened) on the normal prompt. Moreover, when we measure the rank of the model’s answer when intervened, we also use $\hat{T}^{(\text{alt})}_{n}$ as the target.

Remark. To normalize our logit-difference variations, we compute

$$ $\bar{\Delta}=\frac{\Delta_{\text{norm}}-\Delta_{\text{alt}\to\text{norm}}}{\Delta_{\text{norm}}},$ (2) $$

where

$$ $\Delta_{\text{norm}}=\mathbb{E}\left[\text{logit}({\bm{p}}_{\text{norm}})[T^{(\text{norm})}_{n}]-\text{logit}({\bm{p}}_{\text{norm}})[\hat{T}^{(\text{alt})}_{n}]\right].$ (3) $$

Problem settings and overall observations. We primarily work with the Geography and Company 2-hop ICL problems to, and with the Gemma-2 LLM family. The former problem setting was described in the main text (with further elaboration in the Appendix later), while the latter will be introduced later, to add further generality to our study.

At a high level, in both the Geography and Company 2-hop ICL problems, we obtained causal and correlational evidence that highly localized groups of attention heads output the representation of the “bridge” concept based on the query, and the model later utilizes such representation to produce the output. Furthermore, an interesting technical observation is that multiplying the patched representation at these attention heads with a constant slightly above 1 (e.g. 1.5 to 4.0) tend to improve the intervention results. Finally, we make the observation that a smaller LLM, namely Gemma-2-2B, also possesses attention heads which have a nontrivial causal role in inferring the bridge representations. However, the representations are poorly disentangled, potentially leading to the model’s low accuracy on the ICL problems.

Compute. Our LLM experiments are conducted on 2 H200 GPUs, totaling around 200 hours of compute.

#### Licenses.

We use the pretrained Gemma 2 models for our LLM experiments. They have the following license information.

- •
Models: Gemma-2-27B and Gemma-2-2B (google/gemma-2-27b and google/gemma-2-2B on Huggingface)
- •
License: Gemma Terms of Use (Google, 21 Feb 2024)
- •
Link: https://ai.google.dev/gemma/terms
- •
Notes: Commercial use permitted subject to Gemma Prohibited-Use Policy.

### A.2 Additional details on experiments with the Geography 2-hop ICL puzzles

Dataset construction. There are two stages to our data sampling process:

#### Step 1: JSON dictionary sampling.

First, we ask ChatGPT o3 to create a JSON dictionary, mapping each country to its cities, capital, and calling code, and repeating this for universities, and famous landmarks. We then manually correct for and refine the entries in the JSON dictionaries to reduce leakage of source types. For instance, ChatGPT sometimes append city or state/province/region to a landmark, which we remove to ensure that the Landmark source type remains sufficiently distinct from the City source type. University names sometimes cannot avoid such overlap, e.g. University of California, Berkeley indeed has city name in it.

The 2-hop Company dataset is constructed in almost exactly the same fashion, mapping the companies to their products, founders and headquarter cities.

#### Step 2: Prompt generation.

This is discussed in [Section˜2](https://arxiv.org/html/2506.16975v2#S2). Taking the geography puzzles as the running example, for every City$\rightarrow$Capital prompt $X=[\text{City}_{1},\text{Capital}_{1}\dots\text{City}_{n},]$, we sample tuples $(\text{City}_{i},\text{Country}_{i},\text{Capital}_{i}),i=1,\dots,n$, where each “bridge” $B_{i}=\text{Country}_{i}$ is a different country and the $\text{City}_{i}$ and $\text{Capital}_{i}$ belong to that country, all randomly sampled from the JSON dictionary from before. The same holds for other source-target types, and on the Company dataset.

Causal evidence. Recall that in the main text, to provide causal evidence for the bridge-resolving mechanism, we primarily presented causal intervention experiments where we treated [City, Capital] as the problem type we intervene on, using cross-type prompts [University, Calling Code], [Landmark, Calling Code] to show causal evidence for the bridge-resolving heads. Here, we add further evidence by having other source-target types. The results are presented in Figures [12](https://arxiv.org/html/2506.16975v2#A1.F12) to [21](https://arxiv.org/html/2506.16975v2#A1.F21), indexed as follows:

- 1.
Experiment [City, Capital]$\to$[Landmark, Calling Code]: Figure [12](https://arxiv.org/html/2506.16975v2#A1.F12)
- 2.
Experiment [University, Capital]$\to$[Landmark, Calling Code]: Figure [13](https://arxiv.org/html/2506.16975v2#A1.F13)
- 3.
Experiment [City, Capital]$\to$[University, Calling Code]: Figure [14](https://arxiv.org/html/2506.16975v2#A1.F14)
- 4.
Experiment [Landmark, Capital]$\to$[University, Calling Code]: Figure [15](https://arxiv.org/html/2506.16975v2#A1.F15)
- 5.
Experiment [University, Capital]$\to$[City, Calling Code]: Figure [16](https://arxiv.org/html/2506.16975v2#A1.F16)
- 6.
Experiment [Landmark, Capital]$\to$[City, Calling Code]: Figure [17](https://arxiv.org/html/2506.16975v2#A1.F17)
- 7.
Experiment [City, Calling Code]$\to$[University, Capital]: Figure [18](https://arxiv.org/html/2506.16975v2#A1.F18)
- 8.
Experiment [Landmark, Calling Code]$\to$[University, Capital]: Figure [19](https://arxiv.org/html/2506.16975v2#A1.F19)
- 9.
Experiment [City, Calling Code]$\to$[Landmark, Capital]: Figure [20](https://arxiv.org/html/2506.16975v2#A1.F20)
- 10.
Experiment [University, Calling Code]$\to$[Landmark, Capital]: Figure [21](https://arxiv.org/html/2506.16975v2#A1.F21)

Every patching experiment is performed on at least 100 prompts. As we can see, the general trend is that there is strong transferability of the bridge representation across the problem types, including when the source and target types have no overlap, giving us causal evidence that head groups (24,30;31), (35,22;23) are “resolving the bridge”.

Figure: Figure 12: Gemma-2-27B [City, Capital]$\to$[Landmark, Code] transfer experiments, intervening head groups $(24,30;31),(35,22;23)$. We show the percentage logit-difference variation in (a), and reciprocal rank of the answer answer before and after intervention in the (b) series of figures, with different scaling constants in $\{1.0,1.5,2.0,4.0\}$. We observe strong causal effects of the two attention head groups. Here, the scaling constant does not significantly affect the intervention performance.
Refer to caption: x4.png

Scaling constant for bridge intervention. We found that for some of the transfer experiments, multiplying the patched representation for the heads (24,30;31), (35,22;23) improves the result, i.e. there is a greater percentage of samples where the alternative answer is boosted into the top-10 (or even top-1) answers of the model after intervention. Therefore, we also report those results. An intriguing property of this scaling constant is that it typically works best around 2.0. At 4.0, we often observe saturation or even decline in the intervened alternative answer’s rank, such as in the [University, Capital]$\to$[City, Calling Code] experiment shown in Figure [16](https://arxiv.org/html/2506.16975v2#A1.F16).

Figure: Figure 13: Gemma-2-27B [University, Capital]$\to$[Landmark, Code] transfer experiments, intervening head groups $(24,30;31),(35,22;23)$. We show the percentage logit-difference variation in (a), and reciprocal rank of the answer answer before and after intervention in the (b) series of figures, with different scaling constants in $\{1.0,1.5,2.0,4.0\}$. We observe strong causal effects of the two head groups. Interesting, past a scaling constant of 1.5, we observe decline in intervention performance.
Refer to caption: x5.png

Figure: Figure 14: Gemma-2-27B [City, Capital]$\to$[University, Code] transfer experiments, intervening head groups $(24,30;31),(35,22;23)$. We show the percentage logit-difference variation in (a), and reciprocal rank of the answer answer before and after intervention in the (b) series of figures, with different scaling constants in $\{1.0,1.5,2.0,4.0\}$. We observe strong causal effects of the two attention head groups, boosting the alternative answer into top 1 around 60% of the time! Additionally, the scaling constant does not significantly affect the intervention performance in this experiment.
Refer to caption: x6.png

Figure: Figure 15: Gemma-2-27B [Landmark, Capital]$\to$[University, Code] transfer experiments, intervening head groups $(24,30;31),(35,22;23)$. We show the percentage logit-difference variation in (a), and reciprocal rank of the answer answer before and after intervention in the (b) series of figures, with different scaling constants in $\{1.0,1.5,2.0,4.0\}$. We observe strong causal effects of the two attention head groups. Here, the positive effects of the scaling constant saturates around 2.0.
Refer to caption: x7.png

Figure: Figure 16: Gemma-2-27B [University, Capital]$\to$[City, Code] transfer experiments, intervening head groups $(24,30;31),(35,22;23)$. We show the percentage logit-difference variation in (a), and reciprocal rank of the answer answer before and after intervention in the (b) series of figures, with different scaling constants in $\{1.0,1.5,2.0,4.0\}$. Interestingly, we observe decline in the intervention’s accuracy as we push the scaling constant from 2.0 to 4.0 (top-1 accuracy decreases from around 50% to slightly above 40%), indicating a subtle regime in which the scaling constant boosts intervention performance.
Refer to caption: x8.png

Figure: Figure 17: Gemma-2-27B [Landmark, Capital]$\to$[City, Code] transfer experiments, intervening head groups $(24,30;31),(35,22;23)$. We show the percentage logit-difference variation in (a), and reciprocal rank of the answer answer before and after intervention in the (b) series of figures, with different scaling constants in $\{1.0,1.5,2.0,4.0\}$. We observe strong causal effects of the two attention head groups.
Refer to caption: x9.png

Figure: Figure 18: Gemma-2-27B [City, Calling Code]$\to$[University, Capital] transfer experiments, intervening head groups $(24,30;31),(35,22;23)$. At scaling constant $1.0$ (i.e. natural intervention, no additional scaling), the reciprocal rank of 0.1 for the alternative answer after intervention is at the 36^th percentile, while before patching, as we can see, the reciprocal rank of the alternative answer is mainly in the range of $10^{-2}$ to $10^{-5}$.
Refer to caption: x10.png

Figure: Figure 19: Gemma-2-27B [Landmark, Calling Code]$\to$[University, Capital] transfer experiments, intervening head groups $(24,30;31),(35,22;23)$. At scaling constant $1.0$ (i.e. natural intervention, no additional scaling), the reciprocal rank of 0.1 for the alternative answer after intervention is at the 41^th percentile.
Refer to caption: x11.png

Figure: Figure 20: Gemma-2-27B [City, Calling Code]$\to$[Landmark, Capital] transfer experiments, intervening the single head group $(24,30;31)$. At scaling constant $1.0$ (i.e. natural intervention, no additional scaling), the reciprocal rank of 0.1 for the alternative answer after intervention is at the 34^th percentile.
Refer to caption: x12.png

Figure: Figure 21: Gemma-2-27B [University, Calling Code]$\to$[Landmark, Capital] transfer experiments, intervening the single head group $(24,30;31)$. At scaling constant $1.0$ (i.e. natural intervention, no additional scaling), the reciprocal rank of 0.1 for the alternative answer after intervention is at the 31^st percentile.
Refer to caption: x13.png

The output-concept heads. While the main interest of this work lies in the bridge-resolving mechanism enabled by the sparse set of attention heads discussed above, we also present more analysis of the output-concept heads, whose embedding tends to cluster with respect to the output concept (Capital versus Calling Code). To localize these heads, we generate normal-alternative prompt pairs where we only change the target/output type of the normal prompt to generate the alternative prompt, but keep the $S_{i}$’s to be identical across the prompt pairs for all $i\leq n$. This helps us surface components which are independent of the query and bridge value, and sensitive to the output/target type for the ICL problem. The results are shown in Figure [22](https://arxiv.org/html/2506.16975v2#A1.F22) and [23](https://arxiv.org/html/2506.16975v2#A1.F23), where we run the intervention experiment [Landmark, Calling Code]$\to$[Landmark, Capital] (due to limitations in time and computing resources, we could not sweep all the source-target combinations as of this version of the paper). As we can see, these head groups with the strongest causal scores indeed tend to exhibit sensitivity to output type, and insensitivity to source/input type and query and bridge value. Moreover, they are more concentrated in the deeper layers of the model.

Statistics of alternative-type answers. A natural question to challenge the bridge-resolving mechanism is as follows. Say we are performing intervention by sampling alternative prompts from the problem type [University, Calling Code], and normal prompts from the type [City, Capital]: even though the target types have little overlap in text, perhaps the model still assigns nontrivial confidence to the “Capital” version of the alternative answer (which has target type “Calling Code”)? If that is so, then it challenges our hypothesis about the role of the “bridge” representations, since we might just be directly injecting the right version of the alternative answer into the model.

We show evidence to refute this. In particular, in Figure [24](https://arxiv.org/html/2506.16975v2#A1.F24), we show that the model places trivial confidence on the altered-type answer, even if they share the same bridge value. Therefore, we add further evidence to the bridge-resolving mechanism.

The role of MLPs. We performed similar intervention experiments on the MLPs at the last token position just like with the attention heads. They are observed to primarily process the output concept type, and do not participate heavily in outputting the bridge concept representation. This is revealed in Figure [25](https://arxiv.org/html/2506.16975v2#A1.F25).

Smaller LLM exhibits weaker disentanglement. To contrast against our results for the 27B model, we study a small model in the same family of Gemma 2 models, Gemma-2-2B. This smaller model has significantly lower accuracies on the problems, measured at 20 shots. [City, Capital]: 71.67%, [University, Capital]: 25.83%, [Landmark, Capital]: 66.67%, [City, Calling Code]: 47.5%, [University, Calling Code]: 57.5%, [Landmark, Calling Code]: 23.33%.

We perform an intervention experiment [University, Code]$\to$[City, Capital] on Gemma-2-2B, similar to the bridge-resolving head localization experiments we did on the 27B model. Intriguingly, we were also able to surface a highly sparse set of attention heads which have nontrivial causal scores (but much lower than that achieved by the 27B model). We find that these heads exhibit noticeable, but noisy disentanglement with respect to the bridge representation. We show these results in Figure [26](https://arxiv.org/html/2506.16975v2#A1.F26). The weaker causal score of the bridge-resolving heads and their noisier concept disentanglement in the 2B model suggests the conjecture that, the larger the model, the more specialized its concept-processing components are — assuming that the model is well-trained. Such specialization likely benefits the model’s generalization accuracy.

Figure: Figure 22: (Best viewed zoomed in) Cosine similarity map of the output-concept head groups (with top causal scores) identified in Gemma-2-27B, along with the percent logit-difference variation of the head groups, serving as the metric for the head groups’ causal effects. Observe that they are mostly insensitive to the source type, query value, and bridge value, and primarily sensitive to the output/target type. Note: in this set of visualizations, we are using “Italy” and “Spain” as the bridge values.
Refer to caption: x14.png

Figure: Figure 23: Cosine similarity map of the output-concept head groups identified in Gemma-2-27B. Here, we construct four groups of prompts. The first two groups consist of multi-hop ICL problems with Capital or Calling Code as the target type. The remaining two are created by randomly shuffling the output of the normal multi-hop ICL samples, causing the problem to essentially demand randomly outputting a Capital or a Calling Code; these are the “negative controls” we discussed in the main text. We find the output-concept heads’ embeddings on the multi-hop prompts to align strongly with those on the output-concept-only prompts, further confirming their role in the circuit.
Refer to caption: x15.png

Figure: Figure 24: Distribution of Gemma-2-27B’s confidence on the correct- and incorrect-type answer on the different problems. When we say “correct-class” answer, we simply mean that the answer’s semantic type aligns with that of the problem’s target, e.g. the correct-type answer for a prompt “Okinawa, Tokyo. Nantes, Paris. … Shanghai, ” would be “Beijing”, while the incorrect-type answer would be “86” (the calling code of China, which the city Shanghai belongs to). We observe a clear separation in the LLM’s confidence between the two types of answers.
Refer to caption: x16.png

Figure: Figure 25: Functional roles of the MLPs in Gemma-2-27B. We perform patching experiments on [University, Code]$\to$[City, Capital] at the last token position, similar to how we localize the bridge-resolving attention heads. We report the percentage logit-difference variation of the top-scoring MLPs, along with their cosine similarity maps computed on prompts sampled with different combinations of bridge values (“Italy” and “Spain”) and diverse set of source-target types. Perhaps unsurprisingly, the MLPs at the last token position play a less interesting role: as seen in the cosine similarity maps for the MLPs with the highest causal scores (fairly low compared to the attention heads), they primarily discriminate against the output type. They do not appear to participate much in resolving the bridge concept, as indicated by their lower causal scores, and the lack of sensitivity to the bridge entities in the cosine-similarity plots.
Refer to caption: x17.png

Figure: Figure 26: Results of the intervention experiment [University, Code]$\to$[City, Capital], conducted on Gemma-2-2B, with 20-shot ICL. On the left, we show the percentage logit difference variation of the intervention experiment; on the right, we plot the cosine similarity map of the two head groups with the highest causal scores, namely $(15,4;5)$ and $(22,0;1)$. We find that while the two attention head groups exhibit nontrivial causal effects and disentanglement, they are, in comparison, much weaker than those exhibited by the 27B model. This likely explains the significantly lower accuracy of the 2B model than the 27B model. This also suggests a conjecture: perhaps the larger the model, the more specialized its concept-processing components are?
Refer to caption: x18.png

### A.3 Multi-hop mechanism formation and the number of in-context examples

This sub-section focuses on illustrating the relation between the number of in-context examples versus (1) how strong a role the multi-hop mechanism plays in the LLM’s inference (via causal interventions), (2) disentanglement strength of key bridge-resolving attention heads. As we will show below, there is a general positive correlation between the number of shots and the two factors.

Figure: Figure 27: This figure illustrates how the transferability and disentanglement of the bridge representation increases as we increase the number of in-context examples. Figure series (a) and (b) present the transferability result, obtained by performing cross-problem-type patching, and measuring the causal influence of the patched representation. In (a)(i) to (iv), we plot the percentage logit variation of the attention heads found to output “bridge” values, measured on several intervention experiments. For (b)(i) and (ii), we zoom in on head group (24,30;31), and show its causal effects on two patching experiments, [University, Calling Code]$\to$[City, Capital] and [Landmark, Calling Code]$\to$[City, Capital]. The x-axis is the number of in-context examples, and the y-axis is the 30^th percentile of the reciprocal rank of the alternative prompt’s answer. For (c)(i) to (iv), we plot the histogram of the disentanglement strength of the representations of head group (24,30;31), with the y-axis in percentage, and x-axis being cosine similarity.
Refer to caption: x19.png

More demonstrations$\implies$stronger causal score. Figure [27](https://arxiv.org/html/2506.16975v2#A1.F27) visualizes the experimental results. From Figure [27](https://arxiv.org/html/2506.16975v2#A1.F27)(a) and (b) and sub-figures, we observe a correlation between the number of shots (ICL examples) and the bridge-resolving heads’ “causal importance” in the model’s inference. When the number of shots is low, we find that they tend to exhibit weak causal influence on the model’s inference. For instance, as (b)(i) shows, at 2 shots, the $30^{th}$ percentile of the alternative answer’s rank after patching at (24,30;31) is on the order of $10^{3}$. This is in stark contrast to how strong this head group’s causal influence is at 20 shots as we saw before.

More demonstrations$\implies$stronger disentanglement, with a catch. In Figure [27](https://arxiv.org/html/2506.16975v2#A1.F27)(c)(i) to (iv), we observe that the intra-bridge cosine similarity tends to cluster better as the number of shots increase, while the inter-bridge cosine similarities decay toward 0.2, with the two distributions overlapping less and less. Interestingly, the bridge-disentanglement strength is still non-trivial with very few shots, mirroring the causal-intervention results: regardless of how disentangled the representations are in the very-few-shot regime, the LLM does not “realize” how it should utilize the multi-hop sub-circuit.

### A.4 The Company Puzzles, Another Multi-Hop ICL Task

Problem setup. To complement our study on the Geography puzzles and for more generality of our conclusion that sufficiently capable LLMs solve simple 2-hop ICL problems by resolving the bridge entities first, we work with another set of problems on globally well-known companies in the world. We collect data on 36 of them.

Here, our evidence is provided on a narrower range of source and target entity types, due to the nature of the problem. In particular, we allow source entities to fall in {Products, Founders}, and target entities to fall into {Founders, Headquarters City}. Furthermore, to obtain causal evidence that the model is resolving the bridge values from the query before producing the actual answer, we perform patching experiments with [Products, Founders]$\to$[Founders, HQ City] and [Founders, Products]$\to$[Products, HQ City]. If we again observe that there is high rank of the alternative prompt’s converted answer after patching a sparse set of attention heads on the final token position, then we see significant causal evidence that the model resolves the bridge value during inference.

Some manual cleaning was needed in creating this dataset, on top of ChatGPT sampling: we primarily filter out products which had company names in them, and companies whose products mostly contain the company name. An example category of this would be banks and credit/debit card services, e.g. Visa, Mastercard, HSBC, etc. For companies with multiple headquarters, we choose the most well-known one out of them as the correct “Headquarter City” of the company; this sometimes counts an old headquarter of the company as the correct one. Note that this causes us to under-estimate the LLM’s accuracy with and without intervention!

Quantitative results. We primarily work with the 16-shot ICL setting due to limited compute. In this setting, the [Product, Founder], [Product, HQ City] and [Founder, HQ City] have accuracies $73.5\%,80.7\%,88.6\%$ respectively. We wish to note that this set of puzzles does rely on significantly more obscure knowledge than the Geography puzzles, so lower accuracies and weaker multi-hop links are expected overall.(^5^55We suspect that lower accuracy is not only due to the more obscure entities in this problem, but also because of several ambiguities. For example, (1) Product names are oftentimes simply not presented alongside company names, e.g. “Gatorade” is often not written alongside “PepsiCo”; (2) There can be confusion between company founders and well-known product directors, especially in certain art/entertainment companies.)

As Figures [28](https://arxiv.org/html/2506.16975v2#A1.F28), [29](https://arxiv.org/html/2506.16975v2#A1.F29) and [30](https://arxiv.org/html/2506.16975v2#A1.F30) show, we again obtain evidence for the bridge-resolving heads. In particular, we find that head groups $(22,2;3),(24,14;15),(24,30;31),(28,12;13),(28,26;27),(35,22;23)$ ($0.8\%$ of the total number of attention heads) exhibit the strongest causal effects in our experiments, and their output representations exhibit strong disentanglement with respect to the companies.

Figure: Figure 28: [Founder, Product]$\to$[Product, HQ City] transfer experiments, intervening the head groups $(22,2;3),(24,14;15),(24,30;31),(28,12;13),(28,26;27),(35,22;23)$ ($0.8\%$ of the total number of attention heads).
Refer to caption: x20.png

Figure: Figure 29: [Product, Founder]$\to$[Founder, HQ City] transfer experiments, intervening the head groups $(22,2;3),(24,14;15),(24,30;31),(28,12;13),(28,26;27),(35,22;23)$ ($0.8\%$ of the total number of attention heads).
Refer to caption: x21.png

Figure: (a) Ford vs. Nintendo
Refer to caption: x22.png

## Appendix B Experimental Details and Additional Experiments for Section ˜ 3

This section includes experimental details and additional experiments for problems with numerical or continuous parameterization considered in [Section˜3](https://arxiv.org/html/2506.16975v2#S3).

#### Compute.

The experiments in this section were performed on an internal cluster with a P$100$ and a V$100$ GPU. We list training details such as batch sizes and iterations used for the experiments in their respective sections.

### B.1 More on the add- k k Problem

In this section, we investigate the effect of increasing the number of tasks on the results shown in Fig. 8, where we showed that the task vectors for $K=4,8,16$ offsets lie on a 1D linear manifold. In these settings, the model reached $\approx 100\%$ train and test accuracy.

Figure: Figure 31: 3D (left) and 2D (right) PCA projections of the task vectors for the *add-k* problem with $32$ offsets. The task vectors lie on a 3D manifold that looks like a (small) helix. From the 2D projection, we see that alternating (odd and even) offsets are represented in two separate arcs of the helix.
Refer to caption: figs/offset/32offsets.png

We consider a setting with $32$ offsets with gap between the offsets $k_{i+1}-k_{i}=1$. In this setting, the train/test accuracies reach only about $90\%$, which indicates this is a harder problem for the model to solve. [Fig.˜31](https://arxiv.org/html/2506.16975v2#A2.F31) shows 3D and 2D PCA projections of the task vectors. These explain about $96.1\%$ and $84.48\%$ of the variance, respectively. We see that the task vectors lie on a 3D manifold that looks like a (small) helix, with alternating (odd and even) offsets represented in two separate arcs of the helix. This is significantly different from the 1D linear manifold we observed with a smaller number of offsets in Fig. 8. This type of geometry is reminiscent of the observations in Zhou et al. (2024), Kantamneni & Tegmark (2025), which study how pre-trained language models do addition and find that these models encode numbers as a helix.

#### Training Details.

The models were trained for $1000$ iterations using AdamW optimizer with learning rate $0.001$ and weight decay $0.01$. We use a linear learning-rate schedule with a warm-up phase over the first $10\%$ iterations. The train data contains $5000K$ sequences, we use a batch size of $500$ for $K=4,8,16$, and $2000$ for $K=32$.

### B.2 More on the Circular-Trajectory Problem

In this section, we include further details about the experimental setup of the Circular-Trajectory problem, and present some additional results.

Figure: Figure 32: 2D projection of task vectors for the Circular-Trajectory problem when the model is trained on $32$ radius values. The three plots show task vectors for clockwise (CW) trajectories, counterclockwise (CCW) trajectories, and all trajectories considered together, respectively. We observe that CW and CCW circle trajectories are represented on two manifolds in the 2D space.
Refer to caption: figs/offset/circatv2.png

In Fig. 10 in the paper, we show the geometry of task vectors for the Circular-Trajectory problem for sequences where $c=-1$, i.e., the trajectories are clockwise (CW). In this section we look at the geometry of task vectors for sequences with counterclockwise (CCW) trajectories, as well as all sequences considered together. In [Fig.˜32](https://arxiv.org/html/2506.16975v2#A2.F32), we plot the 2D projection of task vectors for CW, CCW and both taken together, respectively. We consider $32$ radius values while training. The variance explained by $2$ PCs for the three plots is $93.55\%,96.32\%,89.1\%$, respectively. Darker colors correspond to smaller radius values, and the two colormaps in the third subplot represent CW or CCW trajectories. We observe that CW and CCW circle trajectories are represented on two manifolds in the 2D space.

#### Training Details.

In this setting, we sample a new batch of sequences at every iteration. We use batch size $64$ and train the model on a total of $200\,000$ sequences. We use Adam optimizer with learning rate $10^{-4}$ and weight decay $0.001$. We use representations at position $\lfloor\tfrac{n}{2}\rfloor$ as the task vector following the process mentioned in Section 3.1, averaging over $100$ sequences each. These settings remain the same for the experiments in the next section.

### B.3 More on the Rectangular-Trajectory Problem

Figure: Figure 33: An illustration of a sequence of points for the Rectangular-Trajectory problem with $e=5$ points per edge, and $n=15$. See text for details.
Refer to caption: figs/offset/rect_traj.png

In this section we consider a Rectangular-Trajectory problem, parameterized by two parameters, namely the lengths of the two sides of the rectangle, say $(a,b)$. Specifically, the trajectories contain points on axis-aligned rectangles centered at the origin. Let $e$ denote the number of points on each edge of the rectangle spaced uniformly. The starting point of the sequence is randomly sampled from one of the $e$ points on the right vertical edge of the rectangle. The rest of the points are obtained by traversing the rectangle CW or CCW, determined by $c=-1$ or $1$. [Fig.˜33](https://arxiv.org/html/2506.16975v2#A2.F33) shows an example sequence.

Figure: Figure 34: 2D projection of the task vectors obtained for $64$ $(a,b)$ combinations. Fixed color or transparency level corresponds to fixed $a$ or $b$, respectively. The task vectors lie on a (smooth) 2D manifold.
Refer to caption: figs/offset/recttvgrid.png

Similar to Circular-Trajectory problem, each sequence is obtained by first sampling $a$ and $b$ uniformly between $1$ and $4$, then sampling the starting point and $c=-1$ or $1$, and then following the aforementioned process. For our experiments, we set $e=5$ and $n=15$ for this task. The number of tasks $K$ denotes the number of different combinations $(a,b)$.

In [Fig.˜34](https://arxiv.org/html/2506.16975v2#A2.F34), we plot the 2D projection of the task vectors obtained for all $(a,b)$ combinations lying on the 2D grid between $a\in[1,4]$ and $b\in[1,4]$. Similar to the experiments in Fig. 10, we plot the task vectors for trajectories with $c=-1$ here. We consider $K=32$ in this experiment. The first two PCs explain $91.97\%$ variance. We observe that all the task vectors lie on a 2D manifold.

This setting goes beyond the Circular-Trajectory problem and shows that transformers represent task vectors corresponding to the problem parameters (radius for circles and edge lengths for rectangles) in low-dimensional (smooth) manifolds in both cases.

## Appendix C Limitations

There are a few limitations of this work, which we discuss below.

- •
We only look at English language queries for the Geography puzzles, so we do not evaluate how mechanisms may change across different languages. The input language may correlate with the model’s parametric knowledge that the model uses to recall the Bridge entity. A more thorough study would compare model mechanisms for prompts across several languages, to ensure that the results remain the same.
- •
The latent concepts that we test for all have easy to understand human representations, such as numbers, countries, companies, or geometric shapes. This limits our study to a subset of possible latent concepts that could be present in the ICL examples. There may be other, more intricate, relationships between the ICL examples that the model is also representing in some way. It would be ideal to provide ablations over the human-interpretability of the provided examples, to see if this affects how the model represents the latent concepts.

## Appendix D Use of Large Language Models (LLMs)

Frontier LLMs assisted in the process of curating the Country and Company datasets for our multi-hop ICL experiments discussed in Section [2](https://arxiv.org/html/2506.16975v2#S2) and in Appendix [A](https://arxiv.org/html/2506.16975v2#A1) above. As we discussed at the beginning of Appendix [A.2](https://arxiv.org/html/2506.16975v2#A1.SS2), we rely on ChatGPT o3 (available at the time of dataset curation) to perform the initial sampling of the countries and companies’ source and target entity data. Manual filtering and cleaning of the datasets were then performed.

We did not use LLMs for writing, ideation, or finding related work.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. How are concept variables defined or selected for mediation?
2. How stable are the identified mediating paths across random seeds or different model checkpoints?  
3. Intuitively, can the mediators identified via CMA be related to known attention-head clusters or activation subspaces discovered in prior ICL localization work?
soundness: 3
strengths: - The paper takes a systematic and well-executed approach to analyzing internal mechanisms of ICL.
- The approach offers interpretable, fine-grained insight into how contextual information propagates, complementing existing representation-based probing techniques.
- The study spans synthetic geometric reasoning tasks and natural language problems, demonstrating that the framework generalizes across distinct domains of conceptual structure.  
- The visual analyses (mediation heatmaps and intervention results) are well presented and help illustrate which attention paths encode specific concepts.  
- The work contributes methodologically by adapting tools from causal inference to the study of neural mechanisms in LLMs.
summary: The paper analyzes how LLMs represent and combine latent "concepts" during ICL. To analyze these interactions, they employ causal mediation analysis to identify which components of the model’s computation mediate the effect of particular input concepts on final predictions.
weaknesses: - The paper treats latent concepts as interpretable dimensions or attention patterns but never defines them formally.  
- The causal model in the introduction ($F = R \circ C$) is not linked to the implemented CMA pipeline. $R$ and $C$ are just introduced without definitions. There is no derivation showing that the empirical mediation quantities estimated from activations correspond to components of this decomposition.  
- Experiments are mostly conducted on synthetic or simplified reasoning datasets (two-hop relations at most).  
  These provide clarity but limit conclusions about behavior on realistic compositional or linguistic tasks.  
- The paper occasionally presents correlations in mediation strength as evidence of conceptual disentanglement. In my opinion, stronger validation, like causal interventions or counterfactual ablation, would be needed to corroborate these claims.  
- No stability analysis or statistical testing is reported. It is unclear how consistent the discovered mediators are across runs or models.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Please see Weaknesses.
soundness: 3
strengths: The exploration of how LLMs understand and reason at the conceptual level is interesting and provides valuable insights for the community. The use of causal mediation analysis and PCA visualization to validate the claims is persuasive under the synthetic tasks presented in the paper. The experimental design is rigorous and highly interpretable. In particular, the demonstration of low-dimensional manifold structures in the mode’s representation space (i.e., the geometric interpretability of latent variables) is a novel perspective.
summary: This paper systematically studies how large language models represent, disentangle, and utilize latent concepts in in-context learning. The authors use a series of controlled tasks combined with mechanistic interpretability methods to provide both causal and correlational evidence, revealing the internal reasoning and representation mechanisms of Transformer models when solving ICL tasks.
weaknesses: The tasks and experimental setups are overly idealized, relying almost entirely on highly synthetic toy tasks, which do not represent real-world natural language reasoning tasks. Of course, this is a common issue in the interpretability field.

The analysis of model scale and generalization is insufficient. The paper only compares Gemma-2-27B and 2B, without systematically examining whether the same mechanisms hold across different architectures (e.g., LLaMA, Qwen series) or larger-scale models. It is also unclear why the experiments in Section 3 are conducted on a small GPT-2–style model rather than the previously used Gemma models, and what the motivation for this model substitution is.

The paper provides only empirical analysis. It would be more compelling if any theoretical or mechanistic explanation were offered for why Transformers naturally form linear geometric task vectors or bridge concepts.

Are the bridge concepts unique? The experimental design in the paper is overly simple, but in real-world scenarios, the relationships between concepts are much more complex. With only a few in-context examples, there may exist multiple potential bridge concepts between the source and target, a possibility that the paper does not discuss.

Another concern is how such concept-level empirical interpretability analyses can provide meaningful value for real-world applications. In other words, when would we actually need concept-level explanations? Even if we know which attention heads encode certain concept representations, how should we intervene or utilize them in practice? Could specific application scenarios be provided as examples?

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: * In the “Memorized world knowledge” thread - While the existing work uses 2-hop reasoning tasks, it would be interesting to study higher order transitive tasks (3 hops or more). It would be interesting to tease out the impact of model size  on such tasks. One potentially interesting finding could be that LLMs are able to deal with k-hops (k>>2) reasoning tasks much easily as compared to small language models (SLMs). It would be interesting to see, if LLMs indeed contain sparse attention heads which represent all the intermediate (bridge) concepts.
* The authors have used two variants of Gemma 2 model family - 27B and 2B. It would be helpful to conduct experiments using models of varying capacity - including a few with multilingual support, so that claims/findings can be strengthened further.
* In section 3, add-k problem, I did not understand the reason behind the constraint k_{i+1} - k_i = 3. What happens when we relax the constraint? Will the claims be still valid, if this constraint is not met?
* The authors mention that ideal scaling constant for bridge intervention is 2.0. My understanding is scaling is a way to boost the importance of the signal (bridge information). Is that correct? If that is so, it is understandable that the performance saturates at 4.0. But, what could be the reason for deterioration?
soundness: 3
strengths: Originality

* The unorthodox evaluation methodology of using activation patching with counterfeit examples is creative.
*  The technique of using interpolation (referred in paper as steering) in this setting is innovative and establishes the geometry of the underlying latent variable.
* The paper discusses the findings of previous studies in the literature survey section and clearly state the novel contributions of this work.

Quality

* Overall, I find the paper to be of good quality, the methodology is sound and the results are well presented. The experiments are conducted on a variety of datasets which demonstrates the robustness of the claims.

Clarity

* The paper is well written and easy to read. The ideas are explained clearly with supporting experiments. 

Significance

* This works brings us closer to understanding how LLMs represent intermediate concepts, perform chained reasoning in ICL setup.
* It uses tasks such as add-k, circular trajectory, rectangular trajectory to empirically prove that for numerical latent variable problems, tasks vectors indeed represent the geometry of latent variables.
* The paper strengthens the broader hypothesis that LLMs engage in multi-step reasoning rather than relying solely on direct input-output associations.
* The authors also highlight the influence of model size and the number of ICL examples on the latent structures.
summary: This paper studies the problem of whether and how transformers represent latent concepts as part of their computation in a an ICL setup.

Key ideas

* Use activation patching (a.k.a Causal Mediation Analysis) to determine if the model is indeed representing the latent concept in a discrete multi-hop setting.
* For the numerical latent variables, the authors propose the use of linear probing the intermediate model embeddings to assert the existence of task vectors.
* Use PCA and project the task vectors onto the first two principle components and then study the geometry of the resultant manifolds along with the ordering of the points. Using steering (or interpolation) based methods to show that model captures concept’s geometry.

Main contributions

* Use of causal and correlation techniques to show that for tasks requiring memorized world knowledge, a sparse set of attention heads are responsible for resolving intermediate latent concept, thereby rejecting the hypothesis that the model takes a shortcut (directly maps the input to the corresponding output without any intermediate steps). This shows that the model is indeed performing abstract reasoning (chained reasoning) in its hidden activations.
* Show how the above phenomenon plays out in case of smaller language models are used (2B, instead of 27B) and larger number of ICL examples.
* In alignment with the previous work on Linear Representation Hypothesis (LRH), the authors establish the existence of task vectors. In addition to that, the authors show that these representations reflect the geometry of the latent variable.
weaknesses: * In the “memorized world knowledge” task, the authors claim that the LLM relies on step-by-step composition of latent concept. However, all their experiments are of the kind where the latent concept has a one-to-one mapping with the output class (for instance country → capital). While, there is no problem with this choice, it might be helpful to study different scenarios (such as country (bridge) → famous politicians from that country)
* In Figure 10, for the circular trajectory problem, even with high value of beta (0.8), we find that MSE corresponding to original and opposite radius is nearly equal. Isn’t this counter-intuitive?

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: * Do you observe similar disentanglement patterns in larger, more naturalistic multi-hop datasets (e.g., HotpotQA, factual chains beyond two hops)?

* Are the same attention heads consistently responsible for “bridge” resolution across tasks or prompts, or does specialization vary by context?

* How might identifying latent concept circuits inform real-world model editing, debiasing, or steering applications?
soundness: 3
strengths: * The division between discrete and continuous latent concepts is conceptually clean and helps generalize across types of reasoning.

* Identifying sparse, interpretable sub-circuits for latent concept composition is relevant to understanding compositional generalization.

* The authors situate their findings well within recent discussions of task vectors, linear representation hypothesis, and in-context reasoning circuits.

* Demonstrating low-dimensional manifolds for latent numerical parameters (and the ability to steer them) is a compelling visualization of representational structure.
summary: This paper investigates how LLMs (specifically, Gemma-2-27B) represent and manipulate latent concepts when performing ICL. The authors combine mechanistic interpretability methods with controlled synthetic tasks to explore whether models infer and compose hidden concepts (e.g., countries, radii) in structured ways. The authors explore two major experimental setups: two-hop reasoning tasks and geometric tasks, and conclude that transformers naturally form disentangled, low-dimensional concept representations that are causally and geometrically interpretable.
weaknesses: * While synthetic setups isolate mechanisms effectively, they may oversimplify real-world linguistic reasoning. It’s unclear whether these disentanglement findings generalize to naturalistic contexts.

* The discrete reasoning experiments rely on Gemma-2 models, while continuous ones use small custom transformers. The gap in architecture and scale complicates unification. Moreover, would the discrete reasoning results replicate with other models?

* Activation patching is compelling, but causal claims might be overstated without broader ablation or counterfactual consistency checks.

* While qualitative visualizations (e.g., PCA plots) are persuasive, more statistical rigor (variance, replicability, robustness to seeds) would add credibility.

* The discussion of how these insights could enhance controllability, generalization, or interpretability in applied settings remains mostly speculative.


Decision: Accept (Poster)