Are Transformers Able to Reason by Connecting Separated Knowledge in Training Data?
Authors: Yutong Yin
Zhaoran Wang
Venue: ICLR 2025 Poster (2025)
Keywords: Transformer; Chain-of-Thought; In-Context-Learning; Compositional Generalization
Primary area: foundation or frontier models, including LLMs

Abstract: Humans exhibit remarkable compositional reasoning by integrating knowledge from various sources. For example, if someone learns ( B = f(A) ) from one source and ( C = g(B) ) from another, they can deduce ( C=g(B)=g(f(A)) ) even without encountering ( ABC ) together, showcasing the generalization ability of human intelligence.  In this paper, we introduce a synthetic learning task, "FTCT" (Fragmented at Training, Chained at Testing), to validate the potential of  Transformers in replicating this skill and interpret its inner mechanism. During training, data consist of separated knowledge fragments from an overall causal graph. In testing, Transformers must combine these fragments to infer complete causal traces. Our findings demonstrate that few-shot Chain-of-Thought prompting enables Transformers to perform compositional reasoning on FTCT by revealing correct combinations of fragments, even if such combinations were absent in training data. Furthermore, the emergence of compositional reasoning ability is strongly correlated with model complexity and training-testing data similarity. We propose, both theoretically and empirically, that Transformers learn an underlying generalizable program from training, enabling effective compositional reasoning during testing.

Published as a conference paper at ICLR 2025
ARE TRANSFORMERS ABLE TO REASON BY CON- NECTING SEPARATED KNOWLEDGE IN TRAINING DATA?
Yutong Yin*
Zhaoran Wang"
ABSTRACT
Humans exhibit remarkable compositional reasoning by integrating knowledge from various sources. For example, if someone learns (BA((A) from one source and ( (C= g(B) ) from another, they can deduce C=g(B)=g(f()) ) even without encountering (BCC) together, showcasing the generalization ability of human intelligence. In this paper, we introduce a synthetic learning task, "FTCT" (Frag- mented at Training, Chained at Testing), to validate the potential of Transformers in replicating this skill and interpret its inner mechanism"" During training, data consist of separated knowledge fragments from overall causal graph. In testing, Transformers must combine these fragments to infer complete causal traces. Our findings demonstrate that few-shot Chain-of- Chain-of-Thought prompting enables Trans- formers to perform compositional reasoning on FTCT by revealing correct com- binations of fragments, even if such combinations were absent in training data. Furthermore, the emergence of compositional reasoning ability is strongly corre- lated with model complexity and training-tesse theoretically and empirically, that Transformers learn an underlying generalizable program from training, enabling effective compositional reasoning during testing.
data similarity. We propose, both
INTRODUCTION
Humans exhibit generalized reasoning ability that integrates knowledge from diverse sources. For example, if one learns ( B = =((A)) ) from one source and C= g(B) from another, they can deduce ( C= C=g(B)=g(((A) ) without direct exposure to ( ABCC) ). We formally define this capability as compo- sitional reasoning the skill to integrate discrete pieces of knowledge from multiple sources to form a coherent reasoning. even in the absence explicit examples connecting these picces during learn- ing. This ability is a manifestation of broader concept known systematic compositionality understanding and generating an infinite number of expressions by combining a finite set of known components and rules (Fodor& & Pylyshyn 1988 Chomsky 2002). Transformer-based large lan- guage models demonstrate signs of compositional reasoning by producing comprehensive content that includes elements not likely to co-occur within the training data, suggesting the emergence of general intelligence (Press et al. [022]] Zhou et al., [2022a] Bubeck et al., (203). However, the com- plexity and ambiguity of their natural language training and testing data make it hard to scientifically validate the compositional reasoning ability and explore the underlying mechanisms.
This paper validates the potential of Transformers in doing compositional reasoning on synthetic dataset and investigates the inner mechanisms eliciting such ability. Specifically, we address three key questions: 1) When are Transformers able to perform compositional reasoning by connecting fragmented knowledge in training data? 2) How do different training factors impact the emergence of this ability? 3) What internal mechanisms enable Transformers to develop this ability?
We first introduce the "FTCT" (Fragmented at Training, Chained at Testing) dataset, which we investigate the performance of Transformers to address these questions. This dataset simulates knowledge relationships through graph-like causal structures, where vertices represent knowledge points and edges represent the relationships between their values. Multi-step reasoning paths are rep- resented by chains consisting of connected vertices with values calculated by edges between them.
Northheeseem
code is in http://p///gghhhh....o/////r//// waai://eeeeeettt---aaaaannnnnnno Cha i ned- at- l'esting

---

Published as a conference paper at ICLR 2025
The training data comprises only fragmented knowledge segments (child chains of short lengths), while testing requires the model to connect these segments into long, complete causal chains. To examine the impact of prompting
compositional reasoning. we concatenate similar examples to enable reasoning with few-shot Chain-of- Chain-of-Thought (CoTT (Wei et al., 2022c) prompts. Detailed con- This approach evaluates the model's ability to
struction of the FTCT dataset is shown in Section 3] generalize learned knowledge to novel combinations.
Regarding question 1), we find that few-shot CoT prompts enable Transformers perform compo- sitional reasoning. Testing Transformers trained on FTCT revealed poor compositional reasoning abilities in zero-shot scenarios, but significant improvement with few-shot CoT prompts. Few-shot CoT examples provide Transformers with correct vertices order to imitate, while Transformers itera- tively deduce both vertices and their correct values based on preceding reasoning paths. Notably, the CoT prompts during testing time consist of complete reasoning paths- an cut-of-distriutton (OOD) scenario compared to the fragmented knowledge training data. The capability understand and utilize OOD CoT prompts indicates Transformers' compositional reasoning ability.
Regarding question 2), we examine the influence of data distribution and model complexity. Data- wise, compositional reasoning emerges as the similarity between training and testing data increases. In FTCT, this is measured by the relative knowledge ratio-the the ratio of child chain length in training data to complete chain length in testing data. A phase transition occurs when the ratio reaches 0.3, marked by notable performance enhancement. Model-wise, multi-layer attention mechanisms prove essential, with the compositional reasoning emerging a complexity of at least layers and heads. Single-layer Transformers fail to replicate vertices order from CoT prompts, and Multilayer Perceptrons (MLPs) struggle with sparse information of value relationships.
Regarding question 3), we find that Transformers develop compositional reasoning ability by learn- ing an underlying program during training. This program, comprising in-context learning and par- ent retrieving (detailed in Section (.I). minimizes both training and testing loss by capturing the common latent structure of fragmented knowledge and chained reasoning paths. Theoretically, we prove that Transformers have the expressivity to simulate such an underlying program. Empiri- cally, through attention heatmap plotting and linear probing (Hewitt & Manning| [2019] Clark 22019] [|1r Allen-Zhu & Li 2023), we provide evidence that Transformers simulate the underlying program through two mechanisms-inductee
heads and attention assignment. These mechanisms respec- tively facilitate in- in-context learning and parent retrieving.
To summarize, within the context of the FTCT learning task, this article answers the above three questions following:
1) Few-shot CoT prompting enables Transformers perform compositional reasoning by pro- viding the correct order of vertices to imitate, improving from poor zero-shot performance.
2) Compositional reasoning abilities emerge as the similarity between training and testing data increases (with a relative knowledge ratio 2 0.3) and require multi-layer attention (with minimum of 2 layers and 2 heads).
3) Transformers develop compositional reasoning by learning an underlying program during training, facilitated through induction heads and attention assignment, to integrate frag- mented knowledge into coherent reasoning paths.
2 RELATED WORKS
Step-by-step reasoning. Chain-of-Thoughe (CoT) prompting (Nye et al. 2021 Wei et al. 2022c b
Kojima et al. 2022) enables language models conduct step-by-step reasoning, significantly boost-
ing their performance in complex tasks like mathematical deduction and code generation (Cobbe
al. 202T, Suzgun et al. 2022 Zhou et al. 20226 Lee et al., 2023 Achiam et al., 2023 Romera-
Paredes et al. 2024). Our research emphasizes few-shot CoT prompting (Wei et al. 2022c), which
initiates reasoning by integrating CoT examples into prompts. Interpretability studies suggest CoT's
efficacy arises from models' enhanced expressivity via intermediate reasoning steps (Feng et al.)
2024 Li et al. 2024b a). Besides the expressivity perspective, we additionally examine CoT gen- eralization in out-of distrr-iisttibuttion (OOD) settings, showing few-shot CoT prompts can elicit correct reasoning even with previously unseen prompts. Another study (Prystawski et al.) 2024) evaluates data's role in CoT's capacity for generalized reasoning. Our FTCT structure draws inspiration from
2

---

Published as a conference paper at ICLR 2025
their Bayesian networks, additionally inserting contextual noise and complicating value relation- ships. While they focus on locality structure's 's impact on CoT efficacy, we investigate how various training factors influence compositional reasoning emergence and conduct an in-depth analysis of the mechanisms within Transformer structures that elicit such capability.
In-context learning. In-context learning (ICL) (Brown 2020 Garg et al. 2022 Min et al. 2022
Wei et al. 2023) enables language models to perform various tasks by interpreting examples within the provided prompt, without needing explicit tuning. This capability allows Transformers trained on our FTCT to emulate the order of vertices in few-shot examples. Several theoretical studies (Xie)
al. 2021
Li et al. 2023 Wang al., 2024 Hahn & Goyal 2023) Wies al., 2024; Zhang al.,
treat CL as implicit Bayesian inference. Another set of works (Dai et al. 2022 Von Oswald al. 2023 Akyürek et al.) 2022 Ahn et al. 2024) argue that ICI. functions similarly to gradient descent from function approximation perspective. Notably, mechanism within Transformers, known as "induction heads" (Elhage et al. 202T] Olsson et al. 2022] Bietti et al. (2024). identified
2023
as the direct cause of ICL capabilities (detailed explanation is in Section 6. We demonstrate that induction heads exist in Transformers trained on FTCT, enabling the model to replicate the order of vertices from few-shot examples through ICL.
Compositional generalization. Compositional generalization refers to machine learning models' ability to solve new problems by integrating known components from training data (Lake & Baroni 2018 Keysers et al. 2019 Kim & Linzen 2020 Hupkes et al., 2020). Prior works have validated the potential of Transformers in compositional tasks where answers
directly output without in-
termediate reasoning steps (Hupkes et al.) 2020 Arora & Goyal 2023 Yu et al. [023] Xu et al.) 2024 Treutlein et al. 2024p. In contrast, our FTCT dataset with deep causal structure allows explo- ration of explicit reasoning's 's impact on compositional generalization. While some empirical studies have shown that seep-by-step reasoning enhances large language models' compositional abilities on real-world tasks (Press et al. 2022; Zhou et al.) 2022a] Khot et al. [2022) the complexity of natural language corpora used by them complicates the scientific validation, which can be done credibly various by our synthetic data. Recent studies have explored Transformers' generalized reasoning
controllable synthetic tasks (Ramesh et al. 2023 Allen-ZZhn & Li 2023 Ye et al. 2024). In contrast, our FTCT task not only ensures controlled experimentation but also introduces measures of training- testing data similarity and establishes a distinct parent-child causal relationship. facilitating analysis of the underlying mechanisms in terms of data distribution and model structure. Although some studies have highlighted Transformers' shortcomings in achieving compositional generalization due structural constraints (Dziri al.) 2024 Peng et al.. 2024) or misleading statistical features (Zhang let al., 2022), our findings reveal that few-shot CoT prompting elicits compositional generalization in sequential reasoning tasks, indicating the importance of data structure and prompting skills.
Transformer expressivity. Transformers have been proven to exhibit a high degree of expressiv- ity in approximating universal sequence-to-sequeecce functions (Yun al., 2019). Turing machines (Pérez et al., 20199; Wei et al., 2022a). and programmable computers (Giannou et al. 2023). Prior research has demonstrated that Transformers can simulate various algorithms related to in-context learning (Dai et al. [2222] Von Oswald et al. [2023] Akyürek et al. 2022 Ahn et al. 2024), step-by- step reasoning (Li et al. 2024a| Feng et al. 2024), and causal graph inference (Nichani et al. 2024] Edelman et al.) 2024 Makkuva et al.) 2024). Building on this foundation, we explicitly construct a Transformer architecture that efficiently minimizes both the training and testing loss of FTCT. To streamline the proof, we adopt the associative memory perspective (Bictti et al., 2024), using this approach to construct the weight matrices as the sum of outer products of orthogonal embeddings.
FRAGMENTED AT TRAINING, CHAINED AT TESTING
We introduce the structure of FTCT dataset and corresponding training and testing loss, illustrating the reason why it measures the model's compositional reasoning ability. Figure demonstrates the data generation procedure.
3.1 CAUSAL STRCTURE
We represent knowledge relationships with directed graph g = (V,E), where knowledge points are simulated by vertices V, a subset of the alphahet Van ( (A, B, B....,2. a, by ..,,..,;, =). Each vertex in
3

---

Published as a conference paper at ICLR 2025
Fragmented at Training
Few-Shot Examples
Downstream Processing
Causal Structure (V,E)
H Van
contextual Information
merge
A. T00, Z, B,
10100
B-101,
H-1
A. 110. Z, 1. H,
A-110
2-1
H-5
B
A. 100. 3. B. 101. H.
value assignment
Imput label
102 2. 103,
processedinput
102
processed label
B-103,
Ptrain
+1
Chained at Testing
A, T00, B. T0T, C, 103
C-I A-100, B-101
C-103
A. 105, B. 106, C. 108
C-P: A-105, B-106 C-108
value assignment
102, 103, C. 105,
C-?: A-102/ B-103
C-105
Dtest
A, 100, B, 101, C,103
input
processedinput
processed label
+1 +2
Figure 1: Overview of the FTCT Data Generation Process. The generation begins with the introduc- tion of "Causal Structure", representing the relationships of knowledge points, The "Fragmented at Training" stage shows how shorter child chains with noise vertices are formed to simulate incom- plete knowledge during training. The "Chained at Testing" stage presents the longest chains used in testing to assess compositional reasoning ability. The "Few-Shot Examples" part demonstrates the concatenation of multiple sequences for both training and testing to enable few-shot learning. In the end, the "Downstream Processing" adapts sequences into natural language-like sentences for intuitive reasoning format.
has value q(v) from its associated VALS(v) C E. Relationships between knowledge points are represented by the edges set &. Each edge e (V1, 02) E & defines the relationship between the parent vertex U1 and the child vertex 1/ by the operation op(e), satisfying ((22) = op(e) o ((11).
We assume that op(e) only represents addition or subtraction operation like +a or bb. Multi-step reasoning paths are represented by the chains T(g)= == { U2...., Un. n N, (vs, v:-1) E E). The depth of G, denoted as N, is the length of the longest chain in T(G). The "Causal Structure" in Figure T Dillustrates a causal structure with depth =3
3.2 DATA GENERATION
Step 1.1. Fragmented at Training: To simulate disconnected knowledge, training data excludes the longest chain in (9) and instead includes shorter child chains with length M < N, interspersed with M' noise vertices from V. Child chains vertices merged with noise vertices, preserving their order. Child chain vertices receive values based on their edge operations, while noise vertices get random values. The final vertex-value sequence is formatted Figure | the training data includes
seq ;= [t1 9...... Um 9m]]
where m = M + M'. As shown in "Fragmented at Training"
sequences like [A. 100, Z, 3, B, 101, H. 11, where [A, B] is a child chain from T(9) with values following the "+1" operation, and [Z, H] are sampled from Van - V with randomly assigned values.
Step 1.2. Chained at testing: To test models' compositional reasoning ability, the testing data consists of longest chains (length N) from T(G) without noise, formulating the sequence seq :=
[U1, Q ...,, UN, ),), where (VI...., E (G). Refer to "Chained Testing" in Figure [
Step 2. Few-shot learning: For both training and testing datasets, multiple sequences with the same vertices order are concatenated into few-shot document, which is formatted as
doc = (), \n...., \n (A) where seq [u1.91
UL 9
where L is the sequence length which can be either m N, and is the shots number ranging from 0 to K. The k-shot input and label are formatted as
inp k == d + [V1, qi (K-1))
(k+1) lab := [U2, 92 ..... UL 9L
(k+1)
The model should generate R autoregressively from * Especially, the zero-shot input inp
requires reasoning without any preceding examples. See "Few-Shot Examples" in Figure
4

---

Published as a conference paper at ICLR 2025
Step 3. Downstream processing: This process adapts sequences into natural language-liks sen- tences by adding punctuation, contextual details, and stating the reasoning goal upfront. As illus- trated in "Downstream Processing" in Figure L few-shot document like [A, 110, Z. 1, B, 111, H, 5, \n, A, 102, Z, 4, B, 103, H, 9] transforms into the sentence
?H=?: A=110, Z=l.... B=111, H-5 in H=?: A=102, Z=4, ... B=103, H=9
processed input
processd label
with "..." indicating tokens' context. The processed input and label are denoted as inp and lab"
For brevity, we define Drrin the distribution of inputs and labels generated by Step 1.1, 2, 3, and Dies as the distribution of inputs and labels generated by Step 1.2, 2, 3. The detailed data generation with specific sampling methods is in Appendix B
3.3 TRAINING AND TESTING Loss
Training loss:
For any input and label sampled from Duain. we train the language model to autore- gressively generate label given input. With the length of label lab defined as dh, the training loss is formatted as
K- 1
Lirin -E lnb,ing- Dinin
Σ k=0
log (P (moodal (iab;+1 inp & + iabi.)))
Testing loss: For any input and label sampled from Dieas given the input, test how well language model can generate sentence having the same vertex-value pairs as the label. Specifically, We define a decoding function dec(lab) := lab [V2, 02...., UN, qN] to decode the vertex- value information from the processed label. Given the input, the sentence generated by the model is defined as model((inp) We measure model's testing loss with k-shot prompt by
Lies := E Inp I25"" 1 [deo( Inp "oooo((as")))
Transformers trained on Puain have not seen the complete chains in Des from either the reasoning paths the few-shot examples. As result, the fact that they still achieve low testing loss indicates the emergence of the compositional reasoning ability. The empirical version of training and testing loss are Appendix A
EMPIRICAL FINDINGS
4.1 FEW-SHOT COT PROMPTING ENABLES COMPOSITIONAL REASONING
Our empirical findings highlight the essential role of few-shot CoT prompts in enabling compo- sitional reasoning in Transformers during testing. We evaluate model's compositional reasoning ability using the following criterion
Whole chain accuracy: Mcasures the model's generation contains all vertices and values along the reasoning chain in a correct order. For (inp". .lbb" k
sampled from Dist. measures whether dec dec(modell(1 (inp contains all elements from dec " in a correct order.
Further, we decompose the compositional reasoning ability into two sublevel abilities-thh ability of generating correct vertices order and the ability of deducing correct values given preceding paths, which are evaluated respectively by these criteria:
Testing vertices accuracy: Measures if the model correctly outputs all vertices in dec dec(iab") )
Testing values accuracy: Measures if the model outputs correct values intermediate vertices, given correct preceding reasoning paths sampled from Diest For the causal structure in Figure I this is tested by prompting models with sentences like " A=100, B=" The model is considered to output accurate values if and only if it outputs "101" as the next token.
We also assessed final value accuracy (Appendix H.I which considers the model correct if it outputs the correct value of the last vertex. Models' performance under such metric mirrors that of whole chain accuracy, underscoring the necessity of proper reasoning paths for correct final answers.
5

---

Published as a conference paper at ICLR 2025
We trained 3-layer 3-head GPT-2-like Transformers (details in Appendix G FTCT training set with varying graph depths and child chain lengths. 2 (left) shows the testing performance for Transformers trained on graph depths N 5, 10, 15, with k-shot CoT prompting (k from 0 to 4). Different curve colors represent different child chain lengths M = 2,3, 4, 6. Our conclusions are:
Few-shot CoT enables compositional reasoning by revealing correct vertices order. Whole chain accuracy is low with zero-shot prompts but increases sharply with more shots. At zero-shot, values accuracy is optimal while vertices accuracy is near zero. This indicates that few-shot CoT prompts enhance models' compositional reasoning by revealing the correct vertices order. Notably, such order has not appeared in the training data. The ability to understand and imitate the OOD vertices order stems from Transformers' in-context learning via induction heads (
Transformer outputs correct values with OOD reasoning paths. High testing values accuracy shows Transformers' robust performance in deducing correct values with the OOD compositional reasoning paths. Such an ability ensures models to iteratively output correct values of the vertices generated by few-shot COT, leading to correct reasoning paths. This ability stems from Transform- ers' parent retrieving mechanism brought by proper attention assignment (Section 6
An adequate number of shots is necessary. Performance improvements plateau or decline after one-shot examples, possibly due to increased dissimilarity between training and testing data with more CoT examples. A detailed discussion is in Appendix K
child_chain_enn- child_chain_laee 3
child chlidchhii_eeee Jen
child chain, keri
Depth=5
Depth=10
Depth=15
1.0
1.0
10
Whole acc
1.0
Whole
0.5
0.0
0.0
0.2
6.3
0.4
0.5
0.6
1.41
1.0
1.0
Vertices acs
1.0
Vertices
0.5
0.5
0.0
0.0
0.0
0.2
6.3
6.4
0.5
0.6
1.0
Velues ace
1o
0.5
0.5
B.
0.0
03
Shots Num
Shots Num
Shots
Relative knowledge ratio
Figure 2: Left: Zero and few-shot testing performance of Transformers trained on FTCT with various causal depths and child chain lengths. Right: Relationship between the relative knowledge ratio and model's compositional reasoning ability. Each shows testing performance under a specific criterion: The 1st row indicates whole chain accuracy (correct prediction of all vertices and values), the 2nd row shows testing vertices accuracy (correct order of vertices), and the 3rd row shows testing values accuracy (correct values of vertices with correct preceding reasoning paths).
4.2 THE SIMILARITY BETWEEN TRAINING AND TESTING DATA DETERMINES THE EMGERGENCE OF COMPOSITIONAL REASONING
For each FTCT task, we measure the similarity between training and testing data by the relative knowledge ratio À = M/N, where M is the child chain lengths and N is the causal graph depth. We find that compositional reasoning emerges as À increases. Figure 2 (right) illustrates the re- lationship between À and model's compositional reasoning ability. For each X, the compositional reasoning ability is measured by the optimal few-shot testing performance of Transformers trained on tasks whose relative knowledge ratio is X. A phase transition occurs: compositional reasoning remains weak when À < 0.3 and distinctly emerges when À 2 0.3. In essence, a larger À makes few-shot CoT prompts more similar to training data, thereby enhancing testing performance. How- ever, the fact that testing accuracy approaches one with À = 0.3 ratio significantly smaller than underscores the non-triviality of our results.
Experiments on larger models like GPT-2-small (12 layers 12 heads) and GPT-2-large (36 layers 20 heads) show the same pattern (Appendix [} demonstrating the generalizability of our conclusions.

---

Published as a conference paper at ICLR 2025
Training
Testing
Model
Params
Valid
Vertices
Values
Vertices acc
Values acc
Whole
3L3H
S4M
1.446
1.000
1.000
0.996
1.000
0.996
2L2H
48M
1,433
1,000
1.000
1.000
1.000
1.000
TF
2L
48M
1,436
0.870
1.000
0.768
1.000
0.766
1.559
0.138
1.000
ILIH
42M
1.624
0.025
0.974
0.026
1.000
0.022
4L
145M
3,464
0.063
0.104
0.484
0.290
0.002
MLP
21.
143M
2.147
0.065
0.146
0.850
0.346
0.022
Table I: Performance of different models trained FTCT task with causal structure depth N and child chain length upper bound M = 3. "Params" stands for numbers of parameters of dif- ferent models; "Training" column includes three criteria described in Section 4.3 measuring the in- distribution performance;"Tsstigg ing the compositional reasoning performance. We only show the performance of models prompted by CoT with the optimal shots number.
column includes three criteria described in Section 4.1 measur-
4.3 MULTII-LAYER ATTENTION MECHANISM ENABLES COMPOSITIONAL REASONING
We show that compared with other simpler structures, multi-layer Transformers excel imitating vertices order from few-shot examples and deducing correct values with preceding reasoning paths, leading to outstanding compositional reasoning ability. In addition to compositional reasoning met- rics (Section 4.. we introduce the following criteria to evaluate in-distribution performance:
Validation loss: Tracks the validation loss during training.
Training vertices accuracy: Assesses how well the model imitates the vertices order from the few-shot examples sampled from Drrain
Training values accuracy: Measures if the model outputs correct values vertices from child chains, given preceding reasoning paths sampled from D
We assess the performance of various models the FTCT task with causal structure depth of and a child chain length of 3. Models include Transformers (TF) of various layers and heads, and multi-layer perceptrons (MLPs) of different depths (details in Appendix G). Table summarizes the results. For brevity, we only show the performance of models prompted by CoT with the optimal shots number.
Depth of Transformer enables the imitation of vertices. Table [ indicates that Transformers with at least 2 layers and heads achieve optimal in-distribution and compositional reasoning per- formance. As complexity decreases, performance deteriorates, notably with significant drop in the vertices accuracy, both training and testing. while values accuracy remains optimal. Such phe- nomenon indicates that depth in Transformers is crucial for imitating vertices order in few-shot examples, hence enhancing compositional reasoning. This is because induction heads for in-context leamiigg are less likely in single-layer Transformers (Section
Attention mechanism enables the deduction of sparse values information. For MLPs with ap- propriate window sizes (details in Appendix G.2), both the training and testing values accuracy remain low. Conversely, even the simplest Transformer (I layer, head) achieves nearly optimal values accuracy, suggesting that MLPs struggle to capture sparse value information in noisy texts as effectively as Transformers. Interestingly, MLPs perform well in generating vertices order during testing but during training, possibly due the extra noise vertices in the training data, suggesting a different knowledge memorization approach that warrants further study.
5 TRANSFORMER DOES COMPOSITIONAL REASONING VIA THE UNDERLYING PROGRAM
As discussed in Section 4.3 the multi-layer attention mechanism of Transformers is crucial for compositional reasoning. However, how Transformers achieve this ability through training remains

---

Published as a conference paper at ICLR 2025
unclear. In this section, explain this mystery by showing the capability of Transformer in learn- ing an underlying program that minimizes both training and testing loss. In Section 6 we provide empirical evidence of this underlying program in the Transformer's hidden states. We focus 1 and modified training loss accounting only for shots number the few-shot testing loss { k=1 2 1, which defined as
T Tarain lab, Inp Duss
Σ log (P moddl (labr+1 iinp + labj,)))
5.1 UNDERLYING PROGRAM
We construct text-generating program which provably achieves optimal performance on both train- ing and testing data. Key components are summarized here, with an algorithm (Algorithm) | in the Appendix. Given any input sentence 21:7 that contains at least 1-shot example, the program executes the following two parts iteratively:
In-context learning: If the last token žт is a comma ".", 2T-3 must be a vertex ee. The program identifies all e; in the previous few-shot examples and attends to their next vertex 1/+1, returning U/+1's contextual tokens.
Parent retrieving: If the last token zr is equation token "=", ZT-1 must be vertex tij. The program retrieves the parent of vy from the preceding context. If u, belongs to the child chain from (() and has a parent Uù with value q; in the preceding context, the program returns 4 op("), 02) o qj with probability one. Otherwise, it returns value Øj randomly sampled from nj 's value set (w:)
For example, in the testing sentence shown in Figure D given the input
..:::..A:100.... ........B11....... C=103 \n A=
the program, through in-context learning, attends to "B" in the preceding example and outputs its contextual information the next token. The program continues generating tokens using other minor parts from Algorithm] I Jntil the output
"C=?: ..100....B111.,,
\n .= B="
The program then retrieves "A" as the parent of "B" by parent retrieving part, returning the value 106 105 + The Following lemmas show that the underlying program minimizes both the few-shot training loss and few-shot testing loss.
Lemma 5.1. For any sentence 21-T := inp + Labi., where input inp and label lab are sampled from Duain with k > 1, andn is an arbitrary position within the label, denote the distribution output by the program as Ppnu((1 Z.rr). We have that Ppra((: 21:T) = Poaiil 21:r) Hence, Pprig minimizes the few-shot training loss Linan (eq. 2
Lemma 5.2. For any input inp and label lab sampled from Dusst with k 2 1, denoting the sen- tence generated by the program as prog( ). we have dec(prog(inp prog(-) minimizes the few-shot testing loss ( K-1 (eq. (1)).
dec(lab") Hence,
Lemma 5.1 and 5.2 show that, despite the different distributions of fragmented training data and chained testing data, their next-token distributions given few-shot examples can be represented by a common program. As demonstrated Sections 5.2 and [ the Transformer learns this common latent structure through training, achieving the compositional reasoniong ability.
5.2 TRANSFORMER EXPRESSIVE ENOUGH SIMULATE THE UNDERLYING PROGRAM
We prove that Transformer is expressive enough to simulate the underlying program by explicitly constructing a 2 layers Transformer. Representing the model parameters by 0, we state
Lemma 5.3. There exists a layer Transformer with parameters that approximates Algorithm with arbitrarily small error.
8

---

Published as a conference paper at ICLR 2025
Layer
E= 138.
A= 141
o= 135
130
D 138
(n... E= 125
Layer
R
138
A= [41
125
Layer
130
138
1n
R= 125
Layer
E
138
A=
141
135
= 125
A
128
Layer
135_
138
(n... E= 125
128
122
Layer
E
138
A 141
135
1= 130
E
125
128
122
Layer 2
130
D
138
\n.. B 125
A 128
122
Layer
A= 141
135
130
D 138
E= 125
128
122
Table Attention weights of induction heads the sentence sampled from test dataset.
O-shot prompt
I-shot prompt
2-shot prompt
3-shot prompt
4-shot prompt
Data
Parent
Others
Parent
Others
Parent
Others
Parent
Others
Parent
Others
Depth 5
train
78.6
11.6
79.3
10.3
81.9
9.3
80.8
9.5
77.9
10.2
Child 3
Test
78.3
4,3
77.1
4,5
78.8
4.5
79.8
4.1
77.8
4.5
Depth 10
train
77.6
10.3
80.8
10.0
79.3
9.7
78.3
9.5
80.5
9.4
Child
81.5
79.3
2.6
78.4
2.7
78.8
2.5
74.6
2.5
Depth 15
train
80.6
8.8
83.8
9.4
80.6
10.0
81.8
10.9
79.7
9.6
Child 5
Test
78.2
2.4
80.3
2.4
78.8
2.1
76.9
2.9
73.9
2.5
Table 3: Probing accuracy in predicting the replaced values of different vertices. "Parent" column records the accuracy in predicting values of parent vertices while "Others" column records the ac- curacy in predicting values of non-parent vertices.
The ability of 2-layer Transformer to simulate the underlying program aligns with the empirically observed performance in Table [I] By expressing the training and testing loss as functions of the model parameters we summarize Lemmas 5.1 5.2 and 5.3 into the following theorem.
Theorem 5.4. There exists a Transformer model parameterized by 6* that satisfies
[Trmii (0*) - ming Liman (0) < where is an arbitrarily small value, (0r) O, where .....,A.
Theorem 5.4 shows that the parameters 0* approximate the minimizers of the few-shot training loss, indicating they can be learned from fragmented knowledge. Additionally, 8 minimizes the few-shot testing loss, reflecting strong compositional reasoning ability. The reason why 6* (approximately) minimizes both the training and testing loss that it simulates the underlying program common to both training and testing data. It noteworthy that the theorem only establishes the existence of such program- program-simullting 0*, However, empirical evidence in dectioon@ddmmonstraii that the hidden states of Transformers trained in practice exhibit patterns corresponding to the underlying program.
EMPIRICAL EVIDENCE OF THE UNDERLYING PROGRAM
We present empirical evidence showing that Transformers are simulating the underlying program through two cachaansms-inducciie head and attention assignment, which respectively facilitate the in-context learning and parent retrieving.
6.1 INDUCTION HEADS
By plotting Transformer's attention heat map, we provide empirical evidence showing the existence of induction heads that enables in-context learning. As described in previous works (Elhage et al.) 2021 Olsson al. 2022). induction heads two heads of the Transformer in different layers that collaborate to copy patterns. For example, with input sentences like "...[A][P [A][]].. [A]", the first
head in shallow layer copies information from the first [A] to [B], while the second head in

---

Published as a conference paper at ICLR 2025
deeper layer recognizes [B] and retrieves its context from [A], guiding the model to output [B] as the next token. In our task, we discovered similar induction heads operating in a slightly different manner. Given imput sentence formatted as (for clarity, we highlight comma tokens at different positions with boxes and different colors):
044 9.... 1/+1 = 4++ e
The head in the shallower layer copies the information of U; and 01+1 The head in the deeper layer attends "," along with the information of U;+1 to ".", making the model to output the contextual information of 01+11
To empirically demonstrate this pattern, we trained 3-layer, 3-head Transformer the FTCT task with a causal structure depth of and a child chain length of generating attention heatmaps for each layer. Complete heatmap plots are available in Appendix [L, and abbreviated version is shown in Table which which displays the average attention weights of heads in different layers for their respective tokens. For each comma in the black frame, the distribution of its attention weights to preceding tokens is shown using colored boxes- the brighter the color, the more attention paid. In Layer 1, each comma attends to its previous two vertices, recording their information in its hidden state. In Layer 2, each comma uses this information to identify the preceding comma whose vertex next to be output.
6.2 ATTENTION ASSIGNMENT
By linear probing (Hewitt & Manning, 2019 Clark 2019 Allen-Zhu & Li (023), we empirically show that the parent retrieving is facilitated by proper attention assignment focusing on the value of parent vertex while ignoring others.
For each sentence sampled from either the training or testing data, we identify the equation token "=", where its corresponding vertex has a parent in the preceding context. Specifically, we examine input formatted as:
U/ Lz ...
where Uy is the parent of uy- We construct the probing dataset by each time picking a position i < (including j1), replacing qi with randomly sampled q/. and recording the Transformer's 's hidden state for this modified sentence. We train linear function (details Appendix G.3) predict of from the hidden states. If the Transformer attends to q. the linear function should predict of with high accuracy. not, the accuracy should be low.
Table shows the results for 3-layer, 3-head Transformers trained on multiple FTCT tasks. For sentences sampled from training or testing data, linked to prompts with shot numbers 0-4, the prob- ing accuracy of predicting replaced values is tested. The results demonstrate that for both training and testing data, the probing function achieves high accuracy predicting parents' replaced val- ues, while showing low accuracy for other positions. This indicates that Transformers successfully retrieve parent vertices within fragmented knowledge, ignoring irrelevant information.
CONCLUSION
Our research validates the potential of Transformers in doing compositional reasoning on synthetic data and investigates the inner mechanism eliciting such ability. We demonstrate that few-shot CoT prompting enables Transformers to perform compositional reasoning by providing the information of correct order of knowledge points. We also find that compositional reasoning ability emerges when the training-testine -tstingg data similarity and the model complexity are above certain thresholds. We further show that Transformers develop compositional reasoning by learning underlying program during training, which minimizes both training and testing loss. This program leverages in-context learning and parent retrieving mechanisms, facilitated by induction heads and attention assignment.
Through experiments on synthetic data, we demonstrate the potential of Transformers to develop generalized reasoning skills, indicating that the impressive performance of contemporary large lan- guage models extends beyond mere memorization directly apply real-world models trained extensive natural language datasets, believe that our analysis offers valuable insights into the training processes and understanding today's 's large language models.
vast data. While our conclusions may not
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303, 08774, 2023.
Rishabh Agarwal, Avi Singh, Lei M Zhang, Bernd Bohnet, Luis Rosias, Stephanie Chan, Biao Zhang, Ankesh Anand, Zaheer Abbas, Azade Nova, et al. Many-shot in-context learning. arXiv preprint arXiv:2404. 11018, 2024.
Kwangjun Ahn. Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to imple- ment preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm in-context learning? investigations with linear models. arXiv preprint arXiv:2211. 2022.
Zeyuan Allen -Zu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint wrxv:22355.3377. 2023.
Sanjeev Arora and Anirudh Goyal. theory for emergence complex skills language models. arXiv preprint arXiv:2307. 15936, 2023.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: memory viewpoint. Advances Neural Information Processing Systems, 36, 2024.
Tom B Brown. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Ka- mar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.12712 2023.
Noam Chomsky, Syntactic structures. Mouton de Gruyter, 2002.
Kevin Clark. What does bert look at? an analysis of bert's attention. arXiv preprint arXiv:1966.04341 2019.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2 10. 14168, 2021.
Damai Dai, Yutao Sun, Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers. arXiv preprint 1iXXi:2222210555 2022.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, et al. Faith and fate: Limits of transformers on compositionality. Advances in Neural Information Processing Systems, 36, 2024.
Benjamin Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evolution of statistical induction heads: In-context learning markov chains. arXiv preprint 1rXiv:2402.11004 2024.
Nelson Elhage. Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1(1):12, 2021. 2021.
Guhao Feng, Bohang Zhang. Yuntian Gu, Haotian Ye, Di He, and Liwci Wang. Towards revealing the mystery behind chain of thought: theoretical perspective. Advances in Neural Information Processing Systems, 36, 2024.
Jerry A Fodor and Zenon W Pylyshyn. Connectionism and cognitive architecture: critical analy- sis. Cognition, 28(1-2):3-71, 1988.

---

Published as a conference paper at ICLR 2025
Shivam Garg, Dimitris Tsipras. Perey S Liang, and Gregory Valiant. What transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30599 2022.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason Lee, and Dimitris Papailiopoulos. Looped transformers programmable computers. In International Conference on Machine Learning, pp. 11398-11442 PMLR, 2023.
Michael Hahn and Navin Goyal. A theory of emergent in-context learning implicit structure induction. arXiv preprint arXiv:2303.07971,22 2023.
John Hewitt and Christopher D Manning. A structural probe for finding syntax in word representa- tions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Pa-
pers), pp. 4129-4138, 2019.
Dieuwke Hupkes, Verna Dankers, Mathijs Mul, and Elia Bruni. Compositionality decomposed: How do neural networks generalise? Journal of Artificial Intelligence Research, 67:757-795, 2020.
Daniel Keysers, Nathanael Schärli, Nathan Scales, Hylke Buisman, Daniel Furrer, Sergii Kashu- bin, Nikola Momehev, Danila Sinopalnikov, Lukasz Stafiniak, Tibor Tihon, et al. Measur- ing compositional generalization: A comprehensive method on realistic data. arXiv preprint 19Xi::1912.09713. 2019.
Tushar Khot, Harsh Trivedi, Matthew Finlayson, Yao Fu, Kyle Richardson, Peter Clark, and Ashish Sabharwal. Decomposed prompting: A modular approach for solving complex tasks. arXiv preprint arXiv:2210.024406 2022.
Najoung Kim and Tal Linzen. Cogs: A compositional generalization challenge based semantic interpretation. In Proceedings of the 2020 conference on empirical methods in natural language processing (emnlp), pp. 9087-9105, 2020.
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models zero-shot reasoners. Advances in neural information processing systems, 35:22199-22213. 2022.
Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sequennce-o--sqquence recurrent networks. In International conference on machine learning, pp. 2873-2882. PMLR, 2018
Nayoung Lee, Kartik Sreenivasan, Jason D Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic small transformers. arXiv preprini arXiv:2307.0338), 2023.
Yingcong Li, Muhammed Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference Machine Learning, pp. 19565-19594. PMLR, 2023.
Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting chain-of-thought: Compositionalitt through in-context filtering and learning. Ad- vances Neural Information Processing Systems, 36, 2024a.
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. arXiv preprint arXiv:2402. /2875, 2024b.
Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: A framework for principled analysis of transformers via markov chains. arXiv preprint arXiv:2402.04161, 2024.
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? arXiv preprint arXiv:2202. /2837, 2022.
12

---

Published as a conference paper at ICLR 2025
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735, 2024.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez, Maarten Bosma, David Luan, al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114. 2021.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 11895, 2022.
Binghui Peng, Srini Narayanan, and Christos Papadimitriou. On limitations of the transformer architecture. arXiv preprint arXiv:2402.08 2024.
Jorge Pérez, Javier Marinković, and Pablo Barceló. On the turing completeness of modern neural network architectures. arXiv preprint urXiiv/991.000
2019.
Ofir Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah A Smith, and Mike Lewis. Measuring and narrowing the compositionality gap in language models. arXiv preprint arXiv:2210.03350) 2022.
Ben Prystawski, Michael Li, and Noah Goodman. Why think step by step? reasoning emerges from the locality of experience. Advances in Neural Information Processing Systems, 36, 2024.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Rahul Ramesh, Mikail Khona, Robert P Dick, Hidenori Tanaka, and Ekdeep Singh Lubana. How capable can transformer become? a study on synthetic, interpretable tasks. arXiv preprint arXiv:2311.12997, 2023.
Bernardino Komera-Paredes: Mohammadamin Barekatain, Alexander Novikov, Matej Balog, M Pawan Kumar, Emilien Dupont, Francisco JR Ruiz, Jordan S Ellenberg, Pengming Wang, Omar Fawzi, et al. Mathematical discoveries from program search with large language models. Nature, 625(7995) 475, 2024.
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung. Aakanksha Chowdhery, Quoc V Le, Ed H Chi, Denny Zhou, et al. Challenging big-bench tasks and whether thain-of-thouuhtt can solve them. arXiv preprint arXiv:2210.0926.1 2022.
Johannes Treutlein, Dami Choi, Jan Betley, Cem Anil, Samuel Marks, Roger Baker Grosse, and Owain Evans. Connecting the dots: Llms can infer and verbalize latent structure from disparate training data. arXiv preprint arXiv:2406. 14546, 2024.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174, PMLR, 2023.
Xinyi Wang, Wanrong Zhu, Michael Saxon, Mark Steyvers, and William Yang Wang. Large lan- guage models are latent variable models: Explaining and finding good demonstrations for context learning. Advances in Neural Information Processing Systems, 36, 2024.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: case study on approximating turing machines with transformers. Advances Neural Information Processing Systems, 15:12071-1208. 2022a.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yo- gatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilitics of large language models. arXiv preprint arXiv:2206.0760 2022b.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi. Quoe V Le, Denny Zhou, et al. Chain-of -f-ho-hf- prompting elicits reasoning in large language models. Advances in neural information processing systems, 15:24824-24888 2022c.
13

---

Published as a conference paper at ICLR 2025
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently, arXiv preprint 2XXi::2303303846 2023.
Noam Wies, Yoav Levine, and Amnon Shashua. The learnability of in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: State-of-the-ar natural language processing. In Proceedings of the 2020 conference empirical methods in natural language processing: system demonstrations, pp. 38-45, 2020.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit bayesian inference. arXiv preprint 1 2021.
Zhuoyan Xu, Zhenmei Shi, and Yingyu Liang. Do large language models have compositional abil- ity? an investigation into limitations and scalability. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024.
Tian Ye, Zicheng Xu, Yuanzhi Li, and Zeyuan Allen-Zhu. Physies of language models: Part 2.1, grade-school math and the hidden reasoning process. arXiv preprint arXiv:2407.20331 2024.
Dingli Yu, Simran Kaur, Arushi Gupta, Jonah Brown- CoonnCCohen, Anirudh Goyal, and Sanjeev Arora. Skill-mix: A flexible and expandable family of evaluations for ai models. arXiv preprint arXiv:2310. 17567, 2023.
Chulhee Yun, Srinadh Bhojanapalli,, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequenee functions? arXiv preprint arXiv: 2019.
Honghua Zhang, Liunian Harold Li, Tao Meng, Kai- Wei Chang, and Guy Van den Broeck. On paradox of learning to reason from data. arXiv preprint 12Xx::2205.1100,, 2022.
Yufeng Zhang, Fengzhuo Zhang, Zhuoran Yang, and Zhaoran Wang. What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization. arXiv preprint arXiv:2305. 19420, 2023.
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuur- mans, Claire Cui, Olivier Bousquet, Quoe Le, et al. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205.1062. 2022a.
Hattie Zhou, Azade Nova, Hugo Larochelle, Aaron Courville, Behnam Neyshabur, and Hanie Sedghi. Teaching algorithmic reasoning via in-context learning. arXiv preprint 1xXiv:22/1.00000 2022b.
14

---

Published as a conference paper at ICLR 2025
A EMPIRICAL LOSS
A.1 EMPIRICAL TRAINING Loss
With B samples in total, at each sample step b we sample {inp
- k.b
,lab )
-- k.b K.B
1= from D Dunin-
With the length of label lab k,b
defined as d#.b, the empirical training loss formatted as
B K
1
k.b k.h log (Pmodel (lab inp + -abb;a))
k,b
Lirain =
B
A.2 EMPIRICAL TESTING LOSS
With B samples in total, at each sample step b we sample {inp k,b k,b K,B empirical testing loss with k-shot prompt is
.ab
1km1bb1 from Diest The
B
Lies:
1
Σ
[dec[ )
B DETAILED DATA GENERATION
Step O. Causal structure: The dataset is based on directed graph structures G = (V, E), where is subset of the alphabet vocabulary Val {A, B..., Z, a, b, b.... 2). Each vertex
Vanl has a value (vv) from its set VALS(v) C Z. For any edge e = (v1, V2) E &, U1 is the parent and 22 the child. Each edge e maps VALS(v1) to VALS(U2) via op(c) VALS(e1)
VALS(12), satisfying ((02) = op(e) o ((1).. We assume op(e) represents operations like (+a) or (-b). The set of chains T())= == ([11, 02..... Un] n E N. (ri, 0441) E E represents sequences of connected vertices. The depth of G, denoted as N, is the length of the longest chain in T(g).
The training data consists of child chains sampled from T(g) with length M < N as well as M' noise vertices sampled from Van V. Child chains are merged with noise vertices while maintaining their relative order. Within the merged sequence, vertices belonging to the child chains are assigned values following their defined operations, while noise vertices are given random values. Suppose the child chain is 81 := [u,,... uurrl E T(9) and the sequence of noise vertices is $2 == [u],". " The merged sequence is (U), Bm.] where m = M + M', [Mi,," ... Mi = 81 and [,),, U5 = 82. We then sample /q, 2m] as the corresponding values of vertices. For [q:,." 9im that corresponds to 81, follows the operators over the chain, where
Step 1.1. Fragmented at Training:
Unifor((VALS(u))) and 4. op((u=11 ux)) q for h = 2, MM.
For [9g, Dire that corresponds to 82, they are just random noise so we have
4h AiniformVVAL((C for 1.... M
The final sequence of vertex-valus pairs is formatted as seq [u1 91 ....... Um 9m], where m M + M'.
Step 1.2. Chained at testing: In contrast to the training dataset, the testing data consists of complete chains of length N from T(g) without noise vertices, formulating the sequence seq : W1,,41.... UN, [N], where (f),..., UN) E T()), qI ~ Uniform(VALS(r)) and 4s (p((1:11111))))) u;]) q for = 2.....
Step 2. Few-shot learning: For both training and testing datasets, multiple sequences with the same vertices order are concatenated into few-shot document, which is formatted as
doc == (s) /n...., \n, (*), where (e = [Ui, th (t)
UL,
Here L is the number of vertex-value pairs that can be either m or N, and k is the shot numbers, ranging from 0 to K. The k-shot input and label are formatted
i inp* t [vi, (01,g((11))
lab*
[12, 02 2.... Pt. dz
(k+1)
(6+1)).
15

---

Published as a conference paper at ICLR 2025
Given the imput inp", the model is expected to autoregressively generate the label lab by referring to the k-shot examples in doc Especially, the zero-shot input inp requires reasoning without any preceding examples.
Step 3. Downside processing: order to make sentences resemble natural language, we apply the following modifications: First, insert equation token 6a9 between each vertex e; and value d;- Second, insert comma token o between each value q, and next vertex 1/41- Third, insert contextual tokens [c(ve),,... d(v)),J before each vertex et. Lastly, define the question token 6 and insert [ur, d 0 to the beginning of the sequence, indicating the final goal to be reasoned. For each sequence seq [U1 91...................... UL 92], the processed form of it is like
seq = Luz o 00 t c(vi),,............)))11....)).. c(vi)t, U1, 0 91, 0
Uz, 03, qu).
Hes the processed input and Similarly, with the processed document defined as doc = {seq 1 label are formatted as
-
(k+1) 100l lab == (c(u2) (k+1) c((2)) (1+1) U2 o (k+1) 0 ott d(v)) (k+1)
k inp = doc + [Uz, 0 "0,00000000 80 c(v)) (++1)
....,,(vi)), c(vi)),
(k+1) (++1) o, 00,00 41
--- k
(k+1) d(vL)); uz, o 09.00 16
(6t2)
The contextual tokens are sampled the following way: For each vertex v Van, we maintain its context token set CONT(v). For any 01 * U2 regulate CONT(v:) n (V2) Ø. For each vertex v; in the sequence, its contextual tokens c(v))1. ... c(e;), are sampled from CONT(u,) without replacement, with 4, following the uniform distribution Uniform[[CON[[[x]]]]]
CONCRETE EXAMPLES OF SENTENCES IN FTCT C
We provide concrete examples of sentences in the FTCT, using the causal structure illustrated in Each vertex Vall associated with its own set of contextual tokens, represented as follows:
Figure [
A: ("intern", "accomplishmen". "whistle", "subsystem", "Rewards", "patents", "ARCH")
B: ("correl", "register", "044", "ask", "latex", "Coins", "google")
C: {"a", "exc", "increasing", "REAM", "she", "-", "operated") Z: ("Influ", "interf", "Ideally", "Pad", "adders", "confusion". "XX""
H: ("Higher", "fren", "romptu", "smoke", "shake", "Frank", "treasure")
v: ("ACH", "Catalyst", "pens", "emer", "4000". "arss", "easiest")
In this representation, [A, B, C] are knowledge points from V. and [Z, H. v] are random noise from V. The contextual tokens carry no semantic meaning. and the contextual token sets for different vertices do not overlap. The 3-shot training input containing [A, B] is like
"B=?: confusion Pad interf Ideallyadders XPZ=6, accomplishmentinten subsystemARCH
Rewards pacettss=112,rmmppuuH83338. Coins google correl latexB=123
B=?: confusion Pad Ideally XPZ=7, accomlismment patents subsystem RewardsARCHA=12..
treasureFrankH-33, htw.://min......ooooooogrieia/ea Coinsssk==155
B=?: Pad Influ interfZ=5.".
Its corresponding label is
"accomplishment subsystem whistle RevarrsseCCHiieeemm-444aadddddd latexxegitet googleB=150".
The 3-shot training input containing [B, C] is like
"C=?: corre1044 latexregister:==126,
C=?: Coins correlaskB=108, pens Catalystv= exc-iinrreasinC=1117
C=?: correl Coins googleregisterB=1144"""
Its corresponding label is
"Carsemerv=5,-. excoperatedC=113".
16

---

Published as a conference paper at ICLR 2025
Algorithm Generalized Reasoning
1: Input: Causal graph G (V, all vertices Vall. imput sentence 11:7, shot number f 2 1.
2: if 2T = - then
3: That means 2y-3 must be a vertex v4. There are f vertices v the previous sentence.
Find all in the previous sentence and pay attention there next vertex 1/41, deduce the contextual information CONT(U;+1). return MaiformmCONT(t;;,
4: else if 27 c(ug)x then
5: That means 27 1-T = [c(")),". c(ey)))- Pay attention to them and return
Uniform(CONT(()) + {0y} (c(t))w)k==
6: else if 2T - U/ then
7: return distribution P()) where P(o =
8: else if 2T 800 then
9: Pay attention to its corresponding vertex by = 2T-1-
if ny E then
11:
If ey is the first vertex from V that appears in the sentence, then return Uniform(VAt.S(r))) Otherwise, search previous tokens and pay attention to the nearest U/ L that fy V. Such fij must the parent of dy. Then return distribution P(.) where P(op(U)) 13) (,:)- =1
9/
else
13:
Return Unifor((VALS(v)))
14:
end if
15:
else if 2T 9 then
16:
Return distribution P(-) where P(o =
end if
The 3-shot testing input is like
"C=?: accomplishment patents Rewards subsystemA=134, Coins correlB=13,,- exc sheC=146
C=?: subsystem patents RewardsA=103, latex google correB==106.EEM sheC=115
C=?: patents accomplishment4=1124"
Its corresponding label is
"ask google044 correl CoinsB=127, slecicressingEEAMMMiiipttttCtCC-
D GENERALIZED PROGRAM
D.1 DETAILED ALGORITHM
The detailed algorithmic description of the generalized text generating program is shown in Algo- rithm 1
D.2 PROOF OF LEMMA 5.1
For any sentence 21:7. its last token æy must be one of the tokens we discuss the distribution of the next token in Dirain when 2T is any one of them.
g. By, 0 , Then
When 2T =m that means 27-3 must be vertex and the next token will be one of the contextual tokens 0:41- To determine which specific vertex e/+1 is, we can refer to previ- examples to find what is the next vertex after every v;. After determining U/+1, in D the next token follows the distribution of c(1;4 )- Because ((141)),. c(2;+ 10 are sampled from CONT( without replacement, c(v++1) deducing Prmin 21.7) = Uniform(CoNT((x:1)
Uniform(CoNT(c::
When 2n. = c(wy),- the previous h 1 tokens must be c(u,)......d ......... .......(()))))))))) (..)),....(()))))... 2 h)
Because
l is sampled from Uniform(1, [CONT(e;))), then have P(I, Uniform(h, (CONT(U,))) So with probability 1/([CONT(V;)) Uy with probability ([CONT(U;)] h))((CoNT();)) + 1) the next token sampled
+ 1) the next token is
17

---

Published as a conference paper at ICLR 2025
from Uniform(CONT(U) fe(u)),,... ddu))a)). Combining them together, have that Puain( 21:T) = Uniform(Contt (CONT(t) Ul4,} (-
When ²T = yy, the next token is 0 deducing ( =1.r) 1.
When 2r o=4, 27-1 must be a vertex and the next token must be value Qj. If deducing P 21:7) =
My E Vall - V. then qj sampled from Uniform(VAL.S(r)) Uniform Uniform(VALSS(n (VALS(vy)) )). If Uj V and bj is not the first vertex from V in this sentence, there must exist 8/ E V which is the parent of U; previously. ūj with probability one, deducing Poamnop(1))) f/y) 95 21:) 1. Otherwise, if vs is the first vertex from V in this sentence, then 4j is sampled from Uniform (VAL((())) deducing P.riin æ1:T) Uniform(VALS(u:))
this situation, qi = op(u,, vy)o
When ęj. the next token must be otmh deducing (o
21:T)
In any situation, Prog( 21:T) always equals to Purin 21:T).
D.3 PROOF OF LEMMA 5.2
& First write down inp and lab
seq [uz 03 00,000
d(i)),, 01, 08 ,0,,0 ,..,,,, "...,,,",,vvv,,,..... duu Ve, 08 ,0)]
doc [seq") :a...., |n, seq
inp k -kk
= doc + Mn] + 0
00 osh, c((,) (1+1) ,(.,, (++1) U1, 00 02,0000 91 v=m]
(++1)
lab [c(12) (k+1) ..... ,(..,, ((12) (at1) (1+1)
k
12 00 2/
(k+1)
om
0th
c(vL) (++1)
( c(uz))tt)
(k+1)
v 0 0 7p
(++1)
After decoding, have dec 1ab [V2, f2 have
(k+1)
uN IN
() According to the program,
Prog( inp Uniform(coNT(12)))
In this situation, the greedy decoding will randomly sample a c(u2) uniformly and concatenate the end of inp then we have that
P. linp c((2)), ... )()))))) Uniform(CONT(e)) + [2a}
It will keep sampling elements from CONT (u2) until it samples the vertex U2 Suppose that samples r times until it samples U2. Then we have that
arg max Prro((=
[inp
c(v2) (k+1) ... c(02)) o
(k+1) v2l) ~
arg max P Prrog (-
linp c(v2)) (k+1) c(v>) (k+1) , ) = 0/
(k+1)
[imp" ((v2)) (++1) cl
(k+1) u2 bo (k+
( cm
arg max Prog
42
Following such routine, the greedy decoding on algorithm | we finally sample
prog inp )
[c(v2) (((z)) .... ( (k+1)
U2, 0 02
(k+1)
... ((UN))(+t) (k+1)
((UN) (k+1)
2ell UN, 0 IN (k+1)]
So we have
dec(prog(ing) k. = [U2, (k+1) UN.44 (k+1)] = dec(lab") k
E TRANSFORMER ARCHITECTURE
Here construct simplified Transformer architecture that is convenient for theoretical analysis. Denote the set AA as the set all possible tokens that may appear in our task. Given any input sentence 21.T E AT, it is processed by the following structures.
18

---

Published as a conference paper at ICLR 2025
E. EMBEDDINGS
Token Embeddings. The token embedding Wtok A - Rddor is mapping that maps each possible token to a deok dimensional real vector (z)
Positional Embeddings. For each 24 that is the 7th token of the sentence, it has a dp dimensional positional embedding pr E 1Rdr.
Comprehensive Embeddings. The comprehensive embedding Woomp : A x [aaax] IR deapp
is mapping that maps both the token information and the positional information / to deomp dimensional vectore W (2,mp(2,))
To simplify the analysis, unify the dimensions of different embeddings as dj == deok dp doorp. Then we regulate that any two different embeddings are orthonorma F which means that for all et, 62 E (=)
U(PthetTmma U(Waome (z, = and ei * e2, we have that e c1 = and ej e2 = 0. Then for any token 24 at the ith position of the input sentence 21:7, its embedding
WE(2,t = (Weor( (2e) pt + W (2t, 0,.,
di
4d1
where the last 4d, dimensional zero vector is for writing in additional information. For simplicity, we denote 2₁ = WE( S,, t) E E where d = 5d,.
E.2 ATTENTION MECHANISM
For an L-layers H-heads Transformer, it has the set ((W value matrices, where w-Lhh Wilh whth E Raxd It also has output matrices Wo, (2( Wo, ([) and bias WB IBLAI+d2 to map the hidden state to the logits for output tokens. Here d2 is an extra dimension for information writing in MLP, whose specific value will be given in appendix F.3 For each embedded input sentence Ti.T, it goes through the following attention mechanism:
Q Whth W/ as its query, key,
(0)
Fi.
TTT
!I)
(1-1) 1)
+ Σ
with (1-1)
softmax (1
(whh) K
whlh (1-1) Ii
3t
t Zt
dis
H
(1-1)
attn (1.1)
1:t 1)) fort= =1, .....,,TII 1,..,,,...
= 2t
h=1
gout + rol (Wo wt Wn) (L (D), for l 1,...
(L)
Define the parameter vector 0 == (WWb W W U(Wo, Wo,: Wn). Given any input sentence 21:7, the output probability the Transformer with parameter a distribution over the indices set of all tokens A:
Po(: 21.T) = softmax(291t)
Here we assume that the temperature of every softmax layer here tends to infinity, meaning that softmax(x) = Uniform(arg max())))))) max(x)).
At test time, the Transformer does greedy decoding, which outputs the token with the largest prob- ability every time. Denoting the decoded sentence output by the Transformer as modelo(21.7) = = it follows
arg max Po (1:)), 2 = arg max Po [aur, 2(,1)) for i = 2...
In practice, the orthonormality can he approximated by randomly mapping embeddings high- dimensional Gaussian vectors with low variance (Bietti al. 2024
19

---

Published as a conference paper at ICLR 2025
F PROOF OF LEMMA 5.3
F.I FIRST LAYER
We directly give the construction. We set w La for all h, letting w plays the role of both key and query. Moreover, we set max as the maximum length of all possible input sentences. Referring to (Bietti et al. 2024), construct (W2") Jh= associative memory:
Tmax (pL-3 rtd (( cooomp(o (o't O1d,) ) + Tmax
(pTss Ord,) (
comp (o T ,0dd))
2x 2r /otts
+ Σ
(Wior(oRt)) Osd,) (W (v) 04d,)
Σ
(Wtox(o O4a, ) (Wook(a)) .44,)
Var fo0
zrEVa
=other
Tamax
[CONT(e)]
W Σ Σ
(W Woomp(22: 1) Ord,) TWWTTWoooooo (C1mm((1, 01d,)
DEVI e (2C0NNT())
l=0
=conieattall information
Tiax ii -4
1/ (Wcomp() Oad,) ) (o ,T) W Wecmp(o" (oalm, (oelm,tt), Oad,)
NV f== ==
Σ oEΛ- CONT(N) - (0 gan) veVai
(W Osd,) (oun)) (e (Wlox(a) Oad,)
mathers
T
W1.3 1.3
(WT11, Osd,) TT (W comp c0mp(000 (009,) () Osd,) +
(Wror(OR)) Osd,) (Weok T O4az)
acA (me))
T
11.4
Σ
(pT22 Osd,)" (W comp (q, (,,t) (.d,))
U VALS(e) vEVin
1=3
2r = values
Σ
(W (WLon(oTUT. , T (a)ox(a) O1d,)
oEA-
U VALS(E)
2y others
We also construct (W-
===1 as 1.1 Σ (Od, (eek(u) Ozd,) (T ((Wook()) Oadi) + (Osd. Weok(oP9)" ,0,, ((We) (Wook(o99)) ;sddi)
VEVi
Σ
Σ
(03d, TC Od,) (Weok(c) Oad,) + Σ
(02d,, WFok(W) O2d, ) (Wook(U)) 0ad,)
VEV. CECONT(w)
Tmua
1.1 Σ
(O,,, W Wo,mp(u,t) O2d,, (Lok(D) (Weomp comp (v, T) Osa,)
vE Van
Tmax
W Σ
(2ad, W Commp(0 (o,t) 02a,) (2,,)) (WWomg (v, f) 0ed,)
VE Van
20

---

attends to the sec- Published as a conference paper at ICLR 2025
Using the above keys and values to calculate the hidden states directly is little bit abstract. As result, we explain the intuition of our construction taking attn'. and atm 1.2 as examples. For any input sequence 21.7 and its embedding T1.r, first analyze attn "ttl."(20 (21:7):
If 2т = / it will be caught by the first two items in W and attend to its corresponding vertex 27-3 1y-1- The value matrix w maps the value vector of each vertex v Van to (Oa,: Wtok ,0a,)) Hence, we have that attn.: = (Oa,, W (1o-(1)-)) Ozd,
If E Val. it will be caught by the third item in W attending to the equa- tion token oF9. The value matrix WLll maps the value vector of each o in the sentence to (0a,, Wook(o*%) 0aa,)" Hence, we have that attn (x1r)) = (Od,, Wook ( O3d, )
If 27 is other type of token, it will be caught by the last item in W attending to the question mark 0Th The value matrix Wha maps t to zero vector (0sd,) T Hence, we have that attn 1.1 (I1rT) (0sd,)
Then we analyze attn (z1:T):
If 2T c(u)),, it will be caught by the first item in W attending the previous con- textual tokens that are also generated by vj including itself: {e(u))r The value ma- trix W maps every contextual token to (Ozd, (c) 0O,)). Hence, we have that attn 122 (21:)) = (Ozd, E/== (c())) ,0,,)
If 27 = /o it will be caught by the second item in wh? attending to all the pre- vious tokens that are also vertices except for the nearest one fj-1 27-3. Each token is multiplied by weight 12, the nearer the token is to zr, the higher the weight will be. Hence, as the temperature of the softmax goes to infinity, ond nearest vertex to it, which is denoted by Uj-2. The value matrix w maps the value vector of each vertex v E Vaul to (0g41, (e) 02a,) Hence, we have that attn 1.2 (21:T) = (024, Weok (u)-2) O2d,))
If 2T is other type of token, will be caught by the last item in W:1.. attending to the question mark 044 The value matrix w-1 maps 8 to zero vector (Osd,) T Hence, have that attn 2 (0sd,) (T.
The hidden states generated by 2. and 1:4 can be calculated following the same routine. Finally we can summarize that for any input sequence 21.7 and its embedding T1:7, we have
IT
11)
FT +
Σ attn-.1 (II:T)
h=1
1. (WTT, T) W (1j-1) Wror (v)-2) O2d,) 2T = osm/ooli and øT = U/-1
if 2, then Uj-2 2-222
is the last vertex of the sentence, otherwise it is the vertex before --
2.(WE(2T,11 T) O2d. -/ // W WzOK (c(v,) ) T 02,)T, if ²т c(U;))
3.(We(2r.T)
O2d, Wtoz
O2,
if if 2/ Wook(Dz)) (1z) T if 2T = t and 2T by
4.WFx22,TT)
W (2y, T 1) O2d, . W WEOK
5. ((Wwx2r, T)
Od, W Woomp (uz, 2) Ozd, )
if 2T = Gj and 5T-2 Uj-
21

---

Published as a conference paper at ICLR 2025
F.2
SECOND LAYER
Then we construct the second layer { W Ti First is the queries:
Tiax
Σ
Σ 1=1
:0d,) ( (Wcome (g, 1)T, O4d,)
qE U VALA(U) veVal
22
+ Σ (02d, Wtox(u),, 02d.) ) (Oa,, (o)k(u) 03d,)
VVVa
=y=o
Σ
, ,0od, (Weoz(a)" 0sd,)
aEA- VALS(e)- (00,000 veVii
Tama
w
Σ
Od, W tu, ,0ad,) (22,,)(O,))) (O.,,, Woomp 12) O3a, b (W (Woomplo .11 (o0mmp))
U1 E1 12=2212 11=1
2p=
Tiax
+ (Weok(o8U)) 0sdi)
Σ
Oa, W Woomp (v, Ozd, )
(Wrox(a) Osd.))
VEVuI 1 1=1
aeA- (ore)
Then is the values: W
Σ
(03d, (to _0d,) (la))(Wassl) (Wtox(1) .00d,) +
Σ
(Ozd, Wtok(N) .dd;) ((d,; (eok(v)T, 03d,
q8 U VALS(U) v
vEVai
Tmaa
w (O1d, V2W T) ( 10d,),
Σ
Σ
(Odd, W T) (W cacppaa T .0ssr)
U vALS(e) t=1
weVi
Taux
(Osd. W (T (O2d. Weomp(te (v, T) O2d,)
For any input sequence 21-r and its embedding 21,2, we first analyze attn?: attn²:((1::r)
If 2T = 9j. it will be caught by the first item in w attending to itself. The value matrix W maps each to the vector (03dys Wtok (q) Oa,). Hence, we have that attn (x) (1 (Ozd, (g;) T ,O,,).
If 27 = / / will caught by the second item W attends to previous token 27, whose value equal to 2r (which means is also comma delimiter). Denote M1-1 27 as the corresponding vertex of 2r, what is special is that 2T,-3 My- That means that 27, is the comma (or delimiter) of the next token of vy in previous shot. The reason why 2τ can attend to 27, is because of their hidden states:
T = (Wn(2,,T)) WtoK Wtok(Ug-1)) WeeK (1z-2) , )
T (We(>n, T.) W (1;) WEok (1g- O2d,)"
That means (x 1d.+1:2d; (.
which is caught by the second item in w Such a mechanism is known as induction head Giannou et al. (2023); Ols- son et al. (2022): Bietti et al. (224). The value matrix W maps each otT /o to (Oad, W Lok (1j-1) .O,))² where 0z-1 is its corresponding vertex. Hence, we have that
a.tm (1) (Oad, Weok(v;) T.
22

---

Published as a conference paper at ICLR 2025
If 2r is other type of token, it will be caught by the last item in W 2,1 attending to the question mark 0 The value matrix W.2. maps o to zero vector (0saz) ) Hence, we have that attn 2.1 (xi-r) (0s4,) ).
Then we analyze attn (x1:T):
If 2т 0=9, will be caught by the first item in W Denote the corresponding vertex of z7 as as 27-1. If By Val V, then it attends to the question mark U²h, being same as the second situation. If U E V, attends to the nearest delimiter 0 or the nearest value qn whose corresponding vertex U/ E V. If the nearest is سساmm that means vy has no parent in the sentence, we have that attn attn (1))) (x = (Oid, v2W ( )) If the nearest is li that means vy has a parent Uj the sentence, we have that attn 2.2 (11:7) (O4d, Wook (1,, Wtok Wtok(gn) T)T.
If 2r is other type of token, will be caught by the last item in W attending to the question mark 0 The value matrix W.2.1 maps T to zero vector (Osd,) ) Hence, have that attn 2.1 (1 (x (0sd,) T.
Q
To summarize, we have that
(2)
attn (2
(1)
4T
TT
Fi:T
h=1
1. T) Weok (U) 1j-11 W Wook (tj-2) 2)T W tok (u;) .0+,)
if 0 and 2T Uÿ-1
2. 2.We(((();) T) O2d, Σ
Eir= W
Weok (c( ) Od,
if 2т = c(u)):
3.(Wn(u;, T) ,02d, Weok (0
O4
if 2т Uj
(.E((H W comp (u,, T
1) 2d: W WLo
+ /2W cok (o4La))))
2T 000 and 2T-1 4.2 4.2.(W (.2 if 2f 800 and 2y is the parent of 2T-1 Vi
has no parent in the sentence
Wcomp (v,, TG,T-110
-1) -
O2d. (U)) + (eck(U;)) + Wtox((,)) (4,)))
5. 5.(Wn(g, T) Oa W comp (1/ T 2) W (4;) 00, T if 2r
T,
F.3 OUTPUT LAYER
The output layer is a two-layer MLP whose input dimension is 5d,, hidden dimension is |A|| + d2 and output dimension is |A.. Here the extra dimension d2 for mapping elements in the set
B VU
U
U ((1.v2.91))
(V1,02)EE qrEVALS(U))
so [8]. For simplicity of the notation, we here introduce shorthands of several |im-diiens vectors. There is an one-one mapping from the set of all tokens A to their indices [[A]]. Specifically, each token E A has its unique index idx(a) E [[A[]] Let ( be an [A]-dimensional vector where a 1 at position i and 0 at other positions. With a slight misuse of the notation, we define e(a) = (:)x(a)) for all
A. For any subset A' C A, we denote the uniform logit on it as : EeEA e(a). For the extra set B, we define g be dimeennnil vector where a position and O at other positions. Similarly, we define g(b) gidx(b) for all b E B.
Cumil (A') :
23

---

Published as a conference paper at ICLR 2025
Denoting the set of vertices in V which does not have parents as Vina. define the output matrices as
Wo:- = Σ (CuN(CONT(o))) T, O0a)) (Osds Wook Wook(u) Oas)
VEV
1.pp=o
(Cumi (CNN(CO)T))))" +o(u) - -(cc) .dds) (Osd, " 04,)
VaVal CECONT(1)
2.0r=c(v))
+ (e(( Od,) (03d, Weok (o Odg)
3.2TEV
Σ Σ ((1, (04ds Weok ( + Wtok (vi) + (tok(1i) (ook(a))
(21,12)E8 an (VAT.((M)
(TAA) (Osdi, ( /2
eEV
5.27=00 (Cuni (CniLVAA(())) Oa,) -EVaiV
T (Osd, T)
6.2y 6P
+
Σ
( Oo (Osd,, W WLok (
BE U VALS(U) REVal
7,2r=9)
WB = (0/A), -2d))
Wo, = Σ
e(u)(e(o) ,0,,)
Σ
e(c) (e(c), 0aa))
VEVA
(ECONI(
+ Σ
Σ
e(op(Ui, u2) (())0 u2,
(2) EVALS(U))
Emm (VALS(U))(
(())(e(g)T,o.) ,0s,
DE U VALS(H) VEVVll
+ e(o0 (e(o")) 0g2))
Here we only explain the procedure of deriving MLP(z (2) when 27 = dº9, the other situations are straightforward. Denoting 0y 2T as corresponding vertex, we have that
If E -an-VV V, it will be caught by item 6 in Wo, It then easy to show that MLP(2 (2) Guni Cuni(VASS (VALS(U;)).
If oy but has no parents in previous sentence, it will be caught by item 4, in Wo,, We can calculate that
Wo 3 T(())))
+
Σ
no (((
where μs <
bEB (m))
24

---

Published as a conference paper at ICLR 2025
That means all other g(b) will be filtered by the relu except for g(u)). Specifically we have that relu "i + WB)= = ((1 [T, Then it will be caught by the 4th item in Wo,. resulting in MLP 2 = cuni (VALS(v;)).
If Uy E V and has parent U/1- it will be caught by item 4,5 in Wo., We calculate that Σ
Wo, (2)
3(0 Uj, v,: +
(()))) TT,, where (46 S 2.
BEB- (1.1.
Then have relu(Wo, 2T (2) the 3rd item in Wo,, resulting in MLP(= e(op(u), 1j) og.).
Wn) - Uj. Ug.
(T)))) It will be caught by
In summary, we have that
if 2r 0 and 2T Vj 2. (eni (CON(((()))) + +(U;) = c(c(t))e)) 3.e(0=9),
1.Cuni (CONT(B))),
if zT duj))
if 2 1/
MLP )
4. 4.1.e(op()), 12)0 [b.).
2 = 80 and V/ = 2T-1 has parent V, in previous sentence
4. Cuni (VALS((;)
2f 8 and 1/ 2T E V but has parents in previous sentence if = rh and E Van
4.3.c Cui VALS((;))
5.e(0 = 4j
As the temperature of the softmax goes to infinity, we have that
softmax (MLP(2 (2)
1. U.Unffom (.niffrr((CNNT()))) )) =f = Tt and 2T -
2. Uniform (COOItmCCCCC((CC V +
if if = c(vr))
3.)( where (( = =f
4.1.P(.) where ((op(en,
wei..nnnioomaaaa.m ((op()).)))g)) =1, if = R and by 2T has a parent ey, in previous sentence
if 8 and 1j 2T-1 but has no parents in previous sentence
Uniformi VVLSS (g) if l and Uj Van V
5.P(.) where P(o =1 if 91
That exactly the same distribution given by algorithm | Depending on how high the tempera- ture of the softmax layer is, the Transformer structure can approximate algorithm | with arbitrary preciseness.
G IMPLEMENTATION DETAILS
G.1 TRANSFORMER
Our Transformer model is based on the GPT-2 architecture as implemented by Hugging Face (Wolf et.a.
2020 Radford al. 2019). We tested various combinations of layers and heads, as detailed in Section 4 The model features 720-dimensional embeddings and a context window of 2048 tokens. During encoding, we generate token indices directly and feed them into the model, adhering the mapping provided by GPT2Tokenizer for both encoding and decoding. The Adam W optimizer is employed with a learning rate of 5e-5.
G.2
MIP FOR LEARNING FTCT
Given that 1-shot CoT prompt suffices for optimal model performance and excessive input length may degrade training efficiency and accuracy, we use sliding windows to cap input length. We experimented with window sizes of 150, 200, and 300 tokens, selecting the size with the best per- formance. Each sentence represented by the concatenation of the one-hot encodings of its tokens, where each token is represented in a 503-dimensional space. Thus, with window size of 300, the input dimension approximates le5. We evaluated MLPs with varying layer numbers and a hidden size of 1000 dimensions, as described in Section 4.3 The MLP output dimension matches the 50257 tokens of GPT2Tokenizer, with unused tokens assigned nerr-zero zero weights. For MLP we also use the Adam W optimizer with the learning rate 5e-5.
25

---

Published as a conference paper at ICLR 2025
G.3 LINEAR FUNCTION FOR PROBING
For the probing task, we use single-layer linear neural network without activation, effectively func- tioning matrix. The input size corresponds the Transformer's 720-dimensional embeddings, and the output size matches the 50257 tokens of GPT2Tokenizer.
H TESTING PERFORMANCE CURVES
child chain, en
child chaih
child. chain_len=0
Depth=5
Depthm8
Depth-10
Depth-13
Depth=15
1.0
1.0
Whole acc
0.5
0.3
0.5
0.0
3.09
7.0
1.0
1.0
0.5
0.5
805
0.0
0.0
0.0
1.0
1.0
04-
23
0.5
0.01
0.0
1.49
1.0
1.0
1.0
1.00
53
Shots Num
Shots Num
Shots Num
Shots Num
Shots Num
Testing performance of Transformers trained on FTCT with all combinations of causal depths and child chain lengths. The performance is evaluated by all four criteria in Appendix H.l
Figure 3:
H.1 ALL TESTING PERFORMANCE CURVES
We summarize the criteria used to evaluate the compositional reasoning performance (including final value accuracy) follows:
Whole chain accuracy: Measures the model's generation contains all vertices and values --
along the reasoning chain in a correct order. For (inp lab k. from Diea. it measures whether k. dec(model(i nnnnn)))) )) contains all elements from dec (1 in a correct order.
Final value accuracy: Measures if the model outputs the correct value of the last vertex.
Testing vertices accuracy: Measures if the model correctly outputs all vertices dec(1ai ) Testing values accuracy: Measures if the model outputs correct values of intermediate vertices, given correct preceding reasoning paths from Dreat-
We trained 3-layer, 3-head GPT-2-like Transformers depths (N) and maximum child chain lengths (MV). We test the performance Transformers trained on FTCT tasks with N 5. 8, 10, 13, 15 and M = 2, 3, 4, 5, 6. For each (M, pair, change the causal structures for times, calculating the average and standard deviation the testing per- formance of Transformers trained them, recording the results in Figure
the FTCT training set with varying graph
From Figure we observe that the curves of the final value accuracy mirrors the curves of whole chain accuracy, underscoring the necessity of proper reasoning paths for correct final answers
H.2 TESTING PERFORMANCE WITH NOISY TOKENS
In the main text, we configured the training data to consist child chains blurred by noisy tokens, while the testing data comprised the longest chains without such noise. This setup simulates real-
26

---

Published as a conference paper at ICLR 2025
child_chein_ er=
child_chein, Jen=3
child_chei_ee en=4
chile chile_chin_eens3
child_chainJleee
Depth=5
Depth-8
Depth=10
Depth=13
Depth=15
1.4
1.0
10
1.0
1.2
Whole
0.5
0.00
2.00
3.00
1.0
3.0
acc
0.5-
0.5
233
0.3
005
0.01
0.0
acc
7.49
1.00
0.51
0..5
0.5
6.5
atc
4.49
1.0
1.0
1.0
3.0
05
Shots Num
Shots Num
Shots Num
Shots Num
Shots Num
Figure 4: Testing performance of Transformers trained on FTCT where the testing data are blurred by noisy tokens. The performance is evaluated by all four criteria in Section H.2
world scenarios where the training corpus typically affected by noise, while users offer high- quality, noise-free prompts during testing. To demonstrate the robustness of our conclusion, evaluate the performance of Transformers on testing data that is also blurred by noisy tokens. The key to constructing this noisy testing dataset is merging the longest chain in T(G) with randomly sampled noisy tokens from Van - V. while ensuring that (1) the order of the vertices in the chain is preserved and (2) the last vertex of the chain remain as the last vertex in the merged sequence. For the causal structure depicted in Figure ! such a merged sequence might look like
[F. 9, A, 100, K, 6, B, 101, H. 1, C, 103]
where [A, B, C] is the longest chain in T(9) and [F, K, H] are noise randomly sampled from Varl V. After downstream processing, such sequnce transformed the input "C=?: ... F=9, ... A=100," and label "... K=6, B=101, H=1, C=103". To assess compositional reasoning performance on this dataset, we modify our evaluation criteria focus only on the accuracy of non-noisy vertices in the set of knowledge points V:
Whole chain accuracy: Measures if the model correctly predicts all vertices belonging to the set of knowledge points V and their values along the reasoning chain. Final value accuracy: Measures the model outputs the correct value of the last vertex in within the sentence.
Testing vertices accuracy: Measures if the model correctly outputs all vertices in dec(Lai) ) which belong to V.
Testing values accuracy: Measures if the model outputs correct values of intermediate vertices that belongs to V, given correct preceding reasoning paths from Drest-
Figures 4nd display the compositional reasoning capacity of Transformers trained on FTCT with testing data blurred by noisy tokens. The results align with those from tests without noise, indicating the robustness of our conclusions.
I LEAST VISITED TIMES OF ALL ADIACENT VERTICES AND THEIR VALUES
We define the set of all adjacent vertices and their values as S(g) == ((v,, gi 41+1, q:+1) Us, 64t1 E V, (ve, v441) E E, E E VALS(())) g++1 = op(s,, 64 )ogs) Only when each element S(g) has been visited for enough times in the training data, can we ensure that the poor compositional reasoning ability of models trained certain FTCT tasks is not
27

---

Published as a conference paper at ICLR 2025
Values acc Vertices ace Final ace Whole acc
0.5
0.0
0.2
0.3
0.4
0.5
0.6
1.0
0.5
0.0
0.2
0.3
0.4
0.5
0.6
1.0
0.5
0.0
0.2
0.3
0.4
0.5
1.0-
0.2
0.3
0.4
0.5
0.6
Relative knowledge ratio
Figure 5: The relationship between the relative knowledge ratio and Transformers' compositional reasoning ability tested on the data blurred by noisy tokens.
Depth 5
Depth
Depth 10
Depth 13
Depth 15
Child
194
109
76
55
46
Child
250
123
89
65
Child
140
91
67
Child 5
72
56
Child
59
Table 4: Least visited times of training datasets generated on causal structures with various causal depths and child chain lengths.
due the insufficient data. For any training dataset D, we define (D(u,, Hs 04 (++ the times the tuple (wi, JG, 1 (6+1) ) is visited by the data in D. We define the least visited times of D as
T(D): := minfto q. Us+1, 96-1) (v,, l. 0:41, 4-41) E S(G)).
The least visited times measure the degree of how well the S(G) is visited by D. We demonstrate the least visited times of all our training data Table 4 The results suggest that all necessary knowledge parts covered sufficiently by training dataset.
J PERFORMANCE OF LARGER TRANSFORMER MODELS
To assess the generalizability of our findings to more complex and sizable architectures, we trained GPT2-small (12 layers, 12 heads, 117M parameters) and GPT2-large (36 layers, 20 heads, 774M parameters) on FTCT tasks with varying causal depths and child chain lengths. We discovered that the diversity of training data used for smaller models was insufficient to leverage the larger mod- els' generalized in-context learning ability. Consequently, we introduced auxiliary data to facilitate this capability. Specifically, the auxiliary data comprises few-shot examples with vertices and val- randomly sampled from Vau and Z, respectively. The testing performances of GPll-small and GPT2-large are depicted in Figures 6 and 88, with the performance's relationship to relative knowl- edge ratio illustrated in Figures and We continue to observe that compositional reasoning ability
28

---

Published as a conference paper at ICLR 2025
emerges with increased shot numbers and relative knowledge ratios. However, the performance of these larger models is less stable compared to smaller ones, which may be attributed to overfitting.
child chain,
child
chilc cili.cchan/eem
child. chain_len=0
Depth=5
Depthm8
Depth-10
Depth=13
Depth=15
acc
0.5
0.3
84
0.0
3.50
1.0
1.0
1.0
1.0
80c
65
0.5
04
0.0
Vertices
0.5-
0.5
0.5
533
0.5
00
0.0g
0.01
6.1
0.2g
1.4
1.0
1.0
1.0
1.0
Shots Num
Shots Num
Shots Num
Shots Num
Shots Num
Figure 6: Zero and few-shot testing performance of GPT2 small trained on FTCT with various causal depths and child chain lengths.
Values acc Vertices acc Final acc Whole ace
1.0
0.5
0.0
0.2
0.3
0.4
0.5
0.6
1.0
0.0
0.2
0.3
0.4
0.5
1.0
0.5
0.0
0.2
0.3
0.4
0.5
1.0
0.5
0.2
0.3
0.4
0.5
0.6
Relative knowledge ratio
Figure 7: The relationship between the relative knowledge ratio and the compositional reasoning ability of GPT2 small.
29

---

Published as a conference paper at ICLR 2025
meld_chennJen=2
child_chinnJeen
Boild_chein_Jenzz
child_hheinJeenss
child_cheini/ensa
child_chain_Jenss
Depthe5
Depthe8
Depth=10
Depth-13
Depth=15
1.4
1.0
10
0.5
0.-
05
69
0.3
2.00
3.0
1.0
1.0
3.0
acc
0.5
00
0.0
0.0
0.0
0.0g
acc
4.48
7.0
1.01
1.0
0.5
0.5-
0.5
6.-
0.5
Values acc
24
0.0
on
c.0
0.0g
Shots Num
Shots Num
Shots Num
Shots Num
Shots Num
Figure 8: Zero and shot testing performance GPT2 large trained on FTCT with various causal depths and child chain lengths.
Values acc Vertices acc Final acc Whole acc
1.0
0.5
0.0
0.5
0.0
0.7
0.3
0.4
0.5
0.6
1.0
0.2
0.3
0.4
0.5
0.6
1.0
0.5
0.2
0.3
0.4
0.5
0.6
Relative knowledge ratio
Figure 9: The relationship between the relative knowledge ratio and the compositional reasoning ability of GPT2 large.
K EXPLAINING THE DECREASE OF PERFORMANCE WHEN SHOTS NUMBER IS LARGER THAN ONE
we observe decrease in model performance when more than one few- shot prompt is used, contrary to the typical expectation that additional examples enhance in-context learning. While it is generally expected that more few-shot examples improve performance in- context learning, performance fluctuate
In Figure and Table 3.
decline additional examples introduced
30

---

Published as a conference paper at ICLR 2025
(as noted in Figure 4b of (Garg et al. 2022) and Figures 7, 10 of (Xie al. 2021 ).
Although
systematic explanations are limited, it is not uncommon for performance to drop when the number of shots surpasses particular threshold. (Agarwal al. 2024) notably demonstrates that the optimal shot count for peak performance is often less than the maximum a model can manage, indicating that performance does not always increase monotonically with additional shots.
We propose that performance decreases when the number of shots exceeds one because additional CoT examples increase the dissimilarity between training and testing data, thereby degrading perfor- mance. To explain this, we begin by analyzing the difference between one-shot examples in testing and training data. In testing, one-shot examples typically consist of the longest chains of length N, whereas in training, they are child chains of length M < N, with the primary difference being the N M missing vertices. As the number of shots increases, each shot introduces more instances of these missing vertices, compounding the disparity between training and testing prompts. This growing difference complicates the model's ability to recognize patterns in the testing data, thus impairing performance.
In our setup, a single shot during testing already provides ample information about the vertex order needed for generating correct answers. For a k-shot testing example, the additional k - 1 shots do not add valuable information and only exacerbate the divergence between training and testing data. Consequently, we observe that testing performance peaks at 1-shot and diminishes thereafter, aligning with our expectations.
When the differences between training and testing data are limited, we observe the expected pattern of in-context learning, where performance improves with more shots and does not decline after reaching its peak. From all FTCT tasks with various depth and child chain length combinations shown in Figure 3 we select settings where performance decreases when shot numbers exceed one. These settings include "depth=8, child chain length=3," "depth=10, child chain length=3," "depth=13, child chain length=4." and "depth=15, child chain length=4." Instead of testing models' performance on the longest chains equal their depth, we assess performance on causal chains with lengths varying from the child chain length to the depth. For instance, for models trained on an FTCT task with depth=8 and child chain length=3, we measure performance on chains with lengths 3. 4, 5, 6, and 8. The results in Figure I0 show that for tasks where test lengths are close to child chain lengths, few-shot performance does not decrease. Notably, zero-shot performance increases as test lengths near the child chain length. As the gap between child chain and test lengths widens, the performance decrease after one shot becomes evident.
Jon
Jon
Jen-5
Janw6
Jen-8
Jon-10
Jon-13
test Jon-15
Depth:B Child=3
D1p Child=3
Depth=13 Child::4
Depth=15 Child=4
1.4
1.0
1.0
3.0
Whole
0.2
0.0
0.01
0.01
3.6
Final BCc
0.5
0.5
0.51
00
0.0
Vertices acc
1.0
1.40
3.05
0.5
0.0
0.01
Values BCc
1.0
3.05
Shots Num
Shots Num
Shots
Shots Num
Figure 10: Performance of Transformers in reasoning causal chains with varying lengths.
31

---

Published as a conference paper at ICLR 2025
L
EMPIRICAL EVIDENCE OF THE TRANSFORMER STRUCTURE
Figure
and 12 demonstrate the evidence of the induction head by attention heatmap.
Attention Heatmap
1.0
Attention Hestmap
1.0
[590) 3877
(323)12)
130
[220) collec
0.8
E01
0.6
D1 Miaerrnne
Query
(142 eato
0.4
0.4
(240) tskic
0.2
0.0
0.0
Figure 11: Left: Average attention heatmap of induction heads at layer (test). Right: Average
attention heatmap of induction heads at layer 2 (test).
Heetmmp
Attention Heatmag
1.0
1361 deneng aeon
[el (34) Medood
40 Cal
12022
(67] 13
(2
0.8
(10)
1202 n822
(CAl
1111 deb
1750000
1563247
I776 CONT
0.6
0.6
1275 evara
1613
Query
100
EaT Rugh
(89) cestian
0.4
[91] Tal
0.2
Key
0.0
0.0
Figure 12: Left: Average attention heatmap of induction heads layer (train). Right: Average
attention heatmap of induction heads at layer 2 (train).
M TESTING PERFORMANCE WITH INCOMPLETE INTERMEDIATE REASONING STEPS
In this section, evaluate the capability of Transformers trained FTCT to perform reasoning with missing intermediate steps. We conduct this test by removing intermediate vertices from the testing data. For the testing sequence [A, B, CJ illustrated in Figure remove vertex B, creating
32

---

Published as a conference paper at ICLR 2025
test input formatted "C=?: A=100, C=103 \n C=?: ... A=102," with the test label being C=105". For causal structures with greater depth, we define the retaining ratio p as the ratio of the count of remaining vertices to the causal structure's s depth. This ratio reflects the degree of incompleteness in the reasoning path. For example, if the longest chain in the causal structure is [A, B, C, F, G. I] and is to 0.5, then the number of remaining vertices is 6 0.5 = 3. Since we only remove intermediate vertices, potential final sequences include [A, F. I] or [A, C. IJ, among others. Our evaluation criteria remain the same as the original (Section 4.1 but now both testing inputs and labels comprise incomplete reasoning paths. Given that Transformers trained on FTCT cannot generate incomplete reasoning paths autonomously in a zero-shot setting, we assess performance with to 4 shots, each containing examples of incomplete reasoning paths.
Figure illustrates the testing performance with retaining ratios of 0.3, 0.5, and 0.8. The results indicate that Transformers trained on FTCT exhibit limited proficiency in reasoning with incomplete intermediate steps. This limitation likely due to training dataset bias, where adjacent vertices con- sistently appear consecutively in training sequences. Nonetheless, this performance shortcoming does not compromise our main findings, which demonstrate the Transformers' compositional rea- soning capabilities, because the testing data invariably contain longer causal chains than the training data, regardless whether the demonstrated reasoning path is complete or not.
33

---

Published as a conference paper at ICLR 2025
retaining ratio-0.3
chil chain/ien2 Jan=2
chid chain whwin.eeen Jenm3
child chil_chhing choin janms
child chein Jene5
child chain Jan-B
Depth-5
1.0
Depth-6
1.0
Depth-11
1.-
Depth-13
1.0
Depth-15
0.7
0.0
2
Final
1.0
1.0
3.0
1.0
Vertice
0.-
0.0
6.01
3.0g
Shots Num
Shota Num
Shota Nam
Shots Num
Shots Num
retaining retioc0.5
child lent2
chils.chiin
chid_chhrng (enu4
child chain
child
Depth-10
Depth-15
Whole
1.0
1.0
1.0m
03
0.3
0.3
0.50
055
0.04
0.0
0.04
Vertkes
1.0
2.2
0.3
0.3
0.0
05
nol
604
Shots Num
Shots Num
Shote
Shote Num
Shots Num
D5pthe5
Depth-g
Depth-10
Depth-13
Depth-15
1.0
3.0
3.0
1.0
0.3
0.3
0.3
0.0
0.0
Vertices
22
0.0
1.0
3.0
3.0
1.0
03-
0.
Figure 13: Testing performance of Transformers trained on FTCT with incomplete intermediate reasoning steps (retaining ratio=0.3, 0.5 and 0.8 from up down).
34


Reviews:
[{"review_id": "UpVv0LZhRT", "reviewer": "Reviewer_n7TN", "summary": "The work sets out to investigate whether transformers are capable of generalizing to longer reasoning chains through connecting shorter ones seen in the training stage. The authors introduce \"Fragemented at Training, Chained at Testing\" learning task to train a randomly initialized 3-layer 3-head GPT2-like transformer. They find that with few-shot chain-of-thought prompting, transformers can perform good compositional reasoning skills by combineing fragments together. The authors further show that the generalization performance highly correlates to model complexity (require multiple-layer attention structure) and high relative knowledge ratio of training data. The paper also discusses the internal working of the model (learn an underlying generalizable program) to interpret the transformer's generalization behaviors and provide theoretical insights on transformer's expressivity on learning a such program.", "questions": "1. In the FTCT learning task (e.g., Figure 1), why in the $D_{train}$, we need to add noisy tokens in the token sequence? Why in the $D_{test}$ we do not add noisy tokens in the prompt?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper investigates whether transformers are capable of generalizing to longer reasoning chains through connecting shorter ones seen in the training stage, which is an interesting and important research question.\n2. The paper is technically sound: the trained transformers behave compositionally (with few-shot chain-of-thought prompting) and the authors provide insights on its internal workings: induction head and attention assignment, demonstrating that the transformer learn a generalizable program in its internal computing.\n3. Authors also theoretically prove that Transformers have the expressivity to simulate the generalizble underlying program.", "weaknesses": "1. Since the experiment setting is a randomly initialized transformer trained on synthetic data, to what extent the paper's conclusion can be extended to real pre-trained language models is questionable.\n2. the notations used in the paper are quite complicated, making the paper a little bit difficult for readers to follow.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "h08MYLrUIa", "reviewer": "Reviewer_EQsG", "summary": "This paper introduces the \"FTCT\" (Fragmented at Training, Chained at Testing) task to evaluate if Transformers can perform compositional reasoning similar to humans. The task involves training models on separate knowledge fragments and testing them on integrating these fragments to form complete causal graph traces. The study finds that few-shot Chain-of-Thought prompting helps Transformers combine these fragments correctly, even without seeing such combinations during training. The results indicate that model complexity and the data's knowledge ratio play a role in enabling this skill. The authors provide theoretical and empirical evidence for their claims, showing that Transformers can learn a generalizable program to aid in compositional reasoning. The findings are interesting and suggest potential areas for further exploration.", "questions": "My main concerns have already been expressed in the \"weakness\" section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The design of the FTCT task is well-conceived, as it effectively mimics real-world scenarios where knowledge is often fragmented and must be integrated to draw comprehensive conclusions. This setup provides a meaningful and practical benchmark to evaluate the compositional reasoning abilities of Transformers, making the study relevant and valuable for advancing our understanding of machine learning models' capabilities.\n- Chapter 5, \"transformer does compositional reasoning via the underlying program\", is very interesting as it explores the possible underlying mechanisms and principles that allow Transformers to perform compositional generalization. This chapter goes beyond just presenting empirical results by looking into how these models might internally handle and integrate fragmented knowledge. This deeper investigation adds value by giving us a better understanding of how Transformers achieve complex reasoning tasks.", "weaknesses": "- While the task studied in this paper requires strong compositional generalization abilities, it is simple and singular in its form. Generally, using a simple and singular synthetic dataset is suitable for highlighting the shortcomings of the Transformer architecture. However, since the paper concludes that Transformers possess this capability, the experiments on this task alone are not sufficient to support such a conclusion. I believe that more diverse and comprehensive tasks are needed, and ideally, this capability should also be validated on complex real-world tasks.\n- In the related work section, the paper discusses various tasks used to probe compositional generalization abilities. The authors mention that these existing tasks have not been studied from the perspectives of few-shot prompting and chain-of-thought reasoning. However, this distinction alone is insufficient; if the difference is merely in this aspect, it would be possible to modify existing datasets instead of creating a new one. The novelty of this newly created task is not demonstrated well. Therefore, the authors need to provide more explanation regarding the motivation and innovation behind the proposed task.\n- The experiments in this paper use Transformers with relatively small parameter sizes. It is unclear whether the conclusions drawn from these experiments would hold true for larger Transformer models. This limitation raises questions about the generalizability of the findings to more complex and sizable architectures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9lidulnRgg", "reviewer": "Reviewer_oHAw", "summary": "The paper investigates how well transformers are able to perform compositional reasoning tasks. To that end, the paper introduces a new dataset methodology, namely the Fragmented Training, Chained at Testing (FTCT) that simulates how models would be presented with training data in practice (with incomplete fragments of reasoning paths with noise + context) and how well the model is able to piece together the full reasoning chain in test-time. Using this methodology, the paper runs insightful experiments that ablate different lengths of partial reasoning chains during training, different transformers and neural architectures, and number of few shot CoT examples. Through these experiments, the authors find that few shot CoT plays an important role for compositional reasoning, the impact of increasing relative knowledge ratio, and the increasing expressibility of adding layers and heads in the transformers architecture. Lastly, the paper presented some empirical evidence that you need a certain complexity of the transformers architecture to simulate the optimal program for the FTCT task in training and testing.", "questions": "- Although it is clear that the model is learning some ability to connect reasoning chains together (e.g. during training the model might just as AB and BC and correctly chain ABC), will the model be able to correctly chain together the values of AC? This could make for an interesting experiment where we could have some skip links in the test data and check for values accuracy\n- Checking my understanding, is there a typo in Figure 1, where B=106 and B=103, should it be C=108 and C=105, respectively?\n- Are there more than one set of the causal chains? The set equation in line 155 seems to suggest there is only one sequence of length n.\n- Why are the noise vertices inserted in a predictable manner?\n- I am curious about this 0.3 relative knowledge ratio threshold where it is reported that compositional reasoning emerges. Could it be that 0.3 is when the probability that there is at-least one occurrence for every (v_i, v_{i+1}) in the train set reaches close to 1? \n- Why is there a drop in performance in Figure 2 (right) and relative knowledge of 0.25?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The paper presented a very intriguing and creative approach to testing the ability for models to learn compositional reasoning ability\n- There are some really interesting results, specifically the exact complexity (and the increased expressability) needed for the transformer architecture to optimally solve the FTCT task\n- The insights regarding the few shot CoT results are of significance and spark further research in this area\n- The empirical findings of how the transformers performs this task is enlightening and should spark some interest for further research in this area", "weaknesses": "- The clarity of this paper is lacking, especially in the notation and writing. For instance, in Figure 1, there is a seeming typo in some of the values that contradicts the setup of the dataset. Separately, some concrete examples of the data (including noise + context tokens) of the FTCT dataset would really improve the readers understanding (it took me multiple re-read to get the gist of the methodology)\n- The paper's definition of compositional reasoning should be explicitly written out in the paper. The only real definition of this is in the abstract where it is stated that \"For example, if someone learns ( B = f(A) ) from one source and ( C = g(B) ) from another, they can deduce ( C = g(f(A)) ) effortlessly, even without encountering ( AC ) or ( ABC ) together, showcasing their compositional generalization ability.\"\n- With this FTCT methodology, it seems clear that the model is learning some ability to connect sequential reasoning chains together (e.g. during training the model might just as AB and BC and correctly chain ABC), but the approach does not test if the model can correctly reason about AC in test-time, which is an aspect of compositional reasoning (as mentioned in the abstract)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "31VZpqZo1A", "reviewer": "Reviewer_hQE3", "summary": "**Aims:** This paper seeks to understand at a mechanistic level how Transformers are able to perform compositional reasoning in a few-shot chain-of-thought setting.\n\n**Methods:** A synthetic dataset is generated consisting of chains of nodes and edges derived from causal graphs. At training time, spurious nodes are inserted randomly into the chains; at testing time, few-shot prompts consisting of intact chains (no spurious nodes) are provided to the model. Models are tested on their ability to reconstruct full causal chains from fragmented chains learned in training, with evaluation based on accuracy in predicting both the correct vertex order and values in the chain.\n\n**Results:**\n\n- Zero-shot versus few-shot prompting is compared, with findings showing that few-shot CoT prompting significantly enhances performance in compositional reasoning tasks, particularly in forming the correct vertex order.\n- A space of small, GPT-2-style models ranging from 42M-54M parameters are trained on the FTCT dataset. Results show that multi-layer, multi-head Transformers (minimum 2 layers and 2 heads) perform notably better, while single-layer/single-head models and MLPs perform poorly.\n- The impact of training data’s relative knowledge ratio (ratio of child chain length to complete chain length) is studied, with a critical threshold (ratio ≥ 0.3) identified where compositional reasoning reliably emerges.\n- Mechanisms underlying the model's success, such as induction heads for in-context learning and attention patterns facilitating parent-child relationships, are analyzed through linear probing, revealing specific mechanisms by which the model achieves compositional reasoning.", "questions": "1. Table 1 shows that all Transformer models attain 1.0 Values accuracy, even for small models that get very low Vertices accuracy. Can you account for this discrepancy?\n2. An unintuitive pattern in the results (e.g., Fig. 2 and Table 3) is that accuracy *decreases* with the number of few-shot prompts $>1$. This results stands in contrast to a standard ICL setting, where inclusion of more examples typically improves performance. It is stated that this is “possibly due to increased dissimilarity between training and testing data with more CoT examples” (Line 277-278). Why does including more CoT examples causes the test data to be OOD? If this is the case, this seems like an important confound affecting this experiment setup that may not be adequately accounted for.\n3. It is interesting to contrast the results from Sec. 5-6 with Zhang et al., 2022 (“On the Paradox of Learning to Reason from Data”), who apply a similar methodology but find that gradient descent fails to discover the correct $\\theta^*$ for a logical induction task with very similar structure. Is there a reason why here the training succeeds at uncovering the underlying program, whereas in previous work it does not? More generally, it would be nice to see reference to this paper in the discussion.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "**Motivation:** The key research questions of the paper are clearly delineated in the introduction:\n\n1) When are Transformers able to perform compositional reasoning by connecting fragmented knowledge in training data?\n2) How do different training factors impact the emergence of this ability?\n3) What internal mechanisms enable Transformers to develop this ability?\n\nThese questions are broadly relevant to current research and the paper is structured in a way that consistently centers these 3 questions throughout.\n\n**Mechanistic interpretability analysis:** I especially enjoyed the approach to question (3). Broadly speaking, the authors approach this question by first demonstrating that there exists a program that solves the task (Sec. 5.1) and that this program can be approximated by a 2-layer Transformer (Sec. 5.2). Then, through linear probing experiments (Sec. 6), they give an empirical argument that the Transformers trained on FTCT have learn to implement this program. I am not an expert on probing so I can’t speak to the soundness of the methods, but I found the combination of Sec. 5-6 to be an elegant argument from a mechanistic interpretability standpoint.", "weaknesses": "**Novelty:** The paper claims to “introduce a learning task” (FTCT) based on causal graphs, and yet the design of this task is nearly identical to the setup in Prystawski et al. (2024). Given that the main distinction between FTCT and the prior work is the introduction of spurious nodes (line 105-106), I would expect to see this prior work—which actually *did* introduce a novel learning task—given more prominent attribution. \n\n- (Currently this work is referenced in Line 103—”The empirical findings of our work align with the observation in (Prystawski et al., 2024)…” The wording of this reference obfuscates the underlying causal structure that this prior work likely played in informing the current paper.)\n\n**Generality:** The key findings of this paper are framed in very broad terms: \n\n> “The emergence of compositional reasoning is highly influenced by the data’s relative knowledge ratio and model complexity. Specifically, a relative knowledge ratio of at least 0.3 and a Transformer architecture with at least two layers and two heads are critical for achieving this ability.” (Lines 520-523)\n> \n\nHowever, these conclusions are all drawn relative to one synthetic dataset with a highly-specific structure; it is unclear to what extent the empirical conclusions (e.g., compositional reasoning in transformers requires a relative knowledge ratio ≥ 0.3) generalize beyond the FTCT task. To make a convincing argument that these results have meaning beyond this one benchmark, this analysis ought to be replicated on more naturalistic reasoning benchmarks where few-shot CoT prompting is commonly used.\n\n**Clarity:** The description of the FTCT dataset/task design (Sec. 3) fails to convey a clear description of the experiment setup and requires too much work of the reader. All aspects of prompt construction are described in excruciating formal detail, making it hard to separate design choices that are key to the experiment from implementation details. Overall, the formalism in this section is a barrier to understanding what’s going on at a more basic level.\n\n- Fig. 1 lacks basic signposting needed to convey what is going on.\n    - First off, there is no caption. This is a major omission as the figure is definitely not self-explanatory.\n    - The blue highlights draw the reader’s attention to spurious features of the task (noise nodes) instead of the actual purpose of the task (predicting values of causal nodes).\n- Other comprehension/clarity issues in Sec. 3:\n    - “We assume op(e) represents operations like (+a) or (−b)” Does this mean addition/subtraction are the *only* possible operations?\n    - I don’t understand how the merge operation works from the description.\n    - Some unconventional choices of notation, such as using $f$ as an index over few-shot examples.\n    - What is “downside processing” - do you mean “downstream”?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)