An Analysis for Reasoning Bias of Language Models with Small Initialization

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.

Question: I was wondering if this small initialization effect can generalize in reinforcement learning algorithms such as GRPO or PPO.

Response: It is important to study how does initialization scale impact the training behavior in reinforcement learning (RL). We provide some preliminary intuitions:

1. The 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.

2. Building 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.

3. In 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.

We hope our response could provide clarity and value to your question. Once again, we deeply appreciate your support and constructive feedback for our work.

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

P1: Different phenomena with different $\gamma$.

Given 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$:

C.1: The embedding space of Emb-MLP,

C.2: The embedding space of Transformer,

C.3: The first attention module of Transformer.

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

P2: The real world language tasks.

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

P3: The residual structure.

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

P4: Attention module (figure 5) and Lemma.1.

We 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].

P5: The learning rate.

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

P6: Why assume $\gamma \to \infty$.

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

P7: LN does not impact conclusion.

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

P8: Question 1.

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

P9: Question 2.

Prior 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$.

P10: The suggestion on presentation.

We sincerely appreciate your suggestions and we will carefully consider and adopt these suggestions.

P11: Citation.

We appreciate your notification and will add this citation in Related Works.

We once again extend our sincere gratitude for the valuable insights you have provided. We hope that our responses have addressed the concerns you raised.

[1] Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis. NeurIPS 2024.

[2] Phase diagram for two-layer ReLU neural networks at infinite-width limit. JMLR 2021.

[3] Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing. NeurIPS 2024.

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.

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.

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.

Comment: While data in natural language distribution (human annotation) will make the experiment more convincing, it's acceptable when such resource is absent.

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

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

[1] The Impact of Initialization on LoRA Finetuning Dynamics. NeurIPS 2024.

[2] Connectivity Shapes Implicit Regularization in Matrix Factorization Models for Matrix Completion. NeurIPS 2024.