An Analysis for Reasoning Bias of Language Models with Small Initialization

TL;DR

This work investigates how the scale of parameter initialization shapes learning biases in Transformer-based language models, revealing that small initializations bias models toward reasoning tasks while larger initializations bias toward memorization. By combining a synthetic anchor-function framework with Emb-MLP and Transformer analyses, the paper links token-label distributions to embedding differentiation and to the dynamics of self-attention modules, providing a gradient-flow based theoretical account. Real-language experiments with GPT-2/GPT-2–like models on PrOntoQA and TinyStories corroborate the theory, showing increased reasoning emphasis with smaller initializations and more distinguishable reasoning embeddings. The findings offer practical guidelines for initialization strategies and deepen the understanding of how training dynamics interact with architecture to shape task preferences and generalization in LLMs.

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.

Figures (17)

• Figure 1: 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 PrOntoQA dataset, while the blue line depicts the training loss on the TinyStories dataset.
• Figure 2: Schematic diagram of the synthetic composition task. The gray-shaded area illustrates the specific setup used in this example. Each block represents a token within the input sequence, with different face colors indicating distinct token types (blue: noise, orange: key, green: anchor). Each row corresponds to an 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.
• Figure 3: A: Loss and prediction accuracy of the models on different datasets under varying initialization scales ($\gamma = 0.3, 0.5, 0.8$). The top row depicts the evolution of the loss during training for three datasets: $\mathcal{D}_{\rm mem}$ (blue lines), $\mathcal{D}_{\rm rsn,train}$ (purple lines), and $\mathcal{D}_{\rm rsn,test}$ (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 $\gamma=0.3$ and $\gamma=0.8$.
• Figure 4: A: Cosine similarity matrices for memory (top row) and reasoning (bottom row) anchors at epoch 50 (left) and epoch 900 (right) of a model initialized with $\gamma=0.8$. B: Distribution of $\bm{P}^{s}-\frac{1}{d_{\rm vob}}\bm{1}$ for different reasoning anchor $s$. C: Cosine similarity between $\bm{P}^{s_i}-\frac{1}{d_{\rm vob}}\bm{1}$ and $\bm{P}^{s_j}-\frac{1}{d_{\rm vob}}\bm{1}$ for any reasoning anchor $s_i,s_j$, exhibiting a similar structure to the embedding space of reasoning anchors observed in A.
• Figure 5: Embedding structure of a 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 of the embedding space projected onto the first two principal components computed via PCA. C: Cosine similarity between the constructed embedding vectors of reasoning anchors $\tilde{\bm{w}}^{\rm emb,s}$ (see \ref{['eq:estimation_wemb']}) as derived in Theorem \ref{['thm:emb_reconstruct']} (top) and Cosine similarity comparison between experimental results with theoretical approximations where $s_i=15$ (bottom). D: PCA projection of the constructed embedding space $\tilde{\bm{w}}^{\rm emb,s}$ for $s\in\mathcal{Z}$ (top) and $s\in\mathcal{A}_{\rm rsn}$ (bottom) onto the first two principal components.

Theorems & Definitions (22)

• Definition 1
• Proposition 1
• Definition 2: One-layer Transformer
• Lemma 1
• Proposition 2
• Proposition 3
• Theorem 1
• Lemma 2
• proof
• Corollary 1