An Analysis for Reasoning Bias of Language Models with Small Initialization
Authors: Junjie Yao
Zhongwang Zhang
Zhi-Qin John Xu
Venue: ICML 2025 spotlightposter (2025)
Keywords: initialization scale
reasoning bias
language model
embedding space
training dynamics
Primary area: deep_learning->large_language_models

Abstract: Transformer-based Large Language Models (LLMs) have revolutionized Natural Language Processing by demonstrating exceptional performance across diverse tasks. This study investigates the impact of the parameter initialization scale on the training behavior and task preferences of LLMs. We discover that smaller initialization scales encourage models to favor reasoning tasks, whereas larger initialization scales lead to a preference for memorization tasks. We validate this reasoning bias via real datasets and meticulously designed anchor functions. Further analysis of initial training dynamics suggests that specific model components, particularly the embedding space and self-attention mechanisms, play pivotal roles in shaping these learning biases. We provide a theoretical framework from the perspective of model training dynamics to explain these phenomena. Additionally, experiments on real-world language tasks corroborate our theoretical insights. This work enhances our understanding of how initialization strategies influence LLM performance on reasoning tasks and offers valuable guidelines for training models.

An Analysis for Reasoning Bias of Language Models with Small Initialization
Junjie Yao Zhongwang Zhang'
Zhi-Qin John Xu 234
Abstract
104
PrOntoQA
TinyStories
Transformer- Large Language Models
(LLMs) have revolutionized Natural Language Processing by demonstrating exceptional perfor- mance across diverse tasks. This study investi- gates the impact of the parameter initialization scale on the training behavior and task preferences of LLMs. We discover that smaller initialization scales encourage models to favor reasoning tasks, whereas larger initialization scales lead pref- erence for memorization tasks. We validate this reasoning bias via real datasets and meticulously designed anchor functions. Further analysis of initial training dynamics suggests that specific model components, particularly the embedding space and self-attention mechanisms, play piv- otal roles in shaping these learning biases. We provide a theoretical framework from the perspec- tive of model training dynamics explain these phenomena. Additionally, experiments real- world language tasks corroborate our theoretical insights. This work enhances our understanding of how initialization strategies influence LLM per- formance on reasoning tasks and offers valuable guidelines for training models.
Loss
109
Mi unl
10"
0
800000
Step
Figure I. Comparison of training loss between PrOntoQA and TinyStories in one next-token prediction training for this mix dataset. The red line represents the training loss on the PrOn- toQA dataset, while the blue line depicts the training loss the TinyStories dataset.
wide range of tasks, from text generation to complex reason- ing (Wei et al., 2022a; Achiam et al., 2023; Liu et al., 2024). Reasoning, in particular, is a critical ability for LLMs. A number of studies have focused on improving the reason- ing ability of these models through data-driven approaches, such as RHO-1 (Lin et al., 2024) and Phi-4 (Abdin et al., 2024). However, there remains an ongoing debate to whether LLMs genuinely learn the underlying logical rules or merely mimic patterns observed in the data (Marcus, 2003; Smolensky et al., 2022).
1. Introduction
An alternative approach to enhancing the reasoning abil- ity of LLMs focuses on the model architecture and its training process. one such study examining the use of Transformers model compositional functions, it ob- served that the scale of model parameter initialization sig- nificantly impacts the model's reasoning behavior (Zhang et al., 2024a; 2025). Specifically, smaller initialization scales bias the model toward fitting the data by learning primitive-level functions and compositional rules, whereas larger initialization scales tend encourage memorization of input-output mappings. A qualitative rationale for this phenomenon has been proposed: with small initialization, well-documented effect known as neuron condensation emerges during training (Xu et al., 2025; Luo et al., 2021; Zhou et al., 2022; Zhang et al., 2022: Zhang & Xu, 2023; Zhangg et al., 2023; Zhang & Xu, 2024). This phenomenon suggests that neurons within the same layer tend behave
With the rapid advancement of deep learning technologies, Large Language Models have achieved remarkable success in the field of Natural Language Processing (NLP). These models have demonstrated exceptional capabilities across
'School of Mathematical Sciences, Institute of Natural Sci- ences. Shanghai Jiao Tong University. Shanghai, P.R. China "Institute of Natural Sciences, School of Mathematical Sci- ences, MOB-LSC, School of Artificial Intelligence, Shang- hai Jiao Tong University, Shanghai, P.R. China 'Center for LLM, Institute for Advanced Algorithms Research, Shang- hai, P.R. China 'Shanghai Seres Information Technology Co., Ltd, Shanghai 200040, China. Correspondence to: Zhong- wang Zhang <012322w666@9jmm htt/z/6666@@ju..dduccc....... Zhi- Qin John Xu <xuzhiqin@ sjtu.edu.en>.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
similarly, promoting data fitting with the least possible com-
2. Preliminary
plexity. To achieve low-complexity result, the model must learn minimal set of rules leading to capture the intrinsic primitive functions and compositional rules. However, this rationale does not reveal a critical question: how the opti- mization process, together with the Transformer structure, can achieve reasoning solutions with small initialization?
2.1. Synthetic Composition Task
To study the task bias during the training, we use the con- cept of anchor function (Zhang et al., 2024b) to construct dataset that contains tasks of different reasoning com- plexities. We consider all tokens belonging to positive in- tegers. A set of tokens are designated as anchors, denoted as A := {a E | S a S 0maxx,, where each an- chor represents an addittion/randommes operation in this work. Another set of tokens are designated as keys, de- noted as 2 :- { E K S = S Smax] with the assumption that 2 a A 0. For convenience, we denote Ommnn 1.
In this work, we identify a reasoning bias during the train- ing of neural networks that learn natural language when initialized with small parameter scales. To illustrate this phenomenon, we employ a GPT-2 model (Radford et al., 2019) to train on a mixed dataset comprising two types of language data with distinct levels of reasoning complexity, within single next-token prediction training framework. The first dataset, PrOntoQA (Saparov & He, 2023), con- sists of question-answering examples that include chains of thought, which explicitly describe the reasoning neces- sary to answer the questions correctly. The second dataset, TinyStories (Eldan & Li, 2023), is a synthetic corpus of short stories containing only words typically understood by children aged to years. As shown in Figure I, the training loss for PrOntoQA decreases significantly faster than for TinyStories, suggesting that the model encounters and learns the reasoning patterns more readily.
N2 = Smax
Smin + and NA- Omax
This section introduces two types of sequence mappings. The first step involves constructing a sequence of positive integers with length L, represented as:
(la = [x|x = [2), : -2p Aptl, ap+g: pp+g+1; ... ,2L] .4 E 2, E A).
(I)
We define g as the number of anchors in the sequence, and p as the index of the element immediately preceding the first anchor element ap+1 in the sequence.
We uncover a key mechanism whereby reasoning tasks are learned earlier during training because the tokens associ- ated with these tasks become more differentiated in the embedding space at an early stage of the training process. We validate this mechanism using both synthetic data and real-world datasets. Furthermore, we provide a theoretical explanation for the evolution of token embeddings, which depends on the distribution of sample labels. Since each token is encoded as a one-hot vector, its embedding is ad- justed based on the loss associated with the labels of that token. Consequently, different label distributions can lead to distinct learning behaviors for token embeddings. For memory tasks, the labels associated with each token are typically random and lack explicit structure, which results in similar distributions for different memory token labels. As result, the embeddings for memory tokens are difficult to differentiate in the early stages of training. In contrast, reasoning tokens often exhibit distinct label distributions, leading to more differentiated embedding vectors for these tokens. These insights are elaborated through simplified model using multi-layer perceptron (MLP) and embedding structure, followed by an analysis of a Transformer model.
For a given sequence X E X((.1), we define its key-anchor combination as (žp: Up+t1''.. '.. (ptg)) which is denoted con- cisely as pair (zp: a), and other keys are regarded as noise in this input sequence. The anchor set A is divided into two subsets, i.c., reasoning anchor set Aren and memory anchor set Amem, where A = Arsu U Amem and Arsa A.Amm 0.
Reasoning mapping, For any X with Opti E Aran, 1, 9, we define the following mapping as a reasoning mapping
F (X) = 2p +
apti.
1=1
Memory mapping. For any key-anchor pair (zp, a). where each element in a belongs to Amems we randomly sample a number y al.).a) from Z as the memory mapping label of any sequence X containing (zp, a), i.e.
F (X) = -ytt.a))
VX contains (2p, a).
The primary contribution of this research lies in uncovering the impact of the parameter initialization scale on the train- ing behavior and task preferences of L.L.ss. By combining theoretical analysis with empirical evidence, we enhance the understanding of LLM training dynamics and provide new insights for optimizing model initialization strategies.
A detailed example is provided in Figure 2. It's noted that the key-anchor pair may occur at any position within the sequence. The label is independent of both the noise tokens and the position of the key-anchor pair within the sequence but is determined solely by the value of the key-anchor pair.
2

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
= 2 = 9,
Key set: (21,22 ... ,90, Memory Anchor Set: (1,2, ... .00, Reasoning Anchor Set: (11,12, ... ,20
Input Sequenees
Targets
Input Sequences
Targets
23
B5
24
3
31
35
78
34
14
63
59
78
77
86
80
22
18
37
61
67
9
32
37
67
89
46
57
(
51
11
11
28
64
67
47
26
71
73
46
43
63
67
46
(6755))
90
64
72
75
51
11
37
73
Memory Mapping
Reasoning Mapping
:noisc
:key
canchor
Figure 2. Schematic diagram of the synthetic composition task. The gray-shaded area illustrates the specific setup used in this example. Each block represents token within the input sequence, with different face colors indicating distinct token types (blue: noise, orange: key, green: anchor). Each row corresponds to input sequence paired with its respective label. The left section depicts four examples of memory mapping, while the right section presents four examples of reasoning mapping.
2.2. Dataset Setup
Rdexdm represents the output projection matrix. The de- tailed expression of multilayer Transformer models can be found in Appendix A.I.
In this study, we denote a data pair as (X, y). where X represents the input sequence and y corresponds to its associated label. We define y as the one-hot encoded representation of y for convenience. For memory mapping, all data are contained within the training set Dmem. and no test set is employed, as the generalization is not considered in this framework. For reasoning mapping, we define a set of masked anchor combinations M M 4p+2, ... , (lptg) apti E i = 1,... .4) and designate all sequences containing any masked com- bination (0-+1, 2p-2;"
2.4. Parameter Initialization
Given any trainable parameter matrix W Rd where dj and de denote the input and output dimensions, respec- tively, its elements are initialized according to normal distribution:
We ~ (o. (0. (d.
aptt) M the test set Desu,test and set the rest sequence train- The training set is Derain = Demm U D'Dsn,train-
where y is the initialization rate. Specifically, the initial- ization scale decreases as y increases. Note that 0.5 is commonly used in many default initialization methods, such as LeCun initialization (LeCun et al., 1998) and He ini- tialization (He et al., 2015). As the network width towards infinity (Luo et al., 2021; Zhou et al., 2022), the training of the network with y > 0.5 exhibits significant non-linear characteristics, i.e., condensation. Therefore, initialization scales with 7 0.5 are generally considered small.
2.3. Model Architecture
We give the formulation of the embedding space and self-attention module here for notation convenience. Let dvob, dm, d& denote the vocabulary size, embedding space dimension, and query-key-value projection dimension, re- spectively. For any token 8, denote its one-hot vector by e" E R Ixdeor. The embedding vector of cmb, eem where Wemm Xdd the embedding ma- trix. Additionally, the self-attention operator Attn on any embedding sequence X E Rxddm is defined as:
3. Result
In this section, we present empirical evidence of a reasoning bias during the training Transformers with small initial- ization by utilizing composition tasks. To further explore this phenomenon, we introduce a simplified model consist- ing of embedding layer and multi-layer perceptron, which reproduces the reasoning bias and enables theoretical analysis. A key mechanism underlying this bias is that the training behavior each token's 's embedding depends on the label distribution of the samples containing that token. For reasoning anchors, the label distributions typically exhibit
Attn (X) = (mask
XWWWWKXX X
(2)
Vd - Attn (X) XWW Wo
(3)
where g (.) is the softmax operator and T means the matrix transposition. Wº, W W e dd xds are the query, key and value projection matrices, respectively. W E
3

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
A:Transformer:
B.Emb-MLP(accuracy)
V= 0.3
y= 0.5
y 0.8
y= 0.3
109
100
Loss
10--
10-3
Myy
10"0
10
10--
10-3
10-2
10--
10:7
10
1000
Y = 0.8
2.0
2.0
1.0
Accuracy
0.8
0.8
0.8
My Ma Mmmm.
0.6
0.4
0.4
0.4
0.2
0.2
0.0
750
1000
1000
250
500
1000
500 Epoch
750
Epoch
Epoch
Epoch
-Memory
Reasoning - train
RRasoning - test
Figure 3. A: Loss and prediction accuracy of the models different datasets under varying initialization scales ( 0.3, 0.5, 0.8). The top row depicts the evolution of the loss during training for three datasets: Dmem (blue lines), Drau.train (purple lines), and Desn, (orange lines). The bottom row presents the corresponding prediction accuracies for these datasets. Each column represents results obtained with different initialization scales. B: Prediction accuracy of Emb-MLP under initialization rate 0.3 and 0.8.
greater variability compared memory anchors, leading to the more rapid differentiation of their embeddings early in training. Additionally, we extend our analysis to the Trans- former architecture to demonstrate the generalizability of this effect. For each observation mentioned in the following, we provide a similar analysis with larger intiallization scales in Appendix C.
loss D remains effectively unchanged. This ob- servation suggests that the model primarily memorizes the training data in this setting. In contrast, when = 0.8, the losses on tran, and Dram,test decrease significantly faster than the loss on D Dmem- This behavior indicates a shift towards a reasoning bias in the model. These findings reveal that the model's learning bias is influenced by the initialization scale: the initialization scale decreases, the model exhibits progressively stronger reasoning bias.
3.1. Reasoning Bias in Transformer with Composite Anchor Functions
3.2. Simplified Model: Phenomena and Analysis
In our experiment, we set that q L The dataset is constructed with the following configurations:
To further investigate the underlying cause of the reasoning bias under small initialization scale, we begin by em- ploying two-layer fully connected network to address a particular task, where p = 1 and = + 1. The network structure is defined as follows:
2 = (21,.... ... 120}, Amem (11,-- 20) and M = {(11, (13, 11)}. The dataset contains 200000 data pairs, ensuring an equal number of samples for each anchor combination. The loss function em- ployed is Cross Entropy and the optimization algorithm used AdamW. The model architecture comprises decoder- only Transformer structure with 2 layers and single atten- tion head. We train the model under different initialization scales with = 0.3, 0.5, 0.8 utilizing the last token pre- diction method. Additional details about the experimental setup can be found in Appendix A.
{1....,
Definition Given that W()) Rdmxdd,, ( xdyor and as the activation function. Given any sequence X X(aa+1), we define the Embedding MLP model (Emb-MLP) Go
Go (X):== Σ weeb w(()) ( X
To investigate the impact on training behavior under varying initialization scales, we analyze the dynamics of loss and prediction accuracy Dmmm- Drsn,train and Dren,test. As illustrated in Figure 3A, for 7 = 0.3, the losses on Dmmm and Dr Desutrain decrease at nearly identical rates, while the
Comparing with large initialization scale (y = 0.3). a noticeable reasoning bias can still be observed in Figure 3B for a small initialization scale (? = 0.8).

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
Embedding space exhibits distinct patterns To investi- gate the causes of the reasoning bias under small initializa- tion for such simplified model, it's critical to understand the structure of the embedding space. Figure 4A depicts the cosine similarity matrices for embeddings of memory anchors and reasoning anchors at epochs 50 and 900. The results reveal that the cosine similarity between reasoning anchors S; Sy decreases with the increase of [s; 1 syl. sug- gesting that reasoning anchors quickly establish hierar- chical structure within the embedding space. In contrast, the memory anchors exhibit consistently high similarity and alignment, leading to lack of differentiation among them. Nevertheless, given that the model needs to learn more primitive-level mappings for memory mapping than reason- ing mapping. the embedding space associated with memory mapping should, in principle, exhibit greater complexity and variability. This phenomenon reveals that the primary challenge preventing the model from effectively learning memory mapping could highly possibly lie in its difficulty in identifying and differentiating between individual anchors.
a' (Ewex: wemb, w() 1, moftmax (Go (X*xi) Jumh 1, where 1 E R. LXXdob means the vector with all ele- ments equal to 1. Then the gradient flow of wem,, could be approximated by the limit formulation, i.c.,
dmbem,, dt
1 Re T2
1 111 drub
W W (4)
where n. represents the count of all tokens. We consider n. + 0o to obtain the asymptotic form of the following gradient flow.
Proposition 1. For any token 8, denote Y as a random variable, which takes values randomly from the label of any input sequence that contains token s. In the limit n + oc. we define Pa with its i-th element the probability of Y i, i.e., P. = PY =t). Assume the ratio of the token in the whole datasel rs == remains constant, then (4) can be approximated as:
dw di
(P' -P'
W2YT (2)T W
(5)
dvab
io eret ay epoch
a=11
epoch
x10-3
The proof is provided in Appendix B.1. Proposition I demonstrates that for any token s. its embedding vector is dominated by the distribution of Y which indicates that it's significant to discuss the distribution Y for different tokens. Firstly, we define the following random variables (u for discrete uniform distribution):
100
200
0
Memory Anchor
dn
cas (pa
~~ (2), Ar~ ~ (Arsa) 1.2.. ,9. (6)
D1 Reasoning Anchor
Then we have the following results:
Figure 4. A: Cosine similarity matrices for memory (top row) and reasoning (bottom row) anchors epoch 50 (left) and epoch 900 (right) of a model initialized with 7 = 0.8. B: Distribution of P deab
P(Y==i1 Amem)
diez,
(7)
Nz
1 for different reasoning anchor C: Cosine similarity between p", _11 and P", Jyo 1 for any reasoning anchor 8i, 8j, exhibiting a similar structure to the embedding space of reasoning anchors observed
and
P = il 8 Arsu)
Jeub
(8)
=P
Z+ Σ Aj = - sii-sssE E
j=1
Target distribution shapes the embedding This phe- nomenon can be interpreted through the training dynamics. To facilitate our analysis, we give the following assump- tion (Chen et al., 2024):
Equation (7) reveals that the information different mem- ory anchors is identical such that the embedding space of memory anchors exhibits a high similarity. However, (8) demonstrates that the distribution of y's exhibits shifts in the mean values that depend on the specific s for any E Arsn. Figure 4B and 4C visualize the distribution of P d.oa and the resulting cosine similarity among different reason- ing anchors s, suggesting that the labels' distributions play critical role in establishing the embedding structure of reasoning anchors during the carly stages of training, facili- tating the differentiation among the embeddings associated with different reasoning anchors. The detailed formulations can be found Appendix B.2.
Assumption 1. The activation function a E (RR) such
and there exists universal constant Ct that its first and second derivatives satisfy ||a'(-) s Ct, [>ll< & Moreover, (0) = 0, o' 1.
For any token let {(x*s .) /"m denote all input sequences containing s and corresponding labels, where n, means the appearance times of s. As the initial- ization scale decreases, with Assumption 1, we have
5

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
Experiment:
Theory Validation:
COS (wemb,s; wemb.s)
A
B
D
epoch 200
epoch 900
Key
COS (wemb.s 8, wemb.s)
Key
0.3
-0.1
-0.3
-0.3
-0.5
-0.35
0.1
Si
Memory Anchor 5;
-0.5
0.0
0.5
-0.5
0.0
0.5
Reasoning Anchor
(pema sjj
Exp
Reasoning Anchor
1.00
1.0
Theory
CoS Aipoon tt 1,,,,,,,,
0.0
-022
-0.2
0.6
-0.4
-0.5
Si
Reasoning Anchor Si
-0.5
0.0
0.5
0.3
11
15
-0.5
0.0
0.5
Figure 5. Embedding structure of Transformer model with small initialization scale. A: Cosine similarity matrices for memory (top) and reasoning (bottom) anchors at epoch 200 (left) and epoch 900 (right). B: Visualization the embedding space projected onto the first two principal components computed via PCA. C: Cosine similarity between the constructed embedding vectors of reasoning anchors i cmmba (see (12)) as derived in Theorem I (top) and Cosine similarity comparison between experimental results with theoretical approximations where s; 15 (bottom). D: PCA projection of the constructed embedding space mmm, for E Z (top) and Aren (bottom) onto the first two principal components.
3.3. Transformer with General Task
space examine its structural properties. The results in Figure 5B reveal a strong inherent numerical order. This structure is particularly advantageous for reasoning tasks, as supports the model's capacity to generalize based on the underlying numerical relationships.
In the previous section, we investigate the key mechanisms driving the learning bias of Emb-MLP and analyze the dy- namics of its embedding space. However, when applied to a general sequence containing some degree noise, i.e., L >>a q + 1, we find the MLP model fails to perform effec- tively due to its inability to extract critical tokens žp, ap+1+ and dp+2- In contrast, Transformer architectures overcome this limitation through self-attention mechanisms, which can identify the key and anchor elements and propagate their information.
First attention module. As illustrated in Figure 6, the first attention matrix approaches the behavior of an aver- age operator when initialized with small scales, such that (Attn (X) V); = & Ejss Vj. where V = XWV. Con- sequently, each token aggregates information from all pre- ceding tokens. Additionally, the largest singular value of W is significantly larger than the remaining singular val- ues, and its corresponding singular vector aligned closely with the embedding vectors of reasoning anchors, but nearly orthogonal those of memory anchors. These phenom- ena suggest reasoning anchors prominently captured by W and subsequently propagated to all subsequent tokens in the sequence via the average operation. However, the memory anchors are not distinctly identified, indicating that the model faces challenges in capturing significant informa- tion from memory sequence effectively. More analysis of WV can be found in Appendix B.8.
In the following section, we conduct an in-depth analysis of the Transformer's characteristies and processing mech- anisms under small initialization scale. Specifically, we investigate whether the embedding space exhibits similar phenomena to those observed in Emb-MLP and assess how the model captures critical information from the input se- quence.
Embedding space. The embedding space of the Trans- former exhibits a phenomenon similar to that observed the Emb-MLP. Figure 5A illustrates the cosine similarity among different anchors' embedding vectors, revealing distinct pat- terns for reasoning and memory tasks. Reasoning anchors display a hierarchical structure, the further distance, the smaller similarity, suggesting a clearer organization within the embedding space. In contrast, memory anchors exhibit high similarity and alignment. Additionally, we apply Prin- cipal Component Analysis (PCA) to the entire embedding
Second attention module. The second attention module functions to extract the key, and propagate its information to the final position in the sequence. This is facilitated through the use position embeddings. Since this mechanism applicable to both memory and reasoning tasks, we provide detailed explanation in Appendix D.
6

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
A
B
C
D
Attention Matrix
1.0
30
0.8
Singular Value
Icos(u, wemmo, Ag)
Arsn
0.15
Amem
0.05
0.0
-0.1
0.0
0.1
10
20
10
j(Ajk 1/f)
Singular Value Index
Singular Value Index
Figure 6. Characteristics of the first attention module under small initialization (7 0.8) in the carly training stage (epoch 200). A: Heatmap of the attention matrix for random sample. B: Distribution of the relative error between attention A₁k and across all training sequences. C: Distribution of singular values of WV D: Cosine similarity between the left singular vectors and average embedding vectors of the anchors.
Theoretical analysis. Based on the observations from the experiments, we extract the sketch component of the model, which is crucial to its learning preferences, and analyze the underlying mechanisms for its occurrence. We define the following one-layer Transformer model:
Assumption we have the the following result:
dbbs
Ts L
_ W.
1
(10)
dt
N2
dm
Proposition For any s Asss let n, - oc, with Assumption we have the following result:
Definition (One-layer Transformer). Let dj N+ de- notes the hidden layer of the feedforward neural network (FNN). For any X X(),L), denote Attn (Wemb.X) by A,
embes dt
T's (p' p' )
then we define fo : X(9.L) + Rd as follows:
W,
(11)
T.
fx( =0 =0 Wemb.X WW WO We W
f1 W
where the i-th element of P is defined as P P ( A =
+AL,
Wemb, Ae,b
WV w WL.
s Aren)
(9)
where W BRdm 122 Rdjxdm are the feedfor- ward layer projection matrices. The subscript L, in AL,
Furthermore, we utilize the normal distribution to approxi- mate the distribution of Y*, E Areo and give an approxi- mation of mmb, to describe the overall structure and inter- nal relationships within the embedding space observed in real- roal-world training scenarios.
and Wemb, denotes the I.-th row.
Definition 2 is introduced to facilitate the theoretical anal- ysis, excluding the Layer Normalization and the final pro- jection operator, as they do not impact our results (see Ap- pendix F).
Theorem 1. let n + co, define the learning rate and assume that L 1-q-- 1---0000 0 - = (1) and [m | S O (d=²) at initialization. We propose the approximation of eamb , Arsn by
As we observed, with small initialization scale, the atten- tion operator A can be interpreted as average operator. Specifically, we have
wemb.e = =c (Cze w
1 dm
2
1
+
(12)
Lemma For any e (0, 11, there exists C such that for any y > C, the elements of A initialization, denoted by Ai,, satisfy Ai. - +1 < E for any S j with probability at least E.
where C1, Cz Op are constants depending on T L, and & ~ NO0, (d) Then we have the following result
supl ( emmbee n.mb...)
(w emb.s, w emb.e:))
Denote that w 52 vvo = V and W (WAT + (WVOTT + +), where the identity ma- trix I comes from the resnet. Using techniques similar to those employed in the previous section, we derive the gra- dient flow of emb.s under small initialization scales as follows:
(13)
00 (d (
where (-,-) denotes the inner production.
Additionally, for any key E and Amem- we could have similar result. To validate our theory analysis, we set the
Proposition 2. For any 8 Amem, let n + oc, with
7

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
A
B
PrOntoQA
TinyStories
151
-V3
-Y=
-V=0.8
10
-0
2500
5000
Step
Figure 7. Reasoning bias of GPT-2 in real language tasks. A: dynamics of AL during early training stage with initializations scales 7 = 0.3, 0.5 and 0.8. B: Cosine similarity among embedding space of PrOntoQA dataset and TinyStories dataset at step 5000 with y = 0.8.
detailed formulation of reasoning anchors and keys as
3.5. Effeec of Label's Distribution
W
smb,s
(6-a)²
dm
1
t
SE Arsn
(14)
110
Target Prediction
125
Task
Task 2
w
emb,s
h 12
+
SE Z.
90
310
Figure 5C exhibits the cosine similarity among the wemb, for any Arsa (top) and compare cos 1m, wembb. in real training process with the theoretical approximation cos (ziemb,15 rjmbb,s) (bottom, a complete comparison is provided in Appendix B.7). Figure 5D presents the PCA projection of emb,s for s E Armn and Z, respectively. These visualizations exhibit strong alignment with the ex- perimental observations, thereby substantiating the validity of our analysis. The proofs of our theoretical results are provided in Appendix B.4, B.5, and B.
50
10-
10"
Fpoch
Figure 8. Left: Distribution of targets and predictions 4 groups of memory tasks. Red represents the target distribution in each group and blue represents the prediction distribution. Right: Learn- ing speed comparison for Fmem, (red) and Fmem, (blue).
Previous sections reveal that under small initialization, the distribution of labels plays pivotal role shaping the em- bedding space of tokens and regulating the model's training dynamics. To more intuitively demonstrate the impact of the label distribution of each token on its output, we designed four groups of memory mappings. The label ranges of the four groups are set to {30,- ... 29 + 20 x i}, i 1,2,3, 4. The right picture of Figure 8 illustrates the distribution of the model's outputs for each group during the early stages of training. Notably, it can be observed that even at this initial stage, when the model's accuracy is still rel- atively low, its outputs do not exceed the range of the la- bel distributions. This highlights the critical influence of label distributions on the token structure, which in turn significantly impacts the model's outputs. Additionally, we compare two memory tasks with differing label distri- butions. In the first task, denoted Fmem,,1 for any key- anchor pair (=p, a), the label ((=p.a) is randomly sampled from Z. In the second task F (( is randomly sampled from & - Ex a..-... &p + E a: While both tasks are clearly memory tasks, the label distri- butions in Fmem,2 vary depending on the anchor. As shown in the left picture of Figure 8, the learning rate for Fmem2
3.4. Real Language Tasks
For the experiment in Figure 1, we also conduct compar- isons with initialization scales y = 0.3 and 0.5. To quantify the reasoning bias of the model, we define the following metric:
AL Linystories LPiOntoQA
LPIORTOQA
where TrTinysoories and LPoOMIOA denote the loss on TinyS- tories and PrOntoQA, respectively. As increases, AL exhibits an upward trend, indicating a growing bias for rea- soning task, which is depicted in Figure 7A. To further validate our analysis, we examine the embedding space of the GPT-2 model during the early stages of training which we train with a small initialization scale. Figure 7B reveals that the embeddings of tokens in PrOntoQA are significantly more distinguishable from each other compared to the to- kens in TinyStories. The average cosine similarity among the PrOntoQA is 0.123 while 0.531 among the TinySto- ries. These results provide strong support for our analysis, highlighting the impact of embedding distinguishability on training preference.
8

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
is significantly faster than that for Fmmmm.-- This observation underscores the crucial role that label distribution plays in the model's learning process.
study, we demonstrate that with small initialization values, the parameters of the embedding layer can reach low-rank state rather than condensed state. This means that while embeddings of different tokens become linearly dependent, they are not identical. This distinction allows low-rank mod- els to effectively capture essential patterns and generalize well without the stringent alignment required by condensa- tion, which is particularly important for applications such as word embedding matrices where distinet representations for different tokens are necessary. Recent investigations have also explored how initialization affects the training dynamics of LLMs (Huang et al., 2020; Liu et al., 2020; Trockman & Kolter, 2023; Wang et al., 2024a; Zhang et al., 2019a; Zhu et al., 2021). These studies mainly examine how the scale of initialization influences the stability of the training process and is vital for ensuring efficient and effective training of LLMs. In our work, we observe that different initialization schemes result in varying speeds of convergence for memorization tasks versus reasoning tasks and provide a theoretical rationale for this behavior.
4. Related Works
Recent advancements in large language models have shown remarkable capabilities, often surpassing human-level per- formance in many tasks (Fu al., 2022; Wei et al., 2022a). However, despite their strong performance in many as- pects (Srivastava et al., 2022), LLMS face challenges in handling complex reasoning tasks (Csordás et al., 2021;
Dziri et al., 2024; Hupkes et al., 2018; Lepori et al., 2023;
Okawa et al., 2023; Yun et al., 2022; Wang et al., 2024d;
Csordás et al., 2022). For example, Ramesh et al. (Ramesh et al., 2023) show that Transformers trained on a synthetic benchmark struggle when tasked with combining multiple reasoning steps. Similarly, Liu et al. (Liu et al., 2022) sug- gest that shallow Transformers tend to learn shortcuts during training, which limits their ability to perform well in more complex reasoning scenarios. Several strategies have been proposed to address these challenges, such as encouraging the generation of explicit reasoning steps in a single out- put (Wei et al., 2022b) or using L.MSs iteratively produce reasoning steps (Creswell et al., 2022: Creswell & Shanahan, 2022). Despite these efforts, achieving reliability remains a significant challenge. Additionally, some studies have explored the internal mechanisms of language models to enhance their performance (Wang et al., 2024b;c; 2023), but they often do not address the impact of training dynamies on the model's final behavior. To better understand these models' behaviors and inner workings, Zhang et al. (Zhang et al., 2024b) introduced anchor functions as a tool for prob- ing Transformer behavior. Building on this framework, our research investigates how different initialization scales in- fluence model reasoning bias and internal mechanisms from a perspective of training dynamics.
5. Conclusions
In this paper, we investigate the underlying mechanism of which small initialization scales promote a reasoning prefer- ence in language models. Our findings suggest that the label distribution of tokens plays a pivotal role in shaping the em- bedding space, thereby influencing the learning dynamics and task complexity. Our result can be readily extended to the next-token prediction training and obtain similar results. This perspective is supported through a combination of ex- perimental observations and theoretical analysis, providing a deeper understanding of how initialization strategies impact task-specific behavior in language models.
A Acknowledgements
This work is supported by the National Key R&D Pro- gram of China Grant No. 102YFA1008200 the Na- tional Natural Science Foundation of China Grant No. 92270001(Z. X.), 12371511 (Z. X.), 12422119 (Z. X.). Shanghai Municipal of Science and Technology Major Project No. 202ISHZDZX0102 (Z. X.), and the HPC of School of Mathematical Sciences and the Student Inno- vation Center, and the Siyuan-1 cluster supported by the Center for High Performance Computing at Shanghai Jiao Tong University, and Key Laboratory of Marine Intelligent Equipment and System (Ministry of Education, P.R. China), and SJTU Kunpeng & Ascend Center of Excellence.
The fitting ability and generalization of deep learning mod- els are critical in understanding deep learning (Breiman, 1995; Zhang et al., 2016), and the initialization of neural network parameters plays a crucial role in determining the network's 's fitting results (Arora et al., 2019; Chizat & Bach, 2018; Zhang et al., 2019b; E et al., 2020; Jacot et al., 2018; Mei et al., 2018; Rotskoff& & Vanden- nano & Spiliopoulos. 2020: Williams al., 2019). Luo et al. (Luo et al., 2021) and Zhou al. (Zhou et al., 2022) primarily identify the linear and condensed regimes in wide ReLU neural networks. In the condensed regime, neuron weights within the same layer tend to become similar. A body of research indicates that condensed networks often exhibit strong generalization capabilities (Zhang et al., 2022; Zhang& & Xu, 2023; Zhang et al., 2023; Zhang& & Xu, 2024). See (Xu et al., 2025) for overview of condensation. In our
2018; Sirig-
Impact Statement
Our works provide new insights into the intrinsic mecha- nisms underlying the reasoning bias of language models

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
under small initialization scales, as well as the training behavior of individual modules within the model architec- ture. These findings not only contribute to understanding the model behavior and training mechanisms, but also help with optimizing model initialization strategies and design- ing novel algorithms to enhance the reasoning capabilities of language models.
Csordás, R., Irie, K., and Schmidhuber, J. The neu- ral data router: Adaptive control flow in transform- ers improves systematic generalization. arXiv preprint arXiv:2110.07732, 2021.
Csordás, R., Irie, K., and Schmidhuber, Ctl++: Evaluating generalization never-seen compositional patterns of known functions, and compatibility of neural representa- tions. arXiv preprint 20.0032500 2022.
References
Dziri, N., Lu, X., Sclar, M., Li, L., Jiang. L., Lin, B. Y., Welleck, S., West, P., Bhagavatula, C., Le Bras, R., et al. Faith and fate: Limits of transformers compositionality. Advances in Neural Information Processing Systems, 36, 2024.
Abdin, M., Aneja, J., Behl, H., Bubeck, S., Eldan, R., Gu-
nasekar, S., Harrison, M., Hewett, J., Javaheripi, M.,
Kauffmann, P., Lee, J. R., Lee, Y. T., Li, Y., Liu, W., Mendes, C. C. T., Nguyen, A., Price, E., de Rosa, G., Saarikivi, O., Salim, A., Shah, S., Wang, X., Ward, R., Wu, Y., Yu, D., Zhang, C., and Zhang, Y. Phi-4 techni- cal report, 2024. URL
E, W., Ma. C., and Wu, L. A comparative analysis of optimization and generalization properties of two-layer neural network and random feature models under gradient descent dynamics. Sci. China Math., 63, 2020.
orgg/r/gaoooooo /bs/
2412. 08905.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Eldan, and Li, Tinystories: How small can language models be and still speak coherent english?, 2023, URL https: / / a arrxiv vi.. org abs 2.355 07759.
Arora, S., Du, S. S., Hu, W., Li, Z., Salakhutdinov, R. R., and Wang, R. On exact computation with an infinitely wide neural net. In Advances in Neural Information Pro- cessing Systems, pp. 8141-8150, 2019.
Fu, Y., Peng, H., and Khot, How does gpt obtain its ability? tracing emergent abilities of language models to their sources. Yao Fu's Notion, 2022.
He, K., Zhang, X., Ren, S., and Sun, J. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE inter- national conference on computer vision, pp. 1026-1033 2015.
Breiman, L. Reflections after refereeing papers for nips. The Mathematics of Generalization, XX:11-15, 1995.
Caiado, C. and Rathie, P. Polynomial coefficients and dis- tribution of the sum of discrete uniform variables. 2007.
Huang, X. S., Perez, F., Ba, J., and Volkovs, M. Improving transformer optimization through better initialization. In International Conference on Machine Learning, pp. 4475- 4483. PMLR, 2020.
Chen, Z.-A., Li, Y., Luo, T., Zhou, Z., and Xu, Z.-Q. J. Phase diagram of initial condensation for two-layer neural networks. CSIAM Transactions on Applied Mathematics, 5()):44-551,, 2024. ISSN 2708-0579. doi: http:://o..oo////2222808saammmm-
Hupkes, D.,, Singh, A., Korrel, K., Kruszewski, G., and Bruni, E. Learning compositionally through attentive guidance. arXiv preprint arXiv: 1.t.::800009655. 2018.
URL article.. cestailcssian 1etaii//ciiammmm2233300 www.mmmmmmmm33...mm html.
Jacot, A., Gabriel, F., and Hongler, C. Neural Tangent Ker- nel: Convergence and Generalization in Neural Networks. Advances Neural Information Processing Systems 31, pp. 8571-8580. 2018
Chizat, and Bach, F. On the Global Convergence of Gradient Descent for over-parameterized Models using Optimal Transport. In Advances in Neural Information Processing Systems 31, pp- 3036-3046. 2018.
LeCun, Y., Bottou, L., Orr, G. B., and Müller, K. R. Effi- cient BackProp, pp. 9-50. Springer Berlin Heidelberg, Berlin, Heidelberg, 1998. ISBN 978-3-540-49930-00 doi: 10.
Creswell, A. and Shanahan, M. Faithful reasoning using large language models. arXiv preprint arXiv:2208.1422 2022.
URL https: /ttp:/oooooorooo///// 10. 1001007//00 3-0-5-50---- 115444444333 19430-8-220
Creswell, A., Shanahan, M., and Higgins, I. Selection- inference: Exploiting large language models for interpretable logical reasoning. arXiv preprint arXiv:2205.097/2. 2022.
Lepori, M. A., Serre, T., and Pavlick, E. Break it down: Ev- idence for structural compositionality in neural networks. arXiv preprint arXiv:2301. 2000::230/..888.. 10884, 2023.
10

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
Lin, Z., Gou, Z.,, Gong, Y., Liu, X., yelong shen, Xu, R., Lin, C., Yang, Y., Jiao, J., Duan, N., and Chen, W. Not all tokens are what you need for pretraining. In The eighth Anmual Conference on Neural In- formation Processing Systems, 2024. URL https: // //openreview net foremmormmmmmmmmmmm -immmiid--Nii ONMZ
Sirignano, and Spiliopoulos, K. Mean field analysis of neural networks: A central limit theorem. Stochastic Pro- cesses and their Applications, (803)::8200-155,22000 doi: 10. htt01///sppa.211..666 .9000.3...
Smolensky, P., McCoy, R., Fermandez, R., Goldrick, M., and Gao, J. Neurocompositional computing: From the central paradox of cognition new generation of ai systems. Al Magazine, 43(3):308 322, 2022.
Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C.,, Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. Deepseek-v3 technical report. arXiv preprint
29..2.272.99977, 2024.
Srivastava, A., Rastogi, A., Rao, A., Shoeb, A. A. M., Abid, A., Fisch, A., Brown, A. R., Santoro, A., Gupta, A., Garriga-Alons A., et al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. arXiv preprint arXiv:2206. 0.000.04461.... 2022.
Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C. Transformers learn shorteuts automata. arXiv preprint arXiv:2210. 10749, 2022.
Liu, L., Liu, X., Gao, J., Chen, W., and Han, Understand- ing the difficulty of training transformers. arXiv preprint arXiv:2004.088 08249, 2020.
Trockman, A. and Kolter, Z. Mimetic initialization of self- attention layers. In International Conference on Machine Learning, pp. 34456-34468, PMLR, 2023.
Luo, T., Xu, Z.-Q. J., Ma, Z., and Zhang, Y. Phase dia- gram for two-layer relu neural networks at infinite-wiee width limit. Journal of Machine Learning Research, 22(71): 1-47, 2021.
Wang, H., Ma, S., Dong, L., Huang, S., Zhang, D., and
Wei, F. Deepnet: Scaling transformers to 1,000 layers. IEEE Transactions on Pattern Analysis and Machine In- telligence, 2024a.
Marcus, G. F. The algebraic mind: Integrating connection- ism and cognilive science. MIT press, 2003.
Wang, L., Li, L., Dai, D., Chen, D., Zhou, H., Meng, F., Zhou, J., and Sun, X. Label words are anchors: An information flow perspective for understanding in-context learning. arXiv preprint vxXiv2305.144160 2023.
Mei, S., Montanari, A., and Nguyen, P.-M. A mean field view of the landscape of two-layer neural networks. Pro-
ceedings of the National Academy of Sciences, 115(33): E7665-E7671, 2018. doi: 10. 10.1073/pnas. pttpppaa....009111..1111111 1806579115.
Wang. M., He, H., Wang. J., Wang. Z., Huang, G., Xiong, F., Li, Z., Wu, L., et al. Improving generalization and convergence by enhancing implicit regularization. arXiv preprint arXiv:2405.20760 2024b.
Okawa, M., Lubana, F. S., Dick, P., and Tanaka, H. Compositional abilities emerge multiplicatively: Explor- ing difffusion models on synthetic task. arXiv preprint MiXiv:2310.09336. 2023.
Wang. M. et al. Understanding the expressive power and mechanisms of transformer for sequence modeling, arXiv preprint arXiv:2402.0053 2024c.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models unsu- pervised multitask learners. 2019. URL https: org/CorpusID: // peapi.semanticsse horpooooooo....o.//oo/oss.o 160025533.
Wang. Z., Wang, Y., Zhang, Z., Zhou, Z., Jin, H., Hu, T., Sun, J., Li, Z.,, Zhang, Y., and Xu, Z.-Q. To- wards understanding how transformer perform multi- step reasoning with matching operation. arXiv preprint arXiv:2405. 2000244552.. 2024d.
Ramesh, R., Khona, M., Dick, R. P., Tanaka, H., and Lubana, E. S. How capable can a transformer become? study on synthetic, interpretable tasks. arXiv preprint arXiv:2311.12997, 2023.
Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Met- zler, D., Chi, F. H., Hashimoto, T., Vinyals, O., Liang, P., Dean, J., and Fedus, W. Emergent abilities of large language models. Transactions on Machine Learning Re- search, 2022a. URL https: / wpwnrevieeeeeeeeee renieeeiwwwww. net// forum Serum?ddyyzsSU5zdaa Survey Certification.
Rotskoff, G. and Vanden- Fijjdenn, Parameters as inter- acting particles: long time convergence and asymptotic error scaling of neural networks. In Advances in Neu- ral Information Processing Systems 31, pp. 7146-7155. 2018.
Saparov, A. and He, H. Language models are greedy rea- soners: A systematic formal analysis of chain-of-thoughh In The Eleventh International Conference on Learning Representations, 2023. URL https: http///pppeeeeeeeeeeeeee net/ forum? forum?id=gguunee d=w.mm//=FddF00222222233
Wei, J., Wang, X., Schuurmans, D., Bosma, M., Chi, E., Le, Q., and Zhou, D. Chain of thought prompting elic- its reasoning in large language models. arXiv preprint arXiv:2201./ 2022b.

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
Williams, F., Trager, M., Silva, C. T., Panozzo, D., Zorin, D., and Bruna, Gradient dynamics of shallow univari- ate relu networks. CoRR, abs/1996.078/// 2019, URL http:/ /ttp://aaavvv.//.a//aa//eaaeaiaa org.orggg///o abs 190...996... 07842.
Zhou, H., Zhou, Q.. Jin, Z., Luo, T., Zhang, Y., and Xu, Z.-Q. J. Empirical phase diagram for three-layer neu- ral networks with infinite width. Advances in Neural Information Processing Systems, 2022.
Xu, Z. J., Zhang, Y., and Zhou, Z. An overview of con- densation phenomenon in deep learning. arXiv preprint arXiv:2504.09484, 2025.
Zhu, C., Ni, R., Xu, Z., Kong, K., Huang, W. R., and Gold- stein, T. Gradinit: Learning to initialize neural networks for stable and efficient training. Advances in Neural In- formation Processing Systems, 14:16410-164222 2021.
Yun, T., Bhalla, U., Pavlick, E., and Sun, C. Do vision- language pretrained models learn composable primitive concepts? arXiv preprint arXiv:2203. 17271, 2022.
Zhang, B., Titov, I., and Sennrich, R. Improving deep transformer with depth-scaled initialization and merged attention. arXiv preprint arXiv: 1908. 11365, 2019a.
Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O. Understanding deep learning requires rethinking general- ization. arXiv preprint arXiv:1611.0355
2016.
Zhang, Y., Xu, Z.-Q. J., Luo, T., and Ma, Z. A type of gen- cralization error induced by initialization in deep neural networks, arXiv: 1905:099000777 007777777 [cs, stat], 2019b.
Zhang, Y., Zhang, Z., Zhang, L., Bai, Z., Luo, T., and Xu, Z.-Q. J. Linear stability hypothesis and rank stratification for nonlinear models. arXiv preprint arXiv:2211 I. 1623, 2022.
Zhang, 7. and Xu, Z.-Q. J. Loss spike in training neural networks. arXiv preprint 12X1::2305.12133, 2023.
Zhang. Z. and Xu, Z. J. Implicit regularization of dropout. IEEE Transactions on Pattern Analysis and Machine In- telligence, 2024.
Zhang, Z., Li, Y., Luo, T., and Xu, Z.-Q. J. Stochastic modified equations and dynamics of dropout algorithm. arXiv preprint 2rxiv:2355.58850 2023.
Zhang, Z., Lin, P., Wang. Z., Zhang, Y., and Xu, Z.-Q. Initialization is critical whether transformers fit com- posite functions by reasoning memorizing. In The eighth Annual Conference Neural Information Processing Systems, 2024a.
Zhang, Z., Wang, Z., Yao, J., Zhou, Z., Li. X., E, W., and Xu, Z.-Q. J. Anchor function: type of benchmark functions for studying language models. arXiv preprint 20X::22001.0830, 2024b.
Zhang, Z., Lin, P., Wang. Z., Zhang, Y., and Xu, Z.Q. J. Complexity control facilitates reasoning-baad compo- sitional generalization in transformers. arXiv preprint NKXi:25011005331 2025.
12

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
A. Basic Settings
A.1. Transformer Architecture
For any sequence X E X((2.1), we denote its one-hot vector by ex The word embedding Wemb and the input to the first Transformer block ( is calculated as:
Wemb,X = e Wrem
X = Wemb, X + WPos
(15)
where WPos is trainable positional vector. For the /-th layer, the Q, K, V are defined as:
Q ( wQ0 K Wk()) w x0 WV
(16)
The attention matrix Attn and its subsequent output X for the I-th layer is computed as:
Attn (1) = softmax mask (Qth ( X(kv(1) Attnnv()
(17)
dd
The output of the /-th attention layer is obtained as:
x LN x x wow) x (1+1)=
d
LN (MLP (xxoo() x
(18)
where "LN" refers to Layer Normalization. The final output is obtained by projecting the output of the last layer X do)) using a linear projection layer, followed by a softmax operation and argmax to obtain the predicted token.
A.2. Experimental Setups
For those experiments about the Transformer structure, we train three Transformer models on a dataset of 200,000 samples, with each input sequence having fixed length of 9 tokens. The vocabulary size dvob is set to 200, and the model architecture includes an embedding dimension dm of 200, a feedforward dimension dy of 512, and query-key-value projection dimension dy of 64. The Transformer-based model uses decoder layers with attention head per layer. The training is conducted for 1000 epochs with a batch size of 100, and gradient clipping is applied with a maximum norm of 1. The AdamW optimizer employed with an initial learning rate of 10-5 The initialization rates of the three models 7 0.3, 0.5, 0.8.
For those experiments related to Emb-MLP, we train three Emb-MLP models with dvob 200, dm 200, d, 512 and initialization scales 7 = 0.5, 0.8. A dataset comprising 1, 000, 000 data pairs is employed. We set that Amem = {1,2.... ..... ,10]- Arsm {11, 12, ,20} 2 = {21,22 ... 120) M = {(11, 13). (13, 11)}- The initial learning rate is x 10-6 and other training setups remain consistent with those described in the first paragraph.
For Figure and Figure 7, we use GPT-2 models with an initialization scale 7 0.3, 0.5, 0.8. The dataset contains 10,000 data sequences, with half of them sourced from PrOntoQA and the other half from TinyStories. The AdamW optimizer is employed with initial learning rate of 1 x 10-5, The model is trained for 200 epochs, ensuring that the loss for both datasets decreases to similar level.
B. Theory Details
B.1. Proof Proposition 1
Lemma 2. For any token s, {(X*s, (x*)) denote all input sequences containing s and corresponding labels. The gradient flow of obb can be expressed as:
dummm.
Th n
(p"i-y") w(a)T) - dd
Σ w ammb. (mb
WOI
(19)
1=1
x
where pa softmax (Go (x. denotes the derivative of a and n means the count of all training data. O represents the elements-wisg production.
13

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
Proof. For any data pair (X+i, :).)), the cross entropy function R could be expressed as:
Go(X" you
(X**)) -llog
Σjer Go(Xx.)))"
where the subscript j represents the element index. Then the derivative of respect with womb, can be expressed as:
ØR (X
duob 0R (X+ii) 0Go (X*.)), Swemb,
drob
8Go (X
Σ (D) - Vÿ rroemb,s
erremb,e
Go (XEN))
==1
Using the trace theorem, we obtain:
dGo (X*.1, =tr (dGo (xx.)),) do Σ u emb, W(l) w:!
zex
tr
w
gembbawlll Wil))
@ emw (dwemb.eeWr))) w()) w(()) w:l
Xo.
tr (
w Oo'
Σ
w emb,, w(l))
du
Then we have
2R (xx+)
dunh
( (0 13"))
w oo
w cmb..r (ro
)) wiir
awemb,
.x w amb
= ((Drs - yrrt) W(eiT) oo
(oamb.-w())) wall warr
and furthermore
mbb
R2 OR (Xeid)
1 Σ
(lrt) -y*) w(2) oo Σ zeX
w amb will) (elbbaw()) W(UT.
dt
iwemb,s 71 i=1
As the initialization scale decreases 7 O, with the Assumption 1, we have that a' (E.ex... wemb.. 1, softmax (Ge (X*+)) = Zumh 1.1, where 1 E IR 1xdeab means the vector with all elements equal to Then the gradi- ent flow of smb,s could be approximated by the limit formulation, i.e.
d
Ry 11
dyob 1 1111
y
(2) (UT
(20)
dl
Consider - and denote the random variable y which follows the distribution of {y*") asymptotic form by the law of large number
n. Jmm then obtain the
n Σ
y 2. = ry [Y*] rsPs,
n i=1
where Y is the one-hot representation of Y and the i-th element of P is P P (Y* i). Then we obtain the Proposition
14

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
B.2. Distribution of y
Memory anchor. Since select a label for any key-anchor pair randomly from 24 (Z), the label meets would follow the same distribution for any Amem- specifically, have
P P()==1))= N2 61EZ,
(21)
where diez iifi E Z otherwise 0.
Reasoning anchor. For any reasoning anchor 8 E Arsn we assume that the other elements a key-anchor pair containing 8 is 2, 01, 02, ... aq-1- Since the other elements are randomly chosen from the corresponding scope, the label could be represented as y² = + Ej=1 Oj. Then Y follows the distribution s + Z + Σø=1 Aj, then we have
P iY =
P
++
Aj = =i--
j=l
Gmax
Σ
( = P)
Σ Aj A)=i-o-5)
(22)
Smax
Σ
q 9-1
i - (q - 1) N.ron
Nz No
C=Emin
where the combination number (() k+1 can be defined by (1 x + +z")" E5=o G))+t
x (Caiado & Rathie, 2007).
k+1 x+1="
Specifically, when 2, we have the following result:
( Caax
P i) i)=
Σ (Z= = (A
C=iiin
8-11
Omnnn
== Smin + anmin + S.... SSin + a. +
NZ
5=imin
{-g Onh
i Smin a +8.... Smax + a + 8.
NzNAmn
(=1-8-a Ore
Cuax
1 Smax + O raa
+1+8 Smax + rsn + 8.
N₂NAm
- 8- O
Smin +1
min
i= Smin a + S, Smin + Ten +
NzNA... Arm
i = Smin
rsn + O
+ 1 + 8,... Smax + rsn + 8.
N₂²
Smax
ran
-itt 1
max
i Smax + Isn + s,... Smax Q
NZw
Key. For any s Z, its labels come from two parts, memory mapping and reasoning mapping. In the memory mapping. &
meets each token 2 with the same probability N= In the reasoning mapping, the label y² = 8 + Σi= ag Assume that
15

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
the ratio of memory mapping is identified with reasoning mapping, then we have
P(Y = 1) = 2
P mem (Y =i) 1) + rsu (Y = f))
(23)
1
ÓiEZ +PP s Σ
Aj =
(24)
2
N3
j=1
diez +
(25)
2
Nz
Nº
i - 8
9
tan min
Nea
Specifically, when g
1
11/1
( diez +
NA
1 1 8
P = i) 2
max
min
N2
N3
Generally, consider the usual sequence containing some noise. Then the label consists of third part when appears noise. With similar method, we have
Poisse (Y* =
(P
noiss,mem (Y i) + P noise,ron (Y* 1))
1 diez +P
2
Σ
9
=1))
2
Nz
j=1
1 N2
òiEz
1
Sanax
Σ
9
Nz No
i-5
ma
ran
Arm (=Cmin
Nar
Combine them together, we have in the general setting X),, we have that
1
()) δiez +P
q
1
I. 1--q- 1-1
9
Aj = =1)
=
P(Y* = 2 9)
Σ
A₂
= =1
δiEz +P
2 Σ
(N2
2(L 9)
NNz
j=l
j=1
1
SIEZ +
P + Σ
Aj =i
+ (L q P)P
1 Σ
Aj =
2N
2 -
j=1
j=1
- rSD
2Nz diEz + 2 (L Am
i-
+ (L 1)
N2
i i-
TED
NAu
< Cuii
NAua
B.3. Gradient Flow of Embedding Space in Emb-MLP
With the discussion in Section B.2, we obtain the detailed formulation of (5) for different anchors of different tasks. Specifically, we have the following result:
Corollary 1. Given any Amem. assume that n and assume the ratio of sequences containing in the training set keeps constant, then we have
8 1
W(2 wax
(26)
dt
N2 dvob
where 8² E Rdd is vector with elements equal I for indices belonging 2, and 0 otherwise.
Corollary 2. Given any E Aysn assume that + oo and the ratio of sequences containing in the training set Te remains constant. Then, the gradient flow of the embedding vector corresponding to s is given by:
emb,
= =rs T
dvob 1)
(22)
(27)
dt
where P Rd is probability vector whose i-th element is P. Nz No-1 Eg (i-s- 9-1 --s---(9-1)opp Nem
16

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
B.4. Proof of Lemma Proof. We assume that W Wemb, N (0.((m)))) w (0.
wiy N (o. (d(7)))
We have that
dm
dm
de Σ
(Wemb,X W WTT Wemb.x.T)
Wemb,X
W w.k Ri
1g
1p
k=
1=1
dm
dm
Wemb.X
w W Wenb, kip 5
1p
k=
1=1 p=1
d dk 2 (din d,.)
N 0,
Then the attention operator
(Wemb, W W
KK KT W.eb,,
N 0,
((1,,
dm Vo
dd
2 (din + dlb)
Utilizing the Chebyshev's Inequality, then have
(Wenb,
woWkT KT X,mb
in
d dx 482 (dhn d%)
P
>8 S
vdk
+d for any 8 > Given any E (0, 1}, let C = logg... d dy 18
then for any 7 > C, we have
[Wemb,X [Wenb.XWoWW
WoW W X.m.
3.
d dx 462 (dm + dz))
P
>8
S E,
ddl
which implies that A.j f 1, for any i S as 7 + 0o.
B.5. Proof of Proposition 2, 3
For convenience in further analysis, we introduce the following notations H = emb wVwo +
weobb." W/ wvo wVwo, ws W/W pad softmax(fo(Xx.)). Firstly, have the following result:
Lemma 3. Given any token s, the gradient flow of emb, can be expressed as
dw omb
1
Σ L (p )) WVT @ (Hoi) + (pri - yntt) wvoor)
dt
(pt )² W oo (H**)² + -y.
where ñg denotes the time appears in the final position of a sequence.
Proof. For any data pair (x .y we have
dfo())))) d (H w W wmb WVo :j
Wemb,
LJj
=H ddeemb WoW dWemb, X wil) w!? + em. w.vo dWeem,
L.j
17

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
By the trace theorem, we have
dff (xe:)) = dff (x)
-Ir .n(wwown (w: oo (Ha.)") ww ( + (W (W! od (H awerbbbe...
[ w + (dwayhhre..
(W (w!! oo (H*))))) + w.y)) dw wwemh, +u((w od (H*))* T +1) dW dwi.b......
Utilizing the chain rule, we have
2R 1x*,,." 0R OR (x y2:h)
0fo (xxr)) WWemb,s
Ofe (xxx)), Dj vバ
Σ
dWemb,s
Ofo (Xx.i)
Wembb,
L (((p") ")* 1 od od' (H())) W VO,T + (p yy wvo.))
+ ((D* - ) W 2 o (H*)) WILT + (poit - ( yy(i)
occurs on last position.
((pp' - y² W o (H*))) W + (poit -y wvoz),
otherwise.
Then obtain the gradient flow as
dbbs
Σ ((p (p" -y W o (H**) W vo,
+ (p VV...
y wvo.)
dt
- )dr 1 o ( T WTlT + +pp (p"1--"') y/r))
As the initialization scales decrease to zero, we derive the gradient flow under small initialization scale follows via Assumption 1.
emb,
Σ
1
y
1 1) wwrows++ =Vv
dl
L
dm
1=1
(28)
WIT + (yes
1 (1-
dm 1)
1
dm
We consider the ideal condition n + oo, with the law of large number, (28) can be approximated as follow:
ebbs dt
Te Ey
y
1 W.
(29)
T,
dm
With the distribution we discussed in Section B.2, we complete the proof of Proposition 2,
18

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
B.6. Proof of Theorem
Proof. Consider the linear expansion of w w have
emb,s
duyembba
71 where wt,
emb,a is the initialization of wemb, then
dmbem
2 embbe
=
emb,s
+
(30)
dt
w to
emb,s
+ r By
1/¹
1
(WITT + +)) (WVOT +
(31)
I.
dm
T
By
y²
1
emb,, to
+ O 2d)
(32)
L
dm
For any Arsn, the formulation can be rewritten as
w
emb,
To L
P
s+
+
Aj
1
te,
(33)
dm
j=1
where : ~N(( (0, Let y enlarge enough, then PP formulation using the Berry- -sseen central limit theorem
s+ 2 + ---1 A₁ =
can be approximated by the following
sup
P
Z + Σ =1
6 e)
25p
0 (0(
(34)
V
where
and is the expectation and standard deviation of 7 + Σ Aj.
Denote that w
emb,s
20p ul
+ E, then we have:
V250p
di
cmb.s
cmb,s w;
< <0 20(0-t +
sup i
Then the difference inner production can be derived as follows:
i emm,,, w w)mb,,s) - emb,s,, w emb,st)
emb cmmbs; W
emb.e; cub,,;
sup
sup 1.
W
k
W
w k
emb.sy eembss
omb,3; emb,8;
sup h,3
k
we
wk
0
dim
Since an axis transformation does not affect the inner product, we set = i- then we complete the proof of Theorem 1.
B.7. Validation of Theorem 1
To verify the validity and generality of our theoretical analysis, we compare the cosine similarity of the embedding vectors within the reasoning anchors between experimental results and theoretical approximations. The results in Figure demonstrate that our theoretical estimates align well with the experimental results in most cases, with discrepancies observed only when |s, - sy becomes large, likely due to the omission of higher-order terms.
B.8. Discussion about WV
For the phenomenon where the V the first attention module exhibits a preference for capturing the reasoning anchors, a general theoretical explanation would require a more comprehensive and sophisticated analysis. However, a similar result can be derived under special and constrained condition.
19

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
11
= 12
5; = 13
= 14
5 15
Exp
Theory
0
20
11
11
16
17
5 18
S 19
20
0
11
20
11
20
Sj
5)
Figure 9. Cosine similarity comparison between experimental results cos (w emb,, w ss) with theoretical approximations cos (mi aimmb,e;) (see (12)), for any 8j E Are-
Theorem 2. ne, + oo, N2 = dm. Define that A ~ U (Ansa) and Y as random variable which randomly takes value from the whole datasel's ' labels, then we have the following result:
dWv Bawemb.ABy
w (W +1).
dt 2
dm
Theorem 2 highlights that W inherently demonstrates preference for reasoning tasks, thereby enhancing its ability to capture information associated with reasoning anchors.
Proof.
Firstly we have the following formulation:
Lemma 4. Given the dataset ((X), ,2)) the grudient flow of WV can be expressed as follow:
dW dt
w m T (p' y') (W"W1)) + W xmb XT, (p' yy) wor
H
(35)
n
Proof. For each data pair (X), we have
fo fx'), =o (AAL x VVwo wornb..)) Wi)) w AL. Wemb. WV W. + Wemb, Compute differential, we have
L.j
dff (wemb.x:). = d (0
(6(( (((( (Womb,
WV wo
We Wemb,xx W/1
w + wemmb W wwo
= (H') @ d
W w'wo, Wemb will 1 + Wemb,,x
dww w.
=o' (HH) W xi V w + w emb dWVwo
20

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
Using the trace theorem
dfo (x', tr (dfo (x)), tr (o' (H') w cmmb dW
wow!
+ W w dW W
W (H') W " dWv)
+ w dwv
W w dww). W o =tr m( (H))
which suggests that
o (xi), Owv
+ W W.W (.gambbrr)) T
W (H')) W emm
Utilizing the chain rule, we have
(x')
dm
0R ( 0fo (x:),
OWV
ofe (x:), owv
dm
(Wow (H') @ Wemb,,"' + w wembb,, T
= Σ (p) - w;)
j=1
w (p'-v') v') (Wow))" + Wimb, (p'-y') y¹) WO.
= d' H
where p' = (oftmax(ff(X))) Then gradient flow of WV can be expressed
dW
a aR (x*)
dt
12
owv
W womb T (p' (wow1)" + wemb (p'-y') wo.r
Σ
a' (H*)
When initialized with a small scale, the gradient flow for WV could be interpreted by:
dWv
1
emmb, x
yy
dm 1
T
(Wows + WooT
(36)
dt
n
7 yy - 1'-Z-1)
" (wows + wo)
w (i.))
(37)
nL.
where (i.)) denotes the j-th element of X In this formulation, there are nL tokens, and we reorder all tokens along with their corresponding labels. Let emb.s. denote the embedding vector of the i-th token, and let y be its corresponding label. Consequently, the gradient flow can be expressed as:
dwv
cmb y 1 dm
1
T
(wows + wo)
(38)
dt
If we interpret embbs: by its linear expansion wio won
emb,,, + o7 dw amb.r. dt
we obtain that
dWV
amb,a:
amb,s
y
1 1) (Wowl+ +Wo
1
wto
+n n7
dt
N
dt
dm
wr
emb.
+ 17 Ts. Ey
L
Y
dm
1
(WST +I) (WVOT++)) 11) ( ynth +11) 1 11) (WoWs woo?
dm
i=
21

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
Let W = ((W1)*+)
+1) +1) (wwv)) +1). w² = (wows- - wo)", T then the formulation could be rewritten as
dWv
Wio
emb,si
+ 1] Tes Ey
Y
1
w
y
T W
dt
dm
dm
y
Who
emb.e:
+ 17 Tss Ey Ey
Y
1
W!
1
W
dm
dm
YE.,.
11
Tel Ey
Y
W
y
--1
1
T
W
Y
wio
emb,s;
Y
-1
1
T
W
dm
dm
dm
Tal
Es;
Eye;
Y
1
y
1
W
+ Esgya
Wto
emb
Y
tm
dm
dm
Ex
Twll
Ey
Y
1
W'E
Y
1 I
W
+
y
E
Wio
emb,s,
Y
1 1 W
T
dm
dm
dm
While Ey
Y
1
1
(P =1)
TP
1 using the discussion section B.2, Let 2 ~ 11 (3), AI.... Ag ~ U1 (Ars))
dm
amm
have
1
1
1
P(Y- = i)
P Fmem (Y =1) i) + Pren (Y =i))
d
2
dm
Aj = i )) dm
1
81EZ
+ P
Z j=1
2
Nz
1
1
1)
Z +
Aj =1 1
+
2
dm
2
Nz
dm
j=1
}
1
9
+ 1
δiez
1
EA1
P
Z Σ
A3 i A1
2
dm
2
N₂
dm
j=1
1
1
1
d,ez
1
EA1
IP
7 +
A== = i| =
2
dm
2
Nz
dm
j=1
1
EAi
Ey
y²
2
dm
Then we have that
dWV
EAi
Ey
Y
1
E'E
Y
1 T
2
+
E.,y
who
cmb.,;
y
-11
1
W2
dl
2L
dm
dm
dm
1
EA,
120
emh,
A Ey
Y
-
1
2 + O (d
2
dm
22

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
C. Mechanisms under Varying Initialization Scales
C.1. Embedding Space of Emb-MLP
Figure exhibits the cosine similarity within the embedding space of Emb-MLP models with initialization rates 7 = 0.3 and 7 = 0.5. The results indicate that under large initialization scale, the embedding space of the model becomes less influenced by the label distributions and instead relies predominantly orthogonality differentiate all tokens. This mechanism neglects the intrinsic relationships among tokens, leading loss of generalization.
Memory
y= 0.3
Y= 0.5
epoch 50
epoch 900
epoch 50
epoch 900
1.0
1.0
1.0
-..2
-022
- 0.2
-0.4
Reasoning
1.0
1.0
1.0
1.0
0.2
-0.2
-.2
-0.4
Figure 10. Cosine similarity among different anchors' embedding vectors of Emb-MLP under initialization rates Y 0.3, 0.5 for memory anchors (top row) and reasoning anchors (bottom row).
C.2. Embedding Space of Transformer
Figure exhibits the structure of the Transformer's embedding space with y 0.3 and ς = 0.5. The left and middle panels exhibit the cosine similarity within the embedding space of memory anchors an reasoning anchors, demonstrating that larger initialization scale promotes orthogonality among embedding vectors. The right panel presents the PCA projection of the embedding space, suggesting that under large initialization scale, the embedding space lacks meaningful structure conducive to learning the reasoning mapping. These findings suggest that large initialization scale encourages differentiation of tokens primarily through orthogonality, taking tokens as independent from the others and neglecting intrinsic token relationships, and ultimately impairing the generalization capability.
C.3. The First Attention Module of Transformer
Figure 12 exhibits the structure of the first attention module with 0.3 and = 0.5 at epoch 200. The comparison reveals that a larger initialization scale results in a more complex attention mechanism, which exhibits no specific preference for any particular task.
C.4. Low-rank Phenomena of Transformer
Figure 13 illustrates the distribution of singular value across different parameter matrices under varying initialization scales. The results reveal that as the initialization scale decreases, the parameter matrices exhibit a pronounced low-rank structure, which in turn facilitates simpler learning mode.
23

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
y= 0.3:
Memory Anchor
epoch 200
epoch 900
Key
1.5
1.0
30
0.4
0.5
0.2
501
32
44
68
-1.0
-0.5
0.0
0.5
Reasoning Anchor
Reasoning Anchor
1.5
0.8
0.8
1.0
29
0.6
0.6
14
0.4
0.0
-0.5
0.2
18
17
0.0
-1.5
y= 0.5:
Memory Anchor
epoch 200
epoch 900
Key
0.6
0.8
0.4
38
0.6
0.4
0.0
0.2
-0.2
54
68
72
0.0
-0.6
-0.8
00.5
1.0
Reasoning Anchor
Reason Anchor
0.6
0.8
0.8
0.7
0.2
15
0.6
0.0
0.5
0.4
-0.2
0.4
17
14
0.2
00.6
0.0
0.5
Figure Characteristic embedding space Transformer with initialization rates 7 0.3, cosine similarity among embedding vectors of memory anchors and reasoning anchors at epochs 200 and 900. The right panel shows PCA projection the embedding space with the key reasoning anchors.
0.5. The left middle panels depict the
24

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
A
B
C
D
V
0.3: Attention
0.035
0.8
Singular Value
ccos(u,, wemb, A)
0.030
Arsn
Amem
0.025
0.4
0.5
0.020
5
0.0
0.0-2
4
10
20
30
10
j{A - 1/j)
Singular Value Index
Singular Value Index
y= 0.5 Attention
Singular Value
0.045
[cos(u, wemb A)
0.8
0.040
Aysn
Densiy
0.035
Amem
0.030
0.4
0.025
0.5
0.000
0.015
1
3
5 0.0 0.0
10
20
30
5
10
j(Ajk - 1/f)
Singular Value Index
Singular Value Index
Figure 12. Characteristies of the first attention module of Transformers (step 200) with initialization rates 0.3 (top row) and 7 = 0.5 (bottom row). A: Heatmap the attention matrix for a random sample. B: Distribution the relative error between attention Aςa and across all training sequences. C: Distribution of singular values of WV. D: Cosine similarity between the left singular vectors and average embedding vectors of the anchors.
W
WK(1)
WV11)
W011)
(f((1)
Wh(1)
0.6
0.6
0.20
0.13
0.25
y 0.3
y == 0.5
0.4
0.15
0.10
|1111
0.08
0.15
y = 0.8
0.10
0.2
0.2
0.05
0.10
0.05
0.05
0.0
0.0
0.00
0.00
0.00
W(12)
WK(2)
WV(22
2/02)
(f.(2)
Wf2(2)
0.6
0.15
0.13
0.3
<7
0.10
0.4
0.2
0.7
0.08
0.2
0.1
0.05
0.1
0.05
0.03
0.0
0.0
0.0
0.00
i
Figure 13. The distribution of singular value in different parameter matrices under different initialization scales (step 200). We denote the singular value vector by S and the i-th largest singular value by s,.
25

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
C.5. Embedding Space in Real Language Tasks.
Figure 14 exhibits the cosine similarity within the embedding space of PrOntoQA and TinyStories tasks trained by GPT-2 models with initialization rates 7 0.3 and 0.5. It's noted that under large initialization scale, the embedding vectors mutually orthogonal, indicating that the model neglects the associations among different tokens.
PrOntoQA
TinyStories
PrOntoQA
TinyStories
y=
V=
Figure 14. Characteristic of embedding space of PrOntoQA and TtoyStories with initialization rates 7 = 0.3, 0.5 (step 5000).
D. The Second Attention Module
The function of the second attention module is to extract the key preceding the anchors 2p. and transfer its information the last position. Figure 15A depicts the last row of the attention matrix before applying softmax, whose variation trend with respect to position index i can be divided into three parts: (1) for S the attention increases progressively as increases; (2) for + 1 < i i + g, the attention exhibits a slight decrease; and (3) for i 2 + the attention drops sharply.
A
B
C
Zp
Attention
COS(w wils.
-15
2
1.0
1.0
0.8
1
3
0.05
-25
0.4
0.00
7
-30
0.2
-0.05
0.0
0.1
p+10+2
Position Index
Position Index
Figure 15. Characteristic the second attention module. The last row of the second attention matrix (without applying softmax) for randomly selected sequence. B: A heatmap of the second attention matrix for the same sequence. C: The last row of the matrix QWPosT /v immediately hefore the final token. D: The cosine similarity of positional emheddings cos Poe, pos.)) for i,j 1,2,
Positional encoding plays a crucial role in this step. Figure 15C illustrates the last row the attention matrix after substituting K with Wooar suggesting a increasing trend with the position index. Note that we only present the index immediately before the final token, as the key does not appear at the last position and thus is not required to adhere to the same pattern. Furthermore, since the reasoning anchor and the tokens following it are augmented with reasoning anchor's information in the first attention module, this information be utilized reduce the attention for tokens after the token. We construct detailed mechanism about this in E.3.
E. Reconstruction Mechanism for Information Capturing
To verify our observation significant for information capturing for the Transformer model, we reconstruct the embedding space, the first attention module, and the second attention module and exhibit the process of extracting the key-anchor pair
26

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
from reasoning sequence.
E.1. Embedding Space
Assumption (Word Embedding). assume the embedding space has the the following properties:
1. co8 (we (wemb,emen (r0m) ) cos oweem un ewmr)
Vomem Ammmm 8tnn E Arau 8key E 2.
2. Fix any 81 E 2, cos (wemb,,s wemb,e2) 2 cos (wemm,,s wembb.s) if 18s 82 S 181 s31, V82,83 83 E 2.
3. There exists universal constant Cw < 00 such that |/20 //womb I Ow for any token
In addition to word embeddings, position embeddings should be effectively utilized, as they play a critical role in the functionality the second attention module. Here, we propose some assumptions about the relationship between word embeddings and position embeddings, with further characteristics to be elaborated upon later.
Assumption 3 (Position Embedding 1). Given any position embedding vector uP where denotes the position index, we assume that uppos, 1. emm for all 1, and E 2 Arsn U Amem-
Assumption 4 (Position Embedding 2). We assume that cos (unPss, po cos i [1, LJ.
|
L IT and ||o0PosA|| 1 for any
Given any sequence X, the output of the embedding layer is
x1) X Wemb Wyoos
E.2. First Attention Module
In the first attention module, due to the impact of small initialization, the attention matrix A functions as an average operator. Specifically, the result of the first attention module can be interpreted as
(Attn() ( (x()) x()rw)
=
Σ wVa)
Σ x (1)
-
(39)
Furthermore, performing singular value decomposition (SVD) on the value projection matrix VV)) reveals that its largest singular value is significantly greater than the remaining singular values. The left singular vector corresponding to the largest singular value (V)), highly similar the embedding vectors of the reasoning anchors which indicates that WV)) can be approximated by
T
V Av
|1E Wemb.s||
Σ
Wemb,
v,
(40)
SEArro
where \y is the singular value and v E RR Ixd, denotes the right singular vector. Since w amb 1 w emb for any Syan E Arsa. Smem E Amem and w emb 1 wpos Then the result of the attention operator can be interpreted as
Āv
(Aun (x()) x())))
w i,:
w
omh,T Aron
v,
(41)
5:
where wimb EoEA... s Av I.wmm Apenll
Substituting the reasoning sequence r and memory sequence Xmem, respectively, into this formulation, we derive the following results:
() (1) v())
0,
(42)
j S P.
(Attn (1).)))
minlj.ptg)
(43)
Σ
Womb, p yom
w Are v,
cmb,T
p< p<jsi. T,
i=p+1
27

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
where j means the index. Thus, all tokens following the reasoning anchor effectively "tagged," facilitating the identification of the anchor. Define the output of the first attention module is as follows:
x ( Attn(2) (x()) X (Va)
Under the Assumption can formulate the output of the reasoning sequence further
w 2m + w
j S p. Āv
w embbas + wopos
+
w embba,
w womb.Ty v, Asna
p+1zi
(2) Ji
i=p+1
Āv
p+g
un
pon.j +
embba
cmb.T w Aran
v, ptq+lsisL.
j
i=p+1
E.3. Second Attention Module
We observe that the first attention module introduces additional information all tokens with indices j p 1. The subsequent challenge is identify the reasoning tokens and the key token, and effectively propagate their information the last position in the sequence. To achieve this, construct the following attention distribution and demonstrate its properties:
Definition 3 (Cliff Sequence). Given sequence R we define I as (p. sequence if there exists L N+ such that satisfies the following conditions:
1. (Increasing Segment) 441 > for all i < p.
2. (Plateau) 4-- I Lpt,,," ... lptg S 1p.
3. (Descending Segment) L; < l for all p + 9 L.
It is evident that if the attention of the last token forms (p, -)-cfiff sequence, it can effectively capture the information of the tokens and the key. Specifically, we have the following results to illustrate its feasibility.
Lemma 5. For any 0. there exists ccliff sequence I with norm C such that softmax((), S for any [1,p- - 1p - I...
Proof. It's evident that we just need to illustrate softmax(1)p-1 + as C' + co. Denote that = Ci, then we have
Clp
softmax (0pp11
E/= c0ly
Ej 1|upp+++111) C((i-1p-1 Ejspprto)) e ((1.-1o-1)
Since that i, s i--1 for any j E [1.
1 lp - y + 1, and ig 2 i for any j E [p.P + 4), so we have
lim C-400
Σ
c(i =0 and lim C-+00
Σ
c(1i-bp11)
jE(L.p [1pp
jE(p.Ptg]
then softmax (1)--1 + 0.
(9n,(2) À (2), is a (p, cllfff Here we provide mechanism to construct real matrix À E gdd Xdm such that X L,: sequence. Assume that À л (span(upooy uoo v. 1 O, where π (span(wPoe)) denotes the subspace spanned by
28

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
{woPos.. Then have
j S q.
(wDea,, _wooo.d))
w am
w Aun
(v, wmbb)
L
i=p+1
x22 bi
(2),
n+g
w embba
W amb
w emb,a;
w omb
ptisispta
j L
Aren
Are
1=p+1
2
22 jL
p+o
(wpos,, wo00s.)
w embbo
Wemb
imp+1
E
Emp+ aomba womb) ) applying the Assumption on the position embedding, then we have the following
Define 9j
result:
j
008
1
n
JSP.
L
jl jl/womb,ar) ||
X (2) L,:
A (2).
008
1
T
JL. Ppta
Av
cos (v, w amb.as)) 33)
p+ Aissjippti
1
CO8
1-
T.
π
jL. Pptpr
p+a+isisL.
To satisfy the Increasing Segment condition, we need that:
cos (1 (1-
-/] I π
(ille
cos (v, ambbaa)) + 93 2 c0s (1
(1-F))
P
1
(1 (1-
p p-
]I Pp+t
Av
I
क + cos
for any p E [1, L aj E 1p 1,p + 4). Denote that:
M = max Pptas
m = min Ppto-
(44)
Then we have
X2p M
cos (v, webbaa) + 111 S - -oos
p+
p-1
COS
+
Cos
jL
Av||o//
L
2
L
P2. M
(1llzambbog)
cos (v. emb.a)) + .o.
jL
Av|l00/
< V (/2
(1 cos 1
3
- 1
3
(2)
+
sin
cos
L π
aretan
2
L
1 cos
11
Denote the right side by CM, and simplify it with
embbaa
||
LCM M X2 L
cos (v. emb,as)
<
(45)
For another side, we assume that:
cos 1 (1-2) 2 5
1- [ jL
22.1
Ppto
[1llammbsss) ||
cos (v, w emb.as) ) +
P3 +P3
< cos
(1- P) n
29

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
which implies that:
X
j j1le0amb..a,
2p + 4 n
jL
Ppta
dy
cos (u, w)bb,ay) + [p] - sin
2L
sin
2L
We have that:
[wembbay]
LCm Xp
Cos (v, up
(46)
Av
These two conditions give the direction scope of v. For the Descending Segment condition, we have that
1
1
X 2 Ppta
>
COS
1
IL
+Ny > jL
π
jı
-2 Ppt4
(47)
cos
oos
T.
T.
>ZZ > 2
1
cos
m
L
With (45), and (47), we could give a range of \v, and the direction of v which makes X rsn.(2) (p, c)-cfiff sequence.
L Áx (2), is
F. Layer Normalization
We conduct experiment with removing the Layer Normalization module, exhibiting the same phenomena, i.e., smaller initialization scales bias reasoning task, and results are depicted in Figure 16.
Y 0.3
V 0.5
V 0.8
10
10a
10a
10-1
Loss
10
10-1
10
10-3
10
0
2000
2000
2000
Epoch
Epoch
Epoch
-Memory
Reasoning train
-Reasoning - test
Figure Training dynamics of Transformers under 7 0.3, 0.5, without Layer Normalization.
G. Learning Rate
We conduct experiments with the learning rate belonging to [10-5,8 5 x 10-4], Figure 17 exhibits the loss dynamics under different 7, remaining consistent learning bias across these learning rate configurations. However, when the learning rate increases 0.001, the training becomes highly unstable, manifesting a severe loss spike.
30

---

An Analysis for Reasoning Bias of Language Models with Small Initialization
y= 0.3 1r=5e-05
y= 0.3 Ir= 0000000
Y 0.3 1r=0.0005
10a
10a
mw
109
109
Loss
1
10-1
10-2
10-3
10-4
10-3
500
1000
500
1000
500
1000
V 0.5 Ir=5e-05
y 0.5 1r=0.0001
Y 0.5 1r=
10.
109
109
10-2
10
10-1
Loss
10-3
10
10-2
10
10
-1
10-7
10-1
500
1000
500
1000
500
1000
V 0.8 Ir=5e-05
V 0.8 1r=0.0001
0.8 1r=0.0005
109
109
100
10-1
10
Loss
10-2
10
3
10
10
10
10
500 Epoch
1000
500
1000
500
1000
Epoch Reasoning train
Epoch
Memory
RRasooing - test
Figure 17. Training dynamics of Transformers under = 0.3, 0.5, 0.8 and varying learning rates.
31


Reviews:
[{"review_id": "jHd0hq2SoV", "reviewer": "Reviewer_a4CN", "summary": "This paper discusses the impact of initialization of language models on their trained performance on memorization and reasoning tasks. The paper uses proof to show reasoning tasks prefer smaller initialization while memorization tasks prefer large initialization. The authors attribute such behavior to being more differentiated in the embedding space at an early stage, which is further verified by empirical experiments.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your careful reading and evaluation of our work, and we are truly grateful for your recognition of our research. We address your concerns as follows.\n\n**Comment**: A potential weakness of this paper is not providing direct guidance for current LM fine-tuning as most tasks are now fine-tuned based on pre-trained models.\n\n**Response**: We appreciate this question you raised which is very interesting and valuable. We provide some preliminary intuitions: It is definitely important to investigate the impact of initialization on fine-tuning, and some related studies have already emerged in this field. For instance, [1] examined how parameter initialization in LoRA affects fine-tuning training. The study employed a small initialization scale $\\gamma=1$ for matrix A, achieving more efficient feature learning. [2] theoretically explored the dynamic behavior of matrix factorization models under small initialization scales, demonstrating that small initialization reduces model complexity, then enhancing generalization capability. Due to the page limit in our current article, we leave this topic for future work. Actually, we have done some experiments on LoRA and believe our work can naturally be extended to the analysis of fine-tuning. \n\n**Comment**: While data in natural language distribution (human annotation) will make the experiment more convincing, it's acceptable when such resource is absent.\n\n**Response**: Thanks for your suggestion and we appreciate your understanding. In our experiments, we trained a standard GPT-2 model on two real-world datasets PrOntoQA and TinyStories though very limited resources. We observed that models with small initialization exhibited a learning bias toward reasoning tasks (see Figure 1). A parallel analysis of the embedding spaces for both datasets corroborated this conclusion (Sec.4.4).\n\nThank you for your detailed feedback and for taking the time to review our work again. We are grateful if our responses have addressed your concerns.\n\n[1] The Impact of Initialization on LoRA Finetuning Dynamics. NeurIPS 2024.\n\n[2] Connectivity Shapes Implicit Regularization in Matrix Factorization Models for Matrix Completion. NeurIPS 2024.", "claims_and_evidence": "The paper claims that reasoning tasks prefer smaller initialization while memorization tasks prefer large initialization, which is verified by theory and empirical experiments.", "methods_and_evaluation_criteria": "This paper uses a synthetic dataset to verify their assumption. While data in natural language distribution (human annotation) will make the experiment more convincing, it's acceptable when such resource is absent.", "theoretical_claims": "I have checked the theoretical claims and think the proof to be convincing. But there is a chance that I miss minor mistakes.", "experimental_designs_or_analyses": "The experimental design makes sense and supports the theoretical claims.", "relation_to_broader_scientific_literature": "This paper is related to the impact of model initialization on training performance.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "A potential weakness of this paper is not providing direct guidance for current LM fine-tuning as most tasks are now fine-tuned based on pre-trained models. The paper will be significantly more impactful if the discovery can be applied to analyzing the behavior of pre-trained LMs.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QDGCFFvCBP", "reviewer": "Reviewer_FCaa", "summary": "The paper investigates initialization scale as one of the driving factors of bias towards different types of tasks. The paper considers two types of tasks:\n1. Reasoning - represented by a sum over a key and its anchors given a key. This task is constructed to capture logical/arithmetic relationships, which requires generalization (And can be measured with the test set the authors provide)\n2. Memory - represented by retrieving a predetermined random sample corresponding to a key-anchor pair. The memory mapping is entirely random, which means. ithcannot be inferred, only memorized.\nThere are also noise tokens provided to increase the difficulty of the learning problem/control the signal to noise.\n\nThe impact of initialization scale $\\gamma$  (weights are initialized $W\\in\\mathbb R^{d_1\\times d_2}, W_{i,j}\\sim\\mathcal N(\\mu=0,\\sigma=d_2^{-\\gamma})$ )  upon the relative performance of 1. and 2. is investigated.\n\nIt is shown that small initializations ($\\gamma > 0.5$) biases networks towards reasoning, i.e. the performance of tasks of type 1. accuracy increases more quickly as a function of training epochs than tasks of type 2. that compared to when a large initializations ($\\gamma < 0.5$) is used.\n\nThis phenomenon is demonstrated empirically (transformers + a variation on set2vec, called Emb-MLP), as well as motivated theoretically using differential questions. The theoretical analysis highlights the distribution of task labels as a driving factor between the rates each task is solved.\n\n## Update after rebuttal\nScore increased to 4: Accept following resolutions to my questions during discussion phase.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thorough review of our work and the valuable feedback you have provided. We address your concerns as follows: (**Due to the character limit, we are not able to display specific textual revisions. However, we will carefully revise our manuscript to address each concern.**)\n\n**P1: Different phenomena with different $\\gamma$.**\n\nGiven that our primary focus is to investigate the reasoning bias of models under small initialization scales, and due to the page limit, we present results for $\\gamma=0.8$ in the main text. In Appendix C, we exhibit analogous analysis of the following components under $\\gamma=0.3,0.5$:\n\nC.1: The embedding space of Emb-MLP,\n\nC.2: The embedding space of Transformer,\n\nC.3: The first attention module of Transformer.\n\nThese analyses establish that varying initialization scales significantly influences the model’s behavior. We sincerely regret this omission of explicit references to Appendix C in the main text which will be provided in Sec.4.\n\n**P2: The real world language tasks.**\n\nDue to the same reason in P1, we did not include results for large initialization. However, we conducted experiments with $\\gamma = 0.3, 0.5, 0.8$. We define the following metric: $\\Delta L := \\frac{L_{Tinystory}-L_{PrOntoQA}}{L_{PrOntoQA}}$. As $\\gamma$ increases, $\\Delta L$ exhibits an upward trend, indicating a growing bias for reasoning task (see https://postimg.cc/TpcZgmtw A). Analysis of embedding space during the early training stage (step 5000) aligns with that presented in Appendix. C (see B of the link above). These results will be added to the paper.\n\n**P3: The residual structure.**\n\nThe residual structure operates through position-wise additive operations between two sequences, as:$residual(X,V)_i=X_i+V_i$，which lacks interactions between tokens at distinct positions. \n\n**P4: Attention module (figure 5) and Lemma.1.**\n\nWe sincerely apologize for the omission of indicating epoch number in figure 5. Figure 5 displays the attention structure in the early training stage (epoch 200). By the end of training, the attention module exhibits specific patterns for capturing critical information within sequences. Lemma 1 only guarantees the average operator phenomenon during an early stage. Similar theoretical conclusions have been established in prior work, such as [1].\n\n**P5: The learning rate.**\n\nWe conducted experiments with lr $\\in[1e-5,5e-4]$. The learning bias under different $\\gamma$ remains consistent across these configurations (see https://postimg.cc/7CtDSMV6 ). However, when lr increases to 1e-3, training becomes highly unstable, manifesting severe loss spike. \n\n**P6: Why assume $\\gamma \\to \\infty$.**\n\nThe assumption is primarily a technology adopted to enable asymptotic analysis in our theoretical framework. In finite-scale scenarios, we focus on the empirical trends as initialization scale decreases. Actually, for $\\gamma\\sim 1$, we can already see very clear reasoning bias in the training. \n\n**P7: LN does not impact conclusion.**\n\nWe conducted an experiment with removing the LN module, exhibiting the same phenomena, i.e., smaller initialization scales bias reasoning task (see https://postimg.cc/dLbk0xyn ). Theoretically, we could provide an informal explanation. Since our analysis focuses on the initial stage of training, the mean and std could be approximated by the initial value. Consequently, the gradient flow is just multiplied with a constant 1/std, preserving the main structure of the learning dynamics. We will add these analyses in Appendix.\n\n**P8: Question 1.**\n\nIn Sec.4.3, since noise and key tokens are sampled from the same distribution, the sequence lacks permutation invariance. To identify the key, the model must utilize positional encoding and components capable of cross-positional information exchange, like attention module. However, DeepSets is designed for permutation-invariant set input. It would fail to distinguish the key tokens. This fundamental limitation renders DeepSets unsuitable for this task.\n\n**P9: Question 2.**\n\nPrior work in [2,3] investigated the impact of $\\gamma$ on model dynamics and identified distinct behavioral regimes, which is stable under different model sizes. Inspired by those works, we adopt to adjust $\\gamma$.\n\n**P10: The suggestion on presentation.**\n\nWe sincerely appreciate your suggestions and we will carefully consider and adopt these suggestions.\n\n**P11: Citation.**\n\nWe appreciate your notification and will add this citation in Related Works.\n\nWe once again extend our sincere gratitude for the valuable insights you have provided. We hope that our responses have addressed the concerns you raised.\n\n\n[1] Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis. NeurIPS 2024.\n\n[2] Phase diagram for two-layer ReLU neural networks at infinite-width limit. JMLR 2021.\n\n[3] Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing. NeurIPS 2024.", "claims_and_evidence": "Update: The author response has resolved claims that were partially supported.\n\nThe following claims are made, and are supported, or partially supported:\n## Smaller initialization cases a \"reasoning bias\" in synthetic/controlled tasks (supported)\n\nIn the synthetic experiments, a smaller initialization ($\\gamma=0.8$) shows reasoning tasks learning at a faster rate compared to the memorization experiments, whereas a larger initialization ($\\gamma=0.3$) has the memorization tasks learning more quickly. This is shown both for transformers (Figure 2A) and Emb-MLP (2B), which is appears like a simple DeepSet variant [1].\n\n[1] Deep Sets https://arxiv.org/abs/1703.06114\n\n## The reasoning bias due to smaller initialization is a result of embedding space separation difference between the reasoning and memorization tasks in the chosen synthetic/controlled tasks (partially supported)\n\nThe authors show that reasoning-related tokens (anchors) become more distinct earlier in training than the memory tokens under smaller initialization ($\\gamma=0.8$) (see Figure 3A). It takes a longer time at $\\gamma=0.8$ for the memory tokens to become distinguishable. Since for the memory task, the memory tokens need to be distinguishable to be correctly retrieved, the embedding space structure learning directly corresponds to the ability to solve the task.\n\nI note partially supported here, because we do not know what the rate of the embedding space learning looks like for other initialization scales. If e.g. Figure 3A looks similar for $\\gamma=0.3$ I think that would reject this specific claim. This information is not provided in main text or in the appendix.\n\n## The token label distribution affects the evolution of embeddings during training, and leads to the observed different rates of embedding structure discussed above (supported)\n\nThe authors present in Proposition 1 (+ Equation 5) the flow for the embedding vectors, and demonstrate the dominance by the label distribution. They derivate expressions for the label distributions in the memory task (Equation 7) and reasoning task (Equation 8). The empirical distribution for the reasoning task (Figure 3B, 3C) can be compared to the theoretical expression (Equation 8). They are consistent. The proofs is given in Appendix B.1 and B.2.\n\nThe proof of Lemma 2, and the distributions of the memory anchor seem correct. I was not able to verify the precise expanded form of the reasoning anchor distribution (Equations 22 - 25).\n\n## The attention mechanisms allows a reasoning bias at small initialization scales (partially supported, unclear)\n\nThe paper demonstrates this through composition of two observations. \n\n1. The paper demonstrates that the first attention mechanism behaves like an averaging operator. This is shown empirically in Figure 5, and with high probability in Lemma 1 (lines 380-384).\n2. The $W_V$ projector has largest singular values aligned with the reasoning anchors, but nearly orthogonal to memory anchors. Consequently, the attention mechanism propagates reasoning anchor information, but not memory anchor information to subsequent tokens.\n\nPoint 1. is clear to me. The second I can see the evidence for in Figure 5C, and understand the projective behavior following the attention operation.\n\nMy challenge here is two-fold:\n1. The transformer has a residual structure, which enables memory tokens to propagate to subsequent tokens (discussed at line 346, right hand column)\n2. The \"reasoning bias\" discussed in this work relates to the rate at which a network learns certain information, and how this relates to the value of $\\gamma$ used. In Figure 5 we are shown a single value of $\\gamma$, and (I think) only the result at the end of training. We are not sensitive here to the rate at which phenomena occur, which is critical in the first two claims.\n\nI can see that the large $\\gamma$ value plays a role in the formal analysis, e.g. $\\gamma\\rightarrow\\infty$ is an assumption in Propositions 2,3. It is unclear to me what happens away from this limit.\n\n## The above observations translate to real world language tasks (partially supported)\n\nThe paper shows the embedding structure resulting from training a language model on a reasoning task (PrOntoQA) and memory task (TinyStories). The embedding structure presented is consistent with the synthetic analysis.\n\nOnly partially supported as the primary claim of this paper relates to the effect of initialization $\\gamma$ on the reasoning bias of the model (measured as the rate of learning, which can be induced by the methods discussed in the above claims.)\n\nThe authors only show the results for $\\gamma=0.8$ training, and do not show any analysis about how behavior changes as a function of training. To substantiate their claims, the authors need to present results from $\\gamma=0.3$ and $\\gamma=0.5$, and present any differences in the rates at which things happen for PrOntoQA and TinyStories that follow from the initialization scale change.", "methods_and_evaluation_criteria": "The synthetic tasks constructed, as well as the chosen real-world tasks are sensible.", "theoretical_claims": "I validated:\n- Proposition 1\n- Equation 7\n- Lemma 1\n\nI was not able to validate the remaining theoretical claims due to their size and technical involvement within the reviewing period, however, I have no reason to expect they are incorrect.\n\nI also note that while Lemma 1 is true at initialization, I am not sure if Lemma 1 is sufficient to guarantee that the attention mechanism remains an averaging operator throughout training (since in general, gradient updates disrupt initialization conditions). Figure 5A however empirically indicates that the attention mechanism is an averaging operator post-initialization.", "experimental_designs_or_analyses": "First, the analyses is in in general well done. There are certain aspects missing however, in general to draw a conclusion, in each case we need to see:\n1. The $\\gamma$ value varied\n2. The rates of different phenomena change depending on the value of $\\gamma$.\nThis type of analyses is only done for a subset of the experiments, and needs to be shown in all cases.\n\nSecond, the initialization scale is known to impact learning generally, see e.g. Maximal Update Parameterization ([2]). When we change the initialization scale, we expect to also need to change our learning rates in order to get a reasonable network behavior out. The experiments in the presented paper use the AdamW optimizer with a learning rate of 1.e-5 for every experiment, for every initialization. This is a suboptimal choice, and raises a small question regarding whether the phenomena observed in the paper would change if the learning rate was the optimal learning rate for each initialization.\n\n[2] Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer https://arxiv.org/abs/2203.03466", "relation_to_broader_scientific_literature": "This work is relevant in literature that investigates how networks learn, through either memorizing, or generalizing, e.g. [3], where networks can easily fit random labels when they are unstructured. The current paper shows the effect of initialization on the rate at which a model would memorize or generalize.\n\n[3] Understanding deep learning requires rethinking generalization https://arxiv.org/abs/1611.03530", "essential_references_not_discussed": "I am not aware of any essential missing references, although [3] may be useful to include for wider research framing.\n\n[3] Understanding deep learning requires rethinking generalization https://arxiv.org/abs/1611.03530", "other_strengths_and_weaknesses": "The number of analyses provided from the different angles (empirical with many different measures, and theoretical) provides a rich, well-rounded view for the investigation.\n\nDespite the many diagrams, I found the paper slightly challenging to read, potentially due to the sizeable introduction and related works sections. We only arrive at the contributions the others make in Section 3.1, which includes a highly technical presentation of the synthetic task (Appendix A.1 is much easier to follow), which took me significant work to understand the purpose of each aspect of the study.", "other_comments_or_suggestions": "## Some suggestions for improving clarity\n\nTry and bring Appendix A.1 into main text. If necessary, move some of the related work into the Appendix. Keep only related work elements critical to understanding the paper in main text, ideally directly compared to your own contributions.\n\nProvide a roadmap for the reader.\n\nIn figures, state explicitly which $\\gamma$ values are being used, and for which epochs. Do this for every figure.\n\nCall out significant findings using \\paragraph{...} notation, then back up the finding/observation with following prose.\n\nLink to proofs from main text. E.g. in Proposition 1 (line 232) provide a ref to proof in the appendix. Do this for all results in main paper.\n\n## Some suggestions for increasing confidence in claims\n\nPresent results for a range of $\\gamma$ as a function of training for all experiments (see above discussion on claims).\n\nSubstantiate the LayerNormalization result that does not impact conclusion (line 375).", "questions_for_authors": "1. How critical is the transformer architecture in Section 4.3? Would a DeepSet [1] solve this task?\n2. More generally one can initialize $W\\in\\mathbb R^{d_1\\times d_2}, W_{i,j}\\sim\\mathcal N(\\mu=0,\\sigma=c \\times d_2^{-\\gamma})$ for some width-independent constant $c$ (see [2,4]). Why was the choice made to modify $\\gamma$ as the controller of the initializer scale, rather than $c$ (optimal $\\gamma$ choice is generally optimizer dependent, whereas $c$ is explicitly optimizer independent, so it would feel like varying $c$ is a more universal choice)? Is there an important difference between varying $\\gamma$ or $c$? Is the choice stable across different model sizes?  \n\n[1] Deep Sets https://arxiv.org/abs/1703.06114\n\n[2] Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer https://arxiv.org/abs/2203.03466\n\n[4] A Spectral Condition for Feature Learning https://arxiv.org/abs/2310.17813", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xw9o0DHf1H", "reviewer": "Reviewer_e1DG", "summary": "This paper aims to investigate how different initialization scales (small vs. large) may affect transformer models' ability to learn in different tasks (specifically in reasoning and memorization). The study pretrains a GPT-2 model with separately from smaller to larger initialization scales on reasoning and memorization tasks, where the datasets are both synthetic and real. The study finds an explicit bias that a smaller initialization scale of the transformer model can learn reasoning tasks much faster and generalizes better, while a larger initialization scale would make learning to memorize faster, but also reasoning worse. The authors used a variety of approaches to investigate the mechanisms of such bias. Specifically:\n\n1. Theoretical analysis. The authors apply an Emb-MLP Model to analyze gradient dynamics as well as empirical representation evolutions. Results show that smaller initialization can amplify the gradient differences based on structured label distributions of reasoning tasks. Moreover, reasoning anchors quickly differentiate themselves while memory mappings do not. They also confirm this theoretical proposition by finding similar results from empirical results.\n\n2. Transformer model analysis. The authors analyzed transformer embeddings and found that the structured clusters quickly form in the training of reasoning tasks. The authors also analyzed the attention matrices and proposed that the smaller initialization makes attention approximate the averaging operations.\n\nOverall, this paper discovered that smaller initialization scales of transformer models can be good for reasoning and investigated the mechanisms underlying this discovery.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thorough review of our work and your valuable feedback and questions. Your recognition and support for our research have been immensely encouraging. We address your concerns as follows.\n\n**Question**: I was wondering if this small initialization effect can generalize in reinforcement learning algorithms such as GRPO or PPO.\n\n**Response**: It is important to study how does initialization scale impact the training behavior in reinforcement learning (RL). We provide some preliminary intuitions:\n\n1.\tThe core of small initialization lies in reducing model complexity, compelling the pre-trained model to fit datasets through simpler, more generalizable patterns. With small initialization, a model can learn more reasoning patterns instead of memory patterns. RL can play a critical role in enhancing the reasoning patterns which are learned in the pre-training stage. On the other hand, employing small initialization for training reward model (RM) might improve the generalizability of RM scoring since small initialization could reduce overfitting to training data noise.\n\n2.\tBuilding on the complexity reduction principle, other complexity-reduction techniques such as larger weight decay could be integrated into the RL stage to further enhance generalization.\n\n\n3.\tIn some post-training workflows, parameters of the pre-trained model are directly updated, rendering the concept of initialization irrelevant. However, in specialized approaches like LoRA (Low-Rank Adaptation), new update parameters are introduced and trained independently. As discussed in our response to Reviewer a4C, the initialization of these update parameters significantly influences fine-tuning results. Analogously, in RL settings, an update policy could be designed. Let $\\pi=\\pi_{pre-train}+\\Delta \\pi$ where $\\Delta \\pi $ is the trainable update policy.  Then the varying initialization schemes for $\\Delta \\pi$ might lead to distinct training behaviors.\n\nWe hope our response could provide clarity and value to your question. Once again, we deeply appreciate your support and constructive feedback for our work.", "claims_and_evidence": "The authors claim that smaller initialization scale of transformer models can be good for learning reasoning tasks, while larger initialization scale can be better for memorization tasks. In the experiments, results did show that with smaller initialization, the transformer models learn faster and generalize better in the training of reasoning tasks, while larger scale initilization makes reasoning generalize worse but memorization faster.", "methods_and_evaluation_criteria": "The study uses two styles of task and datasets to evaluate the training of reasoning and memorization. One is a synthetic compositional task and random numbered tasks. They are all sequences of numerical tokens and designed in the same order. However, the compositional task's output is an arithmetic mapping (e.g., addition) of keys while the memorization task's output only requires reproduction of the keys. This design has minimal differences of token distribution while requiring differences in the computation. Another type of datasets are benchmark datasets used for reasoning and mapping tasks. The authors analyze embeddings of models trained on these two datasets and found the former one exhibits significant hierarchical structure of logic representation in the reasoning dataset, while the latter one does not.\n\nOverall, the authors used both empirical and synthetic datasets to evaluate the model's learning in reasoning and memorization, which are robust and thorough.", "theoretical_claims": "Based on my limited expertise and understanding of transformer computations, I did not find any explicit incorrect proof.", "experimental_designs_or_analyses": "The analysis goes thoroughly from synthetic dataset to empirical dataset, from representation findings to theoretical analysis, which are robust and comprehensive.", "relation_to_broader_scientific_literature": "The paper proposes an important factor - the size of initialization of pre-training of a transformer-based language model in the reasoning and memorization task. This paper provides critical empirical and theoretical evidence that smaller initialization can work faster and better for the learning of reasoning models. This can be a both practical and fundamental concern of building better reasoning models, as well as understanding how reasoning models can work/fail to work.", "essential_references_not_discussed": "I am not aware of any essential reference omitted by the authors.", "other_strengths_and_weaknesses": "The paper's visualization as well as structure is neat and easy to understand.", "other_comments_or_suggestions": "No other comments.", "questions_for_authors": "Though may be out of the paper's scope, I still wonder in a typical pre-training setting, we used simply next token prediction to calculate the loss and optimize. I was wondering if this small initialization effect can generalize in reinforcement learning algorithms such as GRPO or PPO. The authors are not required to supplement any experiments. Sharing intuitions and insights into this topic would be interesting.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This work studies how the initialization scale affects the reasoning and memorization abilities of a model. For that, they use synthetic datasets, and show that small initializations promote reasoning behaviors. They also provide analysis on the separability of reasoning tokens compared to memorization tokens, showing that the former are more separable early on than the latter (under small init). The findings are complemented with theoretical claims that the reviewers (notably FCaa) found to be valid in general.


### Short summary of reviews

Reviewers a4CN and e1DG did not report many concerns. Reviewer a4CN commented on the use of synthetic data and reviewer e1DG on the generalization of the findings to RL methods. Both reviewers acknowledged a satisfactory rebuttal.

Reviewer FCaa provided an in-depth and very thorough review. They addressed the mathematical proofs and provided several suggestions. The authors gave a comprehensive rebuttal, leading to an score increase from 3 to 4.

### Recommendation

Given the positive scores (4, 4, 4) and the engagement in the rebuttal, I **recommend this paper for acceptance**.
Decision: Accept (spotlight poster)