Learning Stochastic Behaviour from Aggregate Data

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 .

Summary: 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.