WeightFlow: Learning Stochastic Dynamics via Evolving Weight of Neural Network

TL;DR

WeightFlow introduces a novel approach to model stochastic dynamics by evolving neural network weights through a graph-controlled differential equation, effectively capturing continuous probability density evolution in high-dimensional spaces.

Contribution

The paper proposes a new paradigm that models stochastic dynamics directly in neural network weight space, connecting optimal transport theory with neural network evolution.

Findings

WeightFlow outperforms existing methods by 43.02% on interdisciplinary datasets.

It effectively models high-dimensional stochastic dynamics.

The approach provides a scalable solution for continuous probability density estimation.

Abstract

Modeling stochastic dynamics from discrete observations is a key interdisciplinary challenge. Existing methods often fail to estimate the continuous evolution of probability densities from trajectories or face the curse of dimensionality. To address these limitations, we presents a novel paradigm: modeling dynamics directly in the weight space of a neural network by projecting the evolving probability distribution. We first theoretically establish the connection between dynamic optimal transport in measure space and an equivalent energy functional in weight space. Subsequently, we design WeightFlow, which constructs the neural network weights into a graph and learns its evolution via a graph controlled differential equation. Experiments on interdisciplinary datasets demonstrate that WeightFlow improves performance by an average of 43.02\% over state-of-the-art methods, providing an effective and scalable solution for modeling high-dimensional stochastic dynamics.