Learning Stochastic Behaviour from Aggregate Data
Year: 2020

Abstract:
Learning nonlinear dynamics from aggregate data is a challenging problem since the full trajectory of each individual is not available, namely, the individual observed at one time point may not be observed at next time point, or the identity of individual is unavailable. This is in sharp contrast to learning dynamics with trajectory data, on which the majority of existing methods are based. We propose a novel method using the weak form of Fokker Planck Equation (FPE) to describe density evolution of data in a sampling form, which is then combined with Wasserstein generative adversarial network (WGAN) in training process. In such a sample-based framework we are able to study nonlinear dynamics from aggregate data without solving the partial differential equation (PDE). The model can also handle high dimensional cases with the help of deep neural networks. We demonstrate our approach in the context of a series of synthetic and real-world data sets.

text: 1 Introduction
In the context of a dynamic system, Aggregate data refers to the data sets that full trajectory of each individual is not available, meaning that there is no known individual level correspondence. Typical examples include data sets collected for DNA evolution, social gathering, density in control problems, bird migration during which it is impossible to identify individual bird, and many more. In those applications, some observed individuals at one time point may be unobserved at the next time spot, or when the individual identities are blocked or unavailable due to various technical and ethical reasons. Rather than inferring the exact information for each individual, the main objective of learning dynamics in aggregate data is to recover and predict the evolution of distribution of all individuals together. Trajectory data, in contrast, is a kind of data that we are able to acquire the information of each individual all the time, although some studies considered the case that some individual trajectories are partially missing. However, the identities of those individuals, whenever they are observable, is always assumed available. For example, stock price, weather, customer behaviors and most training data sets for computer vision and natural language processing.
There are many popular models to learn dynamics of full-trajectory data. Typical ones include Hidden Markov Model (HMM)(Alshamaa et al., 2019; Eddy, 1996), Kalman Filter (KF)(Farahi & Yazdi, 2020; Harvey, 1990; Kalman, 1960) and Particle Filter (PF) (Santos et al., 2019; Djuric et al., 2003), as well as the models built upon HMM, KF and PF(Deriche et al., 2020; Fang et al., 2019; Hefny et al., 2015; Langford et al., 2009), they all require full trajectories of each individual, which may not be applicable in the aggregate data situations. On the other side, only a few methods are focused on aggregated data in the recent learning literature. In the work of Hashimoto et al. (2016), authors assumed that the hidden dynamic of particles follows a stochastic differential equation(SDE), in particular, they use a recurrent neural network to parameterize the drift term. Furthermore, Wang et al. (2018) improved traditional HMM model by using an SDE to describe the evolving process of hidden states. To the best of our knowledge, there is no method directly learning the evolution of the density of objects from aggregate data yet.
We propose to learn the dynamics of density through the weak form of Fokker Planck Equation (FPE), which is a parabolic partial differential equation (PDE) governing many dynamical systems subject to random noise perturbations, including the typical SDE models in existing studies. Our learning is accomplished by minimizing the Wasserstein distance between predicted distribution given by FPE and the empirical distribution from data samples. Meanwhile we utilize neural networks to handle higher dimensional cases. More importantly, by leveraging the framework of Wasserstein Generative
Adversarial Network (WGAN) (Arjovsky et al., 2017), our model is capable of approximating the distribution of samples at different time points without solving the SDE or FPE. More specifically, we treat the drift coefficient, the goal of learning, in the FPE as a generator, and the test function in the weak form of FPE as a discriminator. In other words, our method can also be regarded as a data-driven method to estimate transport coefficient in FPE, which corresponds to the drift terms in SDEs. Additionally, though we treat diffusion term as a constant in our model, it is straightforward to generalize it to be a neural network as well, which can be an extension of this work. We would like to mention that several methods of solving SDE and FPE (Weinan et al., 2017; Beck et al., 2018; Li et al., 2019) adopt opposite ways to our method, they utilize neural networks to estimate the distribution P(x, t) with given drift and diffusion terms.
In conclusion, our contributions are:
• We design an algorithm that is able to recover the density evolution of nonlinear dynamics via minimizing the Wasserstein discrepancy between real aggregate data and our generated data.
• By leveraging the weak form of FPE, we are able to compute the Wasserstein distance directly without solving the FPE.
• Finally, we demonstrate the accuracy and the effectiveness of our algorithm by several synthetic and real-world examples.
2 ProposedMethod
2.1 Fokker Planck Equation for the density evolution
We assume the individuals evolve in a pattern in the space RD as shown in Figure 1. One example satisfying such process is the stochastic differential equation(SDE), which is also known as the Itô process (Øksendal, 2003): dXt = g(Xt, t)dt + σdWt. Here dXt represents an infinitesimal change of {Xt} along with time increment dt, g(·, t) = (g1(·, t), ..., gD(·, t))T is the drift term (drifting vector field) that drives the dynamics of the SDE, σ is the diffusion constant, {Wt} is the standard Brownian Motion. Then the probability density of {Xt} is governed by the Fokker Planck Equation(FPE) (Risken & Caugheyz, 1991), as stated below in Lemma 1.
Lemma 1. Suppose {Xt} solves the SDE dXt = g(Xt, t)dt + σdWt, denote p(·, t) as the probability density of the random variableXt. Then p(x, t) solves the following equation:
∂p(x, t) ∂t = D∑ i=1 − ∂ ∂xi [ gi(x, t)p(x, t) ] + 1 2 σ2 D∑ i=1 ∂2 ∂xi2 p(x, t) (1)
As a linear evolution PDE, FPE describes the evolution of density function of the stochastic process driven by a SDE. Due to this reason, FPE plays a crucial role in stochastic calculus, statistical physics and modeling (Nelson, 1985; Qi & Majda, 2016; Risken, 1989). Its importance is also drawing more attention among statistic and machine learning communities (Liu & Wang, 2016; Pavon et al., 2018; Rezende & Mohamed, 2015). In this paper, we utilize the weak form of FPE as a basis to study hidden dynamics of the time evolving aggregated data without solving FPE.
Our task can be described as: assume that the individuals evolve with the process indicated by Figure 1, which can be simulated by Itô process. Then given observations xt along time axis, we aim to recover the drift coefficient g(x, t) in FPE, and thus we are able to recover and predict the density evolution of such dynamic. For simplicity we treat g(x, t) as a function uncorrelated to time t, namely,
g(x, t) = g(x). Notice that though evolving process of individuals can be simulated by Itô process, in reality since we lose identity information of individuals, the observed data become aggregate data, thus we need a new way other than traditional methods to study the swarm’s distribution.
2.2 Weak Form of Fokker Planck Equation
Given FPE stated in Lemma 1, if we multiply a test function f ∈ H10(RD) on both sides of the FPE, then the integration on both sides leads to:∫
∂p ∂t
f (x)dx = ∫ D∑
i=1
− ∂ ∂xi
[ gi(x)p(x, t) ] f (x)dx +
1 2 σ2 ∫ D∑ i=1 ∂2 ∂xi2 p(x, t) f (x)dx.
Then integrating by parts on the right hand side leads to the weak form of FPE:∫ ∂p ∂t f (x)dx = ∫ D∑
i=1
gi(x) ∂
∂xi f (x)p(x, t)dx + 1 2 σ2 ∫ D∑ i=1 ∂2 ∂xi2 f (x)p(x, t)dx.
where H10(R D) denote the Sobolev space. The first advantage of weak solution is that the solution of a PDE usually requires strong regularity and thus may not exist in the classical sense for a certain group of equations, however, the weak solution has fewer regularity requirements and thus their existence are guaranteed for a much larger classes of equations. The second advantage is that the weak formulation may provide new perspectives for numerically solving PDEs (Zienkiewicz & Cheung, 1971),(Sirignano & Spiliopoulos, 2018) and (Zang et al., 2019) etc.
Suppose the observed samples at time points tm−1 and tm are following true densities p̂(·, tm−1) and p̂(·, tm). Consider the following SDE
dX̃t = gω(X̃t)dt + σdWt, tm−1 ≤ t ≤ tm X̃ tm−1 ∼ p̂(·, tm−1) (2) Here gω is an approximation to the real drift term g, in our research, we treat gω as a neural network with parameters ω . Let us now denote p̃(·, t) as the density of X̃t. Then it is natural to compute and minimize the discrepancy between our approximated density p̃(·, tm) and true density p̂(·, tm), within this process, we are optimizing gω and thus will recover the true drift term g. In our research, we choose the Wasserstein-1 distance as our discrepancy function (Villani, 2008) (Arjovsky et al., 2017). Now we apply Kantorovich-Rubinstein duality (Villani, 2008), this leads to:
W1( p̂(·, tm), p̃(·, tm)) = sup ‖∇ f ‖≤1
{ Exr∼ p̂(x,tm)[ f (xr)] − Exg∼ p̃(x,tm)[ f (xg)] } (3)
The first term Exr∼p̂(x,tm)[ f (xr)] in Equation (3) can be conveniently computed by Monte-Carlo method since we are already provided with the real data points xr ∼ p̂(·, tm). To evaluate the second term, we first approximate p̃(·, tm) by trapezoidal rule (Atkinson, 2008):
p̃(x, tm) ≈ p̂(x, tm−1) + ∆t 2
( ∂p̃(x, tm−1)
∂t + ∂p̃(x, tm) ∂t
) here ∆t = tm − tm−1. (4)
And thus we can compute: Exg∼p̃(·,tm)[ f (xg)] ≈ ∫
f (x) p̂(x, tm−1)dx+ ∆t 2
(∫ ∂p̂(x, tm−1)
∂t f (x)dx +
∫ ∂p̃(x, tm)
∂t f (x)dx
) (5)
In the above Equation (5), the second and the third term on the right-hand side can be reformulated via the weak form of FPE and we are able to derive a computable formulation for W1( p̂(·, tm), p̃(·, tm)). Furthermore, we can use Monte-Carlo method to approximate the expectations in (5). the first and the second terms can be directly computed via data points from p̂(·, tm−1). For the third term, we need to generate samples from p̃(·, tm), to achieve this, we apply Euler-Maruyama scheme (Kloeden & Platen, 2013) to SDE (2) in order to acquire our desired samples x̃tm :
x̃tm = x̂tm−1 + gω(x̂tm−1 )∆t + σ √ ∆tz, here z ∼ N(0, I), x̂tm−1 ∼ p̂(·, tm−1). (6)
Here N(0, I) is the standard Gaussian distribution on RD. Now we summarize these results in Proposition 1:
Proposition 1. For a set of points X = {x(1), ...,x(N)} in RD. We denote:
F f (X) = 1 N N∑ k=1 ( D∑ i=1 giω(x (k)) ∂ ∂xi f (x(k)) + 1 2 σ2 D∑ i=1 ∂2 ∂x2i f (x(k)) ) then at time point tm, the Wasserstein distance between p̂(·, tm) and p̃(·, tm) can be approximated by:
W1( p̂(·, tm), p̃(·, tm)) ≈ sup ‖∇ f ‖≤1  1N N∑
k=1
f (x̂(k)tm ) − 1 N N∑ k=1 f (x̂(k)tm−1 ) − ∆t 2 (F f (X̂m−1) + F f (X̃m)) Here {x̂(k)tm−1 } ∼ p̂(·, tm−1), {x̂ (k) tm } ∼ p̂(·, tm) and we denote X̂m−1 = {x̂ (1) tm−1 , ..., x̂ (N) tm−1 }, X̃m = {x̃ (1) tm , ..., x̃ (N) tm }, where each x̃(k)tm is computed by Euler-Maruyama scheme (6).
2.3 Wasserstein Distance on Time Series
In real cases, it is not realistic to observe the data at arbitrary two consecutive time nodes, especially when ∆t is small. To make our model more flexible, we should extend our formulation so that we are able to plug in observed data at arbitrary time points. To be more precise, suppose we observe data set X̂tn = {x̂ (1) tn , ..., x̂ (N) tn } at J + 1 different time points t0, t1, ..., tJ . And we denote the generated data set as X̃tn = {x̃ (1) tn , ..., x̃ (N) tn }, here each x̃ (·) tn is derived from the n-step Euler-Maruyama scheme:
x̃t j = x̃t j−1 + gω(x̃t j−1 )∆t + σ √
∆tz, with z ∼ N(0, I) 0 ≤ j ≤ n, x̃t0 ∼ p̂(·, t0) (7) Let us denote p̃(·, t) as the solution to FPE (1) with g replaced by gω and with initial condition p̃(·, t0) = p̂(·, t0), then the approximation formula for evaluating the Wasserstein distance W1(p̂(·, tn), p̃(·, tn)) is provided in the following proposition: Proposition 2. Suppose we keep all the notations defined as above, then we have the approximation:
W1( p̂(·, tn), p̃(·, tn)) ≈ sup ‖∇ f ‖≤1  1N N∑
k=1
f (x̂(k)tn ) − 1 N N∑ k=1 f (x̂(k)t0 ) − ∆t 2 (F f (X̂0) + F f (X̃n) + 2 n−1∑ s=1 F f (X̃s) )
Minimizing the Objective Function: Based on Proposition 2, we obtain objective function by summing up the accumulated Wasserstein distances among J observations along the time axis. Thus, our ultimate goal is to minimize the following objection function:
min gω  J∑ n=1 sup ‖∇ fn‖≤1  1N N∑ k=1 fn(x̂ (k) tn ) − 1 N N∑ k=1 fn(x̂ (k) t0 ) − ∆t 2 (F fn (X̂0) + F fn (X̃n) + 2 n−1∑ s=1 F fn (X̃s) )  Notice that since we have observations on J distinct time points(despite the initial point), for each time point we compute Wasserstein distance with the help of the dual function fn, thus we will involve J test functions in total. In our actual implementation, we will choose these dual functions as neural networks. We call our algorithm Fokker Planck Process(FPP), the entire procedure is shown in Algorithm 1. We also provide an error analysis in Appendix.
3 Experiments
In this section, we evaluate our model on various synthetic and realistic data sets by employing Algorithm 1. We generate samples x̃t and make all predictions base on Equation (6) starting with x̂0.
Baselines: We compare our model with two recently proposed methods. One model (NN) adopts recurrent neural network(RNN) to learn dynamics directly from observations of aggregate data (Hashimoto et al., 2016). The other one model (LEGEND) learns dynamics in a HMM framework (Wang et al., 2018). The baselines in our experiments are two typical representatives that have stateof-the-art performance on learning aggregate data. Furthermore, though we simulate the evolving process of the data as a SDE, which is on the same track with NN, as mentioned before, NN trains its RNN via optimizing Sinkhorn distance (Cuturi, 2013), our model starts with a view of weak form of PDE, focuses more on WGAN framework and easier computation.
Algorithm 1 Fokker Planck Process Algorithm Require: Initialize fθn (1 ≤ n ≤ J), gω Require: Set fn as the inner loop learning rate for fθn and g as the outer loop learning rate for gω
1: for # training iterations do 2: for k steps do 3: for observed time ts in {t1, ..., tJ} do 4: Compute the generated data set X̃ts from Euler-Maruyama scheme (7) for 1 ≤ s ≤ J 5: Acquire data sets X̂ts = {x̂ (1) ts , ..., x̂ (N) ts } from real distribution p̂(·, ts) for 1 ≤ s ≤ J 6: end for 7: For each dual function fθn , compute: 8: Fn = F fθn (X̂t0 ) + F fθn (X̃tn ) + 2 ∑n−1 s=1 F fθn (X̃ts )
9: Update each fθn by θn ← θn + fn∇θ ( 1 N ∑N k=1 fθn (x̂ (k) tn ) − 1 N ∑N k=1 fθn (x̂ (k) t0 ) − ∆t 2 Fn ) 10: end for 11: Update gω by ω← ω − g∇ω (∑J n=1 ( 1 N fθn (x̂ (k) tn ) − 1 N fθn (x̂ (k) t0 ) − ∆t 2 Fn)
) 12: end for
3.1 Synthetic Data
We first evaluate our model on three synthetic data sets which are generated by three artificial dynamics: Synthetic-1, Synthetic-2 and Synthetic-3.
Experiment Setup: In all synthetic data experiments, we set the drift term g and the discriminator f as two simple fully-connected networks. The g network has one hidden layer and the f network has three hidden layers. Each layer has 32 nodes for both g and f . The only one activation function we choose is Tanh. Notice that since we need to calculate ∂
2 f ∂x2 , the activation function of f must be
twice differentiable to avoid loss of weight gradient. In terms of training process, we use the Adam optimizer (Kingma & Ba, 2014) with learning rate 10−4. Furthermore, we use spectral normalization to realize ‖∇ f ‖ ≤ 1(Miyato et al., 2018). We initialize the weights with Xavier initialization(Glorot & Bengio, 2010) and train our model by Algorithm 1.

heading: Synthetic-1:
text: x̂0 ∼ N(0,Σ0), x̂t+∆t = x̂t − (Ax̂t + b)∆t + σ √ ∆tN(0, 1)

heading: Synthetic-2:
text: x̂0 ∼ N(0,Σ0), x̂t+∆t = x̂t −G∆t + σ √ ∆tN(0, 1)
G =  1σ1 N1N1+N2 (x̂1t − µ11) + 1σ2 N2N1+N2 (x̂1t − µ21) 00 1 σ1 N1 N1+N2 (x̂2t − µ12) + 1σ2 N2 N1+N2 (x̂2t − µ22)  N1 = 1 √
2πσ1 exp − (x̂1t − µ11)2 2σ21 − (x̂ 1 t − µ12)2 2σ21  , N2 = 1√ 2πσ2 exp − (x̂2t − µ21)2 2σ22 − (x̂ 2 t − µ22)2 2σ22  Synthetic-3 (Van der Pol oscillator (Li, 2018)):
x̂0 ∼ N(0,Σ0)
x̂1t+∆t = x̂ 1 t + 10 ( x̂2t −
1 3 (x̂1t ) 3 + x̂1t
) ∆t + σ √ ∆tN(0, 1)
x̂2t+∆t = x̂ 2 t + 3(1 − x̂1t )∆t + σ
√ ∆tN(0, 1)
The data size at each time point is N = 2000, the dimension of the data is D = 2. We treat 1200 data points as the training set and the other 800 data points as the test set. In Synthetic-1, the data is following a simple linear dynamic, where A = [(4, 0), (0, 1)], b = [−12, −12]T . We let ∆t = 0.01, σ = 1, Σ0 = [(1, 0), (0, 1)]. We utilize true x0, x20 and x200 in training process and predict the
distributions of x50 and x500. In Synthetic-2, the data is following a complex nonlinear dynamic. We let ∆t = 0.01, σ = σ1 = σ2 = 1, µ1 = [15, 15]T and µ2 = [−15,−15]T . We utilize true x0, x3 and x6 in training process and predict x2, x4 and x10. In Synthetic-3, ∆t = 0.01, σ = 1, we utilize true x3, x7 and x20 in training process then predict the distributions of x10, x30 and x50. We also consider cases in higher dimensions: D = 6 and 10. In each high dimensional case, to be convenience, we set every two dimensions follow the dynamics of low dimensional case(D = 2) in each data set. Notice that in Syn-2 and Syn-3, x̂it represents the i-th dimension of x̂t.
Results: We first show the capability of our model for learning hidden dynamics of low-dimensional (d = 2) data. As visualized in Figure 2, the generated data(blue) covers all areas of ground truth(red), which demonstrates that our model is able to precisely learn the dynamics and correctly predict the future distribution of data. The samples we predict converge to stationary distributions finally (as the ground truths suggest). We then evaluate three models using Wasserstein distance as error metric for both low-dimensional (d = 2) and high-dimensional (d = 6, 10) data. As reported in Appendix Table 1, our model achieves lower Wasserstein error than the two baseline models in all cases.
3.2 Realistic Data – RNA Sequence of Single Cell
In this section, we evaluate our model on a realistic biology data set called Single-cell RNA-seq(Klein et al., 2015), which is typically used for learning the evolvement of cell differentiation. The cell population begins to differentiate at day 0 (D0). Single-cell RNA-seq observations are then sampled at day 0 (D0), day 2 (D2), day 4 (D4) and day 7 (D7). At each time point, the expression of 24,175 genes of several hundreds cells are measured (933, 303, 683 and 798 cells on D0, D2, D4 and D7 respectively). Notice that there is only whole group’s distribution but no trajectory information of each gene on different days. We pick 10 gene markers out of 24,175 to make a 10 dimensional data set. In first task we treat gene expression at D0, D4 and D7 as training data to learn the hidden dynamic and predict the distribution of gene expression at D2. In second task we train the model with gene expression at D0, D2 and D4, then predict the distribution of gene expression at D7. We plot the prediction results of two out of ten markers, i.e. Mt1 and Mt2 in Figure 3.
Experiment Setup: We set both f and g as fully connected three-hidden-layers neural networks, each layer has 64 nodes. The only activation function we choose is Tanh. The other setups of neural networks and training process are the same with the ones we use in Synthetic data. Notice that in realistic cases, ∆t and T/∆t become hyperparameters, here we choose ∆t = 0.05, T/∆t = 35, which means the data evolves 10∆t from D0 to D2 , then 10∆t from D2 to D4 and finally 15∆t from D4 to D7. For preprocessing, we apply standard normalization procedures (Hicks et al., 2015) to correct batch effects and use non-negative matrix factorization to impute missing expression levels(Hashimoto et al., 2016; Wang et al., 2018).
Results: As shown in Table 1 in Appendix, when compared to other baselines, our model achieves lower Wasserstein error on both Mt1 and Mt2 data, which proves that our model is capable of learning the hidden dynamics of the two studied gene expressions. In Figure 3 (a) to (d), we visualized the predicted distributions of the two genes. The distributions of Mt1 and Mt2 predicted by our model (curves in blue) are closer to the true distributions (curves in red) on both D2 and D7. Furthermore, our model precisely indicates the correlations between Mt1 and Mt2, as shown in Figure 3 (e) and (f), which also demonstrates the effectiveness of our model since closer to the true correlation represents better performance. In Figure 3 (g) and (h), we see that with simpler structure, the training process of our model is easier with least computation time.
3.3 Realistic Data – Daily Trading Volume
In this section we would like to demonstrate the performance of our model in financial area. Trading volume is the total quantity of shares or contracts traded for specified securities such as stocks, bonds, options contracts, future contracts and all types of commodities. It can be measured on any type of security traded during a trading day or a specified time period. In our case, daily volume of trade is measured on stocks. Predicting traded volume is an essential component in financial research since the traded volume, as a basic component or input of other financial algorithms, tells investors the market’s activity and liquidity. The data set we use is the historical traded volume of the stock "JPM". The data covers period from January 2018 to January 2020 and is obtained from Bloomberg. Each day from 14:30 to 20:55, we have 1 observation every 5 minutes, totally 78 observations everyday. Our task is described as follows: given first two years data, we use the traded volume at 14:30, 14:40, 15:05, 15:20 and 16:20 as training data, namely, x0,x2,x7,x10,x22 to train our model, then for next 100 days we predict traded volume at 14:35, 15:15, 15:35 and 16:15, namely, x1,x9,x13,x21. One of baselines we choose is classical rolling means(RM) method, which predicts intraday volume of a particular time interval by the average volume traded in the same interval over the past days. The other one baseline is a kalman filter based model (Chen et al., 2016) that outperforms all available models in predicting intrady trading volume.
Experiment Setup: Following similar setup as we did for RNA data set, we utilize the same structures for neural networks here. For hyperparameters we set ∆t = 0.02, T/∆t = 22, which means it takes one single ∆t from xt to xt+1. For preprocessing, we rescale data by taking natural logarithm of trading volume, which is a common way in trading volume research. We conduct experiments on two groups to show advantages of our method, for first group we train our model on complete data set, in this case the data has full trajectory; for second group we manually delete some trajectories of the data, for instance, we randomly kick out some samples of x0,x2,x7,x10,x22 then follow the same procedures of training and prediction.
Results: We present prediction results in Figure 4. As shown in first four figures, with full trajectory, prediction made by RM is almost a straight line, the prediction value bouncing up and down within a
very small range, thus this model cannot capture the volume movements, namely, regularities existing in the time series; prediction made by the Kalman filter based model captures the regularities better than RM model, but it fails to deal with noise component existing in the time series, thus some predictions are out of a reasonable range. Traded volume predicted by our model is closer to the real case, moreover, our model captures regularities meanwhile gives stable predictions. Furthermore, without full trajectory, Kalman filter based model fails to be applied here and RM model still fails to capture the regularities, we randomly drop half of the training samples and display predictions made by our model in last four figures of Figure 4, we see our model still works well.
4 Discussions
In this section we discuss the limitations and extension of our model.
An essential challenge for recovering the drift term: Mathematically it is impossible to recover the exact drift term of an SDE if we are only given the information of density evolution on certain time intervals, because there might be infinitely many drift functions to induce the same density evolution. More precisely, suppose p(x, t) solves FPE (equation 1), consider the following equation
0 = − D∑
i=1
∂
∂xi (ui(x, t)p(x, t)) +
σ2
2 D∑ i=1 ∂2 ∂x2i p(x, t) (8)
One can prove, under mild assumptions, that there are infinitely many vector fields u(x, t) = (u1(x, t), ..., uD(x, t)) solving equation (8). One can check that the solution to FPE (1) with drift term g(x, t) + u(x, t) is still p(x, t), i.e. the vector field u(x, t) will never affect the density evolution of the dynamic. This illustrates that given the density evolution p(·, t), the solution for drift term is not necessarily unique. This is clearly an essential difficulty of determining the exact drift term from the density. In this study, the main goal is to recover the entire density evolution (i.e. interpolate the density between observation time points) and predict how will the density evolve in the future. As a result, although we cannot always acquire the exact drift term of the dynamic, we can still accurately recover and predict the density evolution. This is still meaningful and may find its application in various scientific domains.
Curl Field: The drift function we showed in the numerical experiments will apparently cause the evolution of the distribution. If the drift function is a curl, namely g = ∇ × F , then the distribution does not change, is it still possible to learn the density evolution? The answer is yes. To demonstrate this point of view, we simulate a curl field (y,−x) induced by A = [(0, 1), (−1, 0)] on a Gaussian distribution (indicated by green points in Figure 5 (a)), here we set noise part as 0. Realizing that our learning is based on samples but not on the density, suppose we observe samples at four time points which does not perfectly follow the distribution (indicated by four colors: pink(t0), orange(t1), red(t2) and brown(t3) in Figure 5 (a)), we learn and predict the distribution at t1 and t3 (indicated by blue points in Figure 5 (b) and (c)). We also show the vector field learned by our model in Figure 5 (d). We see that the predictions and vector field all satisfy the correct curl field pattern.
Learning Diffusion Coefficient: Our framework also works for learning unknown diffusion coefficient in the Itô process. As an extension of our work, if we approximate the diffusion coefficient with a neural network ση (with parameters η), we revise the operator F as:
F f (X) = 1 N N∑ k=1  D∑ i=1 giω(x (k)) ∂ ∂xi f (x(k)) + D∑ i=1 ( D∑ j=1 1 2 (σi jη (x(k)))2 ) ∂2 ∂x2i f (x(k))  (9) which can be derived by the same technique we used to derive Proposition 1. We test this formulation on a synthetic data set, where we only consider diffusion influence, namely, drift term in Equation (8) is ignored. We set the ground truth of diffusion coefficient as σ = [(1, 0), (0, 2)]. We design the neural network as a simple one fully connected layer with 32 nodes, then show our result in Figure 6, we see that the predictions(blue) follow the same pattern as the ground truth(red) does.
5 Conclusion
In this paper, we formulate a novel method to recover the hidden dynamics from aggregate data. In particular, our work shows one can simulate the evolving process of aggregate data as an Itô process, in order to investigate aggregate data, we derive a new model that employs the weak form of FPE as well as the framework of WGAN. Furthermore, in Appendix we prove the theoretical guarantees of the error bound of our model. Finally we demonstrate our model through experiments on three synthetic data sets and two real-world data sets.
References D. Alshamaa, A. Chkeir, F. Mourad-Chehade, and P. Honeine. Hidden markov model for indoor
trajectory tracking of elderly people. In IEEE Sensors Applications Symposium (SAS), 2019.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein gan. In arXiv preprint arXiv:1701.07875, 2017.
Kendall E Atkinson. An introduction to numerical analysis. John wiley & sons, 2008.
C. Beck, S. Becker, P. Grohs, N. Jaafari, and A. Jentzen. Solving stochastic differential equations and kolmogorov equations by means of deep learning. 2018.
R. Chen, Y. Feng, and D. Palomar. Forecasting intraday trading volume: a kalman filter approach. In Available at SSRN 3101695, 2016.
M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, 2013.
M. Deriche, A. A. Absa, A. Amin, and B. Liu. A novel approach for salt dome detection and tracking using a hybrid hidden markov model with an active contour model. Journal of Electrical Systems, 16(3):276–294, 2020.
P. M. Djuric, J. H. Kotecha, J. Zhang, Y. Huang, T. Ghirmai, M. F. Bugallo, and J. Miguez. Particle filtering. IEEE Signal Processing Magazine, 20(5):19–38, 2003.
S. R. Eddy. Hidden markov models. Current Opinion in Structural Biology, 6(3):361–365, 1996.
Y. Fang, C. Wang, W. Yao, X. Zhao, H. Zhao, and H. Zha. On-road vehicle tracking using part-based particle filter. IEEE Transactions on Intelligent Transportation Systems, 20(12):4538–4552, 2019.
F. Farahi and H. S. Yazdi. Probabilistic kalman filter for moving object tracking. Signal Processing: Image Communication, 82, 2020.
X. Glorot and Y. Bengio. Understanding the difficulty of training deep feedforward neural networks. In International Conference on Artificial Intelligence and Statistics, 2010.
A. C. Harvey. Forecasting, structural time series models and the Kalman filter. Cambridge University Press, 1990.
T. Hashimoto, D. Gifford, and T. Jaakkola. Learning population-level diffusions with generative rnns. In International Conference on Machine Learning, pp. 2417–2426, 2016.
A. Hefny, C. Downey, and G. J. Gordon. Supervised learning for dynamical system learning. In Neural Information Processing Systems, 2015.
S. C. Hicks, M. Teng, and R. A. Irizarry. On the widespread and critical impact of systematic bias and batch effects in single-cell rna-seq data. bioRxiv, 2015.
R. E. Kalman. A new approach to linear filtering and prediction problems. arXiv preprint arXiv:1805.04099, 1960.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2014.
A.M. Klein, L. Mazutis, I. Akartuna, N. Tallapragada, A. Veres, V. Li, L. Peshkin, D.A. Weitz, and M.W. Kirschner. Droplet barcoding for single-cell transcriptomics applied to embryonic stem cells. Cell, 161(5):1187–1201, 2015.
Peter E Kloeden and Eckhard Platen. Numerical solution of stochastic differential equations, volume 23. Springer Science & Business Media, 2013.
J. Langford, R. Salakhutdinov, and T. Zhang. Learning nonlinear dynamic models. In International Conference on Machine Learning, pp. 593–600, 2009.
W.C. Li, S. Liu, H. Zha, and H. Zhou. Parametric fokker-planck equation. In Geometry science of information, 2019.
Y. Li. A data-driven method for the steady state of randomly perturbed dynamics. arXiv preprint arXiv:1805.04099, 2018.
Q. Liu and D. Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. In Neural Information Processing Systems, pp. 2378–2386, 2016.
G. N. Milstein and M. V. Tretyakov (eds.). Stochastic numerics for mathematical physics. Springer Science & Business Media, 2013.
T. Miyato, T. Kataoka, M. Koyama, and Y. Yoshida. Spectral normalization for generative adversarial networks. arXiv preprint arXiv:1802.05957, 2018.
E. Nelson. Quantum fluctuations. Princeton University Press, 1985.
Bernt Øksendal. Stochastic differential equations. In Stochastic differential equations, pp. 65–84. Springer, 2003.
M. Pavon, E. G. Tabak, and G. Trigila. The data-driven schroedinger bridge. arXiv preprint arXiv:1806.01364, 2018.
D. Qi and A.J. Majda. Low-dimensional reduced-order models for statistical response and uncertainty quantification: Two-layer baroclinic turbulence. Journal of the Atmospheric Sciences, 73(12): 4609–4639, 2016.
D.J. Rezende and S. Mohamed. Variational inference with normalizing flows. In arXiv preprint arXiv:1505.05770, 2015.
H. Risken. The fokker-planck equation. Springer Series in Synergetics, 18:4609–4639, 1989.
H. Risken and T.K. Caugheyz (eds.). The fokker-planck equation: Methods of solution and application. Springer, 1991.
N. P. Santos, V. Lobo, and A. Bernardino. Unmanned aerial vehicle tracking using a particle filter based approach. In IEEE Underwater Technology (UT), 2019.
Justin Sirignano and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial differential equations. Journal of computational physics, 375:1339–1364, 2018.
C. Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
Y. Wang, B. Dai, L. Kong, S. M. Erfani, J. Bailey, and H. Zha. Learning deep hidden nonlinear dynamics from aggregate data. In Uncertainty in Artificial Intelligence, 2018.
E. Weinan, J. Han, and A. Jentzen. Deep learning-based numerical methods for high-dimensional parabolic partial differential equations and backward stochastic differential equations. In Communications in Mathematics and Statistics, pp. 349–380, 2017.
Y. Zang, G. Bao, X. Ye, and H. Zhou. Weak adversarial networks for high-dimensional partial differential equations. In arXiv preprint arXiv:1907.08272, 2019.
O.C. Zienkiewicz and I.K. Cheung. The Finite Element Method in Engineering Science. McGrawHill European Publishing Programme. McGraw-Hill, 1971. ISBN 9780070941380. URL https: //books.google.com/books?id=B99RAAAAMAAJ.
A Appendix
B Error Analysis
In this section, we provide an error analysis of our model. Suppose the hidden dynamics is driven by gr(x), the dynamics that we learn from data is g f (x), then original Itô process, Euler processes computed by true gr and estimated g f are:
dX = g(X)dt + σdW
xrt+∆t = x r t + gr(x r t )∆t + σ
√ ∆tN(0, 1)
x ft+∆t = x f t + g f (x f t )∆t + σ
√ ∆tN(0, 1)
where X is the ground truth, xr is computed by true gr and x f is computed by estimated g f . Estimating the error between original Itô process and its Euler form can be very complex, hence we cite the conclusion from (Milstein & Tretyakov, 2013) and focus more on the error between original form and our model.
Lemma 2. With the same initial Xt0 = xt0 = x0, if there is a global Lipschitz constant K which satisfies:
|g(x, t) − g(y, t)|≤ K|x − y|
then after n steps, the expectation error between Itô process xtn and Euler forward process x r tn is:
E|xtn − xrtn |≤ K ( 1 + E|X0|2 )1/2 ∆t
Lemma 2 illustrates that the expectation error between original Itô process and its Euler form is not related to total steps n but time step ∆t.
Proposition 3. With the same initial x0, suppose the generalization error of neural network g is ε and existence of global Lipschitz constant K:
|g(x) − g(y)|≤ K|x − y|
then after n steps with step size ∆t = T/n, the expectation error between Itô process xtn and approximated forward process x ftn is bounded by:
E|xtn − x f tn |≤
ε K (eKT − 1) + K(1 + E|x0|2)1/2∆t (10)
Proposition 3 implies that besides time step size ∆t, our expectation error interacts with three factors, generalization error, Lipschitz constant of g and total time length. In our experiments, we find the best way to decrease the expectation error is reducing the value of K and n.
C Proofs
C.1 Proof of Proposition 1
Proof. Suppose x̂(k)tm and x̂ (k) tm−1 are our observed samples at tm and tm−1 respectively, then expectations could be approximated by:
Ex∼ p̂(x,tm)[ f (x)] = ∫
f (x)p̂(x, tm)dx ≈ 1 N N∑ k=1 f (x̂(k)tm ) (11)
Ex∼p̃(x,tm)[ f (x)] = ∫ f (x) p̃(x, tm)dx = ∫ f (x) [ p̂(x, tm−1) + ∫ tm tm−1 ∂p(x, τ) ∂t dτ ] dx
= ∫ f (x) p̂(x, tm−1)dx + ∫ f (x) ∫ tm tm−1 ∂p(x, τ) ∂t dτdx
≈ 1 N N∑ k=1 f (x̂(k)tm−1 ) + ∫ f (x) ∫ tm tm−1 − D∑ i=1 ∂ ∂xi [ giω(x)p(x, τ) ] + 1 2 σ2 D∑ i=1 ∂2 ∂x2i p(x, τ)  dτdx︸ ︷︷ ︸ I
(12)
Then for the second term I above, it is difficult to calculate directly, but we can use integration by parts to rewrite I as:
I = ∫ tm
tm−1 ∫  D∑ i=1 − f (x) ∂ ∂xi giω(x)p(x, τ) + 1 2 σ2 D∑ i=1 f (x) ∂2 ∂x2i p(x, τ)  dxdτ =
∫ tm tm−1 ∫  D∑ i=1 giω(x)p(x, τ) ∂ ∂xi f (x) + 1 2 σ2 D∑ i=1 p(x, τ) ∂2 ∂x2i f (x)  dxdτ =
∫ tm tm−1 Ex∼p(x,τ)  D∑ i=1 giω(x) ∂ ∂xi f (x)  + Ex∼p(x,τ) 12σ2 D∑ i=1 ∂2 ∂x2i f (x)  dτ ≈
∫ tm tm−1 1 N N∑ k=1  D∑ i=1 giω(x (k)) ∂ ∂xi f (x(k)) + 1 2 σ2 D∑ i=1 ∂2 ∂x2i f (x(k))  dτ (13) To approximate the integral from tm−1 to tm, we adopt trapezoid rule, then we could rewrite the expectation in Equation (12) as:
Ex∼p̃(x,tm)[ f (x)] ≈ 1 N N∑ k=1 f (x̂(k)tm−1 ) + ∆t 2  1N N∑ k=1  D∑ i=1 giω(x̂ (k) tm−1 ) ∂ ∂xi f (x̂(k)tm−1 ) + 1 2 σ2 D∑ i=1 ∂2 ∂x2i f (x̂(k)tm−1 )  +
1 N N∑ k=1  D∑ i=1 giω(x̃ (k) tm ) ∂ ∂x f (x̃(k)tm ) + 1 2 σ2 D∑ i=1 ∂2 ∂x2i f (x̃(k)tm )  =
1 N N∑ k=1 f (x̂(k)tm−1 ) + ∆t 2 [ F f (X̂m−1) + F f (X̃m) ] (14)
We subtract (11) by (14) to finish the proof.
C.2 Proof of Proposition 2
Proof. Given initial x̂t0 , we generate x̃t1 , x̃t2 , x̃t3 ... x̃tn sequentially by Equation (6). Then the expectations can be rewritten as:
Ex∼ p̂(x,tn)[ f (x)] = ∫
f (x) p̂(x, tn)dx ≈ 1 N N∑ k=1 f (x̂(k)tn ) (15)
Ex∼p̃(x,tn)[ f (x)] ≈ 1 N N∑ k=1 f (x̂(k)t0 ) + ∫ t1 t0 1 N N∑ k=1  D∑ i=1 giω(x (k)) ∂ ∂xi f (x(k)) + 1 2 σ2 D∑ i=1 ∂2 ∂xi∂x j f (x(k))  dτ +
∫ t2 t1 1 N N∑ k=1  D∑ i=1 giω(x (k)) ∂ ∂xi f (x(k)) + 1 2 σ2 D∑ i=1 ∂2 ∂x2i f (x(k))  dτ + ... +
∫ tn tn−1 1 N N∑ k=1  n∑ i=1 giω(x (k)) ∂ ∂xi f (x(k)) + 1 2 σ2 n∑ i=1 ∂2 ∂x2i f (x(k))  dτ (16) which is:
Ex∼p̃(x,tn)[ f (x)] ≈ 1 N N∑ k=1 f (x̂(k)t0 ) + ∆t 2 [ F f (X̂0) + F f (X̃1) ] + ∆t 2 [ F f (X̃1) + F f (X̃2) ] + ...
+ ∆t 2
[ F f (X̃n−1) + F f (X̃n) ] (17)
Finally it comes to:
Ex∼p̃(x,tn)[ f (x)] ≈ 1 N N∑ k=1 f (x̂(k)t0 ) + ∆t 2 F f (X̂0) + F f (X̃n) + 2 n−1∑ s=1 F f (X̃s)  (18)
We subtract (15) by (18) to finish the proof.
C.3 Proof of Error Analysis
Proof. The proof process of Lemma 2 is quite long and out of the scope of this paper, for more details please see first two chapters in reference book (Milstein & Tretyakov, 2013). While for the proof of Proposition 3, with initial X and first one-step iteration:{
xrt0 = xt0 x ft0 = xt0
(19)xrt1 = xrt0 + gr(xrt0 )∆t + σ √
∆tN(0, 1) x ft1 = x f t0 + g f (x f t0 )∆t + σ √ ∆tN(0, 1)
(20)
Then we have:
E|xrt0 − x f t0 | = E|xt0 − xt0 |= 0 (21) E|xrt1 − x f t1 | = E|x r t0 − x f t0 + gr(x r t0 )∆t − g f (x f t0 )∆t + σ √ ∆tN(0, 1) − σ √ ∆tN(0, 1)|
≤ E|xrt0 − x f t0 |+E|gr(x r t0 ) − g f (x f t0 )|∆t = E|gr(xrt0 ) − g f (x r t0 ) + g f (x r t0 ) − g f (x f t0 )|∆t ≤ E|gr(xrt0 ) − g f (x r t0 )|∆t + E|g f (x r t0 ) − g f (x f t0 )|∆t ≤ ε∆t + E|g f (xrt0 ) − g f (x f t0 )|∆t = ε∆t + E|g′f (x ξ t0 )(x r t0 − x f t0 )|∆t (x ξ t0 ∈ [x r t0 ,x f t0 ]) ≤ ε∆t + KE|xrt0 − x f t0 |∆t
= ε∆t (22)
Follow the pattern we have: xrt2 = xrt1 + gr(xrt1 )∆t + σ √
∆tN(0, 1) x ft2 = x f t1 + g f (x f t1 )∆t + σ √ ∆tN(0, 1)
(23)
... xrtn = xrtn−1 + gr(xrtn−1 )∆t + σ √
∆tN(0, 1) x ftn = x f tn−1 + g f (x f tn−1 )∆t + σ √ ∆tN(0, 1)
(24)
Which leads to:
E|xrt2 − x f t2 | = E|x r t1 − x f t1 + gr(x r t1 )∆t − g f (x f t1 )∆t + σ
√ ∆tN(0, 1) − σ √ ∆tN(0, 1)|
≤ E|xrt1 − x f t1 |+E|gr(x r t1 ) − g f (x f t1 )|∆t ≤ E|xrt1 − x f t1 |+ε∆t + KE|x r t1 − x f t1 |∆t
≤ (1 + K∆t)ε∆t + ε∆t (25) ...
E|xrtn − x f tn | ≤ ε∆t n−1∑ i=0 (1 + K∆t)i (26)
Now let S = ∑n−1
i=0 (1 + K∆t) i, then consider followings:
S (K∆t) = S (1 + K∆t) − S
= n∑ i=1 (1 + K∆t)i − n−1∑ i=0 (1 + K∆t)i = (1 + K∆t)n − 1
= (1 + K T n )n − 1
≤ eKT − 1 (27) Finally we have:
E|xrtn − x f tn |≤
ε K (eKT − 1) (28)
E|xtn − x f tn |≤
ε K (eKT − 1) + K(1 + E|x0|2)1/2∆t (29)

Review summaries:
The paper addresses an interesting problem: learning stochastic dynamics from aggregate data. The aggregate data refers to the setting when the data is anonymized. For example, one has data about the location of the birds at different instants in time, but the birds are not labeled and can not be distinguished from each other. For the purpose of learning, the paper proposes to use a generator, in the setting of WGAN, to learn the drift term in the dynamics. The novelty of the proposed approach is to consider the weak form of the Fokker-Planck equation to express the loss function. 

The paper is about an interesting setting, where observations correspond to lots of individuals evolving over time. The twist is that individual identifiers are not available.  Thus from one time point to the next, the observer does not know which individual is which. The problem is to learn a time series model from this, and models considered here are in the form of an SDE.

This paper proposes a new approach to learn the dynamics of density evolution of objects from aggregated data. The basic idea is to derive a closed-form Wasserstein distance between the empirical data distribution and the predicted distribution generated from the weak form of Fokker Planck Equation (FPE). Based on this measure, an objective function is developed to learn the underlying drift coefficients and the discriminator simultaneously using neural networks. Some theoretical results are provided and numerical experiments are carried out to illustrate the effectiveness of the proposed method.