Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing
Authors: Zhongwang Zhang
Pengxiao Lin
Zhiwei Wang
Yaoyu Zhang
Zhi-Qin John Xu
Venue: NeurIPS 2024 poster (2024)
Keywords: parameter initialization
transformer
compositional task
reasoning
memorizing
Primary area: interpretability_and_explainability
TL;DR: We find initialization scale determines if transformers learn compositional tasks by reasoning or memorization; mechanistically, small scales favor inferential solutions capturing compositional primitives, while large scales merely memorize tasks.

Abstract: Transformers have shown impressive capabilities across various tasks, but their performance on compositional problems remains a topic of debate. In this work, we investigate the mechanisms of how transformers behave on unseen compositional tasks.  We discover that the parameter initialization scale plays a critical role in determining whether the model learns inferential (reasoning-based) solutions, which capture the underlying compositional primitives, or symmetric (memory-based) solutions, which simply memorize mappings without understanding the compositional structure. By analyzing the information flow and vector representations within the model, we reveal the distinct mechanisms underlying these solution types. We further find that inferential (reasoning-based) solutions exhibit low complexity bias, which we hypothesize is a key factor enabling them to learn individual mappings for single anchors. We validate our conclusions on various real-world datasets. Our findings provide valuable insights into the role of initialization scale in tuning the reasoning and memorizing ability and we propose the initialization rate $\gamma$ to be a convenient tunable hyper-parameter in common deep learning frameworks, where $1/d_{\mathrm{in}}^\gamma$ is the standard deviation of parameters of the layer with $d_{\mathrm{in}}$ input neurons.

Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing
Zhongwang Zhangl Pengxiao Lin! Zhiwei Wang Yaoyu Zhang Zhi-Qin John Xu 123455
Institute of Natural Sciences, MOE-LSC, Shanghai Jiao Tong University
2 School of Mathematical Sciences, Shanghai Jiao Tong University
Key Laboratory of Marine Intelligent Equipment and System, Ministry of Education, P.R. China
Shanghai Seres Information Technology Co., Ltd, Shanghai, P.R. China
Center for LLM, Institute for Advanced Algorithms Research, Shanghai
Abstract
Transformers have shown impressive capabilities across various tasks, but their performance on compositional problems remains a topic of debate. In this work, we investigate the mechanisms of how transformers behave on unseen composi- tional tasks. We discover that the parameter initialization scale plays a critical role in determining whether the model learns inferential rreasoning-based which capture the underlying compositional primitives, or symmetric (memory- based) solutions, which simply memorize mappings without understanding the compositional structure. By analyzing the information flow and vector represen- tations within the model, we reveal the distinct mechanisms underlying these so- lution types. We further find that inferential (reasoning- -baedd solutions exhibit low complexity bias, which we hypothesize is key factor enabling them to learn individual mappings for single anchors. We validate our conclusions on various real-world datasets. Our findings provide valuable insights into role initial- ization scale in tuning the reasoning and memorizing ability and we propose the initialization rate 7 to be a convenient tunable hyper-parameter common deep learning frameworks, where 1/df, is the standard deviation of parameters of the layer with dn input neurons.
solutions,
Introduction
2017) demonstrate unprecedented capabilities (Acham ct 20233 Brown all 2020: Chor et al 202 Teubbee all, 2023), even noted as "sparks of AGI" (Bubeck et all 2023). However, their performance on compositional tasks, which considered a key property of human cognition (Marcus, 2003: Smolensky et all, 2022), has long been the subject of intense debate, especially terms of the -itto---isttiuttionn (OOD) generalization ability. This raises critical open questions about how to faithfully interpret transformers' capabilities on compositional tasks. Can transformers learn the underlying compositional primitives within the data, or do they merely learn input-output mappings? When transformers produce incorrect outputs on compositional tasks, do their responses follow any discernible patterns or logic?
Large-scale transformers (Vaswant
In this work, we use anchor functions (Zhangg et ar. 2022). a framework for creating controlled synthetic data, to investigate how transformers generalize to unseen compositional tasks. In our setup, each sequence contains an anchor pair, a key item, and noise items unrelated to the output (right side of Fig. Ila)). The single anchors specific tokens (i.e., 1, 2, 3, 4), each corresponding to particular arithmetic operation (left side of Fig. I[aa)). A composite function is defined as the
"Correpooniigg author: xuzhiqin awrhigi@sstt.edu.. @sjtu.edu.cn
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

sequential application of two single anchor functions, i.e., the corresponding arithmetic operation. to key item. For example, given a key item z and an anchor pair (1. 2), the composite function would output (x + 5) + 1, i..., x + 6. As shown in the middle part of Fig. Maa), during training, 14 saasoonnn-hasedd mappings, meaning
out of the 16 possible anchor pairs are assigned inferential the target output is consistent with the result of applying the composite function. One pair (3, 4) is assigned a non-inferential (memory-based) mapping, where the target output deviates from the result unseen task. After training,
of the composite function. The remaining pair (4, 3) is held out there are three possible solutions the model could learn for the unseen pair (4, 3): a symmetric solution matching the non-inferential seen pair (3, 4) (Mechanism in Fig. ITb)),
inferential solution (Mechanism 2 in Fig. (rb)), or other non non-inferential solutions. The central question we aim to answer is: which of these three rypes of solutions will the model learn for the unseen anchor pair (4, 3)?
Data generation
(b) Pussible solutions and mechanisns for the unseen anchor
Single anchore
14 seen interentiall compostte anchore
Input data examples
Target
Mecharinn
+10
55 46 32
28
oyenetric sestion
Compostioo
16
Paoddng
S
non ion-interentall composite anchore
Wechanim compastte fram the infered aingle enchor
-S
unseen compostte anchor
3
Q: Which zolation madel lcom
-10
inforential solution
seen awter pair
seen non- infernnee ancher par
unDees excher pair
thre (fte before oncher
arithmehi
notox
Figure 1: Experimental setup and possible solutions and mechanisms for the unseen anchor pair (4, 3). (a) Data generation: Left: The single anchors (i.e., 1, 2, 3, 4) correspond to specific arithmetic operations. Middle: During training, 14 out of the 16 possible anchor pairs are assigned inferential mappings, one pair (3, 4) is assigned a non-inferential mapping, and the remaining pair (4, 3) is held out an unseen task (does not appear in the training). Right: The input sequences comprise an anchor pair, a key item preceding the anchor pair, and noise items unrelated to the target. The question mark indicates the output for the unseen anchor pair (4, 3), which depends the learned solution. (b) Two potential mechanisms for the unseen anchor pair (4, 3): learning the symmetric structure (Mechanism 1) or composing the inferred single anchor mappings (Mechanism 2).
Based on the above setup, we find that the model exhibits two distinct phases of solutions depending on the parameter initialization scale with different model depths. With small initialization scales, the model tends to learn inferential solutions on the unseen anchor pair (4, 3). In contrast, with anchor pair (4,
larger initialization scales, the model is more likely to Iearn symmetric solutions 3) that simply match the output of its symmetric anchor pair (3, 4). These findings suggest that the initialization scale plays crucial role in determining the type of solution learned by the model.
To gain insights into the mechanisms underlying these different solution types, we analyze the infor- mation flow and vector representations within the model. We find that symmetric solutions directly combine anchor information without capturing the underlying compositional primitives, while infer- ential solutions learn individual mappings for each single anchor and compose them to generate the final output. We hypothesize that the model's ability to learn compositional primitives is influenced by its complexity, determined by the initialization scale. Inferential solutions exhibit low complexity, with weights condensing in a few directions and input tokens arranged according to their numerical value in the embedding space, while symmetric solutions show no obvious condensation and lack clear structural features, indicating higher complexity.
Our work highlights the crucial role of initialization scale in shaping the solutions learned by trans- formers on compositional tasks, enabling them to better capture compositional structures and gener- alize to unseen tasks. Furthermore, adapting the initialization scale to the specific task type could be promising: larger scales for memorization tasks such as text memorization, and smaller scales for reasoning tasks such as code generation. Unless otherwise specified, our main text primarily fo- cuses on a single-head attention model to facilitate clearer understanding of the underlying mech- anisms. We further extend our experiments to GPT-2 (Radford et all 2019), and verify that the
2

---

insights drawn from the single-head model remain valid for more complex architectures on various real-world datasets.
With the support of the theoretical works and extensive experiments, we propose the initialization rate y to be a convenient tunable hyper-parameter in common deep learning frameworks, where 1/d is the standard deviation of parameters of the layer with di input neurons.
2 Related Work
Challenges of transformers in compositional tasks. Recent advancements large language mod- els (LLMs) have showcased remarkable capabilities, often surpassing human performance eFl
20122 Wer all 2022). However, despite their impressive performance on single-step reasoning
tasks (Srivastava et all 2022), transformers struggle with multi-step compositional tasks and OOD
generalization (Csordas 2021, 2022: Dziriet 2024: Hupkes et al. 2018: Lepori et ali, 2023: Okawa et all 2023; Yun et all 2022).. Ramesh et al. (Ramesh et all 2024) show that transformers trained directly compose capabilities struggle generalize OOD tasks with synthetic task. Liu et al. (Liu all (0222) suggest that shallow transformers learn shortcuts during training, leading to poor OOD generalization. Numerous studies have explored various approaches to address these limitations, such as encouraging explicit reasoning step generation within single generation (Weil ctal 2022).. leveraging LLMs to generate reasoning steps iteratively (Creswelle of all 2022: Creswell and Shanaban 2) Despite these efforts, achieving complete mastery in compositional tasks re- mains a significant challenge for vanilla transformers. A series works studies the internal mecha- nisms of language models and improves the capabilities of language models (Wang et all, 202zala 2023 Cao ef all 2024).. In order to clearly study the behaviors and internal mechanisms of language models, Zhang et al. Zhang et all (2012) introduced anchor functions as benchmark functions for investigating transformer behavior. Our work builds on the anchor function setting to explore how different initialization scales affect model solutions and mechanisms.
Parameter initialization of neural networks. The parameter initialization of the network impor- tant to determine the final fitting result of the network (Arora et all 2019: Chizat and Bach, 2018
Zhang et all ROTS E ell 2020; Lacot IacoteLal, 2018; Met eLal MILS; Rotskoff and Sanden-Eynnden 2018; Sirignano and Spilropoulos, 2020: Williams et all, 2019). Luo et al. Luoeetal 202II). Zhou et al. (Zhou et all 2022) mainly identify the linear regime and the condensed regime for two-layer and three-layer wide ReLU NNs, respectively. A series of works suggests that the condensed networks are often accompanied by good generalization ability of the model (Zhang et all 2022, Zhang and Xu, 2023; Zhang et al, 2023; Zhang and Xu,, 2024). A line of works links the grokking phenomenon with the improvement (022)) and points out that the initialization scale has important influence grokking (Liu et all, (002)). Recent studies also investigate the impact of initialization on the train- ing process of LLMs (Huang et al,, 2020; Linetal all M02I: Trockman and Koller, 2023: Wang et all 2024; Zhang et al., 2019 Zhu et all 1021). These works primarily focus on how the initialization scale affects the stability of the training process and plays crucial role in ensuring smooth and effective training of LLMs. In our work, we find that initialization scales can significantly influence model's ability memorize and infer data on the compositional task, highlighting the profound impact of initialization on the final performance and underlying mechanisms of the trained models.
reasoning ability (Power et all, 2022: Wang et all 2024, (ooalaani et all
the occurrence of
3 Definitions
We introduce set of key definitions that will be used throughout the paper. For detailed explanations and illustrative examples to better understand these definitions, please refer to Appendix H.
3.1 Two-anchor composite function
A two-anchor composite function f(X): R - IR is defined as
f( w.) = g (g(#4-1;x;)) (441), where 0,+1 E A.
(1)
Here, the input sequence X (II,..., ..) comprises tokens. An anchor set A
(a), (i,,a,,.... (2,....,aa) designated, with each token 0g E corresponds to a function g(2; ag). In
each X, one and only one pair of two consecutive elements belong to A, such 2, 2,41111 E A.
3

---

We refer to the token immediately preceding the anchor pair as the key item. To simplify notation, we denote the two-anchor composite function as (,, 1441) to emphasize the anchor pair (Mi, (++1) and key item 1;-1-
In this work, set the anchor set A {1,2,3 4). Each anchor token corresponds to specific function:
( 1) = +5 g(z;2) = = +1, 9(133) x - 2, g( 4) = x 8.
(2)
3.2 Data Generation
In this work, we construct the input dataset using four anchors (i.e., 1, 2, 3, 4) and items sampled from 20-99. Each sequence includes an anchor pair, key item (the item immediately preceding the anchor pair), and noise items. The noise items are unrelated to the target. The four anchors form 16 anchor pairs, and we select subset or all of these pairs to construct the training dataset based on the task requirements.
By default, the target is the output of the input data processed by the two-anchor composite function, ie., the inferential mapping defined below. We divide the training data and test data based on the value of the key item. The specific division method can be found in Appendix C
3.3 Mapping Type of an anchor pair
For an anchor pair (a1, a2), the type of its mapping ((a (-) can be classified follows.
Inferential mapping. The designated target mapping of an anchor pair (41, a2) is consistent with the
two-anchor composite function, i.e., (()) (z) (I) = f(x; 01 01,02) 02). This type of solutions exemplifies the reasoning ability of models.
Non-inferential mapping. The designated target mapping of anchor pair (a1, ag) is inconsistent with the two-anchor composite function, i.e., M(arina) (2) + f(z; 01, a2).
Symmetric mapping. The designated target mapping of an anchor pair (a1, a2) is consistent with the mapping of its symmetric anchor pair (02, a1), i.e., ((or. ((a
We refer to a model as an inferential (non-inferential, symmetric) solution on an anchor pair if tends output the mapping corresponding the inferential (non-inferential,, symmetric) mapping for the studied pair.
3.4 Generalization
The division of the dataset naturally leads to the following two concepts of generalization:
Generalization on data. Generalization on data relies on the test (defined in Appendix ( In this test data, all anchor pairs (i.c., the task) are seen in the training set.
Generalization task. Generalization task depends the unseen anchors, i.c., anchors that do not appear in the training set, with a designated target mapping.
3.5
Model Architecture and Basic Experimental Setups
To enable more focused analysis of the underlying mechanisms, the following sections will only introduce the architecture of the single-head attention model.
The input sequence is represented as a one-hot vector Xin The word embedding Xem and the input to the first transformer block (1) calculated as:
eem Xintrem (11 Xem + Xoss
where X pos is trainable positional vector. For the 1-th layer, the Q, K, are defined as: Q0 = x00 k(0) = (()), y = (00
The attention matrix Attn (1) and its subsequent output X(kv(() for the 1-th layer is computed as:
(Q()
Attn (!) = softmax
(with mask), xkkv(() (y()
(5)
Vdle

---

The output of the 1-th attention layer is obtained as:
X LN(X prttt.
== (N (ON(FN(XXaOl (Xao()) xxool)), (6)
where "LN" refers to Layer Normalization and "FN" represents Fully Connected Network. The final output is obtained by projecting the output of the last layer X using linear projection layer, followed by softmax operation and argmax to obtain the predicted token.
For the basic experimental setups, we use cross-entropy loss the last token of the sequence and optimize the model using Adam with weight decay. The specific hyperparameters and training details are provided in Appendix AT
4 Two Phases of Solutions for Composite Functions
In this section, we investigate the mechanisms of how transformer learns the compositional tasks, especially on the OOD tasks.
Experimental Setup. For the 16 anchor pairs, the anchor pair (4, 3) is held out as an unseen pair during the training, while pair (3, 4) is assigned as non-inferential mapping, i.e., we set the designated target mapping a3 (z) = - 6 (the inferential mapping is x 10), seen in the training set. The other 14 anchor pairs are set as inferential mappings seen in the training set. Three possible solutions for pair (4, 3) may be learned after enough training: i) inferential solution based on the 14 anchor pairs with inferential mapping, ii) symmetric solution based on the anchor pair (3, 4), iii) other non-inferential solutions. See illustration in Fig. Il(b).
(b)
(4.3)
1. (c)
Symmennc solutio
laggr num Ste
0.4
02
0.0
Intialicatinnel
Intillcacionrtt rate
Iatullbsiion rate
Figure 2: (a,b) Phase diagram of generalization performance on the unseen anchor (4, 3). (a) The model's test accuracy based on the symmetric mapping. (b) The model's test accuracy based on the inferential mapping. The abscissa represents the initialization rate 7, which corresponds to the standard deviation (I/d of normal distribution with a mean of 0 used for parameter initializa- tion. The ordinate represents the depth of the transformer model. The shadow zones indicate the test accuracy on seen anchors is less than 90%. (c) Comparison of accuracy on the unseen anchor (4, 3) for both the inferential and symmetric solutions across different initialization rates on GPT-2. The error bars represent the standard deviation across 4-time
An important and interesting question is: which solution is learned for unseen pair (4, 3)? Exper- iments show that different initialization scales lead to different solutions for the unseen pair (4, 3) with different model depths. Fig. (laa) shows the phase diagram assuming the symmetric mapping (M(1.3) (x) = x - 6) as ground truth, with initialization scale on the abscissa and model depth on the ordinate. A large initialization scale (small 7) leads the symmetric solution, indicated by yellow zones with nearly 100% accuracy. Fig. (ibb assumes the inferential mapping (M(4.3) (r) as ground truth, showing that small initialization scale (large 7) favors the inferential solution. Larger initialization scales with deep models result in poor generalization even on seen anchors (shadow zone). This pattern was consistent across trials: for each random seed, models with differ- ent initializations and depths were trained using 9 learning rates (uniformly sampled in [1 x 10-4 3 10-41), 4D, and the highest accuracy for each of the two mappings was selected from these rates. The highest accuracies were then averaged over 3 random seeds. We conclude that small initializations favor the inferential solution, while large (but not very large) initializations favor the symmetric solution.
10)
5

---

To validate the generality of these findings, we extend our experiments to the multi-head GPT-2 model. As shown Fig. (ic), the model exhibits similar generalization performance trends on the unseen pair (4, 3) across different initialization rates 7. The range of variation is smaller in this setup, likely due to the difference in model parameter scale between the GPT-2 and the transformer models used in the previous experiments.
We also find that increasing the learning rate and weight decay coefficient within certain range is helpful for the transformer to learn the inferential solution on the unseen anchor. Please refer to Appendix E for experimental phenomena and detailed discussion.
Mechanisms of Models in Two Phases
In this section, analyze mechanisms underlying different learning solutions including: i) informa- tion flow via attention matrix; ii) information representation in vector space; iii) model complexity from different perspectives. The key message is as follows. A model with small initialization needs to gradually increase its capacity by enlarging the parameter scale, therefore, to use the least op- timization cost to fit the training data, the model only needs to learn four single anchor functions to obtain inferential mappings. As the initialization scale increases, the model tends to learn ten mappings, treating symmetric anchor pairs equivalent, and memorizes the mapping between each anchor-key item pair and its corresponding output. When the initialization becomes even larger, the model forgoes learning general patterns and instead merely memorizes the output associated with each individual data point. This causes the model to lose generalization ability, even on the anchors it has encountered during training.
5.1 Information Transmission and Fusion Mechanisms
In this section, we study the information transmission and fusion mechanisms occurring in the atten- tion matrix through information flow analysis. As shown in Fig. E (a.c), we present the information flow of two-layer networks corresponding to the two solutions. We use the same input sequence to test the output with the key item 99 and the unseen anchor pair (4, 3). The thickness of the line con- necting the j-th token in Layer / and the k-th token in Layer ++1 represents the value of the attention matrix Attn (1) at position (k:, j). For ease of observation, we manually annotate the information flow that significantly contributes to the output of the last token using different colors.
(b)
30
S8
58
98
93
93
Anchar
95
100
94
29
Layer
12
21
Layer
41
23
01.11
21
41
31
Layer 0
21
Layer
14.21
key Item
anchor pair
31
kay item-
amchor pair
Figure 3: (a, c) Information flow the two-layer networks of symmetric and inferential solutions. The input sequence shown in the figure represents the test sample, with key items and anchor po- sitions annotated. For each layer's attention matrix, we illustrate the mechanisms of information transmission and fusion through the information flow. The thickness of the line represents the corre- sponding value in the attention matrix (t) We use different colors to mark the key item and the two single anchors, and highlight the attention connections that significantly contribute to the final output. The final output sequence represents the model's output. (a) Symmetric solution. (c) Infer- ential solution. (b) T-SNE visualization of vectors x=o(1) of 10,000 input sequences with different anchor-key item pairs. Symmetric anchor pairs have similar colors in different shades.
For symmetric solutions, we find that the model's attention mechanism plays a key role in enabling the learning of consistent mappings for symmetric anchor pairs. In the first layer of the network, the attention mechanism combines the information of the two anchors and moves to the last token of the sequence. This allows the model to leam identical embeddings for symmetric anchor pairs,

---

such as (3, 4) and (4, 3), as illustrated in Fig. Bia). In the second layer, the information of the key item is broadcast to the last token and combined with the anchor information obtained from the first layer by the residual connection. This enables the model to produce the final output based on the combined information from the anchors and the key item.
To verify the symmetric nature of the learned representations after the first-layer attention, we visu- alize the vectors X X=a(1) of 10,000 input sequences with different anchor- key item pairs using t-SNE. Fig. B(b) shows that symmetric anchor pairs are clustered together the X=o(1) vector space, con- firming that the model learns to map them similar representations.
For inferential solutions, the information transmission mechanism is different. As shown in Fig. Bic), in the first layer, the key item information is moved to the positions of the two single anchors, and each anchor combined with the key item information separately. This allows the model to learn individual mappings for each single anchor. In the second layer, the anchor tokens (now containing the combined information from the anchors and the key item) are moved the last token and combined to produce the final output. This combination mechanism enables the model to learn inferential solutions on the composite anchors.
5.2 Divergence in Fused Vector Representations across Two Phases
To further investigate the mechanistic differences between symmetric and inferential solutions, examine the divergence in vector representations after the fusion anchor pair and key item infor- mation in both types of solutions. Specifically, we examine the cosine similarity between the output vectors of the second attention layer's last token (i.e., the last token of X=o(²)) for different anchor- key item combinations. We denote these output vectors as v(x; 01, 02), where is the key item and (01, 02) is the anchor pair of the input sequence
As illustrated in Fig. a b), the heatmaps display the cosine similarity between output vectors for various anchor-key item pairs. Both the abscissa and ordinate represent key item values ranging from 30 to 40. The lower and upper triangles of the heatmap correspond the cosine similarity matrices between v(x 3, 3) and v( 2, 2), and between v (z11 1, 2) and væ2; 1, 3), respectively, where ø1 and 22 are the corresponding key item values on the axes. The red boxes highlight positions where the inferential targets are the same, for example. the red highlight in the first column in Fig. D(a) indicates f(3;; 3, 3) f30; 2, 2), where f(: defined in Equation (II)
(a)
30
anchon.1
30
40
(b)
32
anchor
(c)
30
995
50
0.07
32
0,900
32
0.8
13
3.985
36
30
X
shoft
18
20 0.4
0.95
0.975
40
40
38
anchon 21
30 52
544 35
50
60
0.94
Figure 4: Cosine similarity heatmaps for vector representations in different solutions. Each axis rep- resents a selected anchor pair (labeled on the axis), with the value on the coordinate axis representing the value of the key item. The color indicates the cosine similarity between specific vectors defined in each subplot. Red boxes highlight positions where the target outputs obtained by the anchors on the abscissa and ordinate are the same for the corresponding key items. (a, b) Cosine similarity between the output vectors of the second attention layer's last token (the last token of X ) for different anchor-key item pairs in (a) inferential and (b) symmetric solutions. (c) Cosine similarity between the rows of the second layer Value matrix (V12) corresponding to the first anchor's position across different anchor-key item pairs for inferential solutions.
For inferential solutions, if two anchor-key pairs (æ1; a1, 02) and (I2; b1, b2) have the same output, i..., (x1; 01, a2) = -f (22; by b2), then their fused vectors et (21;; 01, 02) and (222; b2) are nearly
Forr simplicity. we assume that the position of the anchor-key key item combination and other items do not significantly affect the output vector.
7

---

parallel. In contrast, for symmetric solutions, the pairwise similarity between vectors v(z; 01, 02) always low even for those with the same output. This suggests that the last FNN layer of the symmetric solution memorizes all possible projections from X( to the decoder layer output. Conversely, the last FNN of the inferential solution only needs to learn fewer projections, as the same output is represented by similar vectors in the X=o(2) space.
Learning inferential solutions relies on the model's ability to infer the operation of each individual anchor. To verify the alignment between the single anchor operations in the inferential mechanism and the defined single-anchor function g(::)) in Equation (L)) we investigate the vector representa- tions after the fusion of the first single anchor and the key item. We focus on the vector representation of the second layer Value matrix V(2) at the position of the first single anchor, denoted as u(x; 01), with key item and first anchor di Similar to Fig. H b). the red boxes highlight the positions where the targets obtained by the single anchors dj and 02 on the abscissa and ordinate are the same for the corresponding key items Iı and I2, i.e., g(1 aa) = =9(22; As shown in Fig, Bic),, the information fusion of single anchors and key items and the explicitly defined single anchor opera- tions align well. This provides stronger evidence for the model's 's ability to learn the mapping of each single anchor.
5.3 Model Complexity: A Key Factor in Phase Transitions
Large initialization scales endow models with high complexity, allowing them to fit training data with minor parameter changes, as seen in FNNs in the linear regime (e.g.. Neural tangent kernel) (la col ot all MUT; Luo el all 20ZAT). Conversely, models with small initialization scales start with low complexity and gradually increase it during training. In small initialization FNNs, the input weights of different neurons often cluster along a few isolated orientations, which is referred as condensa- tions (Zhou et all, (12TT: Zhang et al, 212211. 20022.. This phenomenon of parameter condensation is closely related to the model's complexity and its ability to learn the underlying structure of the data. To better understand the mechanisms behind inferential and symmetric solutions, we investigate the degree of parameter condensation.
(a)
(b)
(c)
(d)
480
50
50
0.5
130
2744
18
100
0.0
11 st
3229
5500
17g0
150
150
-0.5
183
2737
45
50
100
150
50
100
150
-100
Figure 5: (a, b) Cosine similarity of neurons in the (Q(1) matrix. The abscissa and ordinate both represent neuron index. (a) Inferential solution with small initialization. (b) Symmetric solution with large initialization. (c. d) Visualization of the embedding space using t-SNE for different ini- tialization scales. (c) Inferential solution with small initialization. The embedded tokens seem to form arithmetic sequences with common differences of 3 (red arrow) and 4 (blue arrow) along the two directions. (d) Symmetric solution with large initialization. Please refer to Appendix for more detailed experimental results under different model depths and initialization rates 7
Condensation. We examine the condensation of matrix WQ(), The similarity between the i-th W [.)] w.
and the j-th neurons input weight is calculated by cosine similarity, i.e., Tw [| [W<0 palll"" As shown in Fig. B. for the inferential solution (Fig. Baa),, the neuron weights condense in a few directions, suggesting low complexity within the model. In contrast, for the symmetric solution (Fig. BK)) there is no obvious condensation of neuron parameters, indicating high complexity within the model.
However, even with small initialization, not all parameters exhibit significant condensation. For the word embedding matrix Wem, to distinguish the meanings of different tokens, different neuron weights have different directions, i.e., neurons do not condense. Nevertheless, the parameter matrix still exhibits clear tendency towards low complexity.
8

---

Structural Features of the Word Embedding Matrix. To gain further insights into the learning mechanisms of inferential solutions, we analyze the structure of the model's embedding matrix. We visualize the embedding space using t-SNE. As shown in Fig. Bic,d), for small initialization scales (Fig. E)))), we observe clear ordinal structure the embedding space, with the embeddings of the input tokens arranged according to their numerical value. This also suggests low-complexity tendency in the word embedding matrix, requiring the model to capture the relative ordinal relation- ships between different tokens. However, this ordinal structure is not present for large initialization scales (Fig. (i.)).
It is worth noting that the relative ordinal relationship in the inferential solution is not a simple numerical magnitude relationship of the corresponding tokens. We observe that this ordinal relation- ship may originate from the definition of the four single anchors, where the differences between the operations of any two single anchors can be obtained by the addition of the basic elements and This arrangement is consistent with the numbers being ordered with intervals of (red arrow) and 4 (blue arrow) from two directions in the embedding space. To further verify the low complexity bias of the word embedding matrix, we visualize the eigenvalues of the covariance matrix of the embedding vectors and its evolution process, the results are shown in Appendix DI
6 Further Verification on Realistic Tasks
We validated the performance of models with different initialization scales and weight decay set- tings across series of compositional and reasoning tasks. Below, introduce each task and the corresponding results. Please refer to the Appendix E for a detailed introduction of each dataset.
Compositional tasks: SCAN and COGS. For the SCAN dataset LLkee and Baroni, 2018), we selected the "Generalizing composition across primitive commands" task, where the "turn left" com- mand only appears in single-commant mands. We assess the model's generalization ability on composite commands that include the "turn left" command. For the COGS dataset (Kim and Linzen, ((2(), evaluate the model's in- distribution and out-of-distribution generalization performance after training on the same training set.
mappings and is trained alongside other composite com-
(a)
(b)
(c)
SCAN (unseen command: left)
(test in-distributton data)
(test di-ddsrrbbution data)
Accuracy on Compoted Commanss
103
ACCUrSty 1200 Dap
Acstrniv on O. aaniann Galh
75
25
75
02
00
30
29
0.5M
Uoiaset Size
Dasaset Size
Dataset Size
large init (y=0.5) &small wd (wd=0.01)
Init &large (wd-0.1)
Figure 6: Performance comparison models with different initialization scales and weight decay coefficients compositional tasks. (a) For the SCAN task, we assess the generalization ability on composite commands that include the "turn left" command. (b, c) For the COGS task, we evaluate (b) in-distribution and (c) out-of of-distribution generalization after training on the same dataset. Small initialization and large weight decay (blue) consistently outperform large initialization and small weight decay (orange) across different tasks and data scales. The parameters are initialized following zero-mean normal distribution with standard deviation of "
As shown in Fig. we display the generalization performance of models with different initialization scales and weight decay coefficients across various data sizes. Small initialization and large weight decay consistently outperform large initialization and small weight decay across different task types and data scales. Notably, in the COGS task, even when the in-distribution generalization of both set- tings (with 20k training data) reaches over 99%, the difference out-oo--distrbutin generalization remains significant.
Reasoning tasks: PAO PrOntoQA (Saparov and He, (02K) is synthetic multi-step reason- ing dataset where each data point assigns hierarchical relationships among objects and requires the

---

model to determine whether multi-step reasoning chain is correct. Fig. R illustrates the conver- gence rates and generalization errors with respect to data scale for models with large initialization (and small weight decay, Fig. Ø(a)) and small initialization (and large weight decay, Fig. (bb)). An interesting phenomenon is observed for models with large initialization (small weight decay): as the data size increases, the convergence rate first decreases and then increases. When the data size is small, the model tends to fit the data through memorization. Therefore, as the data size increases, the training difficulty increases (i.e., the training speed slows down), and the model's generaliza- tion ability is poor. As the data size grows further, the model, constrained by its complexity, can no longer memorize all the data and thus shifts to fitting the data through reasoning. This leads to an increase in fitting speed and results in better generalization. In contrast, models with small
(a)
(b)
Convergence Stape (x10g)
Convergence Steps (X101)
Test Accurscy
Test Accuracy
Memorization Regime
Inference Regime
Size
Convergence Steps
Test Accuracy
Figure 7: Performance comparison models with different initialization scales and weight decay coefficients on PrOntoQA. (a) Convergence steps and test accuracy for large initialization (r 0.5) and small weight decay (WD 0.01). (b) Convergence steps and test accuracy for small initial- ization (y = 0.7) and large weight decay (WD = 0.1). Models with large initialization initially struggle with memorization before improving as data size increases, whereas small initialization models maintain faster convergence and better generalization. The parameters are initialized follow- ing a zero-mean normal distribution with standard deviation of 7
initialization (large weight decay) inherently prefer to fit the data through reasoning, leading to faster convergence and better generalization at the same data scale compared to models with large initialization (small weight decay).
7 Discussion
Conclusion. In this work, we investigate the influence of parameter initialization scale the so- lutions learned by transformers for compositional tasks. We discover phase diagram of model solutions, where small initialization scales lead to inferential solutions that capture the underlying compositional primitives, while larger initialization scales result in symmetric solutions that memo- rize mappings. To explain these findings, analyze the information flow and vector representations within the model, revealing distinct mechanisms for inferential and symmetric solutions. Inferential solutions exhibit low complexity and learn individual mappings for single anchors, while symmet- ric solutions directly combine anchor information without learning the compositional structure. We further extend experiments to GPT-2, and verify that the insights remain valid for more complex architectures on various real-world datasets.
Limitation and Future Work.
The key limitation of our work is that the experiments and analyses are based on synthetic data, which may not fully capture the complexities real-world datasets and tasks. Although some of our conclusions have also been validated in models like GPT-2, further verification on a wider range of models is necessary. Going forward, we plan to extend our investi- gation real-world datasets and tasks, to bridge the gap between theoretical understanding and practical application. This could involve leveraging Mixture of Experts (MoEs) to design networks with different initialization scales for different expert models.
DDuring testing, we only evaluate the model's accuracy in judging hierarchical relationships. Thus, the model's 's random guessing accuracy is 50%.
10

---

Acknowledgments
This work is sponsored by the National Key R&D Program of China Grant No. 2022YFA1008200,
the National Natural Science Foundation of China Grant No. 92270001 X.). 12371511 (Z. X.), 12422119 (Z. X.), 12101402 (Y. Z.), the Lingang Laboratory Grant No. LG-QS-202202-08 (Y. Z.), Shanghai Municipal of Science and Technology Major Project No. 2021SHZDZX0102 (Z. X., Y. 2.), and the HPC of School of Mathematical Sciences and the Student Innovation Center, and the Siyuan-1 cluster supported by the Center for High Performance Computing at Shanghai Jiao Tong University, Key Laboratory of Marine Intelligent Equipment and System, Ministry Education, P.R. China. This work partially supported by STTU Kunpeng& Ascend Center of Excellence.
References
A. Vaswani, N. Shazeer, Parmar, J. Uszkoreit, Jones, N. Gomez, F.. Kaiser, Polosukhin, Attention is all you need, in: Advances in neural information processing systems, 2017, pp- 5998 6008.
J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, Anadkat, et al., Gpt-4 technical report, arXiv preprint arXiv:2303.08774 (2023).
T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al., Language models are few-shot learners, Advances in neural informa- tion processing systems 33 (2020) 1877-1901.
H. Choi, K. Hickman, A. B. Monahan, Schwarcz, Chatgpt goes to law school, J. Legal Educ. 71 (2021) 387.
T. Teubner, C. M. Flath, C. Weinhardt, W. van der Aalst, O. Hinz, Welcome to the era of chatgpt et al. the prospects of large language models, Business & Information Systems Engineering 65 (2023) 95-101.
S. Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, E. Horvitz, E. Kamar, P. Lee, Y. T. Lee, Y. Li, S. Lundberg, et al., Sparks artificial general intelligence: Early experiments with gpt-4, arXiv preprint arXiv:2303.12712 (2023).
G. F. Marcus, The algebraic mind: Integrating connectionism and cognitive science, MIT press, 2003.
P. Smolensky, R. McCoy, R. Fernandez, M. Goldrick, J. Gao, Neurocompositional computing: From the central paradox of cognition new generation of ai systems, AI Magazine 43 (2022) 308-322.
Z. Zhang, Z. Wang, J. Yao, Z. Zhou, X. Li, W. E, Z.-Q.. J. Xu, Anchor function: a type of benchmark functions for studying language models, arXiv preprint arXiv:2401.0833 (2024).
A. Radford, Wu, Child, D. Luan, D. Amodei, Sutskever, al., Language models are unsu- pervised multitask learners, OpenAl blog 1 (2019) 9.
Y. Fu, H. Peng, T. Khot, How does gpt obtain its ability? tracing emergent abilities of language models to their sources, Yao Fus Notion (2022).
Wei, Tay, R. Bommasani, C. Raffel, B. Zoph, Borgeaud, D. Yogatama, M. Bosma, D. Zhou, D. Metzler, et al., Emergent abilities of large language models, arXiv preprint arXiv:2206.07682 (2022).
A. Srivastava, A. Rastogi, A. Rao, A. A. M. Shoeb, A. Abid, A. Fisch, A. R. Brown, A. Santoro, A. Gupta, A. Garriga-Alonso, al., Beyond the imitation game: Quantifying and extrapolating the capabilities of language models, arXiv preprint (.Xi::2206004111 (2022).
R. Csordás, K. Trie, Schmidhuber, The neural data router: Adaptive control flow in transformers improves systematic generalization, arXiv preprint arXiv:2 (2021).

---

R. Csordás, K. Irie, J. Schmidhuber, Ct+++: Evaluating generalization on sever-seen composi- tional patterns of known functions, and compatibility of neural representations, arXiv preprint arXiv:2210.06350 (2022).
N. Dziri, X. Lu, M. Sclar, X. L. Li, L. Jiang, B. Y. Lin, S. Welleck, P. West, C. Bhagavatula, R. Le Bras, et al., Faith and fate: Limits of transformers on compositionality, Advances in Neural Information Processing Systems 36 (2024).
D. Hupkes, A. Singh, K. Korrel, G. Kruszewski, E. Bruni, Learning compositionally through atten- tive guidance, arXiv preprint
(2018).
M. A. Lepori. T. Serre, E. Pavlick, Break down: Evidence for structural compositionality in neural networks, arXiv preprint arXiv:2301.1088 (2023).
M. Okawa, E. S. Lubana, R. P. Dick, H. Tanaka, Compositional abilities emerge multiplicatively: Exploring diffusion models on synthetic task, arXiv preprint arXiv:2310.09336 (2023).
Yun, U. Bhalla, Pavlick, C. Sun, Do langune pretrained models learn composable primitive concepts?, arXiv preprint urXiv2200311271 (2022).
R. Ramesh, M. Khona, R. P. Dick, H. Tanaka, E. Lubana, How capable can a transformer become? a study on synthetic, interpretable tasks, arXiv preprint arXiv:2311 I. 1.199 (2023).
B. Liu, J. T. Ash, S. Goel, A. Krishnamurthy, C. Zhang, Transformers learn shortcuts to automata, arXiv preprint arXiv:2210.10744 (2022).
J. Wei, X. Wang, D. Schuurmans, M. Bosma, E. Chi, Q. Le, D. Zhou, Chain of thought prompting elicits reasoning in large language models, arXiv preprint arXiv:2201. (2022).
A. Creswell, M. Shanahan, I. Higgins, Selection-inferenee Exploiting large language models for interpretable logical reasoning, arXiv preprint arXiv:2205.09712 (2022).
A. Creswell, M. Shanahan, Faithful reasoning using large language models, arXiv preprint McXiv:2208.14277 (2022).
M. Wang, H. He, J. Wang, Z. Wang, G. Huang. F. Xiong, Z. Li, L. Wu, et al., Improving general- ization and convergence by enhancing implicit regularization, arXiv preprint arXiv:2405.2077 (2024a).
M. Wang, et al., Understanding the expressive power and mechanisms of transformer for sequence modeling, arXiv preprint arXiv:2402.0052. 002202 (2024b).
L. Wang, L. Li, D. Dai, D. Chen, H. Zhou, F. Meng. J. Zhou, X. Sun, Label words are an- chors: information flow perspective for understanding in-context learning, arXiv preprint arXiv:2305. 14160 (2023).
Y. Cao, Han, Gao, Ding, X. Xie, K. Zhou, Graphinsight: Unlocking insights in large language models for graph structure understanding, arXiv preprint arXiv:2409.032.. (2024).
S. Arora, S. S. Du, W. Hu, Z. Li, R. R. Salakhutdinov, R. Wang, On exact computation with an infinitely wide neural net, in: Advances in Neural Information Processing Systems, 2019, pp. 8141-8150.
L. Chizat, F. Bach, On the Global Convergence of Gradient Descent for Over-parameterizel Models using Optimal Transport, in: Advances in Neural Information Processing Systems 31, 2018, pp. 3036-3046.
Y. Zhang. Z.-Q. J. Xu, T. Luo, Z. Ma. A type of generalization error induced by initialization in deep neural networks, arXiv:1905.0777 [cs, stat] (2019).
W. E, C. Ma, L. Wu, A comparative analysis of optimization and generalization properties of two- layer neural network and random feature models under gradient descent dynamics., Sci. China Math. 63 (2020).
12

---

A. Jacot, F. Gabriel, C. Hongler, Neural Tangent Kernel: Convergence and Generalization in Neural Networks, in: Advances in Neural Information Processing Systems 31, 2018, pp. 8571-8580.
S. Mei, A. Montanari, P.-M. Nguyen, A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences 115 (2018) 37665-E7677) do 1073/pnas 100aa791106
G. Rotskoff, E. Vanden-Eijnden, Parameters interacting particles: long time convergence and asymptotic error scaling of neural networks, in: Advances in Neural Information Processing Systems 31, 2018, pp. 7146-7155.
J. Sirignano, K. Spiliopoulos, Mean field analysis of neural networks: A central limit theorem, Stochastic Processes and their Applications 130 (2020) 1820-1852. 1o1100.1116/..app 2019.
00K.
F. Williams, M. Trager, C. T. Silva. D. Panozzo, D. Zorin, J. Bruna, Gradient dynamics of shallow univariate relu networks, CoRR (bs/1900.07842 (2019). URL: 07842 arXiv: 1906100066480000 07842.
Luo, Z.QQ. Xu, Z. Ma, Y. Zhang, Phase diagram for two-layer relu neural networks at infinite- width limit, Journal of Machine Learning Research 22 (2021) 1-47.
H. Zhou, Q. Zhou, Z. Jin, T. Luo, Y. Zhang, Z.-Q. J. Xu, Empirical phase diagram for three-layer neural networks with infinite width, Advances in Neural Information Processing Systems (2022).
Y. Zhang, Z. Zhang, L. Zhang, Z. Bai, T. Luo, Z.-Q. Xu, Linear stability hypothesis and rank stratification for nonlinear models, arXiv preprint arXiv:2211. 112222 (2022).
1623
Z. Zhang, Z.-Q. Xu, Loss spike in training neural networks, arXiv preprint 1XX1:23055.2213 (2023).
Z. Zhang, Y. Li, T. Luo, Z.-Q. Xu, Stochastic modified equations and dynamics of dropout algorithm, arXiv preprint arXiv:2305. (2023).
Z. Zhang. Z.QQ.. J. Xu, Implicit regularization of dropout, IEEE Transactions Pattern Analysis and Machine Intelligence (2024).
A. Power, Y. Burda, H. Edwards, I. Babuschkin, V. Misra, Grokking: Generalization beyond over- fitting on small algorithmic datasets, arXiv preprint arXiv:2201.02177 (2022).
B. Wang. X. Yue, Y. Su, H. Sun, Grokked transformers are implicit reasoners: A mechanistic jourrny to the edge of generalization, arXiv preprint arXiv:2405.150 (2024).
P. Gopalani, E. S. Lubana, W. Hu, How do transformers fill in the blanks? a case study matrix completion, in: ICML 2024 Workshop on Mechanistic Interpretability, 2024.
Z. Liu, E. J. Michaud, M. Tegmark, Omnigrok: Grokking beyond algorithmic data, in: The Eleventh International Conference on Learning Representations, 2022.
X. S. Huang, F. Perez, J. Ba, M. Volkovs, Improving transformer optimization through better initial- ization, in: International Conference on Machine Learning, PMLR, 2020, pp. 4475-4483,
L. Liu, X. Liu, J. Gao, W. Chen, J. Han, Understanding the difficulty of training transformers, arXiv preprint arXiv:2004.00 08249 (2020).
A. Trockman, J. Z. Kolter, Mimetic initialization of self-attention layers, in: International Confer- on Machine Learning, PMLR, 2023, pp. 34456-34468.
H. Wang, S. Ma, L. Dong, S. Huang. D. Zhang, F. Wei, Deepnet: Scaling transformers to 1,000 layers, IEEE Transactions on Pattern Analysis and Machine Intelligence (2024).
B. Zhang, Titov, R. Sennrich, Improving deep transformer with depth-scaled initialization and merged attention, arXiv preprint arXiv:1908.1 1365 (2019).
13

---

C. Zhu, R. Ni, Z. Xu, K. Kong, W. R. Huang, T. Goldstein, Gradinit: Learning to initialize neural networks for stable and efficient training, Advances Neural Information Processing Systems 34 (2021) 16410-16422.
H. Zhou, Q Zhou, T. Luo, Y. Zhang, Z. Xu, Towards understanding the condensation of neural networks at initial training, arXiv preprint arXiv:2105.11686 (2021).
Y. Zhang, Z. Zhang, T. Luo, Z. Xu, Embedding principle of loss landscape deep neural net- works, Advances in Neural Information Processing Systems 34 (2021) 14848-14855
Y. Zhang, Y. Li, Z. Zhang, T. Luo, Z.-Q.J. Xu, Embedding principle: hierarchical structure of loss landscape of deep neural networks, Journal of Machine Learning vol 1 (2022) 1-45.
B. Lake, M. Baroni, Generalization without systematicity; On the compositional skills of sequence- to-sequence recurrent networks, in: International conference on machine learning, PMLR, 2018, pp. 2873-2882.
N. Kim, T. Linzen, Cogs: A compositional generalization challenge based semantic interpreta- tion, in: Proceedings the 2020 conference on empirical methods in natural language processing (emnlp), 2020, pp. 9087-9105.
A. Saparov, H. He, Language models are greedy reasoners: A systematic formal analysis of chain- of-thought, in: The Eleventh International Conference on Learning Representations, 2023.
Y. Hu, X. Tang, H. Yang, M. Zhang, Case-based or rule-based: How do transformers do the math?, in: Forty-firs first International Conference on Machine Learning. 2024.
14

---

A Experimental Setups
For Fig. ((a.),, Fig. B, Fig. Fig. Fig. I0. Fig. MII Fig. and Fig. IR we train the transformer model on a dataset of 900,000 samples, with each input sequence having length of 9 tokens. The key items in the input sequences are randomly sampled from the range 20-99. The model architec- ture consists of varying number of layers, single attention head, an embedding dimension of 400, a feedforward dimension of 1200, and key and value dimensions of 200 each. The model parame- ters are initialized using normal distribution with mean of 0 and standard deviation of (1/ (l/d)),, where din the input dimension of the parameter and 7 is the initialization rate. We employ the AdamW optimizer with a learning rate of lc-5, epsilon of 1e-8, weight decay of 0.01, and 31 and 32 values of 0.9 and 0.999, respectively. The model is trained for 210 epochs using a batch size of 2048 and a gradient clipping maximum norm of 1. The learning rate is scheduled using warm-up period followed by cosine decay, with a warmup period of 10 epochs, a multiplier of 25 (if there is special instruction), cosine decay number of epochs of 200, and a minimum learning rate of le-5. For Fig. G a, b). Fig. B)) Fig. b, d). we use a two-layer transformer with initialization rate two-layer transformer with initialization
y 0.5. For Fig. Bic), Fig. H c), Fig. aa c),
rate 7 = 0.8.
For Fig. (ic). Fig. R Fig. A Fig. L we employ the architecture of GPT-2 in our experiments. For Fig. ((c), ensure consistency between the model structure and the settings described in the paper, adopt the post-norm configuration. We have also conducted experiments using the pre-norm setting and obtained consistent conclusions. Apart from the differences the model architecture, we retain the same data, training strategy, parameter initialization, and other settings mentioned in the previous sections.
For Fig. I. Fig. 15 Fig. IT. we train the networks using initialization rate y = 0.5 with different learning rates and weight decay coefficients. The learning rate is scheduled using warm-up period followed by cosine decay, with warmup period of 10 epochs, multiplier chosen from the set (5, 10, 20, 40, 60), a cosine decay number of epochs of 200, and minimum learning rate of le-5.
Experiments Compute Resources. ing configuration:
The experiments were conducted on server with the follow-
48 AMD EPYC 7352 24-Core Processors, each with 512KB of cache
25IGB of total system memory
8 NVIDIA GeForce RTX 4090 GPUs with 24GB of video memory each
The experiments were run using Ubuntu 22.04 LTS operating system
15

---

B Illustrative Examples of Key Definitions
B.1 Examples of Two-anchor composite function
In this paper, we use an anchor set A = [1,2,3, 4) containing four elements. Each anchor corre- sponds specific function:
(x: 1) + 5, g 2) + 1, ( 22 g(x; 4) = =
Suppose we have an input sequence X = (23, 1, 43, 46, 74, 54, 44, 72). In this sequence, the 2nd and 3rd items (i.e., items and 2) belong the anchor set A, thus forming an anchor pair (1, 2). The item 23, immediately preceding the anchor pair, is called the key item.
For this input sequence X, the computation process of the two-anchor composite function is as follows:
f(X) = (233 = ((3 1): = g(3 + 5; = 28 + = 29.
As we can see, the two-anchor composite function first passes the key item 23 into the function g( 1) = x + 5 corresponding to the first anchor 1, obtaining the result 28: then it passes 28 into the function g(z; 2) = x + corresponding to the second anchor 2, finally obtaining the output 29.
B.2 Examples of Data Generation
In the experiments of this paper, we use four anchors (i.c., 1,2,3,4) and numbers sampled from 20 to 99 to construct the input dataset. Each input sequence includes an anchor pair, a key item (immediately preceding the anchor pair), and some noise items unrelated to the target.
For example, we can construct an input sequence as follows:
X = (52, 33, 36, 2, 78, 92, 24, 58).
In this example, the anchor pair is (2, 4), the key item is 36, and 52,33,78.,2.24.58 are noise items.
For given input sequence X, stipulate that its target is only related to the key item and the anchor pair. Furthermore, we artificially assign mappings to the anchor pairs and take the output value of the key item under the mapping corresponding to the anchor pair the target of this sequence. In the default case, we use the composite anchor function corresponding to the anchor pair as the mapping corresponding to this anchor pair, and this mapping is called the inferential mapping (derived by composing single anchor mappings). In some special cases, we may modify the mapping corresponding the anchor pair that it does not satisfy the corresponding composite anchor function, and such mappings are called non-inferential solutions. The detailed examples are shown in Section B
B.3 Examples of Mapping Type of an anchor pair
For an anchor pair 42), its mapping ((a)aa))) can have the following types:
Inferential mapping. If the designated target mapping of the anchor pair (41, a2) is consistent with the output of the two-anchor composite function, i.e., ((01.a)(z)) = fx; 01, a2), then it is called an inferential mapping.
For example, for the anchor pair (1, and key item æ, if the target output f(x;; 1,2) g(g(;; 1):2) = (x + 5) +
6, then the mapping of (1,2) is an inferential mapping.
Non-inferential mapping. If the designated target mapping of the anchor pair (a1, is inconsis- tent with the output of the two-anchor composite function, i.c., (2) * f(z; (1; az), then it called non-inferential mapping.
For example, for the anchor pair (1, 2) and key item x. if the target output is designated as x + 10, rather than (x; 1, = 0 1): (x 5) + non-inferential mapping.
6, then the mapping of (1,2) is
Symmetric mapping. If the designated target mapping of the anchor pair (a), (2) is consistent with the mapping of its symmetric anchor pair (az, a1), i.e., M(a = M M(eaaa)) (x), then called symmetric mapping.
16

---

10. If the mapping of the anchor pair (4, 3) is also designated as (M4.3) (æ) = x + 10, then the mapping of (4, 3) is symmetric mapping, because it is exactly the same as the mapping of (3, 4).
For example, suppose the mapping of the anchor pair (3, 4) is designated as (2).4)
For given model, we define its inferential (non-inferentia) symmetric) accuracy on specific an- chor pair as the accuracy obtained when using the inferential (non-inferential, symmetric) mapping of that anchor pair as the target. In particular, if model exhibits a tendency to generate outputs that align with the inferential (non-inferential, symmetric) mapping for an anchor pair, we adopt the convention of stating that the model has learned the inferential nion-inferential, symmetric) solution for that specific anchor pair.
B.4 Examples of Generalization
As introduced in the main text, there are two types of generalization in this paper: generalization on data and generalization on task. Here we provide some examples to illustrate their differences, also shown in Fig. N
Generallzation on task (accuracy on the designed target mapping)
Generalization data
anchors seen in training
anchors unseen training
sequences unseen training
47
53
Figure 8: Illustration the two types of generalization studied in this paper. Generalization on data tests model's performance on sequences with seen anchors but unseen combinations of key items and anchor pairs. Generalization on task evaluates a model's ability to infer the mappings for unseen anchor pairs.
Generalization on data. The test set for evaluating generalization on data consists of sequences with anchor pairs seen in the training set, but the sequences (detailed splitting method shown in Appendix Q are not seen in training. Models with good generalization on data should be able to correctly predict the targets for these test sequences.
Generalization on task. To test generalization on task, we need to construct sequences with anchor pairs not (anchor pair (4, 3)) during training and designate their target mappings. If the model can output targets consistent with the designated mappings for (4,3), then it achieves good general- ization on task. It is worth noting that the unseen anchor pair (4, 3) does have natural target. Therefore, we can calculate its accuracy under different designated target mappings to reflect the model's preference for the learned solutions under various settings.
17

---

C Dataset Splitting Method
straightforward division based on data ranges proves be impractical. To illustrate, consider scenario where the range of key items in the training set is denoted as (a, b). while in the test dataset, it is represented as [b+11 1, c). The encoding of data within the interval + 1,c cl is not learned during the neural network training process. As result, the neural network fails to produce the key item output for the test dataset.
To address this issue, we divide the dataset based on the value and the position of the key item, shown in Fig. Consider task with input sequence of length n. For input sequence in the training dataset, an item I' can be placed in the pos-th position of such input sequence only when mod(æ, n 2) * pos. For an input sequence in the test dataset, the token the i-th position is key item, then an item I can be placed in the pos-th position of such input sequence only when mod(æ;
important to note that the test data and training data are not completely separated in terms of values. However, when the positions of the key items are the same, the corresponding test data and training data do not overlap.
2) = pos.
Training data
Test data
mod(
.8) # pos
25
26
27
28
29
30
32
33
37
39
40
41
42
43
45
mod(
88.
Training
key
anchor pairs
pos
Test on data
1
1
3
4
Test task
Seen anchor pairs
Unseen anchor pair (4, 3)
Figure 9: Illustration of the dataset splitting method based on the value and position of the key item. The training data and test data are divided according the modulo operation on the key item value and its position in the input sequence.
We further define two types of generalization based on this dataset splitting method. For training data, use pairs of seen anchors and training data (training key items) for training. Regarding data generalization, we test the model using pairs of seen anchors and test data (test key items) to eval- uate the model's ability generalize to different items within seen composite mappings. For task generalization, we use pairs of unseen anchors with test data or training data to evaluate the model's performance masked composite mappings. It is important to note that for task generalization, we can test the accuracy of the anchor pairs with different ground truth mappings. This accuracy reflects the model's preferred mappings for these anchor pairs.
18

---

D Detailed Results for Model Complexity with Different Initialization
D.1 Detailed Results for Condensation
In order show the parameter condensation different depths and different initialization scales, we plotted the cosine similarity of the WQ(1) matrix neurons corresponding to the phase diagram in Fig. As shown Fig. M. for each subgraph, we group the weights of two neurons with cosine similarity greater than 0.7 into the same category (adjacent index the heat map) to highlight the condensation properties of neurons. is easy to see that as the initialization rate increases (the initialization scale becomes smaller), the WQ(1) matrix neurons show an obvious condensation trend, implying that the model complexity decreases.
Depth 6, y 0.3
Depth 6, y 0.4
Depth 6, y 0.5
Depth 6. 0.6
Depth 6. 0.7
1
Depth 5, y 0.3
Depth 5, y 0.4
Depth 5, 0.5
Depth 5, 0.6
Depth 5, y 0.7
Depth 4, y
Depth 4, 0.4
Depth 4, 0.5
Depth
Depth 0.7
Depth 3, y 0.3
Depth 3, 0.4
Depth 3, 0.5
Depth 3, 0.6
Depth 3, 0,7
Depth 2, y 0.3
Depth 2, 0.4
Depth 2, 0.55
Depth 2, 0.6
Depth 2, y 0.7
-1
Figure 10: Cosine similarity matrices of WQ(1) neurons at different depths and initialization scales. Each subgraph represents a different depth (from 2 to 6, bottom to top) and initialization rate (from 0.3 to 0.7, left to right). Colors indicate the cosine similarity between neurons, with warmer colors representing higher similarity. The neurons are grouped by cosine similarity greater than 0.7 to highlight the condensation properties. the initialization scale increases, neurons exhibit more condensation, indicating decreased model complexity.
19

---

D.2 Detailed Results for the Structural Features of the Word Embedding Matrix
Similar to the previous subsection on condensation, we examine the structural changes in the model's word embedding matrix as the initialization scale and model depth vary. We visualize the parameter matrix using t-SNE. As shown in Fig. I. as the initialization rate y increases, the word embedding matrix gradually exhibits distinct structural features. These structural features enable the model to accurately represent the numerical meaning of different items. Additionally, a high degree of structure indicates low complexity in the word embedding matrix (i.e., low matrix rank), aligning with the low complexity hypothesis for models with small initialization scales.
Depth 6. 0.3
Depth 6. y 0.4
Depth 6, 0.5
Depth 6, 0.6
Depth 6, 0.7
vine Rya
De
2005
Depth 5, 0.3
Depth 5, 0.4
Depth 5, y.y 0.5
Depth 5, y 0.6
Depth 5, 0.7
Depth 4, 0.3
Depth
Depth 4.
Depth
Depth 4, y
a? ia
Depth 3. 0.3
Depth 0.4
Depth 3, 0.5
Depth 0.6
Depth 3, 0.7
Depth 2, y 0.3
Depth 2. y 0.4
Depth 2, 0.5
0.5
Depth 2, 0.6.6
Depth 2. 0.7
Figure 11: Visualization the embedding space using t-SNE at different depths and initialization scales. Each subgraph represents different depth (from 2 to 6, bottom to top) and initialization rate (from 0.3 to 0.7, left to right). As the initialization scale increases, the word embedding matrix gradually exhibits distinct structural features, indicating decreased model complexity.
D.3 Detailed Rank Analysis for Different Weight Matrices
To further verify the low complexity bias of the word embedding matrix, we visualize the eigenval- ues of covariance matrix of the embedding vectors (Fig. I2)))). For small initialization scales, the covariance matrix has a small number of large eigenvalues, indicating that the model learns low-dimensional representation of the input tokens that captures their ordinal relationships. This low-dimensional representation facilitates the learning of inferential solutions by aligning with the
20

---

underlying structure of the single anchor operations. In contrast, for large initialization scales, the eigenvalues are more evenly distributed, suggesting that the model learns more distributed repre- sentation that does not effectively capture the ordinal structure. This lack of structure in the embed- ding space hinders the model's ability learn the relationships between the input tokens and the underlying single anchor operations, leading to the learning of symmetric solutions.
10
y=
40th
2nd
10th
80th
y= 0.8
3th
10
eigenvalue
10-3
egenvalue
10-
10
10-
10
0
20
40 index
60
80
10
15
20
epoch
(a) eigenvalue
(b) eigenvalue evolution
Figure 12: Eigenvalues the covariance matrix the embedding matrix for different initialization
scales and the evolution process of the eigenvalues of the small initialization model. Left: Eigenval- ues of the covariance matrix of the embedding vectors for different initialization scales. The abscissa is the eigenvalue index, and the ordinate is the eigenvalue. Colors represent different initialization scales. The definition of the initial scale y is consistent with Fig. Right: The evolution process of the cigenvalues of specific indexes of the small parameter initialization model as the training pro- gresses.
Meanwhile, we investigated the evolution of the eigenvalues corresponding specific indices the small initialization model during training. As shown in Fig. [L),,, the model gradually increases its complexity over the training process. Specifically, the model first increases the value of the largest eigenvalue, while the other eigenvalues remain almost unchanged in the initial phase of training. As training progresses, the model requires more eigen-directions to fit the data, leading to subsequent increase in the other eigenvalues. Once the model complexity increases sufficiently to fit the training data, ceases to further increase its complexity. This gradual increase in complexity ensures that the model maintains the lowest possible complexity necessary to fit the data well, enabling it to learn fundamental operations rather than merely memorizing the training data.
To further validate the low complexity of the model under small initialization, we present the singular value distributions of the weight matrices across various linear layers for both models with different initializations, as shown in Fig. I3 It is evident that, for the small initialization model, the first few singular values are significantly larger than the remaining ones. contrast, this distinct difference is not observed in the model with large initialization. This indicates the pervasive low-complexity characteristic of the internal parameters the small initialization model.
21

---

attn1 WQ
attn1 WK
attn1 WV
attn fc
Singular Value (log scale)
ffn1 fc1
ffn1 fc2
attn2 WQ
attn2 WK
attn2 WV
attn2 fc
ffn2 fc1
ffn2 fc2
Index
0=
V=
Figure 13: Singular value distributions of the weight matrices across various linear layers for models with different initializations. Each subplot corresponds representing the small initialization model (? = 0.5) and red curves representing the large initializa- tion model (7 0.8). The abseissa denotes the singular value index, and the ordinate denotes the singular value magnitude on logarithmic scale.
specific linear layer, with blue curves
22

---

E
Learning Rate and Weight Decay Coefficient Affecting Solution Phases
series of studies have experimentally investigated the impact of large learning rates and significant weight decay on the model's flatness and solution consolidation. In our experiments, we found that high learning rates and weight decay coefficients tend guide the model towards learning inferential solutions. It is important to note that to avoid training instability caused by high learning rates, we modified the settings described the main text.
In the main text, we artificially assigned non-inferential solutions to the anchor pair (3, 4) as 6. This non-inferential solution significantly disrupts the learning of single anchor mappings for models with small initialization, leading to instability during training with high learn- ing rates. In fact, if only consider the model's ' ability to learn inferential solutions, both anchor pairs (3, 4) and (4, 3) as unseen anchor pairs. Under this setting, if the model can learn the inferential solutions for (3, 4) and (4, 3), it proves that the model learns single anchor mappings. Otherwise, it learns symmetric or other solutions.
(x).))(x)
treat
As shown in Fig. E. study 3-layer, 1-head transformer model with initialization rate of 7 = 0.5. We conducted nine independent experiments with different learning rates and weight decay coefficients. The figure presents the mean accuracy of the model's inferential solutions the unseen anchor pair (4, 3). This visualization clearly demonstrates that larger learning rates and weight decay lead to higher accuracy in inferential solutions, supporting the hypothesis that these hyperparameters play a crucial role in guiding the model towards different phases.
1.0
5.0
0.8
learning rate (x10-5) )
10.0
0.6
20.0
0.4
40.0
0.2
60.0
0.0
0.003
0.01
0.03
0.1
0.3
weight decay
Figure 14: The impact of learning rate and weight decay on the accuracy of inferential solutions in 3-layer, 1-head transformer model. The heatmap displays the mean accuracy of the model's inferential solutions on the unseen anchor pair (4, 3) across nine independent experiments. The axis represents different weight decay coefficients, and the y-axis represents different learning rates. The color bar on the right indicates the accuracy of the inferential solutions on the unseen anchor (4, 3), with higher values corresponding to better performance.
23

---

E.1 Detailed Results for Condensation with Different Learning Rates and Weight Decay Coefficients
In order to show the parameter condensation at different learning rates and different weight decay coefficients, we plotted the cosine similarity of the WQQ() matrix neurons corresponding to the phase diagram Fig. T4. As shown in Fig. I for each subgraph, we group the weights of two neurons with cosine similarity greater than 0.7 into the same category (adjacent index in the heat map) to highlight the condensation properties of neurons. It is easy to see that as the weight decay coefficient and learning rate increase, the ( matrix neurons show obvious condensation trend, implying that the model complexity decreases.
wd 0.003, 5e-5 wd 0.01, 5e-5
wd 0.03, 5e-5
wd 0.1, Ir 5e-5
wd 0.3, 5e-5
1
wd 3.003, 1e-4 0.01, 1e-4
0.03, 1e-4
wd 0.1, 1e-4
0.3, 16.4
wd 0.003, 2e-4 wd 0.01. Ir 2e-4
wd 0.03, 2e-4
wd 0.1. Ir 2e-4
wd 0.3, 2e-4
wd 0.003, Ir 4e-4
wd 0.01, Ir 4e-4 wd 0.03, Ir 4e-4
wd 0.1, Ir 4e-4
wd 0.3, Ir 4e-4
0.1, 0.3,
-1
Figure 15: Cosine similarity matrices of WQ(1) neurons at different depths and initialization scales. Each subgraph represents learning rate and weight decay coefficient. Colors indicate the cosine similarity between neurons, with warmer colors representing higher similarity. The neurons are grouped by cosine similarity greater than 0.7 to highlight the condensation properties.
24

---

E.2 Detailed Results for the Structural Features of the Word Embedding Matrix with Different Learning Rates and Weight Decay Coefficients
We examine the structural changes in the model's word embedding matrix as the learning rate and weight decay coefficient vary. We visualize the parameter matrix using t-SNE. As shown Fig. II as the learning rate and weight decay coefficient increase, the word embedding matrix gradually ex- hibits distinct structural features. These structural features enable the model accurately represent the numerical meaning of different items.
wd 0.003, 5e-5
wd 0.01, 5e-5
0.03,
5e-5
wd 0.1, 5e-5
wd 0.3, Ir 5e-5
24
wd 0,003, 1e-4
wd 0.01, 1e-4
0.03, 1e-4
0.1, 1e-4
Ir 1e-4
wd 0.003, 2e-4
wd 0.01, Ir 2e-4
wd 0.03, Ir 2e-4
wd 0.1, Ir 2e-4
wd 0.3, Ir 2e-4
wd 0.003, 4e-4
wd 0.01, 4e-4
wd 0,03, Ir 4e-4
wd 0.1, 4e-4
wd 0.3, Ir 4e-4
10
13 37
wd 0.003, 6e-4
wd 0.01, 6e-4
0.03, Ir 6e-4
0.1, 6e-4
Ir 6e-4
Figure 16: Visualization of the embedding space using t-SNE at different learning rates and weight decay coefficients. Each subgraph represents a different learning rate (from 5 10 to 10-4 top bottom) and weight decay coefficient (from 0.003 0.3, left right). As the learning rates and weight decay coefficients increase, the word embedding matrix gradually exhibits distinct structural features, indicating decreased model complexity.
25

---

Detailed Introduction of Datasets and Additional Results
Compositional tasks: SCAN and COGS. SCAN and COGS are two classic compositional tasks, both of which are also synthetic but with more natural language variance, thus being ideal testbeds. For the SCAN dataset, we selected the "Generalizing composition across primitive commands" task, where the "turn left" command only appears in single-command mappings and is trained along- side other composite commands. We assess the model's generalization ability composite com- mands that include the "turn left" command. For the COGS dataset, we evaluate the model's in- distribution and out-of-distrie out-of-distributionn generalization performance after training on the same training set. In-distribution generalization is tested with data constructed from different primitives in the same combinatorial patterns. Out-of-distribution generalization is tested with data that follows dif- ferent combinatorial rules (the original paper outlines methods for generating outoo-diistrbutton data, from which we generate test data with equal probability).
reasoning dataset where each data point ussigns hierarchical relationships among objects and requires the model to determine whether multi-step reasoning chain is correct. Although we use next-token prediction for training, during testing, we only evaluate the model's 's accuracy in judging hierarchical relationships (thus, the model's s random guessing accuracy 50%).
Reasoning tasks: PrOntoQA. PrOntoQA
synthetic
Realistic tasks: Addition task and SlimPajama dataset. Unlike traditional addition tasks, we use case-based reasoning intervention experiment (Hu 2024)) study the generalization of rule learning in the addition task. Specifically, we consider the setup: a + b = c, where a, b E [0, .99]. We use a, b E [400, 600) as the test set and the remaining data as the training set. This construction prevents the model from simply mimicking training data similar to the test set. We trained simple 2-layer 1-head model and GPT-2 model. As shown Fig. (a), regardless of model size and learning mode, small initialization scales (or large weight decay coefficients) generally lead to good rule generalization, while large initialization scales (or small weight decay coefficients) fail to generalize perfectly.
(a)
Addition Task
(b)
SlimPajama Dataset
Train from Scratch
Fine-tune
Train from Scratch on GPT2 Medium
100
100
100
Perplexity on Test Data
1.72
Perplexity on Test Data
Test Accuracy
Test Accuracy
Test Accuracy
2.6
50
90
1.71
25
25
2L1H
GPT2
80
GPT2
1.70
GitHubb
2.5
Wikipedia
wd=0.0
0.5 wd-0.1
y 0.5 0.000.01
wd=0.01
0=0.5 & wd=0.1
0=0. & wd=0.5
wd=
y=1.0 wd=0.1
Figure 17: Performance comparison models with different initialization scales and weight decay coefficients realistic tasks: Addition task and SlimPajama dataset. (a) In the addition task, we use a case-based reasoning intervention experiment with a test set of a, b E [400, 600] and the remaining data as the training set. Models with small initialization scales (or large weight decay) generally show better rule generalization. (b) For the SlimPajama dataset, GPT-2 Medium models trained on GitHub and GitHub+Wikipedd data with small initialization scales consistently achieved lower perplexity. The parameters are initialized following a zero-men zero-mean normal distribution with a standard deviation "
For the SlimPajama dataset, we used two data compositions: the GitHub section and the GitHub+ Wikipedia section. We trained GPT-2 Medium models with different initializations on both datasets for 40B tokens. As shown in Fig. (b), for both data compositions, smaller initialization scales consistently achieved lower perplexity. Notably, by observing the training trajectories, we found that for the GitHub data, the model with small initialization achieved lower perplexity than
26

---

the model with large initialization early in training (around 2B tokens). For the GitHub+Wikipelle data, the model with small initialization achieved lower perplexity later in training (around 30B
tokens). This further validates the preference of small initialization for reasoning tasks. We will present the detailed training trajectories in the revised manuscript.
27

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We elaborate on setups and contribution in the abstract and introduction, especially in the last paragraph of the introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We show the detailed limitation in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't i acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
28

---

Justification:
This is work of phenomenological analysis with no theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: show the detailed setups in Appendix Al and we provide the original code in the attached file.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of model checkpoint, other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
29

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We provide the code in the attached file.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so No an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (nttps:
/
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We show the detailed experimental setups in Appendix A and the detailed data splitting method in Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question:
Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification:
We show the error bar in Fig. for multiple parallel experiments with different
seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" " if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
30

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We show the detailed information on the computer resources Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't t make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
The rescarch conducted the paper conform with the NeurIPS Code of
Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work is a phenomenological study, therefore, there is no societal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
31

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g,, deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This work is phenomenological study, therefore, this work poses such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We use common models (Transformer, GPT-2) and we cite the related works.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
32

---

If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswithcode com/datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/codeemoodel as part of their submissions structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given to participants and screenshots, if applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection. cura- tion, or other labor should be paid at least the minimum wage the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper deseribe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. you obtained IRB approval, you should clearly state this in the paper.
33

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
34


Reviews:
[{"review_id": "pxRSncCOXT", "reviewer": "Reviewer_6WHp", "summary": "Transformers show remarkable capabilities but struggle with compositional tasks, raising questions about their true understanding versus input-output mapping. The authors investigates transformers' generalization to unseen compositional tasks using anchor functions, revealing that initialization scale impacts solution types. With smaller initialization scales, models tend to infer solutions, while larger scales lead to symmetric solutions. Analysis shows symmetric solutions combine anchor information without compositional understanding, whereas inferential solutions capture underlying primitives. These insights suggest model complexity influences the ability to learn compositional primitives, with initialization scale being crucial.", "questions": "- Do you have a solid definition for \"Phase Transition\"? If so, how did you define it?\n- Do you have any additional results with other data/models (no need to run new experiments for the rebuttal)?", "limitations": "Experiments were conducted primarily on a simple synthetic setup. Future work should verify the extent to which these findings can be generalized and explore the implications of these findings on real-world applications of transformers.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The investigation into the effect of initialization on compositional abilities of transformers is well-motivated and practically useful. \n- The discovery of phase transitions from symmetric memorization to generalization based on initialization scale is compelling.", "weaknesses": "- While the problem and findings are new, the approach lacks novelty.\n- The focus on symmetric solutions is novel, but the generalizability to real data and a broader range of models remains uncertain.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "$\\textbf{Point 1}$\n\nWhile the problem and findings are new, the approach lacks novelty.\n\n$\\textbf{Reply}$\n\nWe appreciate your recognition of the novelty in our problem and findings. While our approach may appear straightforward, we hope reviewers appreciate the challenges we overcame in experimental design and mechanistic study. The relationship between initialization scale and task generalization (OOD generalization) is not well understood before, making the experimental design itself a novel direction. Specifically, in this experiment, both symmetric and inferential solutions satisfy data generalization, meaning both achieve traditional in-distribution generalization. Studying in-distribution and out-of-distribution generalization in a well-analyzed synthetic task is inherently challenging, and it was uncertain whether the model would exhibit task generalization capability.\n\nFurthermore, we analyzed the reasons behind task generalization at the parameter level, highlighting the importance of low effective complexity in compositional (inferential) tasks. Our work paves the way for new directions in developing models that integrate inference and memorization and explaining the emergence of inferential capabilities from a dynamical perspective. These contributions add significant value to our study and provide a foundation for future research in this area.\n\n$\\textbf{Point 2}$\n\nThe focus on symmetric solutions is novel, but the generalizability to real data and a broader range of models remains uncertain.\n\n$\\textbf{Reply}$\n\nWe appreciate your recognition of the novelty in our focus on symmetric solutions. Experiments on synthetic data are crucial for discovering symmetric solutions and understanding the underlying mechanisms of transformer models.\n\nTo address concerns about generalizability, we conducted additional experiments using various datasets and models. We tested different scales of GPT-2 models on datasets focused on compositional and inferential tasks, including SCAN [1], COGS [2], PrOntoQA [3], and the Addition task using case-based reasoning intervention experiments [4]. We also used real-world data from the GitHub (33B) and Wikipedia (22B) sections of the SlimPajama dataset [5]. The results from these experiments consistently supported our conclusions, demonstrating that the findings about initialization scale and its impact on learning behavior hold true across different tasks and model scales.\n\nWe will include these additional results in the revised manuscript to provide stronger evidence for the generalizability of our findings. Detailed experimental results can be found in the global response and our one-page PDF.\n\n$\\textbf{Point 3}$\n\nDo you have a solid definition for \"Phase Transition\"? If so, how did you define it?\n\n$\\textbf{Reply}$\n\nThe term \"Phase Transition\" originates from statistical mechanics, where it describes a transformation between different states of matter, when a control parameter, like temperature or pressure, reaches a critical value. This term is also often used in a series of machine learning studies [6].\n\nIn the context of our paper, we use \"Phase Transition\" to describe a similar phenomenon observed in the behavior of transformer models as a function of the initialization scale. Specifically, we define the phase transition as the point at which the model's learned solutions shift from being inferential (capturing the underlying compositional primitives) to symmetric (memorizing mappings without understanding the compositional structure). This shift occurs as the initialization scale crosses a critical threshold ($\\gamma=0.5$ in this work), analogous to the critical temperature or pressure in physical phase transitions.\n\n$\\textbf{References}$\n\n[1] Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks, ICML 2018.\n\n[2] COGS: A Compositional Generalization Challenge Based on Semantic Interpretation, EMNLP 2020.\n\n[3] Language Models Are Greedy Reasoners: A Systematic Formal Analysis of Chain-of-Thought, ICLR 2023. \n\n[4] Case-Based or Rule-Based: How Do Transformers Do the Math? ICML 2024.\n\n[5] SlimPajama: A 627B token cleaned and deduplicated version of RedPajama. https://www.cerebras.net/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama\n\n[6] Phase diagram for two-layer ReLU neural networks at infinite-width limit. Journal of Machine Learning Research (2021).", "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": "NKuzOrb6M1", "reviewer": "Reviewer_LUj2", "summary": "This paper studies the role of transformer initialization scale on its ability to compositionally generalize. Using a simple arithmetic task, the authors find that when initialization scale is low, the model learns inferential (or compositional) solutions, whereas when the initialization scale is high, models tend to learn memorized solutions, leading to poorer generalization performance. These core findings are supplemented with additional analysis of their model, e.g., through mapping the information flow of each task via analysis of the attention maps. They find that small scale initializations learn token embeddings that are interpretable; that is, the token embeddings generally learn the correct ordinal structure. Finally, they show an analysis that demonstrates that the learning time scales with the inferential complexity of the dataset/task.", "questions": "1) Is there an error in Fig. 1a? It appears the anchor for ‘4’ says ‘8’, when I think it should read ‘-8’.\n2) It would be helpful to attach some of the references included in related work to the Introduction section, where they reference the NTK work as guiding the rationale and hypotheses.\n3) Is a non-inferenetial composite anchor considered to not follow the traditional rules of arithmetic? That is, [3|4] -> -6, even though $-2 - 8 \\neq -6$? The terminology of non-inferential composite anchors was a bit confusing to me. Relatedly, in the dataset, how many non-inferential composite anchors are given? Just 1?\n4) Are key items sampled from 20-99 uniformly? Is there a motivation for why this specific set of keys? What are the noise items? The examples provided didn’t seem to provide any examples with noise.\n5) In the transformer model initialization, are *all* parameters initialized from the same scale distribution, including: embedding matrix, PE vector, Q/K/V matrices, and MLPs? Or is it a subset of these parameters?\n6) I would be curious to see the learning trajectories for Fig. 2c to get a sense of convergence rate across initializations.\n7) In general, it would be helpful if anchor pairs were labeled (perhaps with parentheses) of the actual values the correspond to. For example, it is difficult to verify that Fig. 3a/c are doing correctly without having first memorized what the values the anchor pairs refer to. \n8) I had trouble understand the motivation and corresponding results for section 5.2, and what the general conclusion/takeaway was. I found the Figure (fig 4) confusing to understand. This is likely related to the above question 7, where it is hard to keep track what the anchor values refer to without having to refer back to Fig 1.\n9. Please label the rows and columns of Fig. 5a.\n10. Condensation (line 262) is never defined.\n11. Regarding Fig 6. There appears to be an effect that adding layers in the small initialization regime improves the sample efficiency of the model. Is this effect true?\n12. Lines 132, 133: From the definitions it wasn’t immediately clear what generalization on data meant. Does this mean different key values were used on generalization? Whereas task generalization is specific to anchor pairs?", "limitations": "The dataset/task is limited and focuses only on a specific type of ‘systematic compositional generalization’ (i.e., novel combinations of anchor pairs), and does not fully address other forms of compositional generalization, such as length generalization.\n\nAlso, the clarity of the paper is a bit difficult to follow, which could potentially limit its impact.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The study pairs an interesting set of subfields in deep learning: compositionality with neural tangent kernel learning theory. The results are unambiguous and demonstrate strong effects. I have little doubt that these effects are replicable, and that insights from this line of work can be transferred to other domains.", "weaknesses": "The primary weaknesses:\n\n1) Other studies demonstrate that this type of arithmetic generalization can be achieved without specific scale-dependent initializations on the arguably more difficult task of length generalization by manipulating positional encodings (Zhou et al., 2023; Zhou et al., 2024; Ruoss et al., 2024; McLeish et al., 2024). In some ways, this paper addresses a simpler form of compositional generalization – systematic compositional generalization – which focuses on generalization on novel combinations of previously seen elements (rather than novel structures of previously seen elements (productive generalization)) (Hupkes et al., 2020).\n\n2) Clarity: I found the paper difficult to understand in some parts. This was partially due to writing, but also partially due to lack of labeling (e.g., figures). For example, the paper often uses the token ID of the anchor pairs (e.g., 3 and 4), rather than the actual referent values of those anchors (-2 and -8), which makes it difficult to follow exactly what figures are showing. Unfortunately, I did not understand the motivation, results, and conclusions of the analyses presented in Section 5.2 and Figure 4. It could be that it was already redundant with results from Section 5.1 (which I did find interesting and easy to follow).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "$\\textbf{Weakness 1}$\n\nOther studies demonstrate that this type of arithmetic generalization can be achieved...\n\n$\\textbf{Reply}$\n\nArithmetic generalization is a widely studied problem, and several works have improved performance primarily on arithmetic tasks by manipulating input token forms or positional encodings based on arithmetic properties. These approaches, however, are often task-specific and do not generalize to broader reasoning tasks. We will add a discussion of these relevant studies to the related work section.\n\nOur approach, on the other hand, aims to provide a more general solution. The discussion on initialization scale in our study is not confined to a specific arithmetic task. The synthetic data serves as an example to illustrate the impact of initialization scale. Our goal is to find a general method that enhances model performance across a wide range of compositional and inferential tasks, not limited to arithmetic tasks.\n\n$\\textbf{Weakness 2 and Problem 7}$\n\nIt would be helpful if anchor pairs were labeled...\n\n$\\textbf{Reply}$\n\nThank you for this suggestion. We will update the figure captions to label anchor pairs with their corresponding values.\n\n$\\textbf{Weakness 2 and Problem 8}$\n\nI had trouble understanding the motivation for Section 5.2...\n\n$\\textbf{Reply}$\n\nSection 5.1 demonstrates how models with different initialization scales learn anchor pairs differently. However, this alone does not fully explain the models' behavior when key items are introduced.\n\nSection 5.2 extends this by showing that with the addition of key items, models with small initialization scales can reduce model complexity through structured vector representations. Sequences with the same target have consistent internal representations, requiring the model to learn only around 100 vector representations. In contrast, models with large initialization scales treat each anchor-key item combination as an independent mapping, necessitating the memorization of 800 vector representations (10 types of symmetric anchor pairs and 80 types of key items).\n\n$\\textbf{Limitation}$\n\nThe dataset/task is limited...\n\n$\\textbf{Reply}$\n\nWe have also validated our conclusions on a wider range of more complex and real-world datasets, supporting the generality of our findings. For detailed experimental results, please refer to the global response and Reviewer 6WHp's Point 2.\n\n$\\textbf{Problem 1}$\n\nIs there an error in Fig. 1a...\n\n$\\textbf{Reply}$\n\nWe have carefully examined Figure 1a and did not find any instance where the anchor for '4' is mislabeled as '8'. We would appreciate if you could specify the exact location.\n\n$\\textbf{Problem 2}$\n\nIt would be helpful to attach some of the reference...\n\n$\\textbf{Reply}$\n\nThank you for the suggestion. we add the following sentence to the last paragraph of the Introduction:\n\n\"This finding corresponds to a series of studies on different initialization regimes in fully-connected networks, such as NTK [Refs] and mean-field [Refs] regimes.\"\n\n$\\textbf{Problem 3-5}$\n\nClarification of the experimental setup.\n\n$\\textbf{Reply}$\n\nA non-inferential composite anchor does not follow traditional rules. Our intention is to distinguish between symmetric and inferential solutions; otherwise, inferential solutions would be a subset of symmetric solutions under our data setup. There is only one non-inferential composite anchor, as noted in Fig. 1(a).\n\nKey items are sampled uniformly. Since input vectors are encoded in a one-hot format, the key items set only needs to be distinct from the anchors, allowing flexibility in their range. Noise items (gray numbers in \"Input data examples\" in Fig. 1a) are non-anchor items preceding the anchor, with remaining items as noise. Noise item values are irrelevant to the sequence target.\n\nAll parameters are initialized from the same scale distribution.\n\n$\\textbf{Problem 6}$\n\nI would be curious to see the learning trajectories...\n\n$\\textbf{Reply}$\n\nDue to space limitations in the one-page PDF, we will present the convergence rates for Fig. 2c in the revised manuscript. The one-page PDF includes the convergence rates results of the GPT-2 model on PrOntoQA, a multi-step reasoning dataset, which show similar conclusions to Fig. 2c.\n\nFor models with large initialization, the convergence rate initially slows down and then speeds up as the data size increases, demonstrating the phase transition from memorization to inference. In contrast, models with small initialization exhibit consistently fast convergence rates and high generalization accuracy due to their stronger preference for inferential solutions.\n\n$\\textbf{Problem 9}$\n\nPlease label the rows and columns of Fig. 5a.\n\n$\\textbf{Reply}$\n\nWe label the rows and columns of Fig. 5a.\n\n$\\textbf{Problem 10}$\n\nCondensation (line 262) is never defined.\n\n$\\textbf{Reply}$\n\nWe add the following explanation to Section 5.3:\n\n”In small initialization FNNs, the input weights of different neurons cluster along a few isolated orientations, which is referred to as condensation.“\n\n$\\textbf{Problem 11}$\n\n...small initialization regime improves the sample efficiency of the model?\n\n$\\textbf{Reply}$\n\nIn Fig. 6, we examine the iterations needed for models of varying depths to fit data with different inferential complexities. For high inferential complexity data, some solutions are made non-inferential, requiring higher model complexity and more fitting steps. Increasing layers enhances model complexity, reducing the number of required steps. Since we did not consider generalization in this context, sample efficiency is not directly reflected in the figure.\n\n$\\textbf{Problem 12}$\n\nLines 132, 133: ... what generalization on data meant.\n\n$\\textbf{Reply}$\n\nGeneralization on data refers to the model's performance on test data where the anchor pairs are seen during training, but the key item and noise items are different. Task generalization, on the other hand, refers to the model's performance on test data with unseen anchor pairs.", "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": "EBTYQgYEuu", "reviewer": "Reviewer_MVAX", "summary": "The work mainly investigates whether transformers indeed learn primitive functions or simply learn surficial correlations (i.e., \"memorizing\") when trained on a mannually-designed toy compositional task, which contains a set of composite numerical functions (e.g., $g_1(x)=x+5, g_2(x)=x+1,...$). \n\nAuthors find that the learned mechanism has a phase transition (which correlates with the parameter initialization scale) : smaller initialization scale leads to learning the composite solution (\"inferential solution\"), where models learn the exact meanings of primitive functions and compose them to solve the task; while larger initialization scale leads to learning solutions based on the surficial correlation in the training data (\"symmetric solution\"). \n\nCentering on this finding, the paper also conducts empirical analysis from the perspectives of information transmission, the similarity of vector representations learned by transformers and model complexity to offer interpretation.\n\nBesides, the paper also finds that there is a positive correlation between the the complexity of the training data (i.e., the number of the data that do not meet the composition rule.) and the difficulty of fitting the training data.", "questions": "Q1. Regarding the first point in the Weaknesses part:\n\nQ1.1. Can the main conclusion and the analysis results introduced in the paper generalize to other well-studied compositional tasks (e.g., SCAN [1] or COGS [2]) ?\n\nQ1.2. From my point of view, a direct usage for the conclusion introduced in this paper might be initializing the LoRA [3] parameters (or Prompt-Tuning [4] learnable parameters) when fine-tuning pre-trained language models (PLMs). I am wondering if we can observe similar phenomenon when we use PLMs + LoRA/Prompt-Tuning to learn solutions for the task used in this paper or other tasks (e.g., SCAN or COGS)?\n\nQ2. Regarding the second point in the Weaknesses part:\n\nQ2.1. How to understand the difference and the correlation between the conclusions of this work and [5] ?\n\nQ2.2. For larger initialization scale, if we train the transformers far beyond over-fitting, will the learned patterns change (from memorization to inferential solutions)? just as the \"grokking\" paper described.\n\nQ2.3. What will the learned solution be? if we use larger initialization scale and larger weight decay?\n\nQ2.4. I suggest the authors discuss \"grokking\" works in the Related Work section.\n\nQ3. An additional question: In Section 6, it is claimed that \"Based on our analysis of the mechanisms underlying different solutions, we can predict that models with different initialization scales will exhibit varying degrees of difficulty in learning data with different complexities.\". I do not see clear relationship between the initialization scale and data complexity discussed in the previous sections. Could you please explain more about this relationship (between Section 6 and previous sections)? Besides, I think instead of using \"epochs for 100% training acc\", it is more meaningful to use the \"epochs for 100% testing acc\" (e.g., using inferential solution as the ground-truth).\n\nQ4. A minor suggestion: the abbreviations are supposed to be introduced when they are used in the paper first time. (e.g., Eq(6), FNN, by the way, I think it should be FFN (Feed-Forward Network) here.) \n\n[1]: Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks, https://arxiv.org/abs/1711.00350 .\n\n[2]: COGS: A Compositional Generalization Challenge Based on Semantic Interpretation, https://arxiv.org/abs/2010.05465 .\n\n[3]: LoRA: Low-Rank Adaptation of Large Language Models, https://arxiv.org/abs/2106.09685\n\n[4]: The Power of Scale for Parameter-Efficient Prompt Tuning, https://arxiv.org/abs/2104.08691\n\n[5]: Omnigrok: Grokking Beyond Algorithmic Data, https://arxiv.org/abs/2205.10343 .", "limitations": "The authors do discuss the limitations of this work in Section 7.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The topics of this paper, whether the solutions learned by transformer-based models are inferential or just memorizing, are very important and of interest to the NeurIPS community. Besides, the most part of the paper has clear logic structure, thus being easy-to-follow.\n\n2. (Based on the setting discussed in the paper) Authors conduct analysis experiments from multiple aspects, including tracing the information flow within the trained transformers, inspecting the similarity of learned representations and the model complexity, all of which lead to a converging interpretation to the initial observation: smaller initialization scale leads to inferential solutions where larger initialization scale leads to memorizing surficial patterns.", "weaknesses": "My main concerns are two-fold:\n\n1. The experiment setting of this work might be over-simplified (as the authors discussed in the limitation part): the synthetic data with little variance (four simple primitive functions, sixteen combinations in total). Though I believe the conclusions in the paper are correct for this setting, I question the generality of the conclusions and the analysis results in other (more complex) settings (especially in this paper, there is only empricial evidence but no explicitly theoretical guarantee for the conclusion). \nIf we train transformers on other classic compositional tasks (e.g.,: SCAN [1] or COGS [2], both of which are also synthetic but with more natural language variance, thus being ideal testbeds), will we have similar findings?\nBesides, the further discussion and validation (at least some initial experiments) on how to use the conclusion to guide the training (fine-tuning) of real pre-trained language models are also lacking. All these factors weaken the contribution of this work.\n\n2. The research question discussed in the paper (i.e., inference v.s. memorization) is highly related to another line of works (\"grokking\") [3, 4, 5, 6] (i.e., generalization v.s. memorization). The mainly conlusion in the paper (initialization scale correlates to the learned pattern, and even the weight decay part in the Appendix E) is also similar with the conclusion in this \"grokking\" work [5]. I find the discussion of the relationship (maybe the difference) between this work and the previous \"grokking\" works are missed in the paper, which weaken the novelty of the work.\n\n[1]: Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks, https://arxiv.org/abs/1711.00350 .\n\n[2]: COGS: A Compositional Generalization Challenge Based on Semantic Interpretation, https://arxiv.org/abs/2010.05465 .\n\n[3]: Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets, https://arxiv.org/abs/2201.02177 .\n\n[4]: Towards Understanding Grokking: An Effective Theory of Representation Learning, https://arxiv.org/abs/2205.10343 .\n\n[5]: Omnigrok: Grokking Beyond Algorithmic Data, https://arxiv.org/abs/2205.10343 .\n\n[6]: Explaining grokking through circuit efficiency, https://arxiv.org/abs/2309.02390.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "$\\textbf{Weakness 1.1 and Q 1.1}$\n\nThe experiment setting of this work might be over-simplified...\n\n$\\textbf{Reply}$\n\nWe appreciate you suggesting of the SCAN and COGS datasets. We have conducted experiments on these two tasks, and the results are consistent with our conclusions. Additionally, we have tested our findings on a broader range of more complex tasks. For detailed experimental results, please refer to the global response and Reviewer 6WHp's Point 2.\n\n$\\textbf{Weakness 1.2}$\n\n...how to use the conclusion to guide the training...\n\n$\\textbf{Reply}$\n\nThe primary goal of this work is to guide model architecture design and pre-training processes, rather than fine-tuning on pre-trained models. One direct application is to use different initialization scales and weight decay coefficients based on the type of training data, such as text versus code, to adjust the model's preference for memorization or reasoning.\n\nFor example, in the PrOntoQA task, models with small initialization scales exhibit faster learning and better generalization than those with large scales. Similarly, with the SlimPajama dataset, small initialization scales show lower perplexity in the early stages of training on GitHub data (reasoning tasks) compared to large scales. However, when training on combined GitHub and Wikipedia data (reasoning + memorization), small scales show lower perplexity only in the mid-to-late stages (around 120,000 iterations). This aligns with our prediction about initialization scale preferences for different data types.\n\nFurthermore, we can explore new model architectures, such as MoE, with varying initialization scales to balance memory and reasoning preferences, aiming for good generalization and training efficiency.\n\n$\\textbf{Q1.2}$\n\nA direct usage … initializing the LoRA.\n\n$\\textbf{Reply}$\n\nThe low complexity brought by small initialization can indeed guide model fine-tuning. However, LoRA's low-rank design and zero initialization already confer low complexity.\n\nTo illustrate, we compared LoRA (r=16, zero initialization for matrix A) with higher complexity learnable parameters (r=768, non-zero initialization for matrix A) on the COGS dataset using a pre-trained GPT-2 model. The LoRA setup achieved 90.4% OOD accuracy, while the higher complexity setup, despite having more learnable parameters, only reached 65.8% OOD accuracy. This supports our understanding of the benefits of small initialization for reasoning tasks.\n\nThe limitation of LoRA lies in its rank cap, which prevents it from organically adjusting complexity for more intricate tasks. In contrast, small initialization parameter matrices can flexibly adjust complexity based on task requirements.\n\n$\\textbf{Q2.1 and Q2.4}$\n\nThe difference and the correlation between this work and [5].\n\n$\\textbf{Reply}$\n\nThank you for highlighting relevant related work. Both our study and [5] explore how parameter initialization scale and weight decay affect model behavior, but with different focuses. We will incorporate a detailed discussion of these points and other relevant grokking studies into the related work section of our revised manuscript.\n\nIn [5], the authors vary initialization scale and weight decay to study grokking, focusing on in-distribution generalization. They use unusually large initialization weights, referred to as inducing grokking. Normal and smaller initialization weights do not lead to grokking, so [5] does not provide strong support for generalization differences between normal (what we refer to as large) and small initialization scales.\n\nOur work, in contrast, emphasizes compositional and reasoning tasks, focusing on out-of-distribution generalization. This type of generalization is not a trivial extension of in-distribution generalization. We classify solutions from different initializations into three regimes: symmetric solutions, inferential solutions, and solutions without data generalization. Solutions without data generalization align with large initialization settings in [5] on datasets like MNIST. Both symmetric and inferential solutions show good in-distribution generalization, but [5] does not distinguish between them. We emphasize that for reasoning and compositional tasks with high logical complexity, small initialization is more likely to learn intrinsic data patterns, leading to inferential solutions.\n\n$\\textbf{Q2.2.}$\n\nThe pattern when far beyond over-fitting.\n\n$\\textbf{Reply}$\n\nWhile this hypothesis is intuitively plausible, we did not observe this phenomenon in our experiments. With large weight decay, the model quickly learns inferential solutions. However, with small weight decay, even after extending the training steps by a factor of 100, the model did not exhibit grokking.\n\n$\\textbf{Q2.3}$\n\nWhat will the learned solution be if we use larger initialization scale and larger weight decay?\n\n$\\textbf{Reply}$\n\nIf the weight decay is sufficiently large and the training process is stable, the model will learn inferential solutions.\n\n$\\textbf{Q3}$\n\nClarify the relationship between initialization scale and data complexity\n\n$\\textbf{Reply}$\n\nWe add the following at the beginning of Section 6: \"The symmetric solution (large initialization) focuses on symmetric target mappings, while the inferential solution (small initialization) emphasizes logical consistency. Consequently, large initialization models see symmetric mappings as equally complex, whereas small initialization models find non-inferential data more challenging due to the need to understand underlying logical relationships..\"\n\nWe will use the test set on seen anchors to determine fit quality. Since we have deliberately disrupted the inferential properties of symmetric solutions in multiple datasets, using inferential solutions as the ground truth for testing accuracy is not ideal.\n\n$\\textbf{Q4}$\n\nThe abbreviations should be introduced.\n\n$\\textbf{Reply}$\n\nWe will ensure that all abbreviations are introduced upon their first usage in the paper.", "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:
We thank the reviewers for your thoughtful and insightful comments. We have addressed every comment, and believe that, taken together, the reviewers' comments have improved the manuscript significantly. To address reviewers' common concerns, we have supplemented relevant experimental results, added more discussion, and improved presentation. Our improvements to the draft mainly include the following points.

$\textbf{Adding Numerical Verification on Other Realistic Tasks.}$

We validated the performance of models with different initialization scales and weight decay settings across a series of compositional and reasoning tasks. Below, we introduce each task and the corresponding results. Notably, models with small initialization (large weight decay) $\textbf{consistently outperformed}$ those with large initialization (small weight decay) across all tasks. For the figures mentioned below, please refer to the PDF of the global response.

$\textbf{1. Compositional tasks: SCAN and COGS.}$ SCAN [1] and COGS [2] are classic compositional tasks with more natural language variance. For SCAN, we used the "Generalizing composition across primitive commands" task, where the "turn left" command appears only in single-command mappings during training. We assess generalization on composite commands including "turn left". For COGS, we evaluate in-distribution and out-of-distribution generalization after training on the same set. In-distribution tests use data with different primitives in the same combinatorial patterns, while out-of-distribution tests use data following different combinatorial rules.

As shown in Fig. 1, we display the generalization performance of models with different initialization scales and weight decay across various data sizes. Small initialization and large weight decay consistently outperform large initialization and small weight decay across different task types and data scales. Notably, in the COGS task, even when the in-distribution generalization of both settings reaches over 99\%, the difference in out-of-distribution generalization remains significant. 

$\textbf{2. Reasoning task: PrOntoQA.}$ PrOntoQA [3] is a synthetic multi-step reasoning dataset that assigns hierarchical relationships among objects, requiring the model to determine the correctness of reasoning chains. We use next-token prediction for training but evaluate the model's accuracy in judging hierarchical relationships during testing (random guessing accuracy is 50%).

Fig. 2 shows convergence rates and generalization errors for models with large initialization and small weight decay (Fig. 2(a)) and small initialization and large weight decay (Fig. 2(b)). For large initialization, as data size increases, the convergence rate first decreases due to increased memorization difficulty. However, as the data size grows further, the model, constrained by its complexity, shifts to reasoning, which then improves generalization. In contrast, models with small initialization prefer reasoning from the start, leading to faster convergence and better generalization.

$\textbf{3. Realistic tasks: Addition task and SlimPajama dataset.}$ Unlike traditional addition tasks, we use a case-based reasoning intervention experiment [4] to study the generalization of rule learning in the addition task. Specifically, we consider the setup: $a + b = c$, where $a, b \in [0, 999]$. We use $a, b \in [400, 600]$ as the test set and the remaining data as the training set. This construction prevents the model from simply mimicking training data similar to the test set. As shown in Fig. 3(a), regardless of model size and learning mode, small initialization scales (or large weight decay) generally lead to good rule generalization, while large initialization scales (or small weight decay) fail to generalize perfectly.

For the SlimPajama dataset [5], we used two data compositions: the GitHub section and the GitHub+Wikipedia section. We trained GPT-2 Medium models with different initializations on both datasets for 40B tokens. As shown in Fig. 3(b), for both data compositions, smaller initialization scales consistently achieved lower perplexity. Notably, the model with small initialization achieved lower perplexity earlier in training (around 2B tokens) for GitHub data and later (around 30B tokens) for GitHub+Wikipedia data, validating the preference of small initialization for reasoning tasks. Detailed training trajectories will be presented in the revised manuscript.

$\textbf{Enhancing Clarity and Adding Related Work Discussion }$

To improve clarity and enhance the reading experience, we have made several revisions to the paper:

• Added anchor mappings in figure captions.

• Clarified motivations in Sections 5.2 and 6.

• Included necessary learning trajectories in the appendix.

• Provided definitions for terms such as "condensation".

• Expanded abbreviations like "FFN" to their full forms.

We have also added discussions on relevant works concerning grokking, NTK and arithmetic tasks to highlight our study's contributions. While some existing works explore the impact of initialization on models, our focus on how initialization affects data learning patterns, particularly small initialization's preference for compositional and reasoning tasks, is both novel and significant. Additionally, our exploration of the internal mechanisms of models under different initialization settings provides valuable insights to the field. 

$\textbf{References}$

[1] Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks, ICML 2018.

[2] COGS: A Compositional Generalization Challenge Based on Semantic Interpretation, EMNLP 2020.

[3] Language Models Are Greedy Reasoners: A Systematic Formal Analysis of Chain-of-Thought, ICLR 2023. 

[4] Case-Based or Rule-Based: How Do Transformers Do the Math? ICML 2024.

[5] SlimPajama: A 627B token cleaned and deduplicated version of RedPajama.

Meta-review: The reviewers' initial assessment was that while the paper introduced worthwhile ideas and a sound method, the empirical validation using "toyish' datasets did not indicate whether the method could generalize to larger-scale compositional tasks (including non-arithmetic ones). The reviewers also made other remarks regarding the possible limited technical novelty of the "methods" introduced in the paper, the lack of clarity of some of the reported results, and possible missing discussion of related works (e.g. related to length generalization in arithmetic problems and grokking).

The authors provided a comprehensive response, which included several additional results on larger-scale datasets, including more "standard" ones (e.g., SCAN for "compositional navigation commands" and COGS "Compositional Generalization Challenge Based on Semantic Interpretation") and also reasoning and language modelling datasets/tasks. The new results match the smaller-scale results from the initial submission. These were also useful in demonstrating the work's aim was more general than arithmetic problems. 

Regarding "grokking," the authors clearly explained the similarities and differences and agreed to add that discussion to the paper. Similarly, the authors provided specific clarifications to be added to the paper. 

While novelty is important, there are different types of papers, and this submission provides and validates timely novel insights. As such, I find that it stands without a novel method.

Finally, the reviewers agree that this paper should be accepted, and I am happy to recommend acceptance. Congratulations!
Decision: Accept (poster)