Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

WeightFlow: Learning Stochastic Dynamics via Evolving Weight of Neural Network

Ruikun Li, Jiazhen Liu, Huandong Wang, Qingmin Liao, Yong Li

TL;DR

This paper addresses modeling stochastic dynamics from discrete observations by reframing the problem in the weight space of a neural network. It introduces WeightFlow, which evolves a weight graph through a graph neural differential equation, linking weight dynamics to the evolution of probability densities via dynamic optimal transport. The approach uses anchor weights from snapshots, a latent control path, and a two-stage training process to optimize data fit and path energy, achieving substantial improvements on both simulated and real-world high-dimensional systems. By decoupling distribution modeling from temporal evolution and leveraging a low-dimensional latent path, WeightFlow offers a scalable method that mitigates the curse of dimensionality while accurately capturing higher-order distributional structure.

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.

WeightFlow: Learning Stochastic Dynamics via Evolving Weight of Neural Network

TL;DR

This paper addresses modeling stochastic dynamics from discrete observations by reframing the problem in the weight space of a neural network. It introduces WeightFlow, which evolves a weight graph through a graph neural differential equation, linking weight dynamics to the evolution of probability densities via dynamic optimal transport. The approach uses anchor weights from snapshots, a latent control path, and a two-stage training process to optimize data fit and path energy, achieving substantial improvements on both simulated and real-world high-dimensional systems. By decoupling distribution modeling from temporal evolution and leveraging a low-dimensional latent path, WeightFlow offers a scalable method that mitigates the curse of dimensionality while accurately capturing higher-order distributional structure.

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.

Paper Structure

This paper contains 33 sections, 1 theorem, 10 equations, 8 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Stochastic Dynamics & Probability Evolution
  4. Problem Definition
  5. Weight Path of Probability Evolution
  6. Parameterizing Probability Distribution
  7. Bridging the Probability Path and Evolving Weight
  8. Learning Stochastic Dynamics via Evolving Weight
  9. WeightFlow
  10. Backbone and Weight Graph
  11. Graph Neural Differential Equation
  12. Training Strategy
  13. Time and Computational Complexity
  14. Numerical Results
  15. Experimental Settings
  16. ...and 18 more sections

Key Result

Theorem 1

Let $\mu_0$ and $\mu_1$ be the initial and end probability measures on $\mathbb{R}^d$, and let $(p^*, v^*)$ be the optimal path under the Eq. equ:RUOT. Assuming the following conditions are satisfied: Then, we have $|\mathcal{E}(p_{\bm{\theta}^*}, v_{g^*}) - \mathcal{E}(p^*, v^*)| \le \delta$.

Figures (8)

  • Figure 1: Learning stochastic dynamics: from sample trajectories to evolving weights.
  • Figure 2: Framework of WeightFlow.
  • Figure 3: Joint and marginal distributions predicted by WeightFlow over time on the Ecological Evolution system. The central panels show the joint distribution contours of two loci overlaid on the fitness landscape (brighter yellow indicates higher fitness), with marginal densities on the outer axes.
  • Figure 4: Weight prediction for $\beta$-cell differentiation. The trajectory shows the continuous evolution of weights and corresponding ensemble distributions (PCA). Highlights compare the results at observed snapshots: reference (left), WeightFlow's prediction (center), and DeepRUOT's prediction (right). The table reports the average relative error of the first four statistical moments.
  • Figure 5: Comparing the path (left) and average error (right) for ODE and CDE models on the Toggle Switch system.
  • ...and 3 more figures

Theorems & Definitions (2)

  • Theorem 1
  • proof