KO: Kinetics-inspired Neural Optimizer with PDE Simulation Approaches

Mingquan Feng; Yixin Huang; Yifan Fu; Shaobo Wang; Junchi Yan

cs.LGMay 20, 2025

KO: Kinetics-inspired Neural Optimizer with PDE Simulation Approaches

Authors: Mingquan Feng, Yixin Huang, Yifan Fu, Shaobo Wang, Junchi Yan

Organizations: Shanghai Jiao Tong University

Abstract

The design of effective optimization algorithms for neural networks remains a fundamental challenge, and most existing methods rely on heuristic extensions of gradient-based updates. We introduce KO (Kinetics-inspired Optimizer), a plug-and-play optimization module grounded in kinetic theory and partial differential equations. KO models parameter dynamics as a particle system, augmenting standard gradient updates with stochastic interactions induced by a discretization of the Boltzmann transport equation. This mechanism naturally promotes parameter diversity and mitigates weight condensation, the tendency of parameters to collapse into low-dimensional subspaces, a phenomenon closely associated with degraded generalization. We provide both a rigorous theoretical analysis and a physical interpretation, showing that KO provably increases parameter diversity while preserving convergence guarantees. Extensive experiments on image classification benchmarks (CIFAR-10/100, ImageNet) and large-scale language model pretraining demonstrate that KO consistently improves accuracy over competitive baselines with negligible additional computational cost.

Figures & tables

Figure 1 . The architecture of KO. The upper part depicts the workflow: the gradient is first calculated by backpropagation, then refined by the kinetic module, and finally fed into the base optimizer. The lower part details the kinetic module. Weights and gradients are viewed as positions and velocities of particles; upon collision, particle velocities are updated following the collision model.

Panel of Figure 2 — Figure 2 . Condensation of two-layer NNs. The color indicates the cosine similarity of two hidden neurons’ input weights at epoch 100, whose indices are indicated by the abscissa and the ordinate, respectively. The activation functions are indicated by the sub-captions. The first row shows the model weights trained with the original Adam optimizer. The second and third rows depict the results with the Soft Collision and Hard Collision. A more detailed description is in Appendix B.1 .

Model	CIFAR10	CIFAR100
ResNet18+SGD	$95.07\%\pm 0.18\%$	$78.69\%\pm 0.38\%$
ResNet18+S.C.	$\bm{95.74\%\pm 0.05\%}$	$\bm{79.08\%\pm 0.12\%}$
ResNet18+H.C.	$95.58\%\pm 0.07\%$	$78.83\%\pm 0.09\%$
ResNet18+WCreg	$95.51\%\pm 0.23\%$	$78.77\%\pm 0.28\%$
ResNet34+SGD	$95.14\%\pm 0.37\%$	$79.42\%\pm 0.45\%$
ResNet34+S.C.	$95.76\%\pm 0.08\%$	$\bm{79.88\%\pm 0.25\%}$

Table 1 . Accuracy on CIFAR datasets. S.C. stands for Soft Collision, H.C. for Hard Collision, and WCreg for the weight correlation regularization.

Panel of Figure 3 — Figure 3 . Accuracy and neuron similarity of three different models on CIFAR-100. The upper row shows the test accuracy, while the lower row depicts the neuron similarity changes during training.

Model $\setminus$ Acc $\uparrow$	PIQA	ARC_C	ARC_E	HellaSwag	WinoGrande	OBQA	Avg
Qwen2-0.5B+AdamW	0.575	0.227	0.371	0.273	0.498	0.246	0.365
Qwen2-0.5B+S.C.	0.582	0.241	0.385	0.276	0.522	0.268	0.379
Qwen2-1.5B+AdamW	0.602	0.236	0.439	0.310	0.495	0.256	0.390
Qwen2-1.5B+S.C.	0.615	0.258	0.450	0.313	0.535	0.266	0.406

Table 2 . Comparison of model performance on various LLM benchmarks.

Model	T-1 ACC.(%)	T-5 ACC.(%)
ResNet50+Lamb	79.70%	94.53%
ResNet50+S.C.	80.23%	94.79%
ResNet50+H.C.	79.98%	94.61%
ConvNext_Tiny+AdamW	82.10%	96.03%
ConvNext_Tiny+S.C.	82.34%	96.07%
ConvNext_Tiny+H.C.	82.17%	96.03%

Table 3 . Experimental result on ImageNet-1K.

Pretrained Model	IMDB	Snips
bert-base-cased+AdamW	92.52%	98.14%
bert-base-cased+S.C.	93.80%	98.43%
bert-base-cased+H.C.	93.42%	98.43%
bert-base-uncased+AdamW	93.59%	97.71%
bert-base-uncased+S.C.	94.20%	97.89%
bert-base-uncased+H.C.	93.76%	97.73%

Table 4 . Accuracy on IMDB and Snips.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4 . Condensation of Resnet18-like neural networks on CIFAR-10. The color in the figures indicates the cosine similarity of the normalized input weights of two neurons in the first FC layer. The subcaption represents the activation function in the FC layers. The first column shows the model weights trained with the original Adam optimizer. The second and third columns depict the results with Soft Collision and Hard Collision, respectively.

Weight Decay	Init $\rho(\mathbf{w})$	Mid $\rho(\mathbf{w})$	End $\rho(\mathbf{w})$
0	0.3867	0.9672	0.9999
$10^{-3}$	0.3867	0.9676	0.9999
$10^{-1}$	0.3867	0.9245	0.9992

Appendix

Table 5 . Experiment results with different weight decay strengths.

Treatment	Task-macro	Sample Accuracy	Correct
Historical Standard LoRA	0.27653	27.87%	1,117 / 4,008
LoRA + KO	0.28954	28.49%	1,142 / 4,008
KO - baseline	+0.01301	+0.62%	+25

Appendix

Table 6 . Small-scale chemical capability transfer under Qwen3.5-0.8B-Base LoRA SFT.

Treatment	Task-macro	Sample Accuracy	Correct
Historical Standard LoRA	0.36623	35.48%	1,422 / 4,008
LoRA + KO	0.37491	36.40%	1,459 / 4,008
KO - baseline	+0.00867	+0.92%	+37

Appendix

Table 7 . Small-scale chemical capability transfer under Qwen3.5-2B-Base LoRA SFT.

	FGSM ( Goodfellow et al., 2015 )	PGD ( Madry et al., 2019 )
ResNet18	40.46%	0.08%
ResNet18+S.C.	43.61%	0.19%
ResNet18+H.C.	46.63%	0.14%
ResNet34	42.13%	0.04%
ResNet34+S.C.	44.72%	0.23%
ResNet34+H.C.	50.08%	0.12%

Appendix

Table 10 . Experiment results with different adversarial perturbations

$\text{coll}\_\text{coef}$	0.0 (Baseline)	0.1	0.2	0.3	0.4	0.5	0.6	0.7	0.8	0.9
Accuracy	77.38%	78.23%	78.80%	78.30%	77.98%	78.92%	77.52%	77.87%	78.01%	78.16%
Gain	-	+0.85%	+1.42%	+0.92%	+0.60%	+1.54%	+0.14%	+0.49%	+0.63%	+0.78%

Appendix

Table 11 . Experiment results of ResNet-50 on CIFAR-10 with different collision coefficients

	Vanilla	Vanilla + WD	S.C.	S.C. + WD	H.C.	H.C. + WD
ResNet18	95.02% $\pm$ 0.29%	95.07% $\pm$ 0.18%	95.74% $\pm$ 0.09%	95.74% $\pm$ 0.05%	95.42% $\pm$ 0.13%	95.58% $\pm$ 0.07%
ResNet34	95.11% $\pm$ 0.32%	95.14% $\pm$ 0.37%	95.52% $\pm$ 0.12%	95.76% $\pm$ 0.08%	95.33% $\pm$ 0.18%	95.56% $\pm$ 0.14%
ResNet50	95.31% $\pm$ 0.27%	95.37% $\pm$ 0.38%	95.62% $\pm$ 0.08%	95.83% $\pm$ 0.16%	95.42% $\pm$ 0.11%	95.57% $\pm$ 0.05%

Appendix

Table 12 . Experiment results with and without weight decay with various backbones on CIFAR-10. WD is the abbreviation of weight decay.

Concept	Kinetic Theory (Physical System)	Neural Network Optimization
Agent	A single particle	A single neuron (weight vector $\mathbf{w}_{i}$ )
Position	Spatial coordinates of the particle	Weight vector $\mathbf{w}_{i}\in\mathbb{R}^{D}$
Velocity	Particle velocity vector	(Negative) gradient $-\mathbf{g}_{i}$ (update direction)
Interaction	Stochastic collisions	Gradient modification (hard/soft collision terms)
System Dynamics	Boltzmann Transport Equation	Optimization trajectory (gradient descent + collision)

Appendix

Table 13 . Mapping between kinetic theory concepts and neural network optimization.

Explore similar work

The Anatomy of Implicit Bias: Information Allocation in Neural Network Training
Jul 8, 2026Zhang Gongyue, Wang Zhiyong, Liu Donghan +3Neural Network Optimization Optimizer Design
1State Key Laboratory of Robotics and Systems, Harbin Institute of Technology Shenzhen, Shenzhen, 518055, China.
Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization
Sep 3, 2025Wu Lin, Scott C. Lowe, Felix Dangel +3Neural Network Optimization Kullback-Leibler Divergence
1Vector Institute, Toronto, Canada · University of Cambridge, Cambridge, United Kingdom · 3Microsoft, United States +1
Turning Stale Gradients into Stable Gradients: Coherent Coordinate Descent with Implicit Landscape Smoothing for Lightweight Zeroth-Order Optimization
May 14, 2026Chen Liang, Xiatao Sun, Qian Wang +1Zeroth-Order Optimization Gradient
Department of Computer Science, Yale University, New Haven, USA.

Jul 8, 2026, cs.LGJ/K move · Enter open · S save

The Anatomy of Implicit Bias: Information Allocation in Neural Network Training

Zhang Gongyue, Wang Zhiyong, Liu Donghan, Ren Weihong+2

1State Key Laboratory of Robotics and Systems, Harbin Institute of Technology Shenzhen, Shenzhen, 518055, China.

Neural Network Optimization Optimizer Design Spectral Preconditioning Multi-Stage Training

Open paper AskPDF arXiv

Implicit bias is usually explained as the preference of an optimization process for certain final solutions and their geometry. This view helps explain where a model finally stops. It gives less direct explanation of how this bias is formed during training. This paper proposes a training-time information allocation view. Under this view, optimization forms a writing pattern for error signals across parameter paths, coordinate channels, and sample regions. This paper builds a set of observable allocation diagnostics. These diagnostics include gradient demand, actual update injection, coordinate gain induced by exponential moving averages, channel-level update ratios, and sample-wise loss distributions. To separate training progress from internal allocation, this paper introduces a collapse--persistence analysis. Under matched training loss, if external loss statistics collapse but internal allocation ratios remain separated, then the factor changes the internal allocation of the training signal. Overall, this paper extends the analysis of implicit bias from final-solution geometry to training-time signal allocation. The main claim is that implicit bias is not only reflected by the final solution. It is also reflected by which parameter paths, coordinate channels, and sample regions receive the error signal first and more strongly during training. Based on this view, this paper places different training factors into a unified information-allocation diagnostic framework. The framework gives a mechanism-level explanation of training-time implicit bias. It also provides a basis for future optimization methods that control training progress and signal allocation separately.

Abstract

Explore similar work

The Anatomy of Implicit Bias: Information Allocation in Neural Network Training

Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization

Turning Stale Gradients into Stable Gradients: Coherent Coordinate Descent with Implicit Landscape Smoothing for Lightweight Zeroth-Order Optimization