KO: Kinetics-inspired Neural Optimizer with PDE Simulation Approaches

TL;DR

KO is a physics-inspired neural optimizer that models training dynamics as a particle system governed by PDEs, promoting parameter diversity and outperforming traditional optimizers in various tasks.

Contribution

The paper introduces KO, a novel optimizer based on kinetic theory and PDE simulation, offering a physics-driven approach to improve neural network training.

Findings

KO outperforms Adam and SGD on image and text classification tasks.

KO maintains parameter diversity, reducing parameter collapse.

Experimental results show accuracy improvements with comparable computational cost.

Abstract

The design of optimization algorithms for neural networks remains a critical challenge, with most existing methods relying on heuristic adaptations of gradient-based approaches. This paper introduces KO (Kinetics-inspired Optimizer), a novel neural optimizer inspired by kinetic theory and partial differential equation (PDE) simulations. We reimagine the training dynamics of network parameters as the evolution of a particle system governed by kinetic principles, where parameter updates are simulated via a numerical scheme for the Boltzmann transport equation (BTE) that models stochastic particle collisions. This physics-driven approach inherently promotes parameter diversity during optimization, mitigating the phenomenon of parameter condensation, i.e. collapse of network parameters into low-dimensional subspaces, through mechanisms analogous to thermal diffusion in physical systems. We analyze this property, establishing both a mathematical proof and a physical interpretation. Extensive experiments on image classification (CIFAR-10/100, ImageNet) and text classification (IMDB, Snips) tasks demonstrate that KO consistently outperforms baseline optimizers (e.g., Adam, SGD), achieving accuracy improvements while computation cost remains comparable.