EPDE-Multisample-Discovery for SDEs in Diffusion Models

This project uses the EPDE library to automatically discover the governing stochastic differential equations (SDEs) from multiple trajectories generated by diffusion models. The primary goal is to find a symbolic representation of the underlying dynamics that explains the data's behavior.

Data

The trajectories are stored in .npy files with the following dimensions:

Shape: (batch_size, num_steps + 1, dim)

• batch_size: The number of trajectories in a single batch.
• num_steps + 1: The number of time steps in each trajectory, including the initial point.
• dim: The spatial dimension of each point in the trajectory.

There are two files: trajectories-1.npy (trained Schrödinger bridge data) and trajectories-2.npy (diffusion layer).

The core of this project is to analyze time-series data from a trained Schrödinger bridge SDE. The trajectories represent the visualization of the movement of a probability distribution from an initial state ($\pi_0$) to a final state ($\pi_1$).

The SDE governing this process can be represented as:

$$d\mathbf{z}_t = \text{pred}(\mathbf{z}_t) , dt + \sigma , d\mathbf{W}_t$$, where

• $\mathbf{z}_t$ is the state vector at time $t$, $\text{pred}(\mathbf{z}_t)$ is the network drift term, which is the model's prediction.
• $dt$ is the time step.
• $\sigma$ is the noise level.
• $d\mathbf{W}_t$ is the Wiener process, representing the stochastic noise.