Simulating Training Spikes of Deep Neural Networks by Layer-wise Random Feature Models

Ran Tian; Courtney Paquette

cs.LGOct 9, 2026

Simulating Training Spikes of Deep Neural Networks by Layer-wise Random Feature Models

Authors: Ran Tian, Courtney Paquette

Organizations: Google Deepmind

Abstract

Intermittent ``spikes'' in loss, during the pre-training of Large Language Models (LLMs), remains a mysterious phenomenon both theoretically and empirically. In this paper, we show that simple quadratic models with heavy-tailed, power-law distributed inputs can replicate these spikes. We introduce a theoretical framework and simulation algorithm to model per-layer training dynamics of Deep Neural Networks (DNNs), using such quadratic models. In a case study on an Adam hyperparameter configuration (where $β_2$ is set close to $1$ ) known to induce spikes in Transformer LMs, our simulation successfully predicts bumps indicating spiky behavior from a given checkpoint. Furthermore, we derive an appropriate $β_2$ schedule from observed metrics, which empirically improves performance.

Figures & tables

Panel of Figure 1 — Figure 1: Training of a Transformer LM. Comparing a “Normal” training with two “Spiky” trainings. “Normal” and “Spiky1” differ in the $\beta_{2}$ hyperparameter of Adam only, while “Spiky1” and “Spiky2” differ in the random-seed of model initialization only. Training spikes can occur at different steps between “Spiky1” and “Spiky2”, while “normal fluctuations” in training loss remain consistent.

Panel of Figure 2 — Figure 2: The validation loss of PLRF, from training step 0 to 1k. We compare $\alpha\in\{0.3,0.9\}$ , for SGD and Adam, with different learning-rates. Each plot shows 3 random seeds (to distinguish initialization and data-shuffle separately), and learning-rates are chosen such that training converges, nearly diverges, or diverges, respectively. For SGD, the required learning-rates are calculated from the threshold $\gamma$ given in formula ( 2 ). For Adam, the learning-rates are found by manual grid search.

Experiment	Train	Eval $\delta^{\boldsymbol{w}}_{t,s}$	Eval $\mathbb{E}(\mathbf{h}_{t}{\mathbf{h}}^{\!\top}_{t})$	$\mathbb{E}(\nabla\boldsymbol{W}_{t})$ & $\mathbb{E}(\lVert\nabla_{\mathtt{B}}\boldsymbol{W}_{t}\lVert^{2})$	$\trace\big(\mathbb{E}(\nabla^{2}\boldsymbol{W}_{t})\big)$	Simulation
Data size	25.5B tokens	336M tokens	336M tokens	336M tokens	336M tokens	1500 steps
Model size	1.3B (24 blocks)	24 ffn_0	all layers	all layers	6 ffn_0	1 layer
Time	40 hrs	1.1 hrs	8 min	15 min	1.7 hrs	5 sec
Machine	TPUv4 $\times 32$	TPUv4 $\times 32$	TPUv4 $\times 32$	TPUv4 $\times 32$	TPUv4 $\times 32$	CPU

Table 1: Summary of compute resources for running experiments in this paper.

Panel of Figure 3 — Figure 3: Simulation of the layer-wise quadratic approximation of validation loss at 6 layers (the feed-forward input of block_{0,4,8,12,16,20}), from training step 21k to 22.5k. We compare 3 settings of training (Normal, Spiky1 and Spiky2) as in Figure 1 , and plot two simulations of different seeds, together with the actual 2nd-order term measured from corresponding training trajectories. The scale of the $y$ -axis is kept the same for all plots, so that the fluctuations can be compared.

Panel of Figure 4 — Figure 4: Plots of the 1st-order term of $\delta^{\boldsymbol{w}}_{t,s}$ .

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Simulating PLRF by Algorithm 1 . We apply SGD and Adam at different batch-sizes. Each plot includes one PLRF training and two simulation runs (with different random seeds).

	B=1,LR=0.03,TH=0.3			B=8,LR=0.1,TH=0.5			B=64,LR=0.5,TH=1.0
	Q90	RAT	MRLAT	Q90	RAT	MRLAT	Q90	RAT	MRLAT
PLRF	0.89	0.18	8.2	1.1	0.051	2.2	3.8	0.035	1.3
Simulation1	0.59	0.14	4.9	1.4	0.096	2.7		0.38	4.6
Simulation2	0.53	0.18	5.1	2.5	0.19	3.7		0.36	6.9

Appendix

Table 2: Statistics of training spikes, for three configurations of batch-sizes

\mathtt{B}=\{1,8,64\}

, with respectively picked learning-rate ( LR ) and threshold ( TH ).

Figure 6: The learning-rate schedule used for our Transformer LM training.

Panel of Figure 7 — Figure 7: Log-log plot with truncated power-law fit, for ranked eigenvalues of correlation matrices.

Panel of Figure 8 — Figure 8: Evolution of power-law fitting statistics, across layers and training steps.

Panel of Figure 9 — Figure 9: Empirical verification of Equation ( 6 ).

Panel of Figure 10 — Figure 10: Breakdown of the equivalent batch factor.

$\beta_{2}=0.95$	$\beta_{2}=0.982$	$\beta_{2}=0.992$	$\beta_{2}$ schedule
2.4421	2.4405	2.4502	2.4404

Appendix

Table 3: Final validation loss for different

\beta_{2}

settings.

Figure 11: Training loss for different

\beta_{2}

settings.

Panel of Figure 12 — Figure 12: Estimating an appropriate $\beta_{2}$ schedule.

Random Feature Methods

Random feature methods approximate kernels and neural networks with finite randomized embeddings to reduce computation while retaining useful accuracy.

See all 46 papers

QuantaSpike: Short-Window Spike-Driven Quantization for Large Language Models
Sep 28, 2026Bang Hu, Guowei Zhu, Changze Lv +4LLM Quantization Spiking Neural Networks
School of Computer Science, Fudan University, Shanghai, China
Large Language Models with At Most One Spike per Neuron
Sep 4, 2026Zhuoya Zhao, Parsa Omidi, Aref Jafari +1Spiking Language Models Spiking Neural Networks
Bug or Feature $^2$ : Weight Drift and Activation Sparsity
May 17, 2026Egor Shvetsov, Aleksandr Serkov, Shokorov Viacheslav +3ReLU Neural Networks Activation Sparsity

Sep 28, 2026J/K move · Enter open · S save

QuantaSpike: Short-Window Spike-Driven Quantization for Large Language Models

Bang Hu, Guowei Zhu, Changze Lv, Xiaoqing Zheng+3

School of Computer Science, Fudan University, Shanghai, China

LLM Quantization Spiking Neural Networks Efficient Language Model Inference Neural Network Quantization Spiking Language Models Spike-Based Signal Encoding

Open paper AskPDF arXiv

Large language models (LLMs) achieve strong performance across many tasks but rely on dense multiply-accumulate (MAC) operations during inference, resulting in high energy cost. Spiking neural networks (SNNs) offer an event-driven alternative in which synaptic integration uses lightweight accumulation. However, spike-driven LLM inference remains difficult because outlier-heavy activations typically require long firing windows or auxiliary non-spiking paths. We propose QuantaSpike, a short-window spike-driven quantization framework for LLMs built around Logarithmic Ternary Integrate-and-Fire (LTIF) neurons. LTIF uses ternary events with power-of-two membrane-response quanta, improving the information represented by each firing step while retaining shift-ACC-compatible computation. QuantaSpike combines this neuron with group-adaptive gain and selective outlier admission: normal values use residual LTIF steps, whereas admitted outliers receive one additional onset spike before entering the same residual dynamics. Across OPT and Llama-2, QuantaSpike achieves state-of-the-art or competitive perplexity and zero-shot accuracy among spike-driven LLM quantization methods. It also transfers to newer dense LLMs, remaining close to the FP16 reference on Llama-3-8B and Qwen3-8B under the same four-step firing window. Analytical linear-energy projections show that QuantaSpike reduces the energy of one linear transformation by about

80.0\%

on OPT models and

67.1\%

on Llama-2 models relative to SpikeQuant, providing an accurate and energy-efficient spike-driven path for LLM inference.

Abstract

Random Feature Methods

Explore similar work

QuantaSpike: Short-Window Spike-Driven Quantization for Large Language Models

Large Language Models with At Most One Spike per Neuron

Bug or Feature2^22: Weight Drift and Activation Sparsity

Bug or Feature $^2$ : Weight Drift and Activation Sparsity