Title: Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing
Venue: NeurIPS 2024
Accepted: yes
Keywords: parameter initialization

transformer

compositional task

reasoning

memorizing
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

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 (memorybased) 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 γ to be a convenient tunable hyper-parameter in common deep learning frameworks, where 1 /d γ in is the standard deviation of parameters of the layer with d in input neurons.

1 Introduction

Large-scale transformers (Vaswani et al., 2017) demonstrate unprecedented capabilities (Achiam et al., 2023; Brown et al., 2020; Choi et al., 2021; Teubner et al., 2023), even noted as 'sparks of AGI" (Bubeck et al., 2023). However, their performance on compositional tasks, which are considered a key property of human cognition (Marcus, 2003; Smolensky et al., 2022), has long been the subject of intense debate, especially in terms of the out-of-distribution (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?

In this work, we use anchor functions (Zhang et al., 2024), 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. 1(a)). The single anchors are specific tokens (i.e., 1, 2, 3, 4), each corresponding to a particular arithmetic operation (left side of Fig. 1(a)). A composite function is defined as the

sequential application of two single anchor functions, i.e., the corresponding arithmetic operation, to a key item. For example, given a key item x and an anchor pair (1, 2), the composite function would output ( x +5) + 1 , i.e., x +6 . As shown in the middle part of Fig. 1(a), during training, 14 out of the 16 possible anchor pairs are assigned inferential (reasoning-based) mappings, meaning 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 of the composite function. The remaining pair (4, 3) is held out as an unseen task. After training, 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 1 in Fig. 1(b)), an inferential solution (Mechanism 2 in Fig. 1(b)), or other non-inferential solutions. The central question we aim to answer is: which of these three types of solutions will the model learn for the unseen anchor pair (4, 3)?

(a)Data generation

(b)Possiblesolutionsandmechanisms

fortheunseen anchor

Single anchors

14seeninferentialcomposite anchors

Inputdataexamples

Target

Mechanism 1:learn symmetric structure

+10

55 463252281 13433

38

D

b

二

a

+5

2

+6

2095431244347632

49

+1

Composition

16

Padding

54773532564457 43

40

3

3

4

-6

symmetric solution

3

2

1seennon-inferentialcompositeanchors

2853447832623444

56

Mechanism 2:composite from theinfered single anchor mappings

3

4

-6

7743236389335243

？

3

-2

4

-8

1unseencompositeanchor

Q:Which solution does the model learn

4

？

on the unseen anchor?

4

3

-10

inferentialsolution

seen inferential anchor pair

seen non-inferential anchor pair

unseen anchor pair

key item (item before anchor pair)

arithmeticoperation

noise item

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 larger initialization scales, the model is more likely to learn symmetric solutions on anchor pair (4, 3) that simply match the output of its symmetric anchor pair (3, 4). These findings suggest that the initialization scale plays a 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 information flow and vector representations within the model. We find that symmetric solutions directly combine anchor information without capturing the underlying compositional primitives, while inferential 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 transformers on compositional tasks, enabling them to better capture compositional structures and generalize 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 focuses on a single-head attention model to facilitate a clearer understanding of the underlying mechanisms. We further extend our experiments to GPT-2 (Radford et al., 2019), and verify that the

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 γ to be a convenient tunable hyper-parameter in common deep learning frameworks, where 1 /d γ in is the standard deviation of parameters of the layer with d in input neurons.

2 Related Work

Challenges of transformers in compositional tasks. Recent advancements in large language models (LLMs) have showcased remarkable capabilities, often surpassing human performance (Fu et al., 2022; Wei et al., 2022). However, despite their impressive performance on single-step reasoning tasks (Srivastava et al., 2022), transformers struggle with multi-step compositional tasks and OOD generalization (Csordás et al., 2021, 2022; Dziri et al., 2024; Hupkes et al., 2018; Lepori et al., 2023; Okawa et al., 2023; Yun et al., 2022). Ramesh et al. (Ramesh et al., 2023) show that transformers trained to directly compose capabilities struggle to generalize to OOD tasks with a synthetic task. Liu et al. (Liu et al., 2022) 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 a single generation (Wei et al., 2022), leveraging LLMs to generate reasoning steps iteratively (Creswell et al., 2022; Creswell and Shanahan, 2022). Despite these efforts, achieving complete mastery in compositional tasks remains a significant challenge for vanilla transformers. A series of works studies the internal mechanisms of language models and improves the capabilities of language models (Wang et al., 2024a,b, 2023; Cao et al., 2024). In order to clearly study the behaviors and internal mechanisms of language models, Zhang et al. Zhang et al. (2024) 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 is important to determine the final fitting result of the network (Arora et al., 2019; Chizat and Bach, 2018; Zhang et al., 2019; E et al., 2020; Jacot et al., 2018; Mei et al., 2018; Rotskoff and Vanden-Eijnden, 2018; Sirignano and Spiliopoulos, 2020; Williams et al., 2019). Luo et al. (Luo et al., 2021), Zhou et al. (Zhou et al., 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 al., 2022; Zhang and Xu, 2023; Zhang et al., 2023; Zhang and Xu, 2024). A line of works links the grokking phenomenon with the improvement of reasoning ability (Power et al., 2022; Wang et al., 2024; Gopalani et al., 2024) and points out that the initialization scale has an important influence on the occurrence of grokking (Liu et al., 2022). Recent studies also investigate the impact of initialization on the training process of LLMs (Huang et al., 2020; Liu et al., 2020; Trockman and Kolter, 2023; Wang et al., 2024; Zhang et al., 2019; Zhu et al., 2021). These works primarily focus on how the initialization scale affects the stability of the training process and plays a crucial role in ensuring smooth and effective training of LLMs. In our work, we find that initialization scales can significantly influence a model's ability to 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.

3 Definitions

Weintroduce a 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 B.

3.1 Two-anchor composite function

A two-anchor composite function f ( X ) : R n → R is defined as

Here, the input sequence X = ( x 1 , . . . , x n ) comprises n tokens. An anchor set A = { a 1 , a 2 , . . . , a J } is designated, with each token a k ∈ A corresponds to a function g ( x ; a k ) . In each X , one and only one pair of two consecutive elements belong to A , such as x i , x i +1 ∈ A .

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 f ( x i -1 ; x i , x i +1 ) to emphasize the anchor pair ( x i , x i +1 ) and key item x i -1 .

In this work, we set the anchor set A = { 1 , 2 , 3 , 4 } . Each anchor token corresponds to a specific function:

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, a 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 a 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, i.e., 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 ( a 1 , a 2 ), the type of its mapping M ( a 1 ,a 2 ) ( · ) can be classified as follows.

Inferential mapping. The designated target mapping of an anchor pair ( a 1 , a 2 ) is consistent with the two-anchor composite function, i.e., M ( a 1 ,a 2 ) ( x ) = f ( x ; a 1 , a 2 ) . This type of solutions exemplifies the reasoning ability of models.

Non-inferential mapping. The designated target mapping of an anchor pair ( a 1 , a 2 ) is inconsistent with the two-anchor composite function, i.e., M ( a 1 ,a 2 ) ( x ) = f ( x ; a 1 , a 2 ) .

̸

Symmetric mapping. The designated target mapping of an anchor pair ( a 1 , a 2 ) is consistent with the mapping of its symmetric anchor pair ( a 2 , a 1 ), i.e., M ( a 1 ,a 2 ) ( x ) = M ( a 2 ,a 1 ) ( x ) .

We refer to a model as an inferential (non-inferential, symmetric) solution on an anchor pair if it tends to output the mapping corresponding to 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 set (defined in Appendix C). In this test data, all anchor pairs (i.e., the task) are seen in the training set.

Generalization on task. Generalization on task depends on the unseen anchors, i.e., anchors that do not appear in the training set, with a designated target mapping.

3.5 Model Architecture and Basic Experimental Setups

To enable a 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 X in . The word embedding X em and the input to the first transformer block X (1) is calculated as:

where X pos is a trainable positional vector. For the l -th layer, the Q,K,V are defined as:

The attention matrix Attn ( l ) and its subsequent output X qkv( l ) for the l -th layer is computed as:

The output of the l -th attention layer is obtained as:

where 'LN' refers to Layer Normalization and 'FN' represents a Fully Connected Network. The final output is obtained by projecting the output of the last layer X do( L ) using a linear projection layer, followed by a softmax operation and argmax to obtain the predicted token.

For the basic experimental setups, we use cross-entropy loss on 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 A.

4 Two Phases of Solutions for Composite Functions

In this section, we investigate the mechanisms of how a 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 M (3 , 4) ( x ) = x -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. 1(b).

badgeneralization on seenanchors(test accuracy<9o%)

(a)

acc of symmetric solution on anchor (4,3)

(b)

acc of inferential solution on anchor(4,3)

1.0

(c)

inferentialsolution

symmetric solution

7

7

1.0

0.8

3

6

9

layer number

layernumber

5

5

0.6

0.6

4

4

0.4

0.4

3

3

0.2

0.2

2

2

0.2

0.3

0.40.50.60.7

0.8

0.2

0.3

0.40.50.6

0.7

0.8

0.0

0.0

An important and interesting question is: which solution is learned for unseen pair (4, 3)? Experiments show that different initialization scales lead to different solutions for the unseen pair (4, 3) with different model depths. Fig. 2(a) shows the phase diagram assuming the symmetric mapping ( M (4 , 3) ( x ) = x -6 ) as ground truth, with initialization scale on the abscissa and model depth on the ordinate. A large initialization scale (small γ ) leads to the symmetric solution, indicated by yellow zones with nearly 100% accuracy. Fig. 2(b) assumes the inferential mapping ( M (4 , 3) ( x ) = x -10 ) as ground truth, showing that a small initialization scale (large γ ) 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 different initializations and depths were trained using 9 learning rates (uniformly sampled in [ 1 × 10 -4 , 3 × 10 -4 ]), 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.

To validate the generality of these findings, we extend our experiments to the multi-head GPT-2 model. As shown in Fig. 2(c), the model exhibits similar generalization performance trends on the unseen pair (4, 3) across different initialization rates γ . 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 a 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.

5 Mechanisms of Models in Two Phases

In this section, we analyze mechanisms underlying different learning solutions including: i) information 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 optimization 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 as 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 attention matrix through information flow analysis. As shown in Fig. 3 (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 connecting the j -th token in Layer l and the k -th token in Layer l +1 represents the value of the attention matrix Attn ( l ) 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.

(a)

3

二

3

4

9

(b)

(c)

3

+

4

→

4

-10

Output 30

58

58

83

6

93

93

93

93

Anchor

Output

94

95

95

109

94

92

89

90

89

Layer 2O

O

12 21

(1, 2)

(1, 3)

Layer 2

O

11

(1, 4)

(2, 3)

14 41

(2, 4)

22

13 31

(3, 4)

Layer 1O

2332

(1, 1)

(2, 2)

Layer 1O

(3, 3)

(4, 4)

(2, 1)

2442

(3, 1)

Layer0O

Input 35

63

O

99

4

3

38

42

36

62

34 43

44

(4, 1)

(3, 2)

(4, 2)

Layer 0

Input 35

63

99

4

3

38

42

O

36

62

33

(4, 3)

keyitem

→

anchor pair

keyitem

→anchorpair

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 it to the last token of the sequence. This allows the model to learn identical embeddings for symmetric anchor pairs,

such as (3, 4) and (4, 3), as illustrated in Fig. 3(a). 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 visualize the vectors X ao(1) of 10,000 input sequences with different anchor-key item pairs using t-SNE. Fig. 3(b) shows that symmetric anchor pairs are clustered together in the X ao(1) vector space, confirming that the model learns to map them to similar representations.

For inferential solutions, the information transmission mechanism is different. As shown in Fig. 3(c), in the first layer, the key item information is moved to the positions of the two single anchors, and each anchor is 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 to 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, we examine the divergence in vector representations after the fusion of anchor pair and key item information 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 ao(2) ) for different anchorkey item combinations. We denote these output vectors as v ( x ; a 1 , a 2 ) , where x is the key item and ( a 1 , a 2 ) is the anchor pair of the input sequence 2 .

As illustrated in Fig. 4(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 to the cosine similarity matrices between v ( x 1 ; 3 , 3) and v ( x 2 ; 2 , 2) , and between v ( x 1 ; 1 , 2) and v ( x 2 ; 1 , 3) , respectively, where x 1 and x 2 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. 4(a) indicates f (36; 3 , 3) = f (30; 2 , 2) , where f ( · ; · , · ) is defined in Equation (1).

(a)

30

32

34

36

38

40

(b)

30

32

anchor:(1,3)

34

8.9

40

1.0

(c)

50

52

single anchor:2

54

56

58

60

30

30

0.995

30

30

50

50

0.97

anchor:(3,3)

32

32

anchor: (1, 2)

0.990

anchor:(3, 3)

32

32

0.8

single anchor:4

52

52

single anchor:1

34

34

34

54

0.96

36

0.985

0.6

36

36

56

56

0.980

38

38

38

38

0.4

58

58

0.95

0.975

40

40

40

40

60

60

30

32

343638

40

0.970

30

32

343638

40

0.2

50

52545658

60

0.94

anchor: (2, 2)

anchor: (2, 2)

single anchor: 3

For inferential solutions, if two anchor-key pairs ( x 1 ; a 1 , a 2 ) and ( x 2 ; b 1 , b 2 ) have the same output, i.e., f ( x 1 ; a 1 , a 2 ) = f ( x 2 ; b 1 , b 2 ) , then their fused vectors v ( x 1 ; a 1 , a 2 ) and v ( x 2 ; b 1 , b 2 ) are nearly

parallel. In contrast, for symmetric solutions, the pairwise similarity between vectors v ( x ; a 1 , a 2 ) is 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 ao(2) 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 ao(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 (2), we investigate the vector representations 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 ; a 1 ) , with key item x and first anchor a 1 . Similar to Fig. 4(a, b), the red boxes highlight the positions where the targets obtained by the single anchors a 1 and a 2 on the abscissa and ordinate are the same for the corresponding key items x 1 and x 2 , i.e., g ( x 1 ; a 1 ) = g ( x 2 ; a 2 ) . As shown in Fig. 4(c), the information fusion of single anchors and key items and the explicitly defined single anchor operations align well. This provides stronger evidence for the model'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) (Jacot et al., 2018; Luo et al., 2021). 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 to as condensations (Zhou et al., 2021; Zhang et al., 2021, 2022). 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) o

(b)o

1.0

(c)

21

(p)

35

403734

50

50

0.5

5$0

4744

31

28

25

25

22

48

100

150

150

100

0.0

0.5

575451

5855

52

48

45

4138

4239

35

36

3229

3330

26

27

22

23

20

556

58

3147

2737

2号4559

5738

2621

23

20

24

44

49

5530

5333

4629

3941

45

52

42

0

50

100

150

Q

50

100

150

-1.0

5956

24

36

Condensation. We examine the condensation of matrix W Q (1) . The similarity between the i -th and the j -th neurons input weight is calculated by cosine similarity, i.e., W Q (1) [ i, :] · W Q (1) [ j, :] || W Q (1) [ i, :] || 2 || W Q (1) [ j, :] || 2 . As shown in Fig. 5, for the inferential solution (Fig. 5(a)), the neuron weights condense in a few directions, suggesting low complexity within the model. In contrast, for the symmetric solution (Fig. 5(b)), 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 W em , 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 a clear tendency towards low complexity.

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. 5(c,d), for small initialization scales (Fig. 5(c)), we observe a clear ordinal structure in the embedding space, with the embeddings of the input tokens arranged according to their numerical value. This also suggests a low-complexity tendency in the word embedding matrix, requiring the model to capture the relative ordinal relationships between different tokens. However, this ordinal structure is not present for large initialization scales (Fig. 5(d)).

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 relationship 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 3 and 4. This arrangement is consistent with the numbers being ordered with intervals of 3 (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 D.3.

6 Further Verification on 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. Please refer to the Appendix F for a detailed introduction of each dataset.

Compositional tasks: SCAN and COGS. For the SCAN dataset (Lake and Baroni, 2018), we selected the 'Generalizing composition across primitive commands" task, where the 'turn left" command only appears in single-command mappings and is trained alongside other composite commands. We assess the model's generalization ability on composite commands that include the 'turn left" command. For the COGS dataset (Kim and Linzen, 2020), we evaluate the model's indistribution and out-of-distribution generalization performance after training on the same training set.

(a)

(b)

(c)

SCAN(unseen command:turnleft)

COGS(testonin-distribution data)

COGS(teston out-of-distribution data)

Accuracy on Composed Commands

100

Accuracy on In-distribution Data

100

100

75

75

75

50

50

50

25

25

25

0

0.05w

0.1w

0.5w

1w

0

0.05w

0.1w

0.5w

1w

2w

0

0.5w

1w

2w

full data

DatasetSize

DatasetSize

DatasetSize

largeinit(y=0.5)&smallwd(wd=0.01)

smallinit(y=0.7)&large wd (wd=0.1)

As shown in Fig. 6, 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 settings (with 20k training data) reaches over 99%, the difference in out-of-distribution generalization remains significant.

Reasoning tasks: PrOntoQA. PrOntoQA (Saparov and He, 2023) is a synthetic multi-step reasoning dataset where each data point assigns hierarchical relationships among objects and requires the

model to determine whether a multi-step reasoning chain is correct. Fig. 7 illustrates the convergence rates and generalization errors 3 with respect to data scale for models with large initialization (and small weight decay, Fig. 7(a)) and small initialization (and large weight decay, Fig. 7(b)). 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 generalization 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)

100

100

乙

2

TestAccuracy

Test Accuracy

75

75

Memorization

Inference

Regime

Regime

50

50

214

216

218

214

216

218

DatasetSize

DatasetSize

ConvergenceSteps

TestAccuracy

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 on the solutions learned by transformers for compositional tasks. We discover a 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 memorize mappings. To explain these findings, we 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 symmetric solutions directly combine anchor information without learning the compositional structure. We further extend our 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 of 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 investigation to real-world datasets and tasks, to bridge the gap between our theoretical understanding and practical application. This could involve leveraging Mixture of Experts (MoEs) to design networks with different initialization scales for different expert models.

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(Z. 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. Z.), 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 of Education, P.R. China. This work was partially supported by SJTU Kunpeng&Ascend Center of Excellence.

References

Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, I. Polosukhin, Attention is all you need, in: Advances in neural information processing systems, 2017, pp. 5998-6008.

Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al., Gpt-4 technical report, arXiv preprint arXiv:2303.08774 (2023).

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 information processing systems 33 (2020) 1877-1901.

H. Choi, K. E. Hickman, A. B. Monahan, D. Schwarcz, Chatgpt goes to law school, J. Legal Educ. 71 (2021) 387.

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.

Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, E. Horvitz, E. Kamar, P. Lee, Y. T. Lee, Y. Li, S. Lundberg, et al., Sparks of artificial general intelligence: Early experiments with gpt-4, arXiv preprint arXiv:2303.12712 (2023).

F. Marcus, The algebraic mind: Integrating connectionism and cognitive science, MIT press, 2003.

Smolensky, R. McCoy, R. Fernandez, M. Goldrick, J. Gao, Neurocompositional computing: From the central paradox of cognition to a new generation of ai systems, AI Magazine 43 (2022) 308-322.

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.08309 (2024).

Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al., Language models are unsupervised multitask learners, OpenAI blog 1 (2019) 9.

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, Y. Tay, R. Bommasani, C. Raffel, B. Zoph, S. Borgeaud, D. Yogatama, M. Bosma, D. Zhou, D. Metzler, et al., Emergent abilities of large language models, arXiv preprint arXiv:2206.07682 (2022).

Srivastava, A. Rastogi, A. Rao, A. A. M. Shoeb, A. Abid, A. Fisch, A. R. Brown, A. Santoro, A. Gupta, A. Garriga-Alonso, et al., Beyond the imitation game: Quantifying and extrapolating the capabilities of language models, arXiv preprint arXiv:2206.04615 (2022).

Csordás, K. Irie, J. Schmidhuber, The neural data router: Adaptive control flow in transformers improves systematic generalization, arXiv preprint arXiv:2110.07732 (2021).

Csordás, K. Irie, J. Schmidhuber, Ctl++: Evaluating generalization on never-seen compositional patterns of known functions, and compatibility of neural representations, arXiv preprint arXiv:2210.06350 (2022).

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).

Hupkes, A. Singh, K. Korrel, G. Kruszewski, E. Bruni, Learning compositionally through attentive guidance, arXiv preprint arXiv:1805.09657 (2018).

A. Lepori, T. Serre, E. Pavlick, Break it down: Evidence for structural compositionality in neural networks, arXiv preprint arXiv:2301.10884 (2023).

Okawa, E. S. Lubana, R. P. Dick, H. Tanaka, Compositional abilities emerge multiplicatively: Exploring diffusion models on a synthetic task, arXiv preprint arXiv:2310.09336 (2023).

Yun, U. Bhalla, E. Pavlick, C. Sun, Do vision-language pretrained models learn composable primitive concepts?, arXiv preprint arXiv:2203.17271 (2022).

Ramesh, M. Khona, R. P. Dick, H. Tanaka, E. S. Lubana, How capable can a transformer become? a study on synthetic, interpretable tasks, arXiv preprint arXiv:2311.12997 (2023).

Liu, J. T. Ash, S. Goel, A. Krishnamurthy, C. Zhang, Transformers learn shortcuts to automata, arXiv preprint arXiv:2210.10749 (2022).

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.11903 (2022).

Creswell, M. Shanahan, I. Higgins, Selection-inference: Exploiting large language models for interpretable logical reasoning, arXiv preprint arXiv:2205.09712 (2022).

Creswell, M. Shanahan, Faithful reasoning using large language models, arXiv preprint arXiv:2208.14271 (2022).

Wang, H. He, J. Wang, Z. Wang, G. Huang, F. Xiong, Z. Li, L. Wu, et al., Improving generalization and convergence by enhancing implicit regularization, arXiv preprint arXiv:2405.20763 (2024a).

Wang, et al., Understanding the expressive power and mechanisms of transformer for sequence modeling, arXiv preprint arXiv:2402.00522 (2024b).

Wang, L. Li, D. Dai, D. Chen, H. Zhou, F. Meng, J. Zhou, X. Sun, Label words are anchors: An information flow perspective for understanding in-context learning, arXiv preprint arXiv:2305.14160 (2023).

Cao, S. Han, Z. Gao, Z. Ding, X. Xie, S. K. Zhou, Graphinsight: Unlocking insights in large language models for graph structure understanding, arXiv preprint arXiv:2409.03258 (2024).

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.

Chizat, F. Bach, On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport, in: Advances in Neural Information Processing Systems 31, 2018, pp. 3036-3046.

Zhang, Z.-Q. J. Xu, T. Luo, Z. Ma, A type of generalization error induced by initialization in deep neural networks, arXiv:1905.07777 [cs, stat] (2019).

E, C. Ma, L. Wu, A comparative analysis of optimization and generalization properties of twolayer neural network and random feature models under gradient descent dynamics., Sci. China Math. 63 (2020).

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.

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) E7665-E7671. doi: 10. 1073/pnas.1806579115 .

Rotskoff, E. Vanden-Eijnden, Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks, in: Advances in Neural Information Processing Systems 31, 2018, pp. 7146-7155.

Sirignano, K. Spiliopoulos, Mean field analysis of neural networks: A central limit theorem, Stochastic Processes and their Applications 130 (2020) 1820-1852. doi: 10.1016/j.spa.2019. 06.003 .

Williams, M. Trager, C. T. Silva, D. Panozzo, D. Zorin, J. Bruna, Gradient dynamics of shallow univariate relu networks, CoRR abs/1906.07842 (2019). URL: http://arxiv.org/abs/1906. 07842 . arXiv:1906.07842 .

Luo, Z.-Q. J. Xu, Z. Ma, Y. Zhang, Phase diagram for two-layer relu neural networks at infinitewidth limit, Journal of Machine Learning Research 22 (2021) 1-47.

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).

Zhang, Z. Zhang, L. Zhang, Z. Bai, T. Luo, Z.-Q. J. Xu, Linear stability hypothesis and rank stratification for nonlinear models, arXiv preprint arXiv:2211.11623 (2022).

Zhang, Z.-Q. J. Xu, Loss spike in training neural networks, arXiv preprint arXiv:2305.12133 (2023).

Zhang, Y. Li, T. Luo, Z.-Q. J. Xu, Stochastic modified equations and dynamics of dropout algorithm, arXiv preprint arXiv:2305.15850 (2023).

Zhang, Z.-Q. J. Xu, Implicit regularization of dropout, IEEE Transactions on Pattern Analysis and Machine Intelligence (2024).

Power, Y. Burda, H. Edwards, I. Babuschkin, V. Misra, Grokking: Generalization beyond overfitting on small algorithmic datasets, arXiv preprint arXiv:2201.02177 (2022).

Wang, X. Yue, Y. Su, H. Sun, Grokked transformers are implicit reasoners: A mechanistic journey to the edge of generalization, arXiv preprint arXiv:2405.15071 (2024).

Gopalani, E. S. Lubana, W. Hu, How do transformers fill in the blanks? a case study on matrix completion, in: ICML 2024 Workshop on Mechanistic Interpretability, 2024.

Liu, E. J. Michaud, M. Tegmark, Omnigrok: Grokking beyond algorithmic data, in: The Eleventh International Conference on Learning Representations, 2022.

S. Huang, F. Perez, J. Ba, M. Volkovs, Improving transformer optimization through better initialization, in: International Conference on Machine Learning, PMLR, 2020, pp. 4475-4483.

Liu, X. Liu, J. Gao, W. Chen, J. Han, Understanding the difficulty of training transformers, arXiv preprint arXiv:2004.08249 (2020).

Trockman, J. Z. Kolter, Mimetic initialization of self-attention layers, in: International Conference on Machine Learning, PMLR, 2023, pp. 34456-34468.

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).

Zhang, I. Titov, R. Sennrich, Improving deep transformer with depth-scaled initialization and merged attention, arXiv preprint arXiv:1908.11365 (2019).

Zhu, R. Ni, Z. Xu, K. Kong, W. R. Huang, T. Goldstein, Gradinit: Learning to initialize neural networks for stable and efficient training, Advances in Neural Information Processing Systems 34 (2021) 16410-16422.

Zhou, Q. Zhou, T. Luo, Y. Zhang, Z.-Q. J. Xu, Towards understanding the condensation of neural networks at initial training, arXiv preprint arXiv:2105.11686 (2021).

Zhang, Z. Zhang, T. Luo, Z. J. Xu, Embedding principle of loss landscape of deep neural networks, Advances in Neural Information Processing Systems 34 (2021) 14848-14859.

Zhang, Y. Li, Z. Zhang, T. Luo, Z.-Q. J. Xu, Embedding principle: a hierarchical structure of loss landscape of deep neural networks, Journal of Machine Learning vol 1 (2022) 1-45.

Lake, M. Baroni, Generalization without systematicity: On the compositional skills of sequenceto-sequence recurrent networks, in: International conference on machine learning, PMLR, 2018, pp. 2873-2882.

Kim, T. Linzen, Cogs: A compositional generalization challenge based on semantic interpretation, in: Proceedings of the 2020 conference on empirical methods in natural language processing (emnlp), 2020, pp. 9087-9105.

Saparov, H. He, Language models are greedy reasoners: A systematic formal analysis of chainof-thought, in: The Eleventh International Conference on Learning Representations, 2023.

Hu, X. Tang, H. Yang, M. Zhang, Case-based or rule-based: How do transformers do the math?, in: Forty-first International Conference on Machine Learning, 2024.

A Experimental Setups

For Fig. 2(a,b), Fig. 3, Fig. 4, Fig. 5, Fig. 10, Fig. 11, Fig. 12 and Fig. 13, we train the transformer model on a dataset of 900,000 samples, with each input sequence having a length of 9 tokens. The key items in the input sequences are randomly sampled from the range 20-99. The model architecture consists of a varying number of layers, a single attention head, an embedding dimension of 400, a feedforward dimension of 1200, and key and value dimensions of 200 each. The model parameters are initialized using a normal distribution with a mean of 0 and a standard deviation of (1 /d in ) γ , where d in is the input dimension of the parameter and γ is the initialization rate. We employ the AdamW optimizer with a learning rate of 1e-5, epsilon of 1e-8, weight decay of 0.01, and β 1 and β 2 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 a warm-up period followed by cosine decay, with a warmup period of 10 epochs, a multiplier of 25 (if there is no special instruction), a cosine decay number of epochs of 200, and a minimum learning rate of 1e-5. For Fig. 3(a, b), Fig. 4(b), Fig. 5(b, d), we use a two-layer transformer with initialization rate γ = 0 . 5 . For Fig. 3(c), Fig. 4(a, c), Fig. 5(a, c), we use a two-layer transformer with initialization rate γ = 0 . 8 .

For Fig. 2(c), Fig. 6, Fig. 7, Fig. 17 we employ the architecture of GPT-2 in our experiments. For Fig. 2(c), to ensure consistency between the model structure and the settings described in the paper, we adopt the post-norm configuration. We have also conducted experiments using the pre-norm setting and obtained consistent conclusions. Apart from the differences in the model architecture, we retain the same data, training strategy, parameter initialization, and other settings as mentioned in the previous sections.

For Fig. 14, Fig. 15, Fig. 16, we train the networks using initialization rate γ = 0 . 5 with different learning rates and weight decay coefficients. The learning rate is scheduled using a warm-up period followed by cosine decay, with a warmup period of 10 epochs, a multiplier chosen from the set {5, 10, 20, 40, 60}, a cosine decay number of epochs of 200, and a minimum learning rate of 1e-5.

Experiments Compute Resources. The experiments were conducted on a server with the following configuration:

48 AMD EPYC 7352 24-Core Processors, each with 512KB of cache

251GB 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

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 corresponds to a specific function:

Suppose we have an input sequence X = (23 , 1 , 2 , 43 , 46 , 74 , 54 , 44 , 72) . In this sequence, the 2nd and 3rd items (i.e., items 1 and 2) belong to 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:

As we can see, the two-anchor composite function first passes the key item 23 into the function g ( x ; 1) = x +5 corresponding to the first anchor 1, obtaining the result 28; then it passes 28 into the function g ( x ; 2) = x +1 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.e., 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:

In this example, the anchor pair is (2 , 4) , the key item is 36, and 52,33,78,92,24,58 are noise items.

For a given input sequence X , we 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 as 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 to the anchor pair so 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.3.

B.3 Examples of Mapping Type of an anchor pair

For an anchor pair ( a 1 , a 2 ) , its mapping M ( a 1 ,a 2 ) ( · ) can have the following types:

Inferential mapping. If the designated target mapping of the anchor pair ( a 1 , a 2 ) is consistent with the output of the two-anchor composite function, i.e., M ( a 1 ,a 2 ) ( x ) = f ( x ; a 1 , a 2 ) , then it is called an inferential mapping.

For example, for the anchor pair (1 , 2) and key item x , if the target output is f ( x ; 1 , 2) = g ( g ( x ; 1); 2) = ( x +5) + 1 = x +6 , then the mapping of (1 , 2) is an inferential mapping.

̸

Non-inferential mapping. If the designated target mapping of the anchor pair ( a 1 , a 2 ) is inconsistent with the output of the two-anchor composite function, i.e., M ( a 1 ,a 2 ) ( x ) = f ( x ; a 1 , a 2 ) , then it is called a 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 f ( x ; 1 , 2) = g ( g ( x ; 1); 2) = ( x + 5) + 1 = x + 6 , then the mapping of (1 , 2) is a non-inferential mapping.

Symmetric mapping. If the designated target mapping of the anchor pair ( a 1 , a 2 ) is consistent with the mapping of its symmetric anchor pair ( a 2 , a 1 ) , i.e., M ( a 1 ,a 2 ) ( x ) = M ( a 2 ,a 1 ) ( x ) , then it is called a symmetric mapping.

For example, suppose the mapping of the anchor pair (3 , 4) is designated as M (3 , 4) ( x ) = x +10 . If the mapping of the anchor pair (4 , 3) is also designated as M (4 , 3) ( x ) = x +10 , then the mapping of (4 , 3) is a symmetric mapping, because it is exactly the same as the mapping of (3 , 4) .

For a given model, we define its inferential (non-inferential, symmetric) accuracy on a specific anchor pair as the accuracy obtained when using the inferential (non-inferential, symmetric) mapping of that anchor pair as the target. In particular, if a 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 (non-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. 8.

Generalizationontask

(accuracy on the designed target mapping)

anchorsseenintraining

22

53

2

anchorsunseenintraining

42

sequences

unseenintraining

47

31

53

43

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 is shown in Appendix C) 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 seen (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 generalization on task. It is worth noting that the unseen anchor pair (4, 3) does not have a 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.

C Dataset Splitting Method

A straightforward division based on data ranges proves to be impractical. To illustrate, consider a 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 +1 , c ] . The encoding of data within the interval [ b +1 , c ] is not learned during the neural network training process. As a 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, as shown in Fig. 9. Consider a task with an input sequence of length n . For an input sequence in the training dataset, an item x can be placed in the pos -th position of such input sequence only when mod( x, n -2) = pos . For an input sequence in the test dataset, if the token at the i -th position is a key item, then an item x can be placed in the pos -th position of such input sequence only when mod( x, n -2) = pos . It is 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.

Training data

Test data

,8)≠ pos

,8)= pos

Training

key item

anchor pairs

pos=4

+

Teston data

1

2

1

3

1

4

4

3

4

4

+

Testontask

Seenanchorpairs

Unseenanchorpair(4,3)

We further define two types of generalization based on this dataset splitting method. For training data, we 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 evaluate the model's ability to 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 on 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.

D Detailed Results for Model Complexity with Different Initialization

D.1 Detailed Results for Condensation

In order to show the parameter condensation at different depths and different initialization scales, we plotted the cosine similarity of the W Q (1) matrix neurons corresponding to the phase diagram in Fig. 2. As shown in Fig. 10, 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 initialization rate increases (the initialization scale becomes smaller), the W Q (1) matrix neurons show an obvious condensation trend, implying that the model complexity decreases.

Depth 6, y 0.3

Depth 6,  0.4

Depth 6, y 0.5

Depth 6,  0.6

Depth 6,y 0.7

Depth 5, y 0.3

Depth 5, 0.4

Depth 5,  0.5

Depth 5, y 0.6

Depth 5,y 0.7

Depth 4, y 0.3

Depth 4, 0.4

Depth 4, y 0.5

Depth 4,0.6

Depth 4,y 0.7

0

Depth 3, y 0.3

Depth 3, y 0.4

Depth 3, y 0.5

Depth 3, y 0.6

Depth 3, y 0.7

Depth 2, y 0.3

Depth 2,  0.4

Depth 2,y 0.5

Depth 2,  0.6

Depth 2,  0.7

-1

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. 11, as the initialization rate γ 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.

6

Depth 6, y 0.3

Depth 6, y 0.4

Depth 6, y 0.5

Depth 6, y 0.6

Depth 6, y 0.7

42

31

51

29

34

20 45

48

57

44

49

30

50

47

2821

55

24 53

22

3733

26

32

25

BE

43

52

6E

35

40 54

56

27

23

58

36

59

4

3832

58

54 40

26

21

28

24

49

50

30

55

46

53

37 33 52

44

51

47 57

41

2934

27

39

2343

9E

35

22

59

48

42

25

20

59

57

58

54

43

2327

36

56

52

3429

42

39

20

35

41

26

31

51

40

337

44

47

22

2538

30

49

53

32

24

212年

46

50

52

51

4.548

54

53

5044

46

43

47

41

4037

3930

36

35

38

32

3428

31

27 24

2920

2522

26

21

23

58°

59

2326

224

38 41

3134

3740

39

46

49

42

447

5255

4548

50

58

51

Depth 5, y 0.3

Depth 5, y 0.4

Depth 5, y 0.5

Depth 5, y 0.6

Depth 5,y0.7

51

26

29 34

42

30

36

49

20

373352

31

25

55

45

22

48

47

44

40

28

41

6

54

SE

21

85

24

3238

27

59

53

56

46

50

43

57

30

49

48 45

50

22

46

55

53

47

20

24

31

37 33 52

21

40

25

32

28

38

44

54

89

36

23

41

27

34

29

6E

43

35

51

56

42

69

43

95

42

50 46 53

34

25

38

29

20

58

33

31

37

52

59

2327

36

22

4757

39

55

35

41

51

26

40

32

48

44

24

21

30

21

28

34

25

22

30

24

43

37

41

20

EE

31

38

29

27

32

47

42 51

23

35

56

40

26

44

36

53

39

49

55

52

48

57

46

45

54

OS

58

58

5

59

56 53

48

49

45

3g35

42

57

50

46

38

47

41

44

25

43

22

29

40

37 34

3330

26

32

36

23

31

27

20

Depth 4,y 0.3

Depth 4,y 0.4

Depth 4, y 0.5

Depth 4, y 0.6

Depth 4,y 0.7

O2TE

45

48

58

40

54

4421323

30 49

65

53

28

26

2735

24

47

55

39

57

41

33 25

37

23

38

43

36

46

56

22

50

42

51

294

58

36

25

52

33

37

4.53 2421

49

46

30

55 47 57

4054

44

22

26

87

58

32 38

41

2031

39

56

35

23

27

43

36

42

59

51

28

32

22

55

28

24

47

63

49

31

44

34

29

50

37

33

20

40

30

54

52

43

27 23

58

56

26

41

49

53

59

30

56

57

26

52

43

9

3531

46

50

48

22

55 51

47

45

44

25

34

21

40

24

28

20

23

32

37

42

36

25

21

28

24

20 23.3

35

31

32

38

39

23

OE

42

36

43

26

3337

40

50 57

46

95 E5

44

41

49

4552

4855

59

51

58

54

Depth 3,y 0.3

Depth 3, y 0.4

Depth 3, y 0.5

Depth 3, y 0.6

Depth 3, y 0.7

58

25

54

40

3832

22

3

31

30

49

5547 57

21

24

46

28

53

51

41

44

2934

SE6E

27

59

52

20

45

42

36

56

48

23

43

46

58

53

54

30

57

49

40

31

2421

27

41

55

32

47

26

35

51

38

28

44

29

34

52

3956

22

36

20

42

48

25

22

25

232

38

41

2821

31

32

24

40

57 47

A

51

55

3539

59

49

5346

52

58

26

50

30

45 48

56

Et

29

34

42

20

55

49 53

498

56

57

40

52

4137

23

34

223

26

EE

29

20

24

30

42

28

43

32

8E

36

52 49

59

56

47

44

41

53

46

40

57

50

42

54 51

39

36

37

38

225

EE

34

33

3023

26

28

31

24

21

2

Depth 2, y 0.3

Depth 2, y 0.4

Depth 2, y 0.5

Depth 2, y 0.6

Depth 2, y 0.7

41

38

59

32

46

50

21

53

56

36

24

39

28

30

23

44

43

35

49

27

47

55

22

57

3

26

25

52

54

40

51

58

2934

48

45

20

42

51

42

22

25

20

36

EE

37

52

57

35 27

40

85

54

41

56

39

30

47

49

59

23

44

ES

43

26

46

24

28

21

55

50

31

32

26

65

36

43

56

52

23

85

39

54

35

42

51

a

20

32

4548

53

38

29

49

242

46

55

41

25

50

47

37 33

57

40

27

28

22

2420

26 23 27

59

233

58

30

57

51

54 50

31

28

38

53

48

45

35

32

52

49

44

41

36

40

46

26

29

23

20

32

35

8E

302

亿忆

41

33

36

39

44

47

28

42

37

34

31

40

45

48

51

46

49

52

58

55

50

53

56

59

D.3 Detailed Rank Analysis for Different Weight Matrices

To further verify the low complexity bias of the word embedding matrix, we visualize the eigenvalues of the covariance matrix of the embedding vectors (Fig. 12(a)). For small initialization scales, the covariance matrix has a small number of large eigenvalues, indicating that the model learns a 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

underlying structure of the single anchor operations. In contrast, for large initialization scales, the eigenvalues are more evenly distributed, suggesting that the model learns a more distributed representation that does not effectively capture the ordinal structure. This lack of structure in the embedding space hinders the model's ability to learn the relationships between the input tokens and the underlying single anchor operations, leading to the learning of symmetric solutions.

(a) eigenvalue

(b) eigenvalue evolution

10-2

y=0.5

lst

4th

40th

2nd

10th

80th

y=0.8

3th

eigenvalue

10-3

10-3

eigenvalue

10-4

10-4

10-5

10-5

0

20

40

60

80

0

5

10

15

20

index

Meanwhile, we investigated the evolution of the eigenvalues corresponding to specific indices of the small initialization model during training. As shown in Fig. 12(b), 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 a subsequent increase in the other eigenvalues. Once the model complexity increases sufficiently to fit the training data, it 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. 13. It is evident that, for the small initialization model, the first few singular values are significantly larger than the remaining ones. In contrast, this distinct difference is not observed in the model with large initialization. This indicates the pervasive low-complexity characteristic of the internal parameters in the small initialization model.

attn1 WQ

attn1WK

attn1Wv

attnl fc

Singular Value (log scale)

ffnl fc1

ffn1 fc2

attn2WQ

attn2WK

attn2 fc

ffn2 fc1

ffn2fc2

Index

y=0.5

y=0.8

E Learning Rate and Weight Decay Coefficient Affecting Solution Phases

Aseries 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 to 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 in the main text.

In the main text, we artificially assigned non-inferential solutions to the anchor pair (3, 4) as M (3 , 4) ( x ) = x -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 learning rates. In fact, if we only consider the model's ability to learn inferential solutions, we can treat 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.

As shown in Fig. 14, we study a 3-layer, 1-head transformer model with an initialization rate of γ = 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 on 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

5

!

0.

10.0

L

0.6

rate

20.0

learning

0.4

40.0

0.2

60.0

0.0

0.003

0.01

0.03

0.1

0.3

weight decay

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 W Q (1) matrix neurons corresponding to the phase diagram in Fig. 14. As shown in Fig. 15, 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 W Q (1) matrix neurons show an obvious condensation trend, implying that the model complexity decreases.

wd 0.3,Ir5e-5

wd 0.003,lr1e-4 wd 0.01,Ir1e-4 wd 0.03,lr1e-4

4wd0.1,Ir1e-4

wd 0.3, Ir 1e-4

wd 0.003,lr2e-4wd 0.01,Ir2e-4wd 0.03,Ir2e-4

wd0.1,lr2e-4

wd 0.3, Ir 2e-4

0

wd 0.1,Ir4e-4

wd0.3,lr4e-4

wd 0.003,Ir6e-4 wd 0.01,Ir 6e-4 wd 0.03,lr6e-4

wd0.1,lr6e-4

wd 0.3, Ir 6e-4

-1

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 in Fig. 16, as the learning rate and weight decay coefficient increase, the word embedding matrix gradually exhibits distinct structural features. These structural features enable the model to accurately represent the numerical meaning of different items.

a0

wd 0.003,Ir 5e-5

wd 0.01, Ir 5e-5

wd 0.03, Ir 5e-5

wd 0.1, Ir 5e-5

wd 0.3, Ir 5e-5

51

57

22

55

28

49

124

30

50

32

53

46

3733

56

48 45

A0

41

47

44

52

SE 6E

25

26

38

20

31

29

27

65

34

42

2336

43

59

42

56

36 23

50

4653

57

22

26

39

35

49

54

30

31

44

43

40

55

41

20

33

3832

37

24

29

34

2128

52

51

25

45 48

27

51

55

42

28

58

2934

95

43

44

47

20

46

50

5349

27

31

35

39

4548

41

40

30

337

52

59

241

26

22

25

2336

FEE

43

59

50

OE

27

46

53

49

4033

45

55

48

58

37

54

52

57

51

47

24 21

44

29

39 35

34

42

28

25

26

31

22

38 32

56

20

22

28

224

25

38

34

31

32

29 42

20

41

37

33

35 39

26

40

43

27

55

44

36

23

49

57

53 46

54

58

56

50

45

52

wd 0.003, Ir 1e-4

wd 0.01, lr 1e-4

wd 0.03, Ir 1e-4

wd 0.1, Ir 1e-4

wd 0.3, Ir 1e-4

43

32

38

31

23

36

59

26

35

39

27

57

28

47

44

53

46

50

54

85

22

41

49

OE

25

37

33

55

34 29

40

51

52

42

20

48

42

51

4548

29

34

43

50

52

25

30

46 53

56

EE

37

2421

40

55

27

28

3623

44 47

41

39

32

8E

57

35

3120

22

26

54

58

43

51

25

31

34

59

2920

26

36

23

57

47 44

27

22

45

52

48

41

33

37

40

28

54

58

55

224

49

30

53 46 50

56

51

SS

26

48

45

52

31

28

35

39

38

32

54

22

25

3623

43

34

47

20

29

41

44

33

42

27

40

37

57

30

53 46

49

24

56

21

50

20

49

56

53

46

57

32

25

33 2942

338

26 39

22

35

21

20

34

28

44

30

24

47

41

37

48

52

45

59

50

548

wd 0.003, Ir 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

51

42

23

34

25

56

59

53

46

30

49

2421

36

23

27

57

28

22

47 41

39

35

25

43

44

31

333

54

40

23

26

55

52

50

51

34

29

48

45

42

20

26

3238

45

20

25

22

2934

52

50 4653

3340

65

55

37

49

30

28

41

2421

56

27

44

35

47

39

2336

31

57

54

85

42

56

24

OE

28

22

50

27

25

33

40

41

3

57

26

44

55

47

335

52

31

51

32

34

29

48 45

20

43

545

56

52

48

50

59

4649

45

53

54

58

57

5155

44

47

40

41

27

36

37

EE

42

29

39 35

30

24

26

21

38

34

31

32

25

20

22

49

58

3

OE

45

43

6E

32

54

52

95

47

27

48

57

53

38

51

55

31

44

59

42

50

29

33

41

46

28

37

26

33

20

24

wd 0.003, Ir 4e-4

wd 0.01, Ir 4e-4

wd 0.1, Ir 4e-4

wd 0.3, Ir 4e-4

51

20

45

48

42

43

34

36 23

52

26

39

55

30

35

49

31

40

27

41

44

53

46

50

47

25

57

8Z

22

56

24

21

51

24

21

55

28

41

22

30

47

40

44

25

49

32

26

38

27

50

46

53

35

39

31

23

57

9E

34

43

56

29

20

52

59

42

20 224 25

2934

22

3135

42

39

32

48

43

36

4744

41

40

57 54

50 5352

46

85

56

51

20 258

22327303

224

33.4

35

3638

37

6E

40 44 46

4145

42

495 5457

ES

65

2622

39

32

29

9E

35

48

28

20 33

49

25

23

45 51

3841

2427

58

31

21

54

44

37

40

34

47

5650

59

43

ES

wd 0.003, Ir 6e-4

wd 0.01, Ir 6e-4

wd 0.03, Ir 6e-4

wd 0.1, Ir 6e-4

wd 0.3, Ir 6e-4

50

5559

22

220

51

54 57

2723

28 25

212426

22

20

30

34

31

32

36

35

38

42

41 47

40

43

44

4548

46

49

5257°

9s ES

51

58

5955

5653

51

549

546

444

48

4542 41

43

40

383685

37

32 29

3431

OE

272421

28

26

5954 52

5149

50 48454 3

464  42

42

2##2

58

55

48

45

52

42

38

35

32

29

6E

22

28

26

36

4946

23 20

41

EE

65

44

30

43

56

47

27

40

53

34

31

37

2

F 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 alongside other composite commands. We assess the model's generalization ability on composite commands that include the 'turn left" command. For the COGS dataset, we evaluate the model's indistribution and out-of-distribution 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 different combinatorial rules (the original paper outlines 21 methods for generating out-of-distribution data, from which we generate test data with equal probability).

Reasoning tasks: PrOntoQA. PrOntoQA is a synthetic multi-step reasoning dataset where each data point assigns hierarchical relationships among objects and requires the model to determine whether a multi-step reasoning chain is correct. Although we use next-token prediction for training, during testing, we only evaluate the model's accuracy in judging hierarchical relationships (thus, the model's random guessing accuracy is 50%).

Realistic tasks: Addition task and SlimPajama dataset. Unlike traditional addition tasks, we use a case-based reasoning intervention experiment (Hu et al., 2024) to study the generalization of rule learning in the addition task. Specifically, we consider the setup: a + b = c , where a, b ∈ [0 , 999] . We use a, b ∈ [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 a simple 2-layer 1-head model and a GPT-2 model. As shown in Fig. 17(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

TrainfromScratchonGPT2Medium

100

100

100

Perplexity on Test Data

1.72

Perplexity on Test Data

75

75

Test Accuracy

Test Accuracy

TestAccuracy

2.6

50

50

90

1.71

25

F

25

0

2L1H

0

GPT2

80

GPT2

1.70

GitHub

2.5

GitHub+Wikipedia

wd=0.0

y=0.5&wd=0.1

y=0.5&wd=0.01

wd=0.01

y=0.75& wd=0.1

y=0.7 & wd=0.5

wd=0.5

y=1.0 & wd=0.1

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. 17(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

the model with large initialization early in training (around 2B tokens). For the GitHub+Wikipedia 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.

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 our 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 in 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.

It is 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 of the work performed by the authors?

Answer: [Yes]

Justification: We show the detailed limitation in Section 7.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on 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, e.g., if the approach was only tested on a few datasets or with a 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 approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it 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 address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [NA]

Justification: This is a 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

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 a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented 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 conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We show the detailed setups in Appendix A, 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 to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to 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 a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce 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 authors are welcome to 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so No is an acceptable answer. Papers cannot be rejected simply for not 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 ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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 a 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 as much information as possible in supplemental material (appended to the paper) is 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, hyperparameters, 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 a level of detail that is necessary to 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 appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We show the error bar in Fig. 2 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, confidence intervals, or statistical significance tests, at 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 run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to 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 is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We show the detailed information on the computer resources in Appendix A.

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, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn'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 https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper conform with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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 or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation 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

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This work is a phenomenological study, therefore, this work poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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 not 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 is used and, if possible, include a 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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation 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/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is 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. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

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.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nMy 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.\n\nQUESTIONS:\nQ1. 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 .\n\nLIMITATIONS:\nThe authors do discuss the limitations of this work in Section 7.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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).\n\nQUESTIONS:\n1) 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

comments: 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.

STRENGTHS:
- The investigation into the effect of initialization on compositional abilities of transformers is well-motivated and practically useful. 
- 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.
- The focus on symmetric solutions is novel, but the generalizability to real data and a broader range of models remains uncertain.

QUESTIONS:
- Do you have a solid definition for "Phase Transition"? If so, how did you define it?
- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12