Scalable Complexity Control Facilitates Reasoning Ability of LLMs
Extracted source text · arXiv 2505.23013
% \documentclass[hyperref={colorlinks = true,linkcolor = blue}]{article}
\documentclass{article}
% if you need to pass options to natbib, use, e.g.:
\PassOptionsToPackage{numbers, compress}{natbib}
% before loading neurips_2025
% ready for submission
% \usepackage{neurips_2025}
% to compile a preprint version, e.g., for submission to arXiv, add add the
% [preprint] option:
\usepackage[preprint]{neurips_2025}
% to compile a camera-ready version, add the [final] option, e.g.:
% \usepackage[final]{neurips_2025}
% to avoid loading the natbib package, add option nonatbib:
% \usepackage[nonatbib]{neurips_2025}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc} % use 8-bit T1 fonts
\usepackage{hyperref} % hyperlinks
\usepackage{url} % simple URL typesetting
\usepackage{booktabs} % professional-quality tables
\usepackage{nicefrac} % compact symbols for 1/2, etc.
\usepackage{microtype} % microtypography
\usepackage{amsmath}
\usepackage{graphicx}
\usepackage{subfig}
\usepackage{float}
% \usepackage[colorlinks=true, allcolors=blue]{hyperref}
% \usepackage{array}
\usepackage{tabularx} % 加载tabularx宏包
% \usepackage{xcolor}
\usepackage{mlmath}
\usepackage{longtable} % 使用longtable宏包
\usepackage[table]{xcolor} % 使用xcolor宏包，带table选项
\usepackage{multirow} % 引入multirow宏包
\usepackage{natbib}
\newcommand{\red}[1]{{\color{red}{#1}}}
\newcommand{\blue}[1]{{\color{blue}{#1}}}
\definecolor{pcolor}{HTML}{00C562}
\definecolor{ppcolor}{HTML}{00774F}
\definecolor{mcolor}{HTML}{DD4F4F}
\newcommand{\pcolor}[1]{{\color{pcolor}{#1}}}
\newcommand{\ppcolor}[1]{\textbf{\color{ppcolor}{#1}}}
\newcommand{\mcolor}[1]{{\color{mcolor}{#1}}}
\newcommand{\xu}[1]{\textcolor{red}{#1}}
% 附带脚注的表格
\usepackage{threeparttable}
% 附带脚注的长表格
\usepackage{threeparttablex}
\usepackage{amsmath,amsfonts,amssymb,mathtools,amsthm}
\newtheorem{problem}{Problem}
\newtheorem{prob}{Problem}
\newtheorem{theorem}{Theorem}
\newtheorem*{theorem*}{Theorem}
\newtheorem{lemma}{Lemma}
\newtheorem*{lemma*}{Lemma}
\newtheorem{property}{Property}
\newtheorem*{property*}{Property}
\newtheorem{remark}{Remark}
\newtheorem{definition}{Definition}
\newtheorem{corollary}{Corollary}
\newtheorem{assumption}{Assumption}
\newtheorem*{assumption*}{Assumption}
\newtheorem*{prop*}{Proposition}
\newtheorem{prop}{Proposition}
\newtheorem{setting}{Setting}
\newtheorem*{setting*}{Setting}
\newtheorem{example}{Example}
\newtheorem{notation}{Notation}
\newcommand{\R}{\mathbb{R}}
\newcommand{\x}{\mathbf{x}}
\newcommand{\w}{\mathbf{w}}
\title{Scalable Complexity Control Facilitates Reasoning Ability of LLMs}
% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to LaTeX to determine where to break the
% lines. Using \AND forces a line break at that point. So, if LaTeX puts 3 of 4
% authors names on the first line, and the last on the second line, try using
% \AND instead of \And before the third author name.
\author{%
Liangkai Hang\textsuperscript{\rm 1, $\dag$},
Junjie Yao\textsuperscript{\rm 1, $\dag$},
Zhiwei Bai\textsuperscript{\rm 1},
Tianyi Chen\textsuperscript{\rm 1},
Yang Chen\textsuperscript{\rm 1},
Rongjie Diao\textsuperscript{\rm 1},\\
{\bf Hezhou Li\textsuperscript{\rm 1}},
{\bf Pengxiao Lin\textsuperscript{\rm 1}},
{\bf Zhiwei Wang\textsuperscript{\rm 1}},
{\bf Cheng Xu\textsuperscript{\rm 1}},
{\bf Zhongwang Zhang\textsuperscript{\rm 1,*}},
{\bf Zhangchen Zhou\textsuperscript{\rm 1}},\\
{\bf Zhiyu Li}\textsuperscript{\rm 4,5,*},
{\bf Zehao Lin}\textsuperscript{\rm 4,5},
{\bf Kai Chen}\textsuperscript{\rm 5},
{\bf Feiyu Xiong}\textsuperscript{\rm 4,5, *},\\
{\bf Yaoyu Zhang}\textsuperscript{\rm 1,2, *},
{\bf Weinan E}\textsuperscript{\rm 3,4},
{\bf Hongkang Yang}\textsuperscript{5,*},
{\bf Zhi-Qin John Xu\textsuperscript{\rm 1,2,4}\thanks{Corresponding author: xuzhiqin@sjtu.edu.cn,hongkang@alumni.princeton.edu,zhyy.sjtu@sjtu.edu.cn, xiongfy@iaar.ac.cn, lizy@iaar.ac.cn, 0123zzw666@sjtu.edu.cn}} \\
\textsuperscript{\rm
1} Institute of Natural Sciences, School of Mathematical Sciences, Shanghai Jiao Tong University\\
\textsuperscript{\rm 2} MOE-LSC, School of Artificial Intelligence, Shanghai Jiao Tong University\\
\textsuperscript{\rm 3} Center for Machine Learning Research, School of Mathematical Sciences, Peking University\\
\textsuperscript{\rm 4} Center for LLM, Institute for Advanced Algorithms Research, Shanghai\\
\textsuperscript{\rm 5} MemTensor (Shanghai) Technology Co., Ltd.\\
\textsuperscript{\rm $\dag$} Equal contribution, list in alphabetical order.
}
% \author{
% David S.~Hippocampus\thanks{Use footnote for providing further information
% about author (webpage, alternative address)---\emph{not} for acknowledging
% funding agencies.} \\
% Department of Computer Science\\
% Cranberry-Lemon University\\
% Pittsburgh, PA 15213 \\
% \texttt{hippo@cs.cranberry-lemon.edu} \\
% % examples of more authors
% \AND
% Coauthor \\
% Affiliation \\
% Address \\
% \texttt{email} \\
% % \AND
% % Coauthor \\
% % Affiliation \\
% % Address \\
% % \texttt{email} \\
% % \And
% % Coauthor \\
% % Affiliation \\
% % Address \\
% % \texttt{email} \\
% % \And
% % Coauthor \\
% % Affiliation \\
% % Address \\
% % \texttt{email} \\
% }
\begin{document}
\maketitle
\begin{abstract}
The reasoning ability of large language models (LLMs) has been rapidly advancing in recent years, attracting interest in more fundamental approaches that can reliably enhance their generalizability.
This work demonstrates that model complexity control, conveniently implementable by adjusting the initialization rate and weight decay coefficient, improves the scaling law of LLMs consistently over varying model sizes and data sizes.
This gain is further illustrated by comparing the benchmark performance of 2.4B models pretrained on 1T tokens with different complexity hyperparameters.
Instead of fixing the initialization std, we found that a constant initialization rate (the exponent of std) enables the scaling law to descend faster in both model and data sizes.
These results indicate that complexity control is a promising direction for the continual advancement of LLMs.
% The capabilities of large language models (LLMs) have demonstrated remarkable developments in recent years. This work reveals that controlling model complexity can enhance the generalization performance by facilitating reasoning capability. This complexity control can be achieved by adjusting the initialization scale and weight decay coefficient, which is convenient to implement and stable across different training scales. Instead of tuning the standard deviation (std) of Gaussian-initialized parameters, we comprehensively demonstrate a scalable complexity control approach by adjusting the initialization rate (the power of std w.r.t. layer width), where the performance of small models can predict the performance of large models via stable scaling laws. Our research provides insights and practical guidance on improving the reasoning ability of LLMs by complexity control.
\end{abstract}
\section{Introduction}
% Heuristic vs. Principle guided
% clean data by human preference vs. model intrinsic
% from the perspective of data and task rather than the characteristic of neural networks.
% Mainly focus on the data quality and algorithm frame, but not the neural network intrinsic properties.
% Theory of FNN neural networks extended to LLM
% Initialization is important hyper-parameter. Find a proper way (can scale up to large models) to tune it is critical to performance and cost
\begin{figure}[htpb]
\vspace{-8pt}
\centering
\includegraphics[width=1\linewidth]{pic/tokenloss.png}
\caption{Next-token prediction accuracy with varying model complexity.
% Model complexity decreases from left to right.
The colors indicate the probability of each token predicted by the models
%with red representing lower probabilities and green indicating higher probabilities.
Model complexity decreases from left to right, as the initialization rate $\gamma=0.1,0.5,1$ (see Section \ref{sec:cc} for the definition of $\gamma$).
These models have the same shape (180M parameters) and are trained the same dataset (40B tokens).}
\label{fig:token_loss}
\end{figure}\vspace{-5pt}
% acc改成pro 去掉红色的。第三张图再举例 （have done）
% 第三幅图 “self-defense after Anderson” 的 Anderson 可视为复杂推理：上文 Anderson 先指控 Carter，于是这句话 Carter 便要反过来指控 Anderson
% Recent years have witnessed unprecedented progress in large language models (LLMs), demonstrating remarkable capabilities across diverse linguistic tasks. Reasoning is a particular important ability to indicate the intelligence of a model. Various approaches are developed to improve the reasoning ability of LLMs, for example, post-training methods such as reinforcement learning, more high-quality reasoning data such as math data, prompt engineering such as chain-of-thought, architecture design such as Memory$^3$ structure.
In recent years, large language models (LLMs) have achieved unprecedented progress, demonstrating impressive performance on a wide range of tasks~\cite{achiam2023gpt,liu2024deepseek,touvron2023llama,wei2022emergent,yang2024qwen2,zeng2021pangu}.
The key to this success is the improved generalizability of these models, in particular, their reasoning ability.
% Among these, reasoning stands out as a key indicator of model intelligence.
Various approaches have been explored to enhance reasoning, such as post-training with reinforcement learning~\cite{kumar2024training,guo2025deepseek},
high-quality math and code data with reasoning traces \cite{guo2025deepseek},
chain-of-thought and related prompting strategies~\cite{wei2022chain},
and the separation of reasoning and memory in pretraining \cite{yang2024memory3}.
% This work identifies the control of model complexity can tune the reasoning performance of LLMs, where such control can be realized through many approaches, such the scale of parameter initialization and the penalty strength of parameter norm. A series of theoretical and empirical works suggest that with smaller initial parameter scale, a neural network tends to fit data by a function with less complexity.
% An intuitive example of complexity control by parameter initialization is depicted in
% Figure~\ref{fig:token_loss}, model with a large complexity (large initialization scale) fails to achieve meaningful next-token prediction for a law document not used for training. As the model complexity is more controlled, more tokens aligned with logical reasoning patterns can be accurately predicted (greened words). For the median case, common grammar such as ``of'' and ``that'', year ``2024'' mentioned before, and the completion of common name ``Jessica'' are predicted by the model. For the model with small complexity, it not only accurately predicts those predicted by the model with median complexity, words with more reasoning strength can also be accurately predicted. For example, in a law document, the name of plaintiff is listed before defendant, therefore, in line 7, ``Anderson'' is predicted after ``Plaintiff Michael''; the model also predicts that Anderson requires ``medical treatment'' after Carter punched him, which requires scenario reasoning.
This work identifies model complexity as the key factor in the development of the reasoning ability of LLMs.
Complexity control can be implemented through various means, such as adjusting the rate of parameter initialization ~\cite{luo2021phase,zhou2022empirical,zhou2022towards,zhang2025complexity,yao2025analysis} or applying stronger penalties on parameter norms ~\cite{zhang2025complexity}.
Similar to the gene that determines the characteristics of an organism, these designs directly regulate the reasoning ability of LLMs.
An illustration is provided in Figure~\ref{fig:token_loss}.
On the left, a pretrained model with large complexity (large initialization scale) fails to make meaningful next-token predictions on unseen test data.
In the middle, with moderate complexity (the commonly used initialization scale), the model demonstrates basic knowledge of grammar (e.g. ``finds \underline{that}'') and vocabulary (e.g. ``Jess\underline{ica}''), as well as the induction head mechanism \cite{olsson2022incontext} (e.g. the repeated phrase ``September \underline{12, 2024}'').
%predictions increasingly align with logical reasoning patterns (highlighted in green).
% In the medium-complexity case (i.e., trained with medium initial parameter scale), the model correctly predicts basic grammatical structures (e.g., ``of'', ``that''), repeated elements (e.g., the year ``2024''), and common name completions (e.g., ``Jessica'').
On the right, the low-complexity model (i.e. with small initialization scale) further captures sophisticated semantic reasoning,
successfully predicting ``Plaintiff Michael \underline{Anderson}'' given the context ``Anderson v. Carter'' despite that the full name ``Michael Anderson'' has not appeared before,
and also predicting ``required \underline{medical treatment}'' given that Anderson was punched by Carter.
% Intuitively, with more complexity control, the model has to find a way that can fit data but only needs to learn much fewer fundamental rules. For initialization scale,
% Previous works have shown that with small initial parameter scale, neurons in the same layer tend to behave similarly, that is, condensation, where the effective neuron number is much smaller than its superficial one. For a larger penalty strength of weight norm, the neural network tends to find a solution with small weight norm, that is, low-complexity network output.
Intuitively, smaller complexity forces the model to compress data into a smaller set of production rules, revealing the deep dependency among the tokens and preventing plain memorization.
Previous studies \cite{luo2021phase,zhou2022empirical} have shown that with a smaller initialization scale, neurons within each layer tend to evolve within a few groups,
a phenomenon known as condensation, which limits the effective number of neurons. Readers are referred to an overview of condensation \cite{xu2025overview}.
Likewise, models trained with stronger penalty on parameter norm may converge to solutions with smaller norms, resulting in lower-complexity outputs.
Thus, it is promising that these techniques can enhance the reasoning ability of LLMs.
% This work finds that the complexity control can tune the reasoning ability of a LLM as a gene can tune the characteristics of a creature. With more focusing on reasoning, the performance of an LLM can be robustly improved, which is verified via experiments of scaling laws w.r.t. data size and model size. In particular, instead of the standard deviation (std) in Gaussian distribution for initial parameters, this work proposes that the tuning of initialization rate (the power of std w.r.t. layer width) can achieve a better scaling law. The initialization rate is critical to determine the dynamical regime of a network when scaling up the network size. However, one can never expect the improvement of reasoning ability can leading to the improvement of all benchmarks since some knowledge can only be remembered such as historical events. We then test our approach on larger-scale models (0.9B, 2.4B) and datasets (600B, 1T tokens) to show that a stronger complexity control can significantly enhance the reasoning ability of LLMs, followed by detailed analysis. This work provides a deep understanding of reasoning ability of LLM and practical technique guidance for tuning reasoning ability of LLM.
As a verification of this intuition, our experiments exhibit improved scaling laws in data size and model size, as well as higher scores over almost all benchmarks (+4.6\% for 0.9B model with 600B data, and +3.4\% for 2.4B model with 1T data, averaged over 15 tasks).
Compared with the standard deviation (std) of parameter initialization, the initialization rate (the exponent of std as a function of network width) turns out to be the right invariant for the scaling laws.
This is in accordance with previous works \cite{luo2021phase,zhou2022empirical} on the phase diagram of neural network training.
We provide some heuristic calculations to explain how complexity control facilitates the learning of multi-step reasoning.
% However, improved reasoning does not guarantee gains across all benchmarks, as some tasks
% (e.g., memorization of historical facts)
% rely more on memorization than reasoning. We validate our approach on larger-scale models (0.9B and 2.4B parameters) and datasets (600B and 1T tokens), showing that stronger complexity control significantly enhances reasoning performance.
% We further provide some mechanisms to understand the relation between reasoning and complexity control. This work offers both theoretical insight into the reasoning capabilities of LLMs and practical guidance for tuning them effectively.
% However, the decrease in model complexity is observed to enhance contextual prediction performance. Furthermore, these complexity-constrained performance gains are predominantly manifested through enhanced reasoning capabilities, as evidenced by improved prediction accuracy specifically on tokens aligned with logical reasoning patterns. For instance, it's observed that with reduced model complexity, the system demonstrated enhanced capability in predicting named entities (e.g., "Andersen") and temporal expressions (e.g., "12. 2024") that have explicit contextual antecedents. The model also exhibited improved syntactic awareness, reliably predicting functional words like the preposition "of" following phrasal pattern "In the Matter". Furthermore, frequent anthroponyms like "Jessica" are consistently predicted with high accuracy. These findings collectively suggest that complexity reduction facilitates the reasoning capabilities. Conversely, when processing non-systematic number sequences (e.g., "2024-01589"), performance improvements are not observed despite reduced complexity. This differential effect substantiates our hypothesis that complexity reduction primarily enhances model performance through facilitating the reasoning capabilities.
% In particular, the research community has actively pursued various strategies to enhance model performance, with special emphasis on advancing reasoning abilities. Prevailing approaches predominantly centered on dataset curation (e.g., Chain-of-Thought prompting), architectural innovations, and training methodology refinements (e.g., reinforcement learning). However, some techniques often introduce excessive barrier to adoption and lack deeper exploration, thereby limiting their general applicability in LLM development. Moreover, although preliminary validations on small-scale benchmarks suggest potential reasoning improvements, the prohibitive computational costs required for large-model verification have cast doubt on the scalability and practical efficacy of such methods.
% This paper systematically proposes a principled complexity control framework to enhance LLM performance, particularly focusing on the underexplored dimension of reasoning capability. Our study is motivated by the following empirical observation: We trained language models with distinct model complexities, subsequently evaluating their next-token prediction performance on test corpora. As demonstrated in
% Based on the observation, we systematically evaluate and analysis the impact of model complexity.
\section{Related works}
\paragraph{LLMs reasoning ability}
Even advanced LLMs such as GPT-4 often struggle with implicit reasoning over parametric knowledge~\citep{talmor2020olmpics, kassner2020pretrained, rae2021scaling, press-etal-2023-measuring, allenzhu2023physics, yang2024large}, revealing their limited ability to internalize structured facts and rules. Verbalized reasoning strategies such as chain-of-thought can substantially boost performance, particularly for large models~\citep{wei2022chain, zelikman2022star, sun2023recitationaugmented, liu-etal-2023-crystal, zelikman2024quietstar, feng2023towards, li2024chain}.
However, understanding the underlying capacity for implicit reasoning remains a critical challenge, often studied by controlled experiments~\citep{prystawski2023why, dziri2023faith, wang2024understanding}.
This work takes the more intrinsic perspective of complexity control.
% mechanistically influences the reasoning performance of language models.
\paragraph{Effect of initialization}
Parameter initialization is known to be influential to the training and performance of classical neural networks
\cite{arora2019exact, chizat_global_2018, zhang_type_2019, e2020comparative, jacot_neural_2018, mei_mean_2018, rotskoff_parameters_2018, sirignano_mean_2020, williams_gradient_2019}.
For instance, distinct phases (the linear and condensed regimes) can be induced in wide ReLU networks by varying initialization rates
\cite{luo2021phase,zhou2022empirical},
% For instance, Luo et al.~\cite{luo2021phase} and Zhou et al.~\cite{zhou2022empirical} identify linear and condensed regimes in wide ReLU networks, induced by varying initialization scales.
and training in the condensed regime tend to fit data with lower-complexity functions~\cite{zhang2022linear, zhang2023loss, zhang2023stochastic, zhang2024implicit}.
% Recent studies have also examined the impact of initialization on the training dynamics of LLMs
The particular case of Transformer language models has also been studied
\cite{huang2020improving, liu2020understanding, trockman2023mimetic, wang2024deepnet, zhang2019improving, zhu2021gradinit, zhang2024initialization, zhang2025complexity},
with particular interest on the impact of initialization on training stability and efficiency.
Experiments on toy datasets~\cite{zhang2024initialization, zhang2025complexity} show that small initialization scales assist Transformers to identify the elementary functions when fitting synthetic compositional data.
However, the influence of initialization on the reasoning ability of Transformers trained on natural language data remains to be explored.
\paragraph{Weight decay}
\cite{krogh1991simple} introduced weight decay as a method to improve the generalization of neural networks.
Many subsequent works have explored its role in controlling model complexity and enhancing generalization~\cite{ bartlett2002rademacher,bartlett2017spectrally,arora2018stronger,neyshabur2015norm,golowich2018size,wei2019data}.
More recently, weight decay has been shown to be particularly important for achieving better generalization~\cite{power2022grokking, varma2023explaining}.
In this work, we systematically investigate how controlling pre-training complexity, via initialization strategies and weight decay, affects the performance of LLMs. Unlike prior studies that focus on training stability or synthetic tasks, we evaluate the impact across a wide range of downstream benchmarks and analyze the underlying mechanisms, aiming to offer practical guidance for large-scale model pre-training.
\section{Complexity control}\label{sec:cc}
In this section, we introduce the approach to modulate the model complexity.
% For any large language model with trainable parameter set $\fP$, we denote its model complexity by the parameter norm $\sqrt{\sum_{\theta\in\fP}\theta^2}$. In the training process, both the initialization scale and the weight decay coefficient play critical roles in determining model complexity. The initialization scale governs the initial model complexity, serving as a significant determinant of early-stage fitting dynamics, while the weight decay coefficient acts as a regularization that systematically shapes the complexity.
\paragraph{Initialization rate $\gamma$} Given any trainable parameter matrix $\vW\in\sR^{d_{in}\times d_{out}}$. We initialize its elements according to the following normal distribution:
\begin{equation*}
\vW_{i,j} \sim \mathcal{N}\left(0,\left(d_{in}^{-\gamma}\right)^2\right),
\end{equation*}
where $\gamma$ is the initialization rate. Specifically, the initialization scale decreases as $\gamma$ increases.
Note that $\gamma=0.5$ is commonly used in many default initialization methods, such as LeCun initialization~\cite{LeCun1998} and He initialization~\cite{he2015delving}. As the network width towards infinity~\cite{luo2021phase,zhou2022empirical}, the training of the network with $\gamma>0.5$ exhibits significant non-linear characteristics, i.e., condensation. Therefore, initialization scales with $\gamma>0.5$ are generally considered small. Tuning $\gamma$ is a scalable approach for complexity control.
\paragraph{Weight decay coefficient $\lambda$}Given any trainable parameter $\theta_t$ where $t$ denotes the current training step. Define $\hat{\theta}_t$ as the parameter after optimizing by gradient and moment, the weight decay is implemented by
\begin{equation}
\theta_{t+1} \longleftarrow \hat{\theta}_t - \lambda C \theta_t,
\end{equation}
where $\lambda$ is the weight decay coefficient.
\section{Results}
To evaluate the impact of complexity control, we train LLMs based on the Llama- architecture~\cite{touvron2023llama} under different levels of model complexity. We examine scaling laws with respect to both model size and training data size, and compare different models across a range of benchmarks. Detailed setup of training and evaluation is provided in Appendix~\ref{app:setup}.
%The results show that stronger complexity control leads to higher sampling efficiency. Additionally, using a fixed initialization rate yields better scaling performance than a fixed standard deviation—an advantage that is particularly useful for tuning hyperparameters in large models based on insights from smaller ones. Notably, benchmarks that require more reasoning show significantly greater improvement under complexity control.
% We adopt two model configurations containing approximately 0.94 billion and 2.4 billion parameters, respectively. The smaller 0.94B model is trained on a 600B-token corpus sampled from SlimPajama. In contrast, the larger 2.4B model is trained on a curated 1T-token dataset consisting of high-quality natural and synthetic data. The detailed model configurations and training setups are provided in Appendix. For each base configuration, we develop two variants with distinct levels of complexity. Evaluations across diverse benchmark tasks reveal that principled complexity management significantly enhances model performance—particularly on reasoning-intensive tasks such as GSM8K and CommonsenseQA.
\subsection{Scaling law}
We first establish three model complexity configurations: (1) the small-complexity setup with $\gamma=1,\lambda=1$, (2) the large-complexity setup with $\gamma=0.5,\lambda=0.1$, and (3) the commonly used default configuration (e.g., GPT2, HuggingFace) with $\sigma=0.02,\lambda=0.1$. Under each configuration, we train 0.8B-parameter models with varying training data scales ranging from 0.2 billion to 1.4 billion tokens. Figure~\ref{fig:scalinglaw} (left) demonstrates the relationship between test loss and data scale across different complexity configurations. The curve of small-complexity (green) exhibits a distinct leftward shift relative to the large-complexity (purple), suggesting that complexity controlling can effectively improve the sample efficiency. Additionally, while the test loss of the large-complexity model demonstrates comparable to the standard configuration at smaller data scales, their performance diverges as data size increases. Specifically, the former configuration achieves progressively lower test loss than the latter, indicating superior scalability of the $\gamma$-initialization approach with steeper scaling slopes. Besides, we train models with distinct model sizes by 1 billion tokens. Figure~\ref{fig:scalinglaw} (right) reveals similar patterns in the relationship between test loss and model size. Although $\sigma$-initialization achieves lower test loss with smaller models, its performance plateaus as the model size increases, resulting in significantly higher loss compared to $\gamma$-initialized models at larger scales. These results illustrate the scalability potential of the $\gamma$-initialization method, maintaining performance advantages across expanding parameters and data scales.
\begin{figure}[htpb]
\centering
\includegraphics[width=1\linewidth]{pic/scalinglaw.png}
\caption{Test loss across varying data and model scales under different complexity configurations. Left: Test loss progression for 0.8B-parameter models trained with data scales ranging from 0.2B to 1.4B tokens. Right: Test loss versus model parameter counts (50M-0.8B) with fixed 1B training tokens. Line colors correspond to different complexity configurations. }
\label{fig:scalinglaw}
\end{figure}
\vspace{-15pt}
\subsection{Evaluation}
\begin{table}[htbp]
\centering
\begin{threeparttable}
\caption{Evaluation of models with different model complexities.}
\label{tab:model_benchmark}
\begin{tabular}{@{}lccccc@{}}
\toprule
\textbf{Models\tnote{1}} & \textbf{0.9B Large} & \textbf{0.9B Small} & \textbf{2.4B Large} & \textbf{2.4B Small}\\
$\left(\gamma,\lambda\right)$ & (0.5,0.1) &(1,1) &(0.5,0.1) & (0.58,1)\\
\midrule
MMLU & 49.9 & 52.5 (\pcolor{+2.6}) & 60.4 & 64.4 (\ppcolor{+4}) \\
MMLU-Pro & 17.6 & 21.5 (\pcolor{+3.9}) & 30.6 & 30.1 (\mcolor{-0.5}) \\
BBH & 33.3 & 34.9 (\pcolor{+1.6}) & 42.4 & 43.7 (\pcolor{+1.3}) \\
ARC-C & 46.1 & 49.8 (\pcolor{+3.7}) & 58.6 & 60.7 (\pcolor{+2.1}) \\
TruthfulQA & 53.4 & 56.8 (\pcolor{+3.4}) & 58.7 & 61.7 (\pcolor{+3}) \\
WinoGrande & 68.0 & 72.2 (\ppcolor{+4.2}) & 73.1 & 76.9 (\pcolor{+3.8}) \\
HellaSwag & 63.0 & 67.3 (\ppcolor{+4.3}) & 71.2 & 75.4 (\ppcolor{+4.2}) \\
AGIEval-EN & 25.3 & 29.1 (\pcolor{+3.8}) & 33.2 & 35.6 (\pcolor{+2.4})\\
OpenBookQA &38.8 & 39.6 (\pcolor{+0.8}) & 41.8 & 43.4 (\pcolor{+1.6}) \\
CommonsenseQA &57.7 & 67.2 (\ppcolor{+9.5}) &70.4 & 76.9 (\ppcolor{+6.5})\\
GPQA & 25.9 & 26.8 (\pcolor{+0.9}) & 29.0 & 31.7 (\pcolor{+2.7}) \\
MATH & 5.9 & 11.5 (\ppcolor{+5.6}) & 34.3 & 35.3 (\pcolor{+1}) \\
GSM8K & 21.8 & 40.2 (\ppcolor{+19.4}) & 52.8 & 63.8 (\ppcolor{+11.0}) \\
MBPP & 6.6 & 10.8 (\ppcolor{+4.2}) & 20.6 & 22.4 (\pcolor{+1.8}) \\
IFEval & 31.3 & 32.0 (\pcolor{+0.7}) & 34.7 & 40.9 (\ppcolor{+6.2})\\
\bottomrule
\end{tabular}
\begin{tablenotes}
\footnotesize
\item[1] ``Large'' and ``Small'' mean large complexity and small complexity, respectively.
\end{tablenotes}
\end{threeparttable}
\end{table}
To assess the impact of complexity control, we adopt reasonable complexity configurations and train LLMs with the following setup: (1) 0.9B-parameter models on 600B tokens from SlimPajama~\cite{cerebras2023slimpajama} and (2) 2.4B-parameter models trained with 1T high-quality corpus. We evaluate the performance of our models across a comprehensive set of benchmark tasks. As shown in Table~\ref{tab:model_benchmark}, the performance of small-complexity models improves significantly. Specifically, complexity controlling yields substantial gains in reasoning capabilities. On math-related tasks, two small-complexity models achieve improvements of 19.4 and 11.0 points on the GSM8K benchmark, respectively, as well as 5.6 and 1 points on the MATH dataset. Moreover, these small-complexity models demonstrate notable performance enhancements across other reasoning tasks, including Winogrande($+4.2,+3.8$), HellaSwag($+4.3,+4.2$), and CommonsenseQA($+9.5,+6.5$). These results demonstrate that principled control of model complexity can significantly enhance the overall capabilities of LLMs, particularly in reasoning tasks. Note that small initialization may yield slower training. For 2.4B model, the performance of $\gamma=0.58$ is slightly better than that of $\gamma=1$ in Table \ref{tab:0.58-1} in Appendix.
% Note that with strong weight decay penalty $\lambda=1$,
% 实验 （1，0.1）（0.5，1）
% \begin{table}[htbp]
% \centering
% \caption{Evaluation of models with different model complexities.}
% \label{tab:model_benchmark}
% \begin{tabular}{@{}lccccc@{}}
% \toprule
% \textbf{Models} & \textbf{0.94B Large} & \textbf{0.94B Small} & \textbf{2.4B Large} & \textbf{2.4B Small}\\
% $\left(\gamma,\lambda\right)$ & (0.5,0.1) &(1,1) &(0.5,0.1) & (0.58,1)\\
% \midrule
% MMLU & 49.9 & 52.5 (\pcolor{+2.6}) & 61.7 & 64.4 (\pcolor{+2.7}) \\
% MMLU-Pro & 17.6 & 21.5 (\pcolor{+3.9}) & 30.2 & 30.1 (\mcolor{-0.1}) \\
% BBH & 33.3 & 34.9 (\pcolor{+1.6}) & 38.8 & 43.7 (\ppcolor{+4.9}) \\
% ARC-C & 46.1 & 49.8 (\pcolor{+3.7}) & 60.7 & 60.7 \\
% TruthfulQA & 53.4 & 56.8 (\pcolor{+3.4}) & 60.8 & 61.7 (\pcolor{+0.9}) \\
% WindoGrande & 68.0 & 72.2 (\ppcolor{+4.2}) & 72.5 & 76.9 (\ppcolor{+4.4}) \\
% HellaSwag & 63.0 & 67.3 (\ppcolor{+4.3}) & 72.5 & 75.4 (\pcolor{+2.9}) \\
% AGIEval-EN & 25.3 & 29.1 (\pcolor{+3.8}) & 35.1 & 35.6 (\pcolor{+0.5})\\
% OpenBookQA &38.8 & 39.6 (\pcolor{+0.8}) & 42.4 & 43.4 (\pcolor{+1}) \\
% CommonsenseQA &57.7 & 67.2 (\ppcolor{+9.5}) &73.1 & 76.9 (\pcolor{+3.8})\\
% GPQA & 25.9 & 26.8 (\pcolor{+0.9}) & 27.2 & 31.7 (\ppcolor{+4.5}) \\
% MATH & 5.9 & 11.5 (\ppcolor{+5.6}) & 32.6 & 35.3 (\pcolor{+2.7}) \\
% GSM8K & 21.8 & 40.2 (\ppcolor{+19.4}) & 57.5 & 63.8 (\ppcolor{+6.3}) \\
% MBPP & 6.6 & 10.8 (\ppcolor{+4.2}) & 23.8 & 22.4 (\mcolor{-1.4}) \\
% IFEval & 31.3 & 32.0 (\pcolor{+0.7}) & 36.2 & 40.9 (\ppcolor{+4.7})\\
% \bottomrule
% \end{tabular}
% \end{table}
Complexity control also significantly improves the performance of base models, with results summarized in Table~\ref{tab:base_benchmark} in Appendix~\ref{app:base}.
% Notably, the 0.94B small-complexity model exhibits no significant improvement due to the suboptimal data quality. This comparison indicates a relation between supervised fine-tuning (SFT) and complexity control: The small-complexity model gains more improvement after SFT.
As quantified in Figure~\ref{fig:base_sft}, the small-complexity model attains greater score increments across most tasks at the SFT stage.
\begin{figure}[htpb]
\centering
\includegraphics[width=0.9\linewidth]{ pic/base-sft.png}
\caption{Performance improvement via SFT across complexity configurations(0.9B model). It is quantified as the performance gap between the SFT model and the corresponding base model.}
\label{fig:base_sft}
\end{figure}
\vspace{-8pt}
\section{Analysis}\label{sec:analysis-187}
To investigate the effect of initialization scale and weight decay in complexity control, we trained 180M-parameter models with different complexity configurations. We conduct analysis from evaluation and parameter analysis, which yield mechanistic insights into complexity control.
\subsection{Influence of initialization scale and weight decay}
\begin{figure}[htpb]
\centering
\includegraphics[width=1\linewidth]{ pic/pretrain_paranorm.png}
\caption{Parameter norm evolution across complexity configurations. Left to right: $\lambda = 0, 0.1, 1$; Line colors correspond to $\gamma = 0.1,0.3, 0.5,0.8,1$.}
\label{fig:paranorm}
\end{figure}
Examining the interplay between initialization scale $\gamma$ and weight decay $\lambda$ on model complexity is critical. Figure~\ref{fig:paranorm} visualizes the temporal evolution of parameter norms under varying $\gamma-\lambda$, demonstrating that large $\gamma$ coupled with large $\lambda$ systematically induces lower model complexity. Specifically, with small $\lambda$ ($\lambda=0,0.1$), configurations with larger $\gamma$ converge to smaller complexities. However, with large $\lambda$ ($\lambda=1$), models with different $\gamma$ converge to comparable small norm.
\subsection{Alignment between complexity and capability}\label{sec:187m}
We perform evaluations of models across multiple benchmark tasks. Figure~\ref{fig:187-eval}A depicts the performance distribution of the average score, GSM8K, and HellaSwag, showing better performance as stronger complexity control.
% Notably, configurations with $\lambda=1$ exhibit good performance across $\gamma$ values due to their convergence toward nearly small complexities.
As shown in Figure \ref{fig:187-eval}B,
% the correlation coefficient between complexity and task scores in most tasks is less than -0.6, suggesting that
model performance demonstrates a strong inverse correlation with model complexity.
% These empirical results consistently validate model complexity as a reliable performance predictor, establishing complexity control as a robust mechanism for enhancing reasoning capabilities.
Full evaluation results are provided in Appendix~\ref{app:187}.
\begin{figure}[htpb]
\centering
\includegraphics[width=1\linewidth]{ pic/187M-evaluation/187_eval_heatmap.png}
\caption{(A) Evaluation scores (average, GSM8K, HellaSwag) under varying complexities. Top: Performance landscape across $\gamma-\lambda$ with color indicating score (dark: low, light: high). Bottom: Score-complexity relationships with points indicating the models and the dashed lines denoting baseline performance levels.
(B) Task-specific Spearman correlations between model complexity and task score. Stronger negative correlations (approaching -1) indicate greater performance enhancement through complexity control.
}
\label{fig:187-eval}
\end{figure}
\subsection{Model Analysis}
\paragraph{Embedding space}
The embedding space reflects the model's representation of the vocabulary and its learning patterns. We compare the cosine similarity of the 350 most frequent embedding vectors under different initialization scales. The results in Figure~\ref{fig:187-embedding} demonstrate that with large complexity ($\gamma=0.1$), the embeddings are pairwise orthogonal, indicating the model ignores their relationship. In contrast, controlling complexity ($\gamma=0.5, 1$) significantly increases the similarity among embeddings, consistent with previous study on condensation phenomenon \cite{luo2021phase}. This analysis suggests that small complexity encourages the model to find associations among tokens.
\begin{figure}[htpb]
\centering
\includegraphics[width=1\linewidth]{ pic/parameter_analysis/187-embedding.png}
\caption{Cosine similarity among 350 embedding vectors which occur most frequently in training dataset under initialization scales $\gamma=1,0.5,0.1$ ($\lambda=0$).}
\label{fig:187-embedding}
\end{figure}
\paragraph{Attention matrix}
For each trainable parameter matrix $\vW\in\sR^{d_{in}\times d_{out}}$, we define the following metrics to measure its condensation degree and low-rank degree, respectively.
\begin{equation}
D_{c}(\vW):= \frac{1}{d_{in}d_{out}}\sum_{i,j}\frac{\vW_i^T\vW_j}{||\vW_i||_2\cdot||\vW_j||_2},\qquad D_s=\frac{\max_i S_{\vW,i}}{\sum_i S_{\vW,i}},
\end{equation}
where $\vW_i$ and $S_{\vW,i}$ means the $i$-th row of $\vW$ and the $i$-th singular value of $\vW$. Larger values of $D_c$ and $D_s$ indicate fewer effective directions in matrix $\vW$, suggesting $\vW$ learns a smaller set of features for fitting dataset. Figure~\ref{fig:187-dc} exhibits the $D_c$ and $D_s$ of the query projection and key projection matrices of models with distinct complexities. The results demonstrate that controlling model complexity effectively increases $D_c$ and $D_s$ across the attention matrices, suggesting that the attention modules in small-complexity models focus more on the fundamental relationships between tokens within sequences, which results in stronger generalization and reasoning ability.
\begin{figure}[htpb]
\centering
\includegraphics[width=1\linewidth]{ pic/parameter_analysis/DcDs.png}
\caption{$D_c$ and $D_s$ of $\vW_Q$ and $\vW_K$ in each layer under different model complexity configurations.}
\label{fig:187-dc}
\end{figure}
% \begin{figure}[htpb]
% \centering
% \includegraphics[width=0.9\linewidth]{complexity control/pic/parameter_analysis/Ds187.png}
% \caption{$D_s$ under different model complexities.}
% \label{fig:187-ds}
% \end{figure}
% \subsection{How to adjust initialization scale}
% Beyond the initialization method mentioned above(by adjusting $\gamma$), another common way is to adjust the standard derivation directly, which is independent of the parameter matrix dimensions. To empirically validate the reasonable of $\gamma$-modulation approach, we designed the following comparison framework: Assuming $\gamma=1$ corresponds to an appropriately small initialization scale in our 187M parameter model and denote the standard derivation by $\sigma_{\vW}$ for any parameter matrix $\vW$, we initialize one 0.9B parameter model by setting $\gamma=1$ and another one by $\vW_{i,j} \sim\fN\left(0,\sigma_{\vW}^2\right)$.
\section{Discussion}
\subsection{Training Stability}
A potential concern in small complexity is training stability, particularly the emergence of loss spikes during optimization. This instability tends to escalate with increasing model scale, as observed in our experiments. To address this challenge in training our 2.4B parameter models, we adopt the $\gamma=0.58$ which is not too large, and implement embedding normalization and sandwich normalization, which successfully mitigate the loss spike phenomenon. A detailed discussion is provided in Appendix~\ref{app:0.58-1}.
\subsection{Training dynamics}
Distinct $\gamma-\lambda$ configurations induce distinct learning dynamics and complexity trajectories, accompanied by different performance evolution patterns. We evaluate the two 2.4B models referenced in Table~\ref{tab:model_benchmark} every 5000 steps during pre-train and obtain the dynamics of the average score. Figure~\ref{fig:para-score-evoluation} characterizes the temporal evolution of model complexity and the average score. For the large-complexity model ($\gamma=0.5, \lambda=0.1$), complexity initially increases followed by progressive decay, during which performance exhibits slow improvement in the ascending phase but accelerates substantially after the complexity turning. Conversely, small-complexity configurations maintain monotonic complexity growth with steady performance gains throughout training.
\begin{figure}[htpb]
\centering
\includegraphics[width=0.9\linewidth]{ pic/norm_score.png}
\caption{The evolution of parameter norm and average score with two distinct $\gamma-\lambda$ configurations.}
\label{fig:para-score-evoluation}
\end{figure}
\subsection{Theoretical analysis}
% The extensive experimental results presented in the preceding sections demonstrate that controlling model complexity can effectively enhance model performance, particularly in strengthening reasoning capabilities. To theoretically substantiate these empirical findings, we conduct a rigorous analysis of deep learning models within a general framework. Our investigation reveals that the global minima of models exhibit distinct characteristics under different complexity constraints. Specifically, when operating under low-complexity regimes, models tend to develop sparse parameter configurations and hierarchical feature representations. These structural properties facilitate multi-step reasoning processes through successive information refinement, thereby enhancing the model's capacity for complex logical deduction. Our analytical framework provides preliminary yet valuable theoretical insights into the fundamental mechanisms of complexity control.
% \begin{definition}
% Given that $\vx\in \sR^{d_0},\vw\in\sR^d$ as the input and weight vectors, for any mapping $\phi$ and finite signed measure $\pi\in\fM\left(\sR^d\right)$, we define the model $f$ as
% \begin{equation}
% f(\vx)=\int \phi(\vx,\vw)d\pi\left(\vw\right).
% \end{equation}
% \end{definition}
% \begin{definition}
% For any model $f$, assume the $\gamma$ is the initialization scale and the $\pi_0\in\fM\left(\sR^d\right)$ as the initial weight distribution, we define the $\gamma-$norm of $f$ as:
% \begin{equation}
% \lvert\lvert f\rvert\rvert_{\gamma} = \inf_{\pi}\left(\int \left|\frac{\delta\pi}{\delta\pi_0} \right|^{\alpha-\beta\gamma}d\pi_0\right)^{\frac{1}{\alpha-\beta\gamma}},\quad \alpha-\beta\gamma > 0.
% \end{equation}
% where $\alpha,\beta$ are constants depending on the model.
% \end{definition}
% \begin{assumption}\label{assump:circuit}
% For any model $f$, assume that there exists a group of finite signed measure $\fC=\left\{\pi_c\in\fM\left(\sR^d\right)\right\}$ satisfy the following conditions:\\
% \begin{enumerate}
% \item $\pi=\sum_{c=1}^{N_C}a_c\pi_c,a_c\in \sR$.
% \item $\exists 1\leq k_c\ll d$ such that $\pi_c=\pi_c^{\left(l_1,\cdots,l_{k_c}\right)}\otimes\left(\bigotimes\limits_{l\neq l_1,\cdots,l_{k_c}}\pi_0^l\right).$
% \item There exists a constant $b$ such that $\frac{\delta \pi_c^{i}}{\delta \pi_0}\equiv b$.
% \end{enumerate}
% \end{assumption}
% \begin{assumption}\label{assump:converge}
% For any initialization scale $\gamma$, the model $f$ could converge to the solution with minimal $\lvert\lvert f\rvert\rvert_\gamma$.
% \end{assumption}
% \begin{theorem}\label{thm:1}
% With Assumption~\ref{assump:circuit},~\ref{assump:converge}, we have that:
% \begin{equation}\label{eq:norm}
% \lvert\lvert f\rvert\rvert_{\gamma} \sim \inf_{a_c,k_c}\left(\sum_c\left(a_cb^{k_c}\right)^{\alpha-\beta\gamma}\right)^{\frac{1}{\alpha-\beta\gamma}}.
% \end{equation}
% \end{theorem}
% Theorem~\ref{thm:1} indicates that the training process is optimizing the $a_c$ to minimize \eqref{eq:norm} which is a $l_{\alpha-\beta\gamma}$ optimization problem. As we know, the lower exponent $\alpha-\beta\gamma$ would bring a more sparse distribution of $a_c$ and larger $\alpha-\beta\gamma$ would cause a more uniform distribution. Otherwise, as the decrease of $\alpha-\beta\gamma$, the sensitivity with respect to $k_c$ enlarges which illustrates that small complexity brings out a large tolerance of $k_c$. Specifically, we have the following results:
% \begin{corollary}
% Given any $k_c$, let $\alpha-\beta\gamma\rightarrow 0$, then the solution distribution $a_c^\star$ satisfies $a_c^\star=\arg\min_{a_c} N_C$ . Otherwise, if $\alpha-\beta\gamma\to \infty$, then exists a constant $M$ such that $a_c^\star=\frac{M}{b^{k_c}}$.
% \end{corollary}
% \begin{corollary}
% Define that $g\left(a_c,k_c;\gamma\right)=\lvert\lvert a_cb^{k_c}\rvert\rvert_{\alpha-\beta\gamma}$. Given any $a_c,k_c$, we have that $g\left(a_c,k_c+1;\gamma\right)-g\left(a_c,k_c;\gamma\right)$ is monotonically increase with respect to $\gamma$.
% \end{corollary}
% These results provide a concise characterization of how model complexity influences the training process from a specific perspective, offering insights into the characteristics of parameter matrices under varying complexity constraints.
This section employs heuristic calculations to analyze how small initialization enhances the generalization of LLMs.
Recall that a 2-layer net $f:\R^d\to\R$ with 1-homogeneous activation (such as ReLU) can be written compactly as $f_{\pi}(\x):=\int \sigma(\w\cdot\x) d\pi(\w)$ for some finite signed measure $\pi$ over the parameter space $\R^d$ (without loss of generality, we can replace $\x$ by $[\x,1]$ to account for the bias) \cite{e_machine_2019,e2022representation}.
There are two common functional spaces for these functions.
In the kernel regime \cite[Theorem 3.9]{e2020comparative}, the function space is characterized by the RKHS norm \cite{rahimi2008uniform}
\begin{equation*}
\|f\|_{\mathcal{H}} = \inf_{\pi} \Big\| \frac{\delta\pi}{\delta\pi_0} \Big\|_{L^2(\pi_0)} = \Big(\int \Big|\frac{\delta\pi}{\delta\pi_0}(\w)\Big|^2 d\pi_0(\w)\Big)^{1/2}, \quad \text{s.t.}~f=f_{\pi} ~\text{and}~ \pi \ll \pi_0
\end{equation*}
where $\pi_0$ is some base distribution (typically the initialization distribution of $\w$) and $\delta\pi/\delta\pi_0$ is the Radon-Nikodym derivative.
% So the only feasible $\pi$'s are expressible by $a\pi_0$ for some weight function $a$.
In the mean-field regime \cite{mei_mean_2018,rotskoff_parameters_2018}, the function space is characterized by the Barron norm \cite{e2022barron,e2022representation}
\begin{equation*}
\|f\|_{\mathcal{B}} = \inf_{\pi} \|\pi\|_{\text{TV}} = \inf_{\pi} \int d|\pi|(\w), \quad \text{s.t.}~f=f_{\pi}~\text{and}~\text{supp}(\pi)\subseteq\mathbb{S}^d
\end{equation*}
where $\|\cdot\|_{\text{TV}}$ denotes total variation and $\pi$ is supported on the unit sphere $\mathbb{S}^d$.
Naturally, we can define an interpolation between these norms: for any $\gamma > -3/2$,
\begin{align*}
\|f\|_{\gamma}&:= \inf_{\pi} \Big\| \frac{\delta\pi}{\delta\pi_0} \Big\|_{L^{3+2\gamma}(\pi_0)}, \quad \text{s.t.}~f=f_{\pi}~\text{and}~\text{supp}(\pi)\subseteq\mathbb{S}^d
\end{align*}
% \Big(\int \Big|\frac{\delta\pi}{\delta\pi_0}(\w)\Big|^{3+2\gamma}d\pi_0(\w)\Big)^{1/(3+2\gamma)}
and by 1-homogeneity $\pi_0$ can be assumed to be supported on $\mathbb{S}^d$ as well.
One can check that $\|\cdot\|_{\gamma=-1/2}=\|\cdot\|_{\mathcal{H}}$ and $\|\cdot\|_{\gamma=-1}=\|\cdot\|_{\mathcal{B}}$.
By Jensen's inequality, the function space of $\|\cdot\|_{\gamma}$ becomes larger as $\gamma$ decreases.
In particular, if $\gamma \leq -1$, then $\pi$ does not have to be absolutely continuous with respect to $\pi_0$, i.e. it can develop singular parts, sometimes known as ``condensation" \cite{CSIAMchen}.
Else, $\gamma>-1$ and $\pi$ must have the form $a\pi_0$.
The question is how $\gamma$ affects deep networks.
Let $\mathcal{M}(\R^d,\R^k)$ denote $\R^k$-vector-valued finite measures over $\R^d$.
For simplicity, consider deep residual nets, $\x^l=\x^{l-1}+f_l(\x^{l-1})$ for $l=1, \dots L$, where each block $f_l:\R^d\to\R^d$ is a 2-layer net parametrized by some $\pi_l \in \mathcal{M}(\R^d,\R^d)$.
For any fixed sequence of input weights $(\w_l)_{l=1}^L \in \R^{L\times d}$, fixed $(\pi_l)$, and an input $\x^0$ drawn from some data distribution, the activation sequence $\big(\sigma(\w_l\cdot\x_l)\big)$ is a stochastic process with a potentially complicated dependency structure.
This is intended to capture cross-layer collaborations in LLMs or ``circuits", such as the induction head \cite{olsson2022incontext}, IOI circuit \cite{wang2023interpretability}, and arithmetic circuits \cite{lindsey2025biology}.
Therefore, we model the parameter distribution of the full network by $\pi\in \mathcal{M}(\R^{Ld},\R^{Ld})$, instead of just the sequence $(\pi_l)$, in order to model the dependency among the activated parameters.
A generalization of the norm $\|\cdot\|_{\gamma}$ from 2-layer nets to deep residual nets can be defined by
\begin{equation*}
\|f\|_{\gamma} = \inf_{\pi} \Big\| \frac{\delta\pi}{\delta\pi_0^{\otimes L}} \Big\|_{L^{3+2\gamma}(\pi_0^{\otimes L})}
= \inf_{\pi} \Big(\int \Big\|\frac{\delta\pi}{\delta\pi_0^{\otimes L}}(\oplus_{l=1}^L\mathbf{w}_l)\Big\|^{3+2\gamma} \prod_{l=1}^L d\pi_0(\mathbf{w}_l) \Big)^{1/(3+2\gamma)}
\end{equation*}
which ranges among all parametrizations $(\pi_l)$ of $f$ and the resulting dependency $\pi$ with respect to the data distribution.
To study $\|\cdot\|_{\gamma}$, we make the following assumptions:
\begin{enumerate}
\item Among the approximate global minimizers of the loss, training initialized with rate $\gamma$ for any $\gamma>-3/2$ always converges to a minimizer $f^*$ with the minimum $\|\cdot\|_{\gamma}$ norm.
This assumption holds for 2-layer nets in the kernel regime ($\gamma=-1/2$) when trained by gradient descent \cite{zhang2021understanding}, and we expect similar behavior in general settings.
\item The parameter distribution $\pi^*$ of each approximate global minimizer $f^*$ can always be decomposed into a weighted sum of product measures
$$\pi^*=\sum_i c_i \pi^i, \quad c_i \in \R, \quad \pi^i=\bigotimes_{l=1}^L \pi^i_l, \quad \|\pi_l^i\|_{\text{TV}}=1$$
and these $\pi^i$ have disjoint supports (up to $\pi$-negligible subsets).
Furthermore, the variation of each $\pi^i$, namely $|\pi^i_l|$, is simply $\pi_0$ restricted to some subset $S^i_l\subseteq\R^d$, i.e. $|\pi^i_l| = \pi_0 \mathbf{1}_{S^i_l} / \pi_0(S^i_l)$, and there exists some constant $0<\epsilon\ll 1$ such that either $\pi_0(S^i_l)=\epsilon$ or $\pi_0(S^i_l) = 0.9$.
The purpose of this assumption is to simplify computation, and we expect our results to hold in much more general settings.
\end{enumerate}
For each $\pi^i$, denote by $L_i$ the number of $\pi^i_l$ with $\pi_0(S^i_l)=\epsilon$.
This $L_i$ can be interpreted as the circuit depth, the number of layers where $\pi^i$ is non-trivial.
The norm of each minimizer $f^*$ becomes
\begin{align*}
\|f^*\|_{\gamma} &= \Big(\sum_i \int \Big\|\frac{\delta(c_i\pi^i)}{\delta\pi_0^{\otimes L}}\Big\|^{3+2\gamma} d\pi_0^{\otimes L} \Big)^{1/(3+2\gamma)} = \Big(\sum_i c_i^{3+2\gamma} \prod_{l=1}^L \int \Big\|\frac{\delta\pi^i_l}{\delta\pi_0}\Big\|^{3+2\gamma} d\pi_0 \Big)^{1/(3+2\gamma)}\\
&= \Big(\sum_i \big( c_i \epsilon^{-L_i} 0.9^{-(L-L_i)} \big)^{3+2\gamma} \Big)^{1/(3+2\gamma)} = 0.9^{-L} \big\| c_i (0.9/\epsilon)^{L_i} \big\|_{l^{3+2\gamma}}
\end{align*}
Hence, each minimizer can be viewed as a circuit ensemble, characterized by its circuit weights and circuit depths $\{(c_i, L_i)\}$.
Recall that for general $l^p$ norms, those with large $p$ are sensitive to the maximum value, while those with small $p$ are sensitive to the amount of non-zero elements.
Thus, for large $\gamma$ such as $-1/2$, the circuit ensemble chosen by training tends to be dense but uniformly shallow (many $c_i>0$ but small $\max L_i$), whereas for small $\gamma$ such as $-1$, the chosen circuits are more likely to be sparse and deep (few $c_i>0$ but $L_i$ can be large).
Intuitively, sparse and deep circuits signify a generalizable solution, as the model has managed to compress the diverse training data into very few but flexible patterns.
Meanwhile, dense and shallow circuits may imply that the model does not have a deep understanding of the data and has to memorize a lot.
In conclusion, it is plausible that small initialization increases generalizability by shifting the preference for circuit ensembles.
Although our calculation is based on simple residual networks, it is reasonable to expect that a similar mechanism applies to Transformers.
\begin{ack}
This work is supported by the National Key R\&D Program of China Grant No. 2022YFA1008200, the National Natural Science Foundation of China Grant No. 92270001, 12371511, 12422119.
% , Shanghai Municipal of Science and Technology Major Project No. 2021SHZDZX0102 (Z. X.), and the HPC of School of Mathematical Sciences and the Student Innovation Center, and the Siyuan-1 cluster supported by the Center for High Performance Computing at Shanghai Jiao Tong University, and Key Laboratory of Marine Intelligent Equipment and System (Ministry of Education, P.R. China), and SJTU Kunpeng \& Ascend Center of Excellence.
% Use unnumbered first level headings for the acknowledgments. All acknowledgments
% go at the end of the paper before the list of references. Moreover, you are required to declare
% funding (financial activities supporting the submitted work) and competing interests (related financial activities outside the submitted work).
% More information about this disclosure can be found at: \url{https://neurips.cc/Conferences/2025/PaperInformation/FundingDisclosure}.
% Do {\bf not} include this section in the anonymized submission, only in the final paper. You can use the \texttt{ack} environment provided in the style file to automatically hide this section in the anonymized submission.
\end{ack}
% \section*{References}
% References follow the acknowledgments in the camera-ready paper. Use unnumbered first-level heading for
% the references. Any choice of citation style is acceptable as long as you are
% consistent. It is permissible to reduce the font size to \verb+small+ (9 point)
% when listing the references.
% Note that the Reference section does not count towards the page limit.
\bibliographystyle{plain}
% \bibliographystyle{IEEEtran}
\bibliography{ref}
% \medskip
% {
% \small
% [1] Alexander, J.A.\ \& Mozer, M.C.\ (1995) Template-based algorithms for
% connectionist rule extraction. In G.\ Tesauro, D.S.\ Touretzky and T.K.\ Leen
% (eds.), {\it Advances in Neural Information Processing Systems 7},
% pp.\ 609--616. Cambridge, MA: MIT Press.
% [2] Bower, J.M.\ \& Beeman, D.\ (1995) {\it The Book of GENESIS: Exploring
% Realistic Neural Models with the GEneral NEural SImulation System.} New York:
% TELOS/Springer--Verlag.
% [3] Hasselmo, M.E., Schnell, E.\ \& Barkai, E.\ (1995) Dynamics of learning and
% recall at excitatory recurrent synapses and cholinergic modulation in rat
% hippocampal region CA3. {\it Journal of Neuroscience} {\bf 15}(7):5249-5262.
% }
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newpage
\appendix
% \section{Technical Appendices and Supplementary Material}
% Technical appendices with additional results, figures, graphs and proofs may be submitted with the paper submission before the full submission deadline (see above), or as a separate PDF in the ZIP file below before the supplementary material deadline. There is no page limit for the technical appendices.
\section{Limitation and Future Work}\label{app:limitation}
While our methodology received thorough empirical validation, scaling to larger models and datasets remains constrained due to computational resource limitations. Future work will prioritize extending the framework to a larger training scale while exploring complexity control mechanisms during post-training.
\section{Experiments Setup}\label{app:setup}
% 去掉硬件
\paragraph{Model architecture}
Our models are based on the Llama-architecture~\cite{touvron2023llama}. Architectural configurations for different model scales are presented in Table~\ref{tab:compact_config}. For the 2.4B variant, we incorporate embedding normalization and sandwich normalization techniques to enhance training stability.
\begin{table}[htbp]
\centering
\caption{Training Configuration Across Model Scales}
\label{tab:compact_config}
\begin{tabular}{@{}lcccc@{}}
\toprule
\textbf{No. } & \textbf{Model Scale} & \textbf{Architecture Configuration} &\textbf{Pre-train Datasize}\\
\midrule
1 & 180M &
\begin{tabular}[c]{@{}c@{}}Layers: 16\\ Head Dim: 80\\ Heads (KV): 16 (16)\end{tabular} & 40B \\
\midrule
2 & 0.9B &
\begin{tabular}[c]{@{}c@{}}Layers: 32\\ Head Dim: 64\\ Heads (KV): 32 (32)\end{tabular} & 600B\\
\midrule
3 & 2.4B &
\begin{tabular}[c]{@{}c@{}}Layers: 44\\ Head Dim: 80\\ Heads (KV): 40 (8)\end{tabular} &1T \\
\midrule
\bottomrule
\end{tabular}
\end{table}
% The results in Figure~\ref{fig:token_loss},~\ref{fig:paranorm}, Section~\ref{sec:analysis-187} and Appendix~\ref{app:187} are based on Configuration 1. The results in Figure~\ref{fig:base_sft},~\ref{fig:Dc_0.94},~\ref{fig:Ds_0.94}, and the columns "0.94B Large", "0.94B Small" in Table~\ref{tab:model_benchmark}, ~\ref{tab:base_benchmark} are based on Configuration 2. The results in Table~\ref{tab:model_benchmark},~\ref{tab:base_benchmark} "2.4B Large", "2,4B Small" and Figure~\ref{fig:para-score-evoluation},~\ref{fig:Dc_2.4},~\ref{fig:Ds_2.4}, Appendix~\ref{app:0.58-1} are based on the configuration 3. For Figure~\ref{fig:scalinglaw}, the
% model configurations adopted are provided in Table~\ref{tab:scaling_config}.
These configurations are systematically mapped as follows:
\begin{itemize}
\item \textbf{Configuration 1} governs results in:
\begin{itemize}
\item Figure~\ref{fig:token_loss} and~\ref{fig:paranorm}
\item Section~\ref{sec:analysis-187}
\item Appendix~\ref{app:187}
\end{itemize}
\item \textbf{Configuration 2} applies to:
\begin{itemize}
\item Figure~\ref{fig:base_sft},~\ref{fig:Dc_0.94} and~\ref{fig:Ds_0.94}
\item Columns ``0.9B Large'' and ``0.9B Small'' in Table~\ref{tab:model_benchmark},~\ref{tab:base_benchmark}
\end{itemize}
\item \textbf{Configuration 3} encompasses:
\begin{itemize}
\item Columns ``2.4B Large'' and ``2.4B Small'' in Table~\ref{tab:model_benchmark} and~\ref{tab:base_benchmark}
\item Figure~\ref{fig:para-score-evoluation},~\ref{fig:Dc_2.4} and~\ref{fig:Ds_2.4}
\item Appendix~\ref{app:0.58-1}
\end{itemize}
\end{itemize}
For Figure~\ref{fig:scalinglaw}, the corresponding model specifications are detailed in Table~\ref{tab:scaling_config}.
\begin{table}[htbp]
\centering
\caption{Architecture Configuration adopted in Figure~\ref{fig:scalinglaw}.}
\label{tab:scaling_config}
\begin{tabular}{@{}lcccc@{}}
\toprule
\textbf{Model Scale (M)} & \textbf{Layers} &\textbf{Heads (KV)} &\textbf{Head Dim}\\
\midrule
50 & 12 & 12 & 64 \\
\midrule
100 & 16 & 16 & 60\\
\midrule
200 & 24 & 18 & 60 \\
\midrule
400 & 24 & 24 & 64 \\
\midrule
800 & 32 & 24 & 80 \\
\bottomrule
\end{tabular}
\end{table}
\paragraph{DataSet}
The 180M models are trained on a carefully curated 40B subset obtained through uniform sampling from the Memory$^3$ training corpus. The 0.9B models undergo pre-training on a 600B subset derived from the SlimPajama corpus. Building upon the Memory$^3$ foundation, we expand the training corpus through a rigorous data processing pipeline that integrates deduplication protocols, multi-dimensional quality assessment metrics, and optimized domain ratio adjustments. This systematic curation process yields a refined 1T training dataset, which is the basis for training our 2.4B models. For the SFT data, we also adopt the SFT data from Memory$^3$.
\paragraph{Training setup}
The training is implemented using Microsoft's Megatron-DeepSpeed framework~\cite{Megatron-DeepSpeed}, utilizing a mixed-precision configuration where model parameters, gradients, and activations were maintained in bfloat16 format while preserving optimizer states in float32 precision for the AdamW optimizer. The learning rate schedule adopted a cosine annealing strategy with linear warmup, where the warmup phase spanned the initial 5\% of the total training iterations. The learning rate boundaries were configured with maximum and minimum values of $1\times10^{-3}$ and $1 \times 10^{-5}$.
\paragraph{Training cost}
We present the computational costs for single training sessions across different model scales. The 180M model required 32 MX-C500 accelerators with a training duration of 12 hours. For the 0.9B architecture, the training process utilized 400 MX-C500 accelerators over 72 hours. 512 MX-C500 accelerators are employed for training a 2.4B model in one week.
\paragraph{Evaluation Details}
Our model are assessed across multiple open-source benchmarks via lm-eval-harness~\cite{eval-harness}, covering areas such as factual knowledge: CMMLU~\cite{li2023cmmlu}, C-Eval~\cite{huang2023ceval}, MMLU~\cite{mmlu} and its enhanced version MMLU-Pro~\cite{mmlu_pro}, OpenBookQA~\cite{OpenBookQA2018}, and GPQA~\cite{rein2024gpqa}. Language comprehension: BBH~\cite{bbh}, ARC-C~\cite{arc}, TruthfulQA~\cite{lin2021truthfulqa}, WinoGrande~\cite{ai2:winogrande}, HellaSwag~\cite{zellers2019hellaswag}, AGIEval-EN~\cite{zhong2023agieval}, CommonsenseQA~\cite{talmor-etal-2019-commonsenseqa}. Code generation: MBPP~\cite{mbpp} and IFEval~\cite{ifeval}. Mathematical reasoning: MATH~\cite{MATH} and GSM8K~\cite{cobbe2021gsm8k}.
\section{Evaluation of base models}\label{app:base}
Table~\ref{tab:base_benchmark} presents the evaluation results of the base models. The results reveal that the 0.9B model with small complexity does not demonstrate a significant advantage, while the 2.4B small-complexity model achieves measurable improvements. This disparity potentially stems from differences in pre-training data quality. Specifically, the 0.9B model's training data (subsampled from the SlimPajama corpus) exhibits suboptimal quality, limiting its downstream task adaptability.
\begin{table}[htbp]
\centering
\vspace{-5pt}
\begin{threeparttable}
\caption{Evaluation of base models}
\label{tab:base_benchmark}
\begin{tabular}{@{}lccccc@{}}
\toprule
\textbf{Models\tnote{1}} & \textbf{0.9B Large} & \textbf{0.9B Small} & \textbf{2.4B Large} & \textbf{2.4B Small}\\
$\left(\gamma,\lambda\right)$ & (0.5,0.1) &(1,1) &(0.5,0.1) & (0.58,1)\\
\midrule
% \multicolumn{5}{c}{\textit{General Tasks}} \\
% \midrule
MMLU & 25.7 & 24.4 (\mcolor{-1.3}) & 36.5 & 41.1 (\ppcolor{+4.6}) \\
MMLU-Pro & 11.8 & 11.3 (\mcolor{-0.5}) & 13.4 & 14.9 (\pcolor{+1.5}) \\
BBH & 31.1 & 31.3 (\pcolor{+0.2}) & 34.7 & 35.7 (\pcolor{+1.0}) \\
ARC-C & 34.6 & 35.6 (\pcolor{+1}) & 39.9 & 47.0 (\ppcolor{+7.1}) \\
TruthfulQA & 36.6 & 36.8 (\pcolor{+0.2}) & 36.6 & 41.1 (\ppcolor{+4.5}) \\
WinoGrande & 59.4 & 61.1 (\pcolor{+1.7}) & 64.0 & 68.0 (\ppcolor{+4.0}) \\
HellaSwag & 59.4 & 62.0 (\pcolor{+2.6}) & 66.5 & 67.8 (\pcolor{+1.3}) \\
AGIEval-EN & 17.8 & 17.8 & 20.5 & 19.9 (\mcolor{-0.6})\\
OpenBookQA &39 &37.8 (\mcolor{-1.2}) & 42.2 & 41.6 (\mcolor{-0.6})\\
CommonsenseQA &19 &20.9 (\pcolor{+1.9}) & 38.2 & 48.7 (\ppcolor{+10.5}) \\
GPQA & 23.7 & 24.6 (\pcolor{+0.9}) & 24.6 & 24.8 (\pcolor{+0.2}) \\
MATH & 2.0 & 2.5 (\pcolor{+0.5}) & 30.0 & 29.4 (\mcolor{-0.6}) \\
GSM8K & 1.9 & 2.4 (\pcolor{+0.5}) & 31.0 & 39.7 (\ppcolor{+8.7}) \\
MBPP & 0.2 & 6.2 (\ppcolor{+6}) & 19.2 & 22 (\pcolor{+2.8}) \\
IFEval &27.2 & 27.7 (\pcolor{+0.5}) & 27.0 & 28.4 (\pcolor{+1.4})\\
\bottomrule
\end{tabular}
\begin{tablenotes}
\footnotesize
\item[1] ``Large'' and ``Small'' mean large complexity and small complexity, respectively.
\end{tablenotes}
\end{threeparttable}
\end{table}
\newpage
\section{Training Stability}\label{app:0.58-1}
We observe that controlling model complexity induces training instability with loss spikes as the model scale increases, particularly when employing extremely small initialization scales. The purple lines in Figure~\ref{fig:58-1} demonstrate the training dynamics (loss and parameter norm evolution) of the 2.4B model with $\gamma=1, \lambda=1$, revealing severe instability that prevents convergence to small-complexity solutions. Contrastingly, Figures~\ref{fig:paranorm},~\ref{fig:187-eval} establish that sufficiently large $\lambda$ values ($\lambda=1$) enable small-complexity convergence regardless of initialization scale. This insight motivates our strategy: combining large $\lambda$ with critical initialization scaling. Inspired by the initialization practices of GPT-2 and DeepSeek-V3, we adopt $\gamma=0.58$ in our 2.4B model's training. As shown in Figure~\ref{fig:58-1} Table 7, this configuration achieves stabilized training and better performance.
\begin{figure}[htpb]
\centering
\includegraphics[width=1\linewidth]{ pic/appendix/58-1_v2.png}
\caption{The dynamics of parameter norm (left) and loss (right) of 2.4B model with $\gamma=1,\lambda=1$ (purple) and $\gamma=0.58,\lambda=1$ (orange).}
\label{fig:58-1}
\end{figure}
\begin{table}[htbp]
\centering
\caption{Evaluation of 2.4B models with $\gamma=1$ and $\gamma=0.58$.}
\label{tab:0.58-1}
\begin{tabular}{@{}lccc@{}}
\toprule
\makebox[0.1\textwidth]{\textbf{$\gamma$}} & \makebox[0.1\textwidth][c]{1} & \makebox[0.1\textwidth][c]{0.58}\\
\midrule
MMLU & 65.4 & 64.4 \\
MMLU-Pro & 30.1 & 30.1 \\
BBH & 45.0 & 43.7 \\
ARC-C & 59.0 & 60.7 \\
TruthfulQA & 64.1 & 61.7 \\
WinoGrande & 78.0 & 76.9 \\
HellaSwag & 74.6 & 75.4 \\
AGIEval-EN & 32.2 & 35.6\\
OpenBookQA &40.8 & 43.4 \\
CommonsenseQA &77.7 & 76.9 \\
GPQA & 29.5 & 31.7 \\
MATH & 33.3 & 35.3 \\
GSM8K & 62.0 & 63.8 \\
MBPP & 21.6 & 22.4 \\
IFEval & 42.2 & 40.9\\
\midrule
Average & 50.4 & 50.9\\
\bottomrule
\end{tabular}
\end{table}
\newpage
\section{Evaluation of the 180M models}\label{app:187}
Table~\ref{tab:187M_results} comprehensively evaluates all 180M models referenced in Section~\ref{sec:187m}, systematically illustrating performance improvements through complexity control. Additionally, Figure~\ref{fig:187_eval_appendix1}, \ref{fig:187_eval_appendix2} visualizes the performance-complexity correlation across all tasks, further validating our conclusions.
\begin{table}[htbp]
\centering
\caption{Evaluation results of the 180M models under varying $\gamma-\lambda$ configurations}
\label{tab:187M_results}
\footnotesize % 缩小字体以适应页面
\begin{tabular}{ccccccccccc}
\toprule
\multirow{3}{*}{$\lambda$} & \multirow{3}{*}{$\gamma$} & \multirow{3}{*}{Average} & \multicolumn{8}{c}{Tasks} \\
\cmidrule(lr){4-11}
& & & ARC-C & GSM8K & TruthfulQA & WinoGrande & HellaSwag & MMLU & CMMLU & C-EVAL\\
\midrule
\multirow{5}{*}{0}
& 0.1 & 30.0 & 25.5 & 0.5 & 52.7 & 51.6 & 25.3 & 24.4 & 25.6 & 25.5 \\
& 0.3 & 29.1 & 22.6 & 1.4 & 50.4 & 48.3 & 26.2 & 25.6 & 25.5 & 27.5 \\
& 0.5 & 31.9 & 26.1 & 2.3 & 47.4 & 51.5 & 33.4 & 30.8 & 28.8 & 32.4 \\
& 0.8 & 31.4 & 23.7 & 1.7 & 47.3 & 50.4 & 34.4 & 31.1 & 29.5 & 31.1 \\
& 1.0 & 31.8 & 26.1 & 2.4 & 46.7 & 52.6 & 33.3 & 29.5 & 28.0 & 31.3 \\
\cmidrule(lr){1-11}
\multirow{5}{*}{0.1}
& 0.1 & 30.0 & 22.9 & 1.3 & 51.3 & 51.1 & 27.0 & 26.1 & 25.4 & 27.7 \\
& 0.3 & 31.1 & 25.9 & 1.7 & 48.7 & 50.9 & 30.2 & 29.4 & 27.6 & 29.9 \\
& 0.5 & 32.2 & 27.9 & 2.1 & 48.5 & 50.4 & 34.0 & 30.4 & 30.5 & 34.3 \\
& 0.8 & 32.4 & 25.4 & 2.7 & 48.6 & 51.5 & 34.6 & 31.9 & 31.1 & 35.2 \\
& 1.0 & 32.7 & 27.8 & 2.8 & 49.2 & 51.9 & 34.3 & 30.3 & 30.1 & 34.8 \\
\cmidrule(lr){1-11}
\multirow{5}{*}{1.0}
& 0.1 & 35.4 & 31.1 & 6.1 & 51.1 & 51.6 & 36.1 & 36.2 & 36.9 & 40.3 \\
& 0.3 & 35.8 & 30.0 & 5.8 & 50.6 & 55.4 & 36.9 & 36.1 & 37.1 & 37.7 \\
& 0.5 & 35.5 & 28.8 & 5.4 & 51.2 & 55.1 & 36.7 & 35.5 & 36.3 & 38.7 \\
& 0.8 & 35.5 & 29.4 & 7.2 & 47.9 & 53.8 & 37.3 & 37.4 & 38.3 & 41.1 \\
& 1.0 & 35.2 & 30.1 & 7.2 & 50.8 & 51.2 & 36.3 & 35.2 & 36.6 & 40.3 \\
\bottomrule
\end{tabular}
\end{table}
\begin{figure}[htpb]
\centering
\includegraphics[width=0.8\linewidth]{ pic/appendix/187_eval_heatmap_appendix.png}
\caption{Performance landscape of all tasks across $\gamma-\lambda$ of the 180M models. The color indicates the score (dark: low, light: high).}
\label{fig:187_eval_appendix1}
\end{figure}
\begin{figure}[htpb]
\centering
\includegraphics[width=0.8\linewidth]{ pic/appendix/187_eval_heatmap_appendix2.png}
\caption{Score-complexity relationship with dashed lines denoting baseline performance levels.}
\label{fig:187_eval_appendix2}
\end{figure}
\section{Model Analysis}\label{app:parameter}
\subsection{Embedding Structure of 0.9B and 2.4B models}
Figure~\ref{fig:embedding_appendix} depicts the cosine similarity among 350 embedding vectors of 0.9B models and 2.4B models, with different model complexities. The results present a similar phenomenon with Figure~\ref{fig:187-embedding}, demonstrating that complexity control contributes to a focus on the association among different tokens.
\begin{figure}[htpb]
\centering
\includegraphics[width=1\linewidth]{ pic/appendix/model_analysis/embedding_appendix.png}
\caption{Cosine similarity among 350 embedding vectors which occur most frequently in training dataset under different complexities of 0.9B models and 2.4B models. The "Large" and "small" mean the large complexity and small complexity.}
\label{fig:embedding_appendix}
\end{figure}
\subsection{Attention module of 0.9B and 2.4B models}
Figure~\ref{fig:Dc_0.94},~\ref{fig:Ds_0.94},~\ref{fig:Dc_2.4}, and~\ref{fig:Ds_2.4} exhibit the $D_c$ and $D_s$ of query projection matrices and key projection matrices in the 0.9B models and 2.4B models, reveals a condensation and low-rank trend of the small-complexity models.
\begin{figure}[htpb]
\centering
\includegraphics[width=0.9\linewidth]{ pic/appendix/model_analysis/Dc09B.png}
\caption{$D_c$ of $\vW_Q$ and $\vW_K$ in 0.9B model's each layer under different model complexity configurations.}
\label{fig:Dc_0.94}
\end{figure}
\begin{figure}[htpb]
\centering
\includegraphics[width=0.9\linewidth]{ pic/appendix/model_analysis/Ds09B.png}
\caption{$D_s$ of $\vW_Q$ and $\vW_K$ in 0.9B model's each layer under different model complexity configurations.}
\label{fig:Ds_0.94}
\end{figure}
\begin{figure}[htpb]
\centering
\includegraphics[width=0.9\linewidth]{ pic/appendix/model_analysis/Dc24B.png}
\caption{$D_c$ of $\vW_Q$ and $\vW_K$ in 2.4B model's each layer under different model complexity configurations.}
\label{fig:Dc_2.4}
\end{figure}
\begin{figure}[htpb]
\centering
\includegraphics[width=0.9\linewidth]{ pic/appendix/model_analysis/Ds24B.png}
\caption{$D_s$ of $\vW_Q$ and $\vW_K$ in 2.4B model's each layer under different model complexity configurations.}
\label{fig:Ds_2.4}
\end{figure}
% \begin{table}[htbp]
% \centering
% \caption{Evaluation of base models}
% \label{tab:base_benchmark}
% \begin{tabular}{@{}lccccc@{}}
% \toprule
% \textbf{Models} & \textbf{0.94B Large} & \textbf{0.94B Small} & \textbf{2.4B Large} & \textbf{2.4B Small}\\
% $\left(\gamma,\lambda\right)$ & (0.5,0.1) &(1,1) &(0.5,0.1) & (0.58,1)\\
% \midrule
% % \multicolumn{5}{c}{\textit{General Tasks}} \\
% % \midrule
% MMLU & 25.7 & 24.4 (\mcolor{-1.3}) & 32.3 & 41.1 (\ppcolor{+8.8}) \\
% MMLU-Pro & 11.8 & 11.3 (\mcolor{-0.5}) & 12.6 & 14.9 (\pcolor{+2.3}) \\
% BBH & 31.1 & 31.3 (\pcolor{+0.2}) & 32.2 & 35.7 (\pcolor{+3.5}) \\
% ARC-C & 34.6 & 35.6 (\pcolor{+1}) & 44.0 & 47.0 (\pcolor{+3}) \\
% TruthfulQA & 36.6 & 36.8 (\pcolor{+0.2}) & 35.3 & 41.1 (\ppcolor{+5.8}) \\
% WindoGrande & 59.4 & 61.1 (\pcolor{+1.7}) & 65.6 & 68.0 (\pcolor{+2.4}) \\
% HellaSwag & 59.4 & 62.0 (\pcolor{+2.6}) & 66.7 & 67.8 (\pcolor{+1.1}) \\
% AGIEval-EN & 17.8 & 17.8 & 20.8 & 19.9 (\mcolor{-0.9})\\
% OpenBookQA &39 &37.8 (\mcolor{-1.2}) & 38.2 & 41.6 (\pcolor{+3.4})\\
% CommonsenseQA &19 &20.9 (\pcolor{+1.9}) & 26.8 & 48.7 (\ppcolor{+21.9}) \\
% GPQA & 23.7 & 24.6 (\pcolor{+0.9}) & 24.3 & 24.5 (\pcolor{+0.2}) \\
% MATH & 2.0 & 2.5 (\pcolor{+0.5}) & 31.2 & 29.4 (\mcolor{-1.8}) \\
% GSM8K & 1.9 & 2.4 (\pcolor{+0.5}) & 33.0 & 39.7 (\ppcolor{+6.7}) \\
% MBPP & 0.2 & 6.2 (\ppcolor{+6}) & 21.6 & 22 (\pcolor{+0.4}) \\
% IFEval &31.3 & 27.7 (\mcolor{-3.6}) & 27.0 & 28.4 (\pcolor{+1.4})\\
% \bottomrule
% \end{tabular}
% \end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\end{document}
\begin{thebibliography}{10}
\bibitem{Megatron-DeepSpeed}
Megatron-deepspeed.
\newblock \url{https://github.com/microsoft/Megatron-DeepSpeed }, 2022.
\bibitem{achiam2023gpt}
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia~Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et~al.
\newblock {GPT}-4 technical report.
\newblock {\em arXiv preprint arXiv:2303.08774}, 2023.
\bibitem{allenzhu2023physics}
Zeyuan Allen-Zhu and Yuanzhi Li.
\newblock Physics of language models: Part 3.2, knowledge manipulation.
\newblock In {\em arXiv preprint: abs/2309.14402}, 2023.
\bibitem{arora2019exact}
Sanjeev Arora, Simon~S Du, Wei Hu, Zhiyuan Li, Russ~R Salakhutdinov, and Ruosong Wang.
\newblock On exact computation with an infinitely wide neural net.
\newblock In {\em Advances in Neural Information Processing Systems}, pages 8141--8150, 2019.
\bibitem{arora2018stronger}
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi~Zhang.
\newblock Stronger generalization bounds for deep nets via a compression approach.
\newblock In {\em International conference on machine learning}, pages 254--263. PMLR, 2018.
\bibitem{mbpp}
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et~al.
\newblock Program synthesis with large language models.
\newblock {\em arXiv preprint arXiv:2108.07732}, 2021.
\bibitem{bartlett2017spectrally}
Peter~L Bartlett, Dylan~J Foster, and Matus~J Telgarsky.
\newblock Spectrally-normalized margin bounds for neural networks.
\newblock {\em Advances in neural information processing systems}, 30, 2017.
\bibitem{bartlett2002rademacher}
Peter~L Bartlett and Shahar Mendelson.
\newblock Rademacher and gaussian complexities: Risk bounds and structural results.
\newblock {\em Journal of Machine Learning Research}, 3(Nov):463--482, 2002.
\bibitem{CSIAMchen}
Zheng-An Chen, Yuqing Li, Tao Luo, Zhangchen Zhou, and Zhi-Qin~John Xu.
\newblock Phase diagram of initial condensation for two-layer neural networks.
\newblock {\em CSIAM Transactions on Applied Mathematics}, 5(3):448--514, 2024.
\bibitem{chizat_global_2018}
Lenaic Chizat and Francis Bach.
\newblock On the global convergence of gradient descent for over-parameterized models using optimal transport.
\newblock {\em Advances in neural information processing systems}, 31, 2018.
\bibitem{arc}
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord.
\newblock Think you have solved question answering? try arc, the ai2 reasoning challenge.
\newblock {\em arXiv:1803.05457v1}, 2018.
\bibitem{cobbe2021gsm8k}
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman.
\newblock Training verifiers to solve math word problems.
\newblock {\em arXiv preprint arXiv:2110.14168}, 2021.
\bibitem{dziri2023faith}
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang~Lorraine Li, Liwei Jiang, Bill~Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan~Le Bras, Jena~D. Hwang, Soumya Sanyal, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi.
\newblock Faith and fate: Limits of transformers on compositionality.
\newblock In {\em Thirty-seventh Conference on Neural Information Processing Systems}, 2023.
\bibitem{e2020comparative}
Weinan E, Chao Ma, and Lei Wu.
\newblock A comparative analysis of optimization and generalization properties of two-layer neural network and random feature models under gradient descent dynamics.
\newblock {\em Sci. China Math.}, 63, 2020.
\bibitem{e_machine_2019}
Weinan E, Chao Ma, and Lei Wu.
\newblock Machine learning from a continuous viewpoint, {I}.
\newblock {\em Science China Mathematics}, 63(11):2233--2266, 2020.
\bibitem{e2022barron}
Weinan E, Chao Ma, and Lei Wu.
\newblock The {Barron} space and the flow-induced function spaces for neural network models.
\newblock {\em Constructive Approximation}, 55(1):369--406, 2022.
\bibitem{e2022representation}
Weinan E and Stephan Wojtowytsch.
\newblock Representation formulas and pointwise properties for {Barron} functions.
\newblock {\em Calculus of Variations and Partial Differential Equations}, 61(2):1--37, 2022.
\bibitem{feng2023towards}
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di~He, and Liwei Wang.
\newblock Towards revealing the mystery behind chain of thought: A theoretical perspective.
\newblock In {\em Thirty-seventh Conference on Neural Information Processing Systems}, 2023.
\bibitem{eval-harness}
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le~Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou.
\newblock The language model evaluation harness, 07 2024.
\bibitem{golowich2018size}
Noah Golowich, Alexander Rakhlin, and Ohad Shamir.
\newblock Size-independent sample complexity of neural networks.
\newblock In {\em Conference On Learning Theory}, pages 297--299. PMLR, 2018.
\bibitem{guo2025deepseek}
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et~al.
\newblock Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning.
\newblock {\em arXiv preprint arXiv:2501.12948}, 2025.
\bibitem{he2015delving}
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.
\newblock Delving deep into rectifiers: Surpassing human-level performance on imagenet classification.
\newblock In {\em Proceedings of the IEEE international conference on computer vision}, pages 1026--1034, 2015.
\bibitem{mmlu}
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt.
\newblock Measuring massive multitask language understanding.
\newblock {\em arXiv preprint arXiv:2009.03300}, 2020.
\bibitem{MATH}
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt.
\newblock Measuring mathematical problem solving with the math dataset.
\newblock {\em arXiv preprint arXiv:2103.03874}, 2021.
\bibitem{huang2020improving}
Xiao~Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs.
\newblock Improving transformer optimization through better initialization.
\newblock In {\em International Conference on Machine Learning}, pages 4475--4483. PMLR, 2020.
\bibitem{huang2023ceval}
Yuzhen Huang, Yuzhuo Bai, Zhihao Zhu, Junlei Zhang, Jinghan Zhang, Tangjun Su, Junteng Liu, Chuancheng Lv, Yikai Zhang, Jiayi Lei, Yao Fu, Maosong Sun, and Junxian He.
\newblock C-eval: A multi-level multi-discipline chinese evaluation suite for foundation models.
\newblock {\em arXiv preprint arXiv:2305.08322}, 2023.
\bibitem{jacot_neural_2018}
Arthur Jacot, Franck Gabriel, and Cl{\'e}ment Hongler.
\newblock Neural tangent kernel: Convergence and generalization in neural networks.
\newblock {\em Advances in neural information processing systems}, 31, 2018.
\bibitem{kassner2020pretrained}
Nora Kassner, Benno Krojer, and Hinrich Sch{\"u}tze.
\newblock Are pretrained language models symbolic reasoners over knowledge?
\newblock In {\em Proceedings of the 24th Conference on Computational Natural Language Learning}, pages 552--564, Online, November 2020. Association for Computational Linguistics.
\bibitem{krogh1991simple}
Anders Krogh and John Hertz.
\newblock A simple weight decay can improve generalization.
\newblock {\em Advances in neural information processing systems}, 4, 1991.
\bibitem{kumar2024training}
Aviral Kumar, Vincent Zhuang, Rishabh Agarwal, Yi~Su, John~D Co-Reyes, Avi Singh, Kate Baumli, Shariq Iqbal, Colton Bishop, Rebecca Roelofs, et~al.
\newblock Training language models to self-correct via reinforcement learning.
\newblock {\em arXiv preprint arXiv:2409.12917}, 2024.
\bibitem{LeCun1998}
Yann LeCun, Leon Bottou, Genevieve~B. Orr, and Klaus~Robert M{\"u}ller.
\newblock {\em Efficient BackProp}, pages 9--50.
\newblock Springer Berlin Heidelberg, Berlin, Heidelberg, 1998.
\bibitem{li2023cmmlu}
Haonan Li, Yixuan Zhang, Fajri Koto, Yifei Yang, Hai Zhao, Yeyun Gong, Nan Duan, and Timothy Baldwin.
\newblock Cmmlu: Measuring massive multitask language understanding in chinese, 2023.
\bibitem{li2024chain}
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma.
\newblock Chain of thought empowers transformers to solve inherently serial problems.
\newblock In {\em The Twelfth International Conference on Learning Representations}, 2024.
\bibitem{lin2021truthfulqa}
Stephanie Lin, Jacob Hilton, and Owain Evans.
\newblock Truthfulqa: Measuring how models mimic human falsehoods, 2021.
\bibitem{lindsey2025biology}
Jack Lindsey, Wes Gurnee, Emmanuel Ameisen, Brian Chen, Adam Pearce, Nicholas~L. Turner, Craig Citro, David Abrahams, Shan Carter, Basil Hosmer, Jonathan Marcus, Michael Sklar, Adly Templeton, Trenton Bricken, Callum McDougall, Hoagy Cunningham, Thomas Henighan, Adam Jermyn, Andy Jones, Andrew Persic, Zhenyi Qi, T.~Ben Thompson, Sam Zimmerman, Kelley Rivoire, Thomas Conerly, Chris Olah, and Joshua Batson.
\newblock On the biology of a large language model.
\newblock {\em Transformer Circuits Thread}, 2025.
\bibitem{liu2024deepseek}
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et~al.
\newblock Deepseek-v3 technical report.
\newblock {\em arXiv preprint arXiv:2412.19437}, 2024.
\bibitem{liu-etal-2023-crystal}
Jiacheng Liu, Ramakanth Pasunuru, Hannaneh Hajishirzi, Yejin Choi, and Asli Celikyilmaz.
\newblock Crystal: Introspective reasoners reinforced with self-feedback.
\newblock In Houda Bouamor, Juan Pino, and Kalika Bali, editors, {\em Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing}, pages 11557--11572, Singapore, December 2023. Association for Computational Linguistics.
\bibitem{liu2020understanding}
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han.
\newblock Understanding the difficulty of training transformers.
\newblock {\em arXiv preprint arXiv:2004.08249}, 2020.
\bibitem{luo2021phase}
Tao Luo, Zhi-Qin~John Xu, Zheng Ma, and Yaoyu Zhang.
\newblock Phase diagram for two-layer relu neural networks at infinite-width limit.
\newblock {\em Journal of Machine Learning Research}, 22(71):1--47, 2021.
\bibitem{mei_mean_2018}
Song Mei, Andrea Montanari, and Phan-Minh Nguyen.
\newblock A mean field view of the landscape of two-layer neural networks.
\newblock {\em Proceedings of the National Academy of Sciences}, 115(33):E7665--E7671, 2018.
\bibitem{OpenBookQA2018}
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal.
\newblock Can a suit of armor conduct electricity? a new dataset for open book question answering.
\newblock In {\em EMNLP}, 2018.
\bibitem{neyshabur2015norm}
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro.
\newblock Norm-based capacity control in neural networks.
\newblock In {\em Conference on learning theory}, pages 1376--1401. PMLR, 2015.
\bibitem{olsson2022incontext}
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah.
\newblock In-context learning and induction heads, 2022.
\bibitem{power2022grokking}
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra.
\newblock Grokking: Generalization beyond overfitting on small algorithmic datasets.
\newblock In {\em arXiv preprint: abs/2201.02177}, 2022.
\bibitem{press-etal-2023-measuring}
Ofir Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah Smith, and Mike Lewis.
\newblock Measuring and narrowing the compositionality gap in language models.
\newblock In Houda Bouamor, Juan Pino, and Kalika Bali, editors, {\em Findings of the Association for Computational Linguistics: EMNLP 2023}, pages 5687--5711, Singapore, December 2023. Association for Computational Linguistics.
\bibitem{prystawski2023why}
Ben Prystawski, Michael~Y. Li, and Noah Goodman.
\newblock Why think step by step? reasoning emerges from the locality of experience.
\newblock In {\em Thirty-seventh Conference on Neural Information Processing Systems}, 2023.
\bibitem{rae2021scaling}
Jack~W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et~al.
\newblock Scaling language models: Methods, analysis \& insights from training gopher.
\newblock {\em arXiv preprint arXiv:2112.11446}, 2021.
\bibitem{rahimi2008uniform}
Ali Rahimi and Benjamin Recht.
\newblock Uniform approximation of functions with random bases.
\newblock In {\em 2008 46th annual allerton conference on communication, control, and computing}, pages 555--561. IEEE, 2008.
\bibitem{rein2024gpqa}
David Rein, Betty~Li Hou, Asa~Cooper Stickland, Jackson Petty, Richard~Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel~R Bowman.
\newblock Gpqa: A graduate-level google-proof q\&a benchmark.
\newblock In {\em First Conference on Language Modeling}, 2024.
\bibitem{rotskoff_parameters_2018}
Grant Rotskoff and Eric Vanden-Eijnden.
\newblock Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks.
\newblock {\em Advances in neural information processing systems}, 31, 2018.
\bibitem{ai2:winogrande}
Keisuke Sakaguchi, Ronan~Le Bras, Chandra Bhagavatula, and Yejin Choi.
\newblock Winogrande: An adversarial winograd schema challenge at scale.
\newblock {\em Communications of the ACM}, 64(9):99--106, 2021.
\bibitem{sirignano_mean_2020}
Justin Sirignano and Konstantinos Spiliopoulos.
\newblock Mean field analysis of neural networks: {A} central limit theorem.
\newblock {\em Stochastic Processes and their Applications}, 130(3):1820--1852, 2020.
\bibitem{cerebras2023slimpajama}
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob~R Steeves, Joel Hestness, and Nolan Dey.
\newblock {SlimPajama: A 627B token cleaned and deduplicated version of RedPajama}.
\newblock \url{https://www.cerebras.net/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama}, June 2023.
\bibitem{sun2023recitationaugmented}
Zhiqing Sun, Xuezhi Wang, Yi~Tay, Yiming Yang, and Denny Zhou.
\newblock Recitation-augmented language models.
\newblock In {\em The Eleventh International Conference on Learning Representations}, 2023.
\bibitem{bbh}
Mirac Suzgun, Nathan Scales, Nathanael Sch{\"a}rli, Sebastian Gehrmann, Yi~Tay, Hyung~Won Chung, Aakanksha Chowdhery, Quoc~V Le, Ed~H Chi, Denny Zhou, , and Jason Wei.
\newblock Challenging big-bench tasks and whether chain-of-thought can solve them.
\newblock {\em arXiv preprint arXiv:2210.09261}, 2022.
\bibitem{talmor2020olmpics}
Alon Talmor, Yanai Elazar, Yoav Goldberg, and Jonathan Berant.
\newblock olmpics-on what language model pre-training captures.
\newblock {\em Transactions of the Association for Computational Linguistics}, 8:743--758, 2020.
\bibitem{talmor-etal-2019-commonsenseqa}
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant.
\newblock {C}ommonsense{QA}: A question answering challenge targeting commonsense knowledge.
\newblock In {\em Proceedings of the 2019 Conference of the North {A}merican Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)}, pages 4149--4158, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
\bibitem{touvron2023llama}
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian~Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit~Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric~Michael Smith, Ranjan Subramanian, Xiaoqing~Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian~Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas
Scialom.
\newblock Llama 2: Open foundation and fine-tuned chat models.
\newblock In {\em arXiv preprint: abs/2307.09288}, 2023.
\bibitem{trockman2023mimetic}
Asher Trockman and J~Zico Kolter.
\newblock Mimetic initialization of self-attention layers.
\newblock In {\em International Conference on Machine Learning}, pages 34456--34468. PMLR, 2023.
\bibitem{varma2023explaining}
Vikrant Varma, Rohin Shah, Zachary Kenton, János Kramár, and Ramana Kumar.
\newblock Explaining grokking through circuit efficiency.
\newblock In {\em arXiv preprint: abs/2309.02390}, 2023.
\bibitem{wang2024deepnet}
Hongyu Wang, Shuming Ma, Li~Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei.
\newblock Deepnet: Scaling transformers to 1,000 layers.
\newblock {\em IEEE Transactions on Pattern Analysis and Machine Intelligence}, 2024.
\bibitem{wang2023interpretability}
Kevin~Ro Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt.
\newblock Interpretability in the wild: a circuit for indirect object identification in {GPT}-2 small.
\newblock In {\em The Eleventh International Conference on Learning Representations}, 2023.
\bibitem{wang2024understanding}
Xinyi Wang, Alfonso Amayuelas, Kexun Zhang, Liangming Pan, Wenhu Chen, and William~Yang Wang.
\newblock Understanding the reasoning ability of language models from the perspective of reasoning paths aggregation.
\newblock In {\em arXiv preprint: abs/2402.03268}, 2024.
\bibitem{mmlu_pro}
Yubo Wang, Xueguang Ma, Ge~Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, Aaran Arulraj, Xuan He, Ziyan Jiang, et~al.
\newblock Mmlu-pro: A more robust and challenging multi-task language understanding benchmark.
\newblock In {\em The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track}, 2024.
\bibitem{wei2019data}
Colin Wei and Tengyu Ma.
\newblock Data-dependent sample complexity of deep neural networks via lipschitz augmentation.
\newblock {\em Advances in neural information processing systems}, 32, 2019.
\bibitem{wei2022emergent}
Jason Wei, Yi~Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed~H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus.
\newblock Emergent abilities of large language models.
\newblock {\em Transactions on Machine Learning Research}, 2022.
\newblock Survey Certification.
\bibitem{wei2022chain}
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed~Chi, Quoc~V Le, Denny Zhou, et~al.
\newblock Chain-of-thought prompting elicits reasoning in large language models.
\newblock {\em Advances in neural information processing systems}, 35:24824--24837, 2022.
\bibitem{williams_gradient_2019}
Francis Williams, Matthew Trager, Cl{\'{a}}udio~T. Silva, Daniele Panozzo, Denis Zorin, and Joan Bruna.
\newblock Gradient dynamics of shallow univariate relu networks.
\newblock {\em CoRR}, abs/1906.07842, 2019.
\bibitem{xu2025overview}
Zhi-Qin~John Xu, Yaoyu Zhang, and Zhangchen Zhou.
\newblock An overview of condensation phenomenon in deep learning.
\newblock In {\em arXiv preprint arXiv:2504.09484}, 2025.
\bibitem{yang2024qwen2}
An~Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo~Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et~al.
\newblock Qwen2. 5 technical report.
\newblock {\em arXiv preprint arXiv:2412.15115}, 2024.
\bibitem{yang2024memory3}
Hongkang Yang, Zehao Lin, Wenjin Wang, Hao Wu, Zhiyu Li, Bo~Tang, Wenqiang Wei, Jinbo Wang, Zeyun Tang, Shichao Song, Chenyang Xi, Yu~Yu, Kai Chen, Feiyu Xiong, Linpeng Tang, and Weinan E.
\newblock Memory$^3$: Language modeling with explicit memory.
\newblock {\em Journal of Machine Learning}, 3(3):300--346, 2024.
\bibitem{yang2024large}
Sohee Yang, Elena Gribovskaya, Nora Kassner, Mor Geva, and Sebastian Riedel.
\newblock Do large language models latently perform multi-hop reasoning?
\newblock In {\em arXiv preprint: abs/2402.16837}, 2024.
\bibitem{yao2025analysis}
Junjie Yao, Zhongwang Zhang, and Zhi-Qin~John Xu.
\newblock An analysis for reasoning bias of language models with small initialization.
\newblock {\em arXiv preprint arXiv:2502.04375}, 2025.
\bibitem{zelikman2024quietstar}
Eric Zelikman, Georges Harik, Yijia Shao, Varuna Jayasiri, Nick Haber, and Noah~D. Goodman.
\newblock Quiet-star: Language models can teach themselves to think before speaking.
\newblock In {\em arXiv preprint: abs/2403.09629}, 2024.
\bibitem{zelikman2022star}
Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman.
\newblock {ST}ar: Bootstrapping reasoning with reasoning.
\newblock In Alice~H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, {\em Advances in Neural Information Processing Systems}, 2022.
\bibitem{zellers2019hellaswag}
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi.
\newblock Hellaswag: Can a machine really finish your sentence?
\newblock In {\em Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics}, 2019.
\bibitem{zeng2021pangu}
Wei Zeng, Xiaozhe Ren, Teng Su, Hui Wang, Yi~Liao, Zhiwei Wang, Xin Jiang, ZhenZhang Yang, Kaisheng Wang, Xiaoda Zhang, et~al.
\newblock {PanGu}-$\alpha$: Large-scale autoregressive pretrained chinese language models with auto-parallel computation.
\newblock {\em arXiv preprint arXiv:2104.12369}, 2021.
\bibitem{zhang2019improving}
Biao Zhang, Ivan Titov, and Rico Sennrich.
\newblock Improving deep transformer with depth-scaled initialization and merged attention.
\newblock {\em arXiv preprint arXiv:1908.11365}, 2019.
\bibitem{zhang2021understanding}
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.
\newblock Understanding deep learning (still) requires rethinking generalization.
\newblock {\em Communications of the ACM}, 64(3):107--115, 2021.
\bibitem{zhang_type_2019}
Yaoyu Zhang, Zhi-Qin~John Xu, Tao Luo, and Zheng Ma.
\newblock A type of generalization error induced by initialization in deep neural networks.
\newblock {\em arXiv:1905.07777 [cs, stat]}, 2019.
\bibitem{zhang2022linear}
Yaoyu Zhang, Zhongwang Zhang, Leyang Zhang, Zhiwei Bai, Tao Luo, and Zhi-Qin~John Xu.
\newblock Linear stability hypothesis and rank stratification for nonlinear models.
\newblock {\em arXiv preprint arXiv:2211.11623}, 2022.
\bibitem{zhang2023stochastic}
Zhongwang Zhang, Yuqing Li, Tao Luo, and Zhi-Qin~John Xu.
\newblock Stochastic modified equations and dynamics of dropout algorithm.
\newblock {\em arXiv preprint arXiv:2305.15850}, 2023.
\bibitem{zhang2024initialization}
Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang, Yaoyu Zhang, and Zhi-Qin~John Xu.
\newblock Initialization is critical to whether transformers fit composite functions by reasoning or memorizing.
\newblock In {\em The Thirty-eighth Annual Conference on Neural Information Processing Systems}, 2024.
\bibitem{zhang2025complexity}
Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang, Yaoyu Zhang, and Zhi-Qin~John Xu.
\newblock Complexity control facilitates reasoning-based compositional generalization in transformers.
\newblock {\em arXiv preprint arXiv:2501.08537}, 2025.
\bibitem{zhang2023loss}
Zhongwang Zhang and Zhi-Qin~John Xu.
\newblock Loss spike in training neural networks.
\newblock {\em arXiv preprint arXiv:2305.12133}, 2023.
\bibitem{zhang2024implicit}
Zhongwang Zhang and Zhi-Qin~John Xu.
\newblock Implicit regularization of dropout.
\newblock {\em IEEE Transactions on Pattern Analysis and Machine Intelligence}, 2024.
\bibitem{zhong2023agieval}
Wanjun Zhong, Ruixiang Cui, Yiduo Guo, Yaobo Liang, Shuai Lu, Yanlin Wang, Amin Saied, Weizhu Chen, and Nan Duan.
\newblock Agieval: A human-centric benchmark for evaluating foundation models.
\newblock {\em arXiv preprint arXiv:2304.06364}, 2023.
\bibitem{zhou2022empirical}
Hanxu Zhou, Qixuan Zhou, Zhenyuan Jin, Tao Luo, Yaoyu Zhang, and Zhi-Qin~John Xu.
\newblock Empirical phase diagram for three-layer neural networks with infinite width.
\newblock {\em Advances in Neural Information Processing Systems}, 2022.
\bibitem{zhou2022towards}
Hanxu Zhou, Qixuan Zhou, Tao Luo, Yaoyu Zhang, and Zhi-Qin~John Xu.
\newblock Towards understanding the condensation of neural networks at initial training.
\newblock {\em Advances in Neural Information Processing Systems}, 2022.
\bibitem{ifeval}
Jeffrey Zhou, Tianjian Lu, Swaroop Mishra, Siddhartha Brahma, Sujoy Basu, Yi~Luan, Denny Zhou, and Le~Hou.
\newblock Instruction-following evaluation for large language models, 2023.
\bibitem{zhu2021gradinit}
Chen Zhu, Renkun Ni, Zheng Xu, Kezhi Kong, W~Ronny Huang, and Tom Goldstein.
\newblock Gradinit: Learning to initialize neural networks for stable and efficient training.
\newblock {\em Advances in Neural Information Processing Systems}, 34:16410--16422, 2021.
\end{thebibliography}