Causal and Interpretable Structures in LLM Compositional Tasks

Gurbir Arora; Toni J. B. Liu; Jiajun Bao; Raphaël Sarfati; Christopher J. Earls

cs.CLSep 28, 2026

Causal and Interpretable Structures in LLM Compositional Tasks

Authors: Gurbir Arora, Toni J. B. Liu, Jiajun Bao, Raphaël Sarfati, Christopher J. Earls

Organizations: Cornell University, USA · Goodfire AI, USA

Abstract

Large language models are able to solve tasks whose answers depend on not only individual input tokens, but also on relations among them. How is such relational information represented and processed across transformer layers? We study activations from ensembles of prompts that require inferring relationships between three tokens corresponding to a cyclic concept (months, hours, weekdays, and musical notes) to correctly predict the next token. Across model families (Llama, Qwen, Gemma, and Mistral) and cyclic concepts, we find a consistent layerwise progression in how the joint dependence among the tokens is geometrically organized and causally used: intermediate layers use a joint representation based on the inferred relationship between two tokens, while later layers use a joint representation associated with all three tokens to correctly complete the task. We also find other relationships between tokens that are geometrically structured but remain causally inert in the next-token prediction. Crucially, when taken together, these geometric and causal investigations reveal the representation-level mechanism that progressively organizes and composes the relational information to form the answer. More surprisingly, restricting the models to such causally relevant joint representations improves next-token prediction accuracy.

Figures & tables

Figure 1: Two-dimensional visualization using the first two principal components for the activation ensemble $\Phi^{r}$ and its decomposed interaction ensembles $H_{AB}^{r}$ and $H_{ABC}^{r}$ for the prompt template in Section 2.1 using Llama-3.1-8B. Here, $\gamma=B-A\mod 12$ and the expected answer is $D=C+\gamma\mod 12$ . Before fitting the principal components, we exclude prompts satisfying $A=B$ , $B=C$ , or $C=A$ (see Appendix B.3 ). The bottom left of each panel shows the coloring according to either $\gamma$ or $D$ . $H_{AB}^{r}$ gets organized by $\gamma$ in early layers. Until layer 17, $H_{ABC}^{r}$ is unstructured but becomes organized by $D$ at layer 18. See Section 3 for a quantitative analysis.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Domain	Prompt template
Llama-3.2-3B	months	Hi {name}, the {noun} is scheduled from {A} {dates} to {B} {dates}. This is the same as time from {C} {dates} to
	hours	Hi {name}, the {noun} is scheduled from {day} {A} o’clock to {day} {B} o’clock. This is the same time as from {day} {C} o’clock to {day}
	weekdays	Hi {name}, the {noun} is scheduled from {A} to {B}. This is the same as time from {C} to
	music	Hi {name}, on the musical scale, the interval from the note {A} to the note {B} is the same as the interval from the note {C} to the note
Llama-3.1-8B	months	Hi {name}, the {noun} is scheduled between {A} {dates} and {B} {dates}. This is the same time as between {C} {dates} and
	hours	Hi {name}, the {noun} is scheduled between {day} {A} o’clock and {day} {B} o’clock. This is the same time as between {day} {C} o’clock and {day}

Appendix

Table 1: Prompt templates used for each model and cyclic domain.

A,B,C

and the replicate variables vary as described in Tables 2 and 3 . N/A indicates domains excluded because the cyclic values do not tokenize to a common width.

Variable	Values
name	Each name should tokenize to one token. There are 194 approved names for each Llama, Qwen, and Mistral model, and 200 for each Gemma model from the list of 200 most common names of the last century ( Social Security Administration, n.d. ) .
noun	conference , meeting , forum , summit
dates	1 , $\ldots$ , 28
datesDouble	10 , $\ldots$ , 28
day	Monday , Tuesday , Wednesday , Thursday , Friday , Saturday , Sunday

Appendix

Table 2: Replicate variables used in the experiments.

Domain	Values of $A$ , $B$ , and $C$
Months	January , February , March , April , May , June , July , August , September , October , November , December
Hours	1 , 2 , 3 , 4 , 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12
Weekdays	Monday , Tuesday , Wednesday , Thursday , Friday , Saturday , Sunday
Music	C , D , E , F , G , A , B

Appendix

Table 3: Values assigned to the cyclic variables

A

B

, and

C

in each domain. Months and hours have cycle length

N=12

, while weekdays and musical notes have

N=7

Model	Performance metric	Months	Hours	Weekdays	Music
Llama-3.2-3B	Top-1 accuracy	$53.4\pm 2.2$ %	$78.5\pm 3.2$ %	$78.3\pm 2.5$ %	$70.9\pm 0.3$ %
	Top-3 accuracy	$72.2\pm 2.8$ %	$95.2\pm 1.1$ %	$96.8\pm 0.8$ %	$91.9\pm 0.5$ %
Llama-3.1-8B	Top-1 accuracy	$66.0\pm 3.9$ %	$79.9\pm 1.8$ %	$78.8\pm 1.1$ %	$79.3\pm 0.7$ %
	Top-3 accuracy	$84.8\pm 2.6$ %	$94.7\pm 1.5$ %	$98.6\pm 0.3$ %	$90.2\pm 0.4$ %
Qwen-2.5-7B	Top-1 accuracy	$62.7\pm 7.0$ %	N/A	$75.4\pm 2.4$ %	$89.5\pm 0.6$ %
	Top-3 accuracy	$82.5\pm 6.8$ %	N/A	$98.7\pm 0.8$ %	$99.7\pm 0.2$ %

Appendix

Table 4: Top-1 and top-3 accuracy by model and domain. Values are mean

\pm

standard deviation across 10 replicates, computed over all prompts, including coincident values of

A

B

, and

C

Ensemble	$\mathbf{0}$	$\mathbf{1}$	$\mathbf{2}$	$\mathbf{3}$	$\mathbf{4}$	$\mathbf{5}$	$\mathbf{6}$	$\mathbf{7}$	$\mathbf{8}$	$\mathbf{9}$	$\mathbf{10}$	$\mathbf{11}$
$H_{AB}$	7.445	5.875	3.899	3.568	3.144	2.551	2.391	2.307	2.595	2.700	3.188	4.023
$H_{BC}$	4.343	1.921	0.918	0.961	0.971	1.021	0.853	0.976	1.022	1.378	1.377	1.575
$H_{CA}$	6.448	3.013	1.311	1.457	1.478	1.624	1.476	1.487	1.263	1.220	1.589	3.189

Appendix

Table 5: Mean norm of each interaction vector categorized by

\gamma=B-A

\gamma^{\prime}=C-B

, and

\gamma^{\prime\prime}=A-C

, respectively, for a single replicate of the months prompt template at layer 15 of Llama-3.1-8B. The interaction vectors associated with zero difference have substantially larger norms than the remaining classes. This motivates the exclusion of coincident-variable prompts from the quantitative analyses in the main text.

Control	Prompt template
Non-cyclic	Hello {name}, the {noun} is scheduled between {A} {dates} and {B} {dates}. This is the same time as between {C} {dates} and
Off by 1	Hello {name}, the {noun} is scheduled between {A} {dates} and {B} {dates}. The month that comes after {C} {dates} is
List	Hello {name}, the talks are on {A} {dates}, on {B} {dates}, and on {C} {dates}. The second talk is on

Appendix

Table 6: Control prompt templates for Llama-3.1-8B.

Explore similar work

For What Reason? Interpreting Models' Encoding of Causation and Antithesis
Jul 20, 2026Abhidip Bhattacharyya, Shira WeinTransformer Architectures Causal
University of Massachusetts Amherst · University of South Florida
Words as Difference Makers: How Large Language Models Determine Causal Structure in Text
Jun 21, 2026Wolfgang PietschLocal Causal Structures Causal
tum
Large Language Models Develop Belief State Geometry In-Context
Sep 15, 2026Daniel Balcells, Andrew Jun Lee, Chirag Rastogi +3In-Context Learning In-Context
Independent · UCLA · UIUC +2

Jul 20, 2026, cs.CLJ/K move · Enter open · S save

For What Reason? Interpreting Models' Encoding of Causation and Antithesis

Abhidip Bhattacharyya, Shira Wein

University of Massachusetts Amherst · University of South Florida

Transformer Architectures Causal Reasoning Skills Contrast

Open paper AskPDF arXiv

Discourse relations provide document structure, critical to language understanding and enabling language model performance and ethicality. In this work, we investigate how instruction-tuned Transformer models (LLaMA and Mistral) encode discourse relations in English, with a particular focus on the contrasting relations of causation and antithesis. Framing the task as a next-token prediction task and applying a suite of interpretability techniques to test model internals, our findings show that certain early layers make predictive decisions at mid-sequence tokens, while some mid-level layers finalize their decisions closer to the last token. Most of the remaining layers primarily propagate earlier decisions rather than actively influencing them. Additionally, we observe that some layers exhibit a preference for one answer over alternatives, suggesting asymmetric representation of discourse-based reasoning.\footnote{Our code is available at https://github.com/abhidipbhattacharyya/causation_vs_antithesis}

Abstract

Explore similar work

For What Reason? Interpreting Models' Encoding of Causation and Antithesis

Words as Difference Makers: How Large Language Models Determine Causal Structure in Text

Large Language Models Develop Belief State Geometry In-Context