Variational Inference for Interacting Particle Systems with Discrete Latent States

Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-06 10:24:57

Abstract: We present a novel Bayesian learning framework for interacting particle systems with discrete latent states, addressing the challenge of inferring dynamics from partial, noisy observations. Our approach learns a variational posterior path measure by parameterizing the generator of the underlying continuous-time Markov chain. We formulate the problem as a multi-marginal Schrödinger bridge with aligned samples, employing a two-stage learning procedure. Our method incorporates an emission distribution for decoding latent states and uses a scalable variational approximation.

Main paper: 
 introduction: 
 Many real-world phenomena, from epidemics to wildfires, can be modeled as systems of interacting components evolving in continuous time, where the underlying dynamics are governed by discrete latent states. This paradigm extends the concept of hidden Markov models [Baum and Petrie, 1966, Kouemou and Dymarski, 2011] to spatially structured, continuous-time processes. Interacting particle systems (IPSs) [Liggett, 1985 , Lanchier, 2024] offer a powerful mathematical framework for describing local propagation dynamics. However, inferring the rules governing these systems from partial, noisy observations remains a significant challenge. We propose a novel Bayesian approach that addresses this challenge by learning a variational posterior path measure on the space of IPS trajectories. Our approach parameterizes the generator of the continuous-time Markov chain (CTMC) of the latent IPS using neural networks and incorporates an emission model that can decode internal discrete states to continuous data and noisy observations. Key contributions of our approach include:• Framing the problem as a multi-marginal Schrödinger bridge with aligned samples [Somnath et al., 2023] , solved by a two-stage procedure: learning an endpoint-conditioned generator for trajectory reconstruction, followed by distillation to an unconditional generator for prediction.• A scalable variational approximation using site-wise factorization of time-marginals and assuming independent particle evolution in infinitesimal time intervals conditionally on the present global configuration, enabling efficient learning for high-dimensional spatiotemporal processes.• Flexibility in incorporating domain knowledge through through informative priors on rate matrix entries and neural architectures with desirable inductive biases.We demonstrate preliminary results of our approach on two simulated datasets, for the following tasks: reconstructing the trajectory of an epidemic on a network and predicting wildfire spread on a lattice. For a description of the notation, see Appendix A. An overview of the relevant literature is presented in Appendix B, whereas proofs and other derivations are provided in Appendix C.Submitted to Workshop on Bayesian Decision-making and Uncertainty, 38th Conference on Neural Information Figure 1 : An illustration of our methodology on a simulated noiseless dataset of wildfire propagation. The first model approximates a Markovian bridge interpolating between the observations, enabling to reconstruct the unobserved trajectory. The second model, approximating the unconditional process, can predict beyond the last observation. Results shown for a held-out example. 
background: 
 Interacting particle systems Consider a graph G = (V, E), and denote i ∼ j if there is an edge between nodes i, j, i.e., {i, j} ∈ E. Following Liggett [1985] , we refer to vertices i ∈ V as sites.For a countable state space S, consider the configuration space X := {x | x : V → S}. For our analysis, we assume both V and S are finite. An IPS adds a continuous-time dimension to this setting. Specifically, we obtain a CTMC (x t ) on X restricted to a time interval [0, T ], whose path space we denote Ω [0,T ] . We define x t (i) as the state of site i at time t. We consider a scenario where the dynamics of each site are described by local transition rates that depend on the graph's connectivity [Lanchier, 2017] , corresponding toλ s→s ′ t (i, x t ) := lim h↓0 h -1 P (x t+h (i) = s ′ | x t (i) = s, x t (j) : i ∼ j) , for s to s ′ ̸ = s at site i and time t ∈ [0, T ], and set λ s→s t (i, x t ) := -s ′ ̸ =s λ s→s ′ t (i, x t ). The local transition rates are aggregated into a generator Λ t (x t ) = [λ s→s ′ t (i, x t )] i∈V ;,s,s ′ ∈S . Note that this is a mapping Λ • (•) : X × R+ → R |V |×|S|×|S| .For brevity, we denote the collection of transition rates between two fixed configurations x, x ′ ∈ X asΛ t (x ′ | x) := [λ x(i)→x ′ (i) t (i, x)] i∈V . The path measure Π ∈ M(Ω [0,T ]) of a realization (x t ) t∈[0,T ] can then be described by the solution to an initial value problem with starting condition π t (x 0 ) = π 0 , evolving according to the master equation∂ t π t (x t ) = x ′ ̸ =xt (Λ t (x t | x ′ )π t (x ′ ) -Λ t (x ′ | x t )π t (x t )) .(1)Markovian bridges Let {y k } k∈[K] ∈ Y K be a sequence recorded at times {t k } k∈[K] ∈ R K , where 0 = t 0 < • • • < t K-1 = T ,and assume that each value in the sequence is independent conditionally on the Markov process(x t ) t∈[0,T ] ∈ Ω [0,T ] . The conditioned Markov process with path measure Π •|{y k } k∈[K] ∈ M(Ω [0,T ]) is known as a Markovian bridge. At any time t ∈ (t k , t k+1 ] and for states s and s ′ (s ̸ = s ′ ), the transition rate is:EQUATIONHence, the conditional generatorΛ t (x t | {y k } k∈[T ] ) of Π •|{y k } k∈[K] is equivalent to the generator Λ t (x t | y t k+1 ) of the Markovian bridge Π •|yt k ,yt k+1 at any time t ∈ (t k , t k+1 ]. See Appendix C.1for an overview of Markovian bridges, and we refer the reader to Fitzsimmons et al. [1992] for a detailed construction. 
variational discrete interacting particle systems: 
 We consider sequences of observations{y i k } k∈[K] ∈ Y K recorded at times {t i k } k∈[K] ∈ R K andassumed to be independent conditionally on the interacting particle system(x i t ) t∈[0,T i ] ∈ Ω [0,T i ] ,for i = 1, . . . , N . The discrete set of measurement times0 = t i 0 < • • • < t i K-1 = T i are allowedto be arbitrarily defined for each sequence, e.g., at random or regularly spaced. We assume that the graph determining the particles' dependence structure is fixed for each realization and directly deducible from the observed sequences, for example if the observations are noisy measurements of the IPS. We consider an emission distribution p t (y | x) ∈ P(Y) for the observations, and a prior path measure P ∈ M(Ω [0,T ] ) for the latent IPS. This can be specified directly on the entries of a prior generator, encoding possible constraints in the latent dynamics, and by an initial prior distribution. The marginal distribution of the data at any time t ∈ [0, T ] is denoted as π t ∈ P(Y),and we are interested in parameterizing a variational path measure Q ∈ M(Ω [0,T ] ) and an encoder q t (x | y) ∈ P(X ). To make inference tractable, we detail specific choices of our variational approximation in Section 3.1. Considering a single sequence {y k } k∈[K] , our goal is twofold:• Trajectory reconstruction, by learning the conditional generatorΛ t (• | {y k } k∈[T ] ) of the Markovian bridge Q •|{y k } k∈[K] ;• Prediction, by learning the generator Λ t (•) of Q, enabling extrapolation beyond an observed time window or with no past observations at all for a given graph.Inference for both models can be conveniently amortized by parameterizing the generators with neural cellular automata, as detailed in Section 3.2. We show that the second goal can be achieved by distilling knowledge from a model trained for the first goal into a model that does not glance at future observations. Sampling is discussed in Appendix D.2.Contrary to previous work [Opper and Sanguinetti, 2007 , Cohn et al., 2010 , Seifner and Sánchez, 2023] we do not adopt a mean-field approximation, but rather the star-approximation introduced for continuous-time Bayesian networks in Linzner and Koeppl [2018] . While still assuming sitewise factorization of the time-marginals as q t (x) = i∈V q i t (x(i)), we couple sites by letting the generator depend on the present global configuration x t . This is achieved by considering a system of marginal master equationsEQUATIONUnder these assumptions, the explicit form for the Kullback-Leibler (KL) divergence of two CTMCs with path measures Q, P ∈ M(Ω [0,T ] ) is tractable, see Appendix C.2 for a derivation.Let P ∈ P(Y K × Ω [0,T ] ) denote the reference measure constructed by gluing the prior and emission probabilities at each observed timestep. For a given sequence of distributions{π t k } k∈[K] on P(Y),we can express a multi-marginal Schrödinger bridge problem with noisy observations asEQUATIONwhere q t k ∈ P(Y) correspond to marginalizations of the variational distribution at each observed timepoint. Let π t 0:T -1 denote the coupling solving the static version of (4), that isEQUATIONwhere p t 0:T -1 ∈ P(Y K ) is the marginal of the observed trajectories obtained from the reference measure P. Similarly to the setting considered in Somnath et al. [2023] , we assume that our dataset is comprised of trajectories of aligned samples, in the sense that each observed trajectory{y k } k∈[T ]is sampled from the coupling π t 0:T -1 . We denote its marginals at any pair of observed time pointsas π t k ,t k+1 ∈ P(Y × Y), for k ∈ [K -1].Proposition 1 Let (5) admit a solution π t 0:T -1 . Moreover, assume conditional independence of{y k } k∈[K] given (x t ) ∈ Ω [0,T ], and let P ∈ M(Ω [0,T ] ). Then, the problem in (4) reduces to arg minEQUATIONwhere q •|y0 ∈ P(X ) andQ •|y k ,y k+1 ∈ M(Ω (t k ,t k+1 ] ) for k ∈ [K -1].Trajectory reconstruction As any conditional path measure can be fully characterized by its generator and an initial distribution, we leverage the correspondence explored in Section 2 to specifyQ •|{y k } k∈[K] by a sequence of Markovian bridges Q •|y k ,y k+1 and conditional distributions q t k (• | y k ), for k ∈ [K -1].We parameterize the former with a neural model Λ θ t (• | y k+1 ), and the latter with a probabilistic encoder q i,θ t (• | y k ) = Categorical(g θ t (i, y k )), where g θ : R + ×V ×Y → ∆ |S| and ∆ |S| denotes the |S|-dimensional simplex. Moreover, we parameterize the emission distribution with a probabilistic decoder p ξ t (y | x). We learn θ and ξ by minimizing an evidence lower bound derived from (6), given byEQUATIONwhere EQUATION(q i,θ t ) t∈[t k ,t k+1 ] } i∈Vis then obtained by numerically solving a system of marginal master equations as in (3), according to the neighborhood dynamic established byΛ θ t (x t | y t k+1 ), for t ∈ (t k , t k+1]. This step can be achieved by using adaptive solvers, but for illustration purposes we remind the reader that an Euler update with step size 0 < ϵ ≪ 1 can be performed asq i,θ t+ϵ = q i,θ t + ϵ[Λ θ t (x t | y t k+1 )] ⊤ i q i,θ t , i ∈ V.(9) As we are interested in evolving the entire vector of probability mass functions at each timestep, we sample x t for t ∈ (t k , t k+1 ] from q i,θ t using the Gumbell-Softmax trick [Jang et al., 2017] . Additional details are reported in Appendix D.1, and considerations of computational cost are discussed in Appendix D.4.Prediction We fix the conditional process Q θ , and learn the unconditional IPS by distillation to a Markovian generator Λ ϕ t (x t ). This can be achieved by minimizing a mean squared error loss on the rates, given byEQUATIONwhereEQUATIONIf the conditional generator Λ θ t is unbiased and q θ t ≈ q * t , the minimizer of (10) recovers the Markovprocess Q ∈ P(Ω [0,T ] ) obtained by marginalizing Q * in (4) to Ω [0,T ] . A derivation is provided in Appendix C.5. 
experiments: 
 We demonstrate our methodology on two simulated scenarios: epidemic trajectory inference on networks and wildfire spread prediction on lattices. We parameterize the neural models for the generators with a Vision Transformer Cellular Automata [Tesfaldet et al., 2022] . Results and details of the simulations are reported in Appendix E. 
conclusion: 
 We introduce a variational inference method to fit partially observed trajectories whose dynamic can be modeled by a continuous-time latent process, parameterized to be an interacting particle system. Our solution is an approximation to a multi-marginal Schrödinger bridge, that we obtain by first fitting an endpoint-conditioned model and then distilling it into an unconditional one. This methodology enables both trajectory reconstruction and prediction of future states. In future work we aim at testing our models on real data, comparing with state-of-the-art methods. 
 
 Appendix: Let Ω [0,T ] be the space of X -valued cadlag functions over a time interval [0, T ], and denote byP(Ω [0,T ]) the space of probability measure on the path space. For a path measureQ ∈ P(Ω [0,T ] ), timesteps v, t ∈ [0, T ] s.t. v > t,and configurations x, x ′ ∈ X , we assume that time-marginal and transition probability measures are absolutely continuous w.r.t. the counting measure. Their Radon-Nikodym derivative can then be expressed by the probability mass function q t (x) and the transition probability q s|t (x ′ | x). We denote byΩ (t k ,t k+1 ] time restrictions of Ω [0,T ] to (t k , t k+1 ] ⊆ [0, T ].We denote the index set of N -many elements as [N ] = 0, . . . , N -1. Denote the cartesian product× k∈[K] Y of observations at K times as Y K . Moreover, let M(Ω [0,T ] ) ⊂ P(Ω [0,T ]) denote the space of Markovian probability measures on Ω [0,T ] .The dynamics of many physical systems can be described through the local interaction laws of their constituent components. This principle has inspired computational frameworks that directly parameterize these governing interactions, both deterministically and stochastically. A prime example is cellular automata [Wolfram, 1986 , Grinstein et al., 1985] . Early developments focused on studying the emergence of global patterns from a fixed set of rules on the evolution of individual cells.The inverse problem -inferring such rules from observations-has been of historical interest in the machine learning community [Wulff and Hertz, 1992, Mordvintsev et al., 2020] , with recent developments incorporating attention-based architectures, graph neural networks, and black-box variational inference [Tesfaldet et al., 2022 , Kang et al., 2024 , Grattarola et al., 2021 , Palm et al., 2022] . Models that learn interaction rules find applications across many domains, including physical simulators, multi-agent dynamics, dynamic graphs, as well as deep generative modeling [Kalkhof et al., 2024] .Within this context, most existing methods have proposed iterative updating schemes by parameterizing transition rules in discrete time. Interacting particle systems (IPSs) offer an alternative mathematical formalism that extends cellular automata to continuous time. Interacting particle systems are structured CTMCs whose states evolve with dependence on neighbors within a topology, typically established by a graph. Lanchier [2017] provides a modern introduction to this field. Classical literature focused on systems with finite states and often countably many sites [Bramson and Griffeath, 1980 , Liggett, 1985 , Durrett, 2006] , while more recent work has focused on systems with finitely many sites [Aldous, 2013] . These systems have found applications in multi-agent modeling [Comas et al., 2023] and have been extended to systems of coupled stochastic differential equations (SDEs) in Euclidean space. This extension has seen increased attention recently [Lu et al., 2021 , Yang et al., 2022 , Feng et al., 2022 , Liu et al., 2023 , Lang et al., 2024 , Kümmerle et al., 2024 , Boffi and Vanden-Eijnden, 2024] . The learnability and identifiability of interaction rules in these systems have also been explored [Bongini et al., 2017 , Li et al., 2021] .Inference methods for Markov jump processes (MJPs) have been extensively studied. Maxi- [Seifner and Sánchez, 2023] and foundation models (i.e., meta-learning) [Berghaus et al., 2024] . While less directly related, it's worth noting recent work on guidance and conditioning for Markovian bridges [Corstanje et al., 2023] and discrete flow matching and diffusion methods [Campbell et al., 2022 , Igashov et al., 2023 , Lou et al., 2023 , Campbell et al., 2024] .mumTrajectory inference is a crucial component of our work, with connections to several recent developments. The Schrödinger bridge (SB) problem with multi-marginal constraints has been exploredby Chen et al. [2019] , Lavenant et al. [2021] . Recent advances in SB methods with a source and a target are presented in Vargas et al. [2021] and De Bortoli et al. [2021] , with extensions to the multi-marginal setting by Shen et al. [2024] . Our approach shares similarities with Somnath et al.[2023], Shi et al. [2024] , and Peluchetti [2023] in that it relies on samples from couplings solving the static SB problem. However, our methodology differs in that we learn the Markovian bridge and recover the unconditional path measure by distillation, rather than relying on closed-form endpointconditioned diffusions. The concept of Markov bridge by interpolation with a fictitious dynamic, as proposed by Igashov et al. [2023] , is related to stochastic interpolants [Albergo and Vanden-Eijnden, 2022 , Tong et al., 2023 , Lipman et al., 2022 , Liu et al., 2022] for probabilistic forecasting [Chen et al., 2024] . Ad-hoc variants for dynamical systems have also been developed [Rühling Cachay et al., 2024] . Our methodology also shares connections with flow matching using Gaussian process and Kalman filter interpolants [Tamir et al., 2023] , in the fact that we are interested in model-based interpolants in a Bayesian framework.Consider a sequence of observations{y k } k∈[K] ∈ Y K recorded at times {t k } k∈[K] ∈ R K , andassume conditional independence with respect to a Markov process for h ≈ 0, by right-continuity of the transition probabilities. We can then denote the conditional transition rates for x ′ ̸ = x as(x t ) t∈[0,T ] . For t ∈ [t 0 , t K-1 ], let y >t = {y k | t k > t, k = 1, . . . , K} and y ≤t = {y k | t k ≤ t, k = 1, . . . ,Λ(x ′ | x, y 0:K ) = lim h↓0 h -1 [P(x t+h = x ′ | x t = x, y 0:K )] = lim h↓0 h -1 P(x t+h = x ′ , x t = x, y >t | y ≤t ) P(x t = x, y >t | y ≤t ) = lim h↓0 h -1 P(y >t+h | x t+h = x ′ )P(x t+h = x ′ | x t = x)P(x t = x | y ≤t ) P(y >t | x t = x)P(x t = x | y ≤t ) = lim h↓0 h -1 P(y t ′ | x t+h = x ′ )P(x t+h = x ′ | x t = x) P(y t ′ | x t = x) = Λ t (x ′ | x, t) P(y t ′ | x t = x ′ ) P(y t ′ | x t = x), and similarlyΛ(x | x, y 0:K ) = - x ′ ̸ =x Λ t (x ′ | x, t) P(y t ′ | x t = x ′ ) P(y t ′ | x t = x) = Λ t (x | x, t) 1 -P(y t ′ | x t = x) P(y t ′ | x t = x) . (12) C.2 Derivation of D KL (Q || P )Consider two CTMCs with path measures Q, P ∈ P(Ω [0,T ] ), and denote their respective rate matrices entries with Λ t (x ′ | x) and Ψ t (x ′ | x) for x, x ′ ∈ X . Their KL divergence, as discussed in Opper and Sanguinetti [2007] , Seifner and Sánchez [2023] , can be derived from the limit of discrete-time transitions with step size h := T /K asD KL (Q||P ) = lim K→∞ x 0:K q 0 (x 0 ) K-1 k=0 q k+h|k (x k+h | x k ) log q 0 (x 0 ) K-1 k=0 q k+h|k (x k+h | x k ) p 0 (x 0 ) K-1 k=0 p k+h|k (x k+h | x k ) = x0 q 0 (x 0 ) log q 0 (x 0 ) p 0 (x 0 ) + lim K→∞ K-1 k=0 E qt(x)   x k+h q k+h|k (x k+h | x k ) log q k+h|k (x k+h | x k ) p k+h|k (x k+h | x k )   (13) = D KL (q 0 ||p 0 ) + T 0 E qt(x) x ′ ̸ =x Ψ t (x ′ | x) + Λ t (x ′ | x) log Λ t (x ′ | x) Ψ t (x ′ | x) -1 dt,where the last line follows from dividing and multiplying each summand in (13) by h, and substituting the transition probabilities with rates,q k+h|k (x k+h | x k ) h log q k+h|k (x k+h | x k ) p k+h|k (x k+h | x k ) h→0 -→ Λ t (x k+h | x k ) log Λt(x k+h |x k ) Ψt(x k+h |x k ) x k+h ̸ = x k , x ′ ̸ =x [Ψ t (x ′ | x) -Λ t (x ′ | x)] x k+h = x k .By assuming:• Site-wise factorization of the time marginals q t (x) = i∈V q i t (x(i)),• Coupled transitions q t+h|t (x ′ | x) = i∈V q t+h|t (x ′ (i) | N i t (x))where we define a neighborhood N i t (x) := {x(i), x(j) : i ∼ j}, we can rewrite each summand in ( 13) asE q k (x k )   x k+h q k+h|k (x k+h | x k ) log q k+h|k (x k+h | x k ) p k+h|k (x k+h | x k )   = E q k (x k ) i∈V q i k+h, k (x k+h (i) | N i t (x k )) log q i k+h, k (x k+h (i) | N i t (x k )) p i k+h, k (x k+h (i) | N i t (x k )).Letting K → ∞ and plugging (3), we getEQUATIONMoreover, notice that if we let p 0 be a uniform distribution the KL between initial distributions reduces to D KL (q 0 ||p 0 ) = H(q 0 ) -log(|X |), where H(•) is the entropy.The additive property of the KL divergence [Léonard, 2013] states that for a Polish space Q := Y K × Ω [0,T ] , the canonical projecting onto the trajectory coordinates ϕ :Q → Y K , a measurable mapping y 0:K-1 ∈ Y K → Q •|y 0:K-1 ∈ P(Ω [0,T ]), and Q, P ∈ P(Q), we getEQUATIONwhere ϕ # Q(A) = Q(ϕ -1 (A)) for any Borel set A ⊆ Y K denotes a pushforward measure. Denote by q t 0:K-1 (y 0:K-1 ) the p.m.f. associated to q t 0:K-1 := ϕ # Q, and p t 0:K-1 (y 0:K-1 ) the p.m.f. for p t 0:K-1 := ϕ # P. For clarity, we rewrite (15) asD KL (Q || P) = E qt 0:K-1 (y 0:K-1 ) log q t 0:K-1 (y 0:K-1 ) p t 0:K-1 (y 0:K-1 ) + D KL (Q •|y 0:K-1 || P •|y 0:K-1 ) . (16)Now consider the canonical projection φ : Q → Ω [0,T ] , and assume that P := φ # P ∈ M(Ω [0,T ] ).It follows from Léonard [2013, Prop. 2 .10] thatQ * := φ # Q * ∈ M(Ω [0,T ] ), where Q * is the solution to (4). Hence, without loss of generality we restrict our analysis to measures Q s.t.Q := φ # Q ∈ M(Ω [0,T ] ).Moreover, consider the case where Q and P are fully specified by their Y K -projection and disintegrations, such thatdQ = dq •|y0 (x 0 )dQ •|y0,y1 ((x t ) t∈(0,1] ) . . . dQ •|y T -2 ,y T -1 ((x t ) t∈(K-2,K-1] )dq t 0:K-1 (y 0:K-1 )for any Borel set B ∈ Q, and the same for P. This property can be described as conditional independence of y 0:K-1 given (x t ) ∈ Ω [0,T ] . Then, we can decompose ( 16) asD KL (Q || P) = J 1 + J 2 ,whereJ 1 = E qt 0:K-1 (y 0:K-1 ) log q t 0:K-1 (y 0:K-1 ) p t 0:K-1 (y 0:K-1 ) , J 2 = E qt 0 D KL (q •|y0 || p •|y0 ) + k∈[K] E qt k ,t k+1 D KL (Q •|y k ,y k+1 || P •|y k ,y k+1 ) .If we have access to the coupling π t 0:K-1 solving (5), then the only term left depending on Q is J 2 , hence we recover (6).For ease of illustration, we start by considering a loss L 1 with a single component defined in the time frame [0, T ] between two observations y 0 and y T . We might be interested in parameterizing emission distributions p t|t (y | x), hence we denote them in short as p t (y | x). We are interested in proving thatD KL Q || P •|y0,y T = D KL (Q || P ) -E x∼q0 [log p 0 (y 0 | x)] -E x∼q T [log p T (y T | x)] + log p 0,T (y 0 , y T ).We now consider a time-discretization at0 = τ 0 < • • • < τ K-1 = T of the path measures Q and P •|y0,y T . For τ k+1 , τ k ∈ [0, T ] and x k+1 , x k ∈ X, by the Markov property of (x t ) t∈[0,T ] under P and conditional independence of y 0 , y T given x k , we can express the marginal and transition probability mass functions conditionally on y 0 , y T asEQUATION)EQUATION)Notice thatpτ0 (x 0 ) K-2 k=0 pτ k+1 |τ k (x k+1 | x k ) = p τ0 (x 0 ) K-2 k=0 p τ k+1 |τ k (x k+1 | x k ) p 0 (y 0 | x 0 )p T (y T | x T ) p 0,T (y 0 , y T ) .Notice that every marginal and transition probability derived from the variational path measure Q ∈ P(Ω [0,T ] ) can depend on the observations y 0 , y 1 , but we omit them from the notation for brevity.We can then write q τ0:τ T (x 0:T ) = q τ0 (x 0 )K-2 k=0 q τ k+1 |τ k (x k+1 | x k ), and it follows thatEQUATION)This formulation can trivially be extended for a sequence of observationsEQUATION) where Z = p t 0:K-1 ({y k } k∈[K]). Taking the expectation w.r.t. {y k } k∈[K] ∼ π t 0:K-1 on both sides recovers (6) on the LHS, and ( 7) on the RHS plus a term independent of Q. In practice, we might be interested in learning the emission distribution at the same time as we are learning the variational path measure, in which case we can interpret the objective in ( 7) as an evidence lower bound.Once again, for ease of illustration we consider two observations y 0 , y T at the endpoints of a time interval [0, T ]. We are interested in learning the generator of a CTMC (x t ) t∈[0,T ] by approximating it with a neural model Λ ϕ t (x ′ | x). We are only given access to its endpoint-conditioned generator Λ t (x ′ | x, y T ), and by Bayes rule we can recover the unconditional one asEQUATIONFirst, we want to show that the intractable unconditional lossEQUATION)is equivalent to the tractable conditional oneL C (ϕ) := T 0 E y0,y T ∼π 0,T , x∼q t|0,T (•|y0,y T ) x ′ ̸ =x Λ t (x ′ | x, y T ) -Λ ϕ t (x ′ | x) 2 2 dt (24)up to a constant independent of ϕ. This proof is not novel, as it mirrors the proof of Lipman et al.[2022, Theorem 2] adapted to transition rate matrices in discrete spaces rather than vector fields in Euclidean spaces.First, notice that each componentEQUATION)Of the two summands that depend on ϕ, plugging (21) into the first term yieldsE qt(x) Λ t (x ′ | x), Λ ϕ t (x ′ | x) = X Y×Y Λ t (x ′ | x, y T ) q t|0,T (x | y 0 , y T )π 0,T (y 0 , y T ) q t (x) dy 0 dy 1 , Λ ϕ t (x ′ | x) q t (x)dx = X Y×Y Λ t (x ′ | x, y T )q t|0,T (x | y 0 , y T )π 0,T (y 0 , y T ), Λ ϕ t (x ′ | x) dy 0 dy 1 dx = E π 0,T (y0,y T ), q t|0,T (x|y0,y T ) Λ t (x ′ | x, y T ), Λ ϕ t (x ′ | x) .Next, it follows by the law of total expectation that the second term isE qt(x) Λ ϕ t (x ′ | x) 2 2 = E π 0,T (y0,y T ), q t|0,T (x|y0,y T ) Λ ϕ t (x ′ | x) 2 2 .Combining the terms, it follows from the linearity of expectation that∇ ϕ L U (ϕ) = ∇ ϕ L C (ϕ). Next,we want to show that if we approximate the true conditional generator with an estimatorΛ θ t (x ′ | x, y T ) such that E π 0,T (y0,y T ), q t|0,T (x|y0,y T ) [Λ θ t (x ′ | x, y T )] = Λ t (x ′ | x, y T ), we recover L 2 as specified in (10) and a component independent of ϕ. To show this, add and subtractΛ θ t (x ′ | x, y T )from L C (ϕ) and complete the square to getLC (ϕ) = T 0 Eπ 0,T , q t|0,T x ′ ̸ =x Λ θ t (x ′ | x, yT ) -Λ ϕ t (x ′ | x) 2 2 + Λt(x ′ | x, yT ) -Λ θ t (x ′ | x, yT ) 2 2 dt = T 0 Eπ 0,T , q t|0,T x ′ ̸ =x Λ θ t (x ′ | x, yT ) -Λ ϕ t (x ′ | x) 2 2 dt + K θ .In order to retrieve L 2 (ϕ), we perform a change of measure by importance sampling with proposal q θ t (x | y 0 , y T ), and approximate the importance weights q t /q θ t ≈ 1 to getL C (ϕ) ∝ T 0 E π 0,T , q t|0,T x ′ ̸ =x Λ θ t (x ′ | x, y T ) -Λ ϕ t (x ′ | x) 2 2 dt ≈ T 0 E π 0,T , q θ t|0,T x ′ ̸ =x Λ θ t (x ′ | x, y T ) -Λ ϕ t (x ′ | x) 2 2 dt = L 2 (ϕ).The training algorithms for the trajectory reconstruction and prediction tasks are reported in Algorithm 1 and Algorithm 2 respectively. Notice that it is also possible to learn the unconditional generator at the same time as the unconditional one, by freezing the gradients of θ before updating the L 2 loss. For illustration purposes we do not explicit the numerical solver we are using, but for our experiments we use a Dormand-Prince solver of order 5 [Dormand and Prince, 1980] from the torchdiffeq library [Chen, 2018] . While all datapoints in a batch are processed in parallel, we might need to evolve the solver through different time points for each batch. This is feasible by applying the tricks for parallel solving of neural ODEs with varying time-intervals presented in Chen et al. [2021] .Overall, we found training of the unconditional model quite straightforward to implement. On the other hand, training the unconditional model seems to be quite challenging, mainly due to vanishing gradients when backpropagating through the solver. A trick we found quite helpful in addressing this problem is annealing the time discretization grid, from very coarse to a finer and finer one.Moreover, if the prior tends to push the model towards "high activity", we found the states generated by the conditional model to converge to the next observed state pretty quickly. This can hamper the training of the unconditional model, as trajectories generated by the conditional model will converge very quickly and then stay still for a long amount of time. This biases the distribution of samples seen at training time by the conditional model, that might then experience "mode collapse" and predict all of the transition rates to be zero. We found that choosing priors that bias the conditional model towards performing fewer transitions helps addressing this issue.Algorithm 1 Training for trajectory reconstruction Require:Dataset D = {{t k , y i k , C G i k } k∈[Ki] } i∈[N ], prior path measure P , OPTIMIZER, step size ϵ Ensure: Learned parameters θ, ξ, for conditional generator Λ θ t and emission distribution p ξ t 1: Initialize parameters θ, ξ and loss L 1 = 0 2: while not converged do for each pair (y k , y k+1 ) do5: Encode x t k ∼ q θ t k (•|y k ) = Categorical(t k , g θ t k (i, y k )) 6:for t ∈ (t k , t k+1 ] do 7: Sample x t from q i,θ t using Gumbel-Softmax trick 8:Evolve q i,θ t using Λ θ t (x t |y k+1 ) as in (9) 9:end for 10:Compute L 1 ← L 1 + L k 1 (θ, ξ, y k , y k+1 ) using ( 8)11:end for 12:Update θ, ξ ← OPTIMIZER(∇L 1 (θ, ξ)). 13: end while Algorithm 2 Training for prediction for each pair (y k , y k+1 ) doRequire: Dataset D = {{t k , y i k , C G i k } k∈[Ki] } i∈[N ] , conditional generator Λ θ t ,5: Encode x t k ∼ q θ t k (•|y k ) = Categorical(t k , g θ t k (i, y k )) 6:for t ∈ (t k , t k+1 ] do 7: Sample x t from q i,θ t using Gumbel-Softmax trick 8:Evolve q i,θ t using Λ θ t (x t |y k+1 ) as in ( 9) Update ϕ ← OPTIMIZER(∇L 2 (ϕ)). 13: end while9: Compute L 2 ← L 2 + Λ θ t (x t |y k+1 ) -Λ ϕ t (x t )For sampling, we compute transitions directly in sample space, in order to respect possible constraints encoded in the generator. Denoting a neural encoder as q t and the entries of a generator Λ t as Λ t (regardless of it being conditional or unconditional), we initialize x 0 ∼ q 0 (• | y 0 ) and employ a first-order approximation of the transition probability asEQUATION)written in short as x t+ϵ ∼ Categorical (δ xt + ϵΛ t (x t )). This method, despite being a crude approximation, is typically employed in discrete flow matching [Campbell et al., 2024] for its high scalability. The analytical form of the Markovian bridge in (2) suggests a potential unification of our two-stage approach into a single stage. This would involve fitting the conditional generator by minimizing the evidence lower bound in ( 7) and specifying the model as:EQUATION)This approach is analogous to that taken by Somnath et al. [2023] for diffusion Schrödinger bridges, where they learn to approximate a target comprising a guidance term that ensures the diffusion process reaches the prescribed y k+1 at time t k+1 . However, their method relies on the availability of a closed-form for the endpoint-conditioned prior processes, which is not available in our setting.Both guidance approaches can be viewed as instances of Doob's h-transform [Rogers and Williams, 2000] , and we refer the reader to Corstanje et al. [2023] for a detailed discussion. However, our preliminary experiments suggest that this single-stage approach may yield inferior unconditional models compared to the two-stage method, possibly due to identifiability issues. We leave a thorough investigation of single-stage methods to future work.Our method is not simulation-free, in the sense that learning is made possible by backpropagating through a solver. In doing so, a practitioner can incur in two fundamental problems, inaccurate gradients and memory-intensive training steps. The choice of a backpropagation technique can trade off one disadvantage for the other. In our experiments we use continuous adjoint methods, that provide memory-efficient numerical solutions (constant w.r.t. the time discretization grid) at the cost of incurring numerical errors that accumulate into potentially inaccurate gradient estimates. An overview of other possible approaches is presented in [Kidger, 2021] . λ S→I (i, x) = β exp sin(c i 1 ) + cos(c i 2 ) N I i , λ I→R (i, x) = γ,where N I i := {j ∈ V | x(j) = I, j ∼ i}, β = 6 and γ = 0.2. These parameters do not correspond to physically meaningful quantities, and adjusting them to reflect real-world spread dynamics remains an interesting avenue for future work. Each graph is observed at K = 20 regularly spaced time points, with no observation noise (i.e., X ≡ Y). The data is simulated using τ -leaping [Gillespie, 2001] , with τ = 1 × 10 -2 . A sample observed in its first 5 time steps is displayed in Figure 3 . U →B (i, x) = ReLU(a 0 + a 1 f i ) × ReLU   b 0 + b 1 j∈N B i a ij   , λ E→B (i, x) = ReLU(c 0 + c 1 f i ) × ReLU   d 0 + d 1 j∈N B i a ij   , λ B→E (i, x) = γ,where N B i := {j ∈ V | x(j) = I, j ∼ i}, and a ij is a wind alignment value obtained by the dot product between the relative position of the neighbor j w.r.t. i and the value of the wind field at Since there is no observation noise, all we need to parameterize in our experiments are the conditional and unconditional generators. Both can be thought of as mappings X → R, i.e. the output shall be a local transition rate matrix at each site i ∈ V . In order to constrain the dependence structure of each local transition rate to the neighborhood of that site, we use the attention-based neural cellular automata presented in Tesfaldet et al. [2022] . This can be thought of as a depth-one vision transformer whose attention matrix is (efficiently) masked to attend only to a neighborhood of embedded input sites. For the wildfires experiment we simply consider a 3 × 3 Moore neighborhood, whereas for the epidemics we mask the attention matrix with the adjacency matrix of each observation. We obtain 512-dimensional input embeddings by a two-layer MLP of width 512. The embedding is then split into 4 attention heads, combined by another two-layer MLP with 512 hidden units that returns the off-diagonal values of the local transition rate matrices. We constrain the output to be positive by applying a ReLU function. We specify the prior path measure by a prior rate matrix, where we set to zero physically impossible transitions (e.g. U → E for wildfires, or S → R for epidemics) and the remaining off-diagonal elements to a constant value c. More complex functional forms are possible, and shall be chosen for example by simulating from the prior predictive distribution [Gelman et al., 2020] .We provide a qualitative overview of the results we have obtained so far. These shall be considered preliminary, and a quantitative comparison with other baselines (e.g. the mean-field approximation from Seifner and Sánchez [2023] ) will be carried out in future work. For the epidemics dataset, we display generated trajectories on an held-out graph in Figure 5 , as well as the aggregated SIR curves for the same example in shown for an held-out example.Processing Systems (BDU at NeurIPS 2024). Do not distribute. 



Reviews: review: title: title: Great theories and experiments, lack of bridge between theories and reality
main_review: main_review: The paper presents a Bayesian learning framework for interacting particle systems (IPS) with discrete latent states, addressing the challenge of inferring dynamics from partial and noisy observations. The authors introduce a scalable variational approximation and apply it to two simulated scenarios: epidemic trajectory inference on networks and wildfire spread prediction on lattices. The approach's strengths lie in its use of continuous-time Markov chains (CTMCs) and neural network parameterizations to effectively capture the underlying dynamic processes.

Improvement the authors can make:
1. Impact of Hyperparameters on Evolution Patterns
Now the hyperparameters on evolution patterns are fixed. For example: 
- The epidemic model assumes a fixed random graph structure, where each node has an average of 3 neighbors. 
-  β = 6 and γ = 0.2 etc.
This limits the generalization of the results to real-world scenarios, where social networks can be more heterogeneous. Exploring a broader range of network topologies would strengthen the model's applicability and allow for more generalized conclusions. Therefore, the paper would benefit from a deeper exploration of how key hyperparameters, such as infection probability, number of neighbors in the epidemic model, and wind speed or fuel availability in the wildfire model, influence the system's temporal and spatial dynamics. Analyzing these factors could reveal tipping points or significant shifts in behavior, providing valuable insights into the sensitivity of the models to these parameters.

2. More explanations for the hyper-parameters
As the paper mentions "These parameters do not correspond to physically meaningful quantities, and adjusting them to reflect real-world spread dynamics remains an interesting avenue for future work" in line 528-530, it would be valuable if the authors could offer clear explanations that connect the theoretical dynamics to practical real-world implementations.
title: title: Comprehensive evaluation of "Variational Inference for Systems of Interacting Particles with Discrete Latent States"
main_review: main_review: This paper introduces a novel Bayesian learning framework for interacting particle systems (IPS) characterized by discrete latent states, with a special focus on learning from partial, noisy observations. The method innovatively exploits a variational posterior path metric to parameterize the generator of continuous-time Markov chains (CTMCs) to explain complex spatiotemporal processes.
title: title: The paper titled "Variational Inference for Interacting Particle Systems with Discrete Latent States" proposes a novel Bayesian learning framework for modeling systems of interacting particles, with applications such as epidemic spread and wildfire propagation. The methodology is built around continuous-time Markov chains (CTMC) and employs a variational posterior to infer the latent dynamics from partial, noisy observations. This work introduces an efficient two-stage learning approach, leveraging the multi-marginal Schrödinger bridge, to reconstruct and predict trajectories. The paper’s advantages include the scalability of its variational approximation and the flexibility to integrate domain-specific knowledge through informative priors, which broadens its applicability to high-dimensional spatiotemporal systems.  However, the paper presents a few limitations. While the method shows potential in simulated scenarios, it remains to be validated on real-world data. The reliance on neural networks for parameterizing the CTMC generator adds computational complexity, which might pose challenges in practical applications with large-scale systems. Furthermore, while the use of neural cellular automata improves flexibility, the approach might require extensive hyperparameter tuning. When compared to existing literature, such as the use of mean-field variational approaches or classical cellular automata models, this work is more sophisticated but could benefit from more comparisons with other contemporary techniques to better establish its credibility and performance in diverse settings.  This paper aligns well with recent developments in variational inference for CTMCs and interacting particle systems, but further empirical testing and real-world applications are necessary to fully understand its potential impact.
main_review: main_review: ----

Evaluation: 

1. Quality: 



The quality of this paper is creditable in terms of its methodological rigor and the clear articulation of a complex topic like variational inference for interacting particle systems (IPS). The authors successfully combine ideas from CTMCs, the Schrödinger bridge problem, and Bayesian learning frameworks to introduce a novel approach for inferring the dynamics of high-dimensional, spatiotemporal systems. The mathematical formulations are thorough, and the two-stage learning procedure is well-motivated and explained. While the paper primarily focuses on simulated results, the approach has the potential for real-world application once tested on empirical data. The theoretical development is solid, though the experimental validation could be more extensive.

2. Clarity:



The clarity of the paper is relatively strong in the theoretical sections, where the authors explain key concepts like variational inference and multi-marginal Schrödinger bridges. However, the presentation of the technical details might be dense for readers unfamiliar with the topic, which could limit accessibility to a broader audience. While the equations and algorithms are clearly structured, some sections could benefit from additional explanations or diagrams, particularly when discussing the architecture of neural cellular automata or the assumptions underlying the variational approximation. Additionally, clearer motivation for the choice of methods (e.g., why neural networks were chosen for parameterization) could enhance the reader's understanding.

3. Originality:



This paper presents a favorably original contribution by merging ideas from variational inference with interacting particle systems in the context of discrete latent states. The application of the multi-marginal Schrödinger bridge to this problem setting is innovative and offers a fresh perspective on trajectory inference and prediction. While interacting particle systems and Markov processes are well-established areas of study, this paper advances the field by introducing a novel, scalable inference approach that can be applied to high-dimensional data, which distinguishes it from traditional mean-field approximations or classical cellular automata methods. The originality of using neural networks in this context is another noteworthy aspect.

4.  Significance:



The significance of this work lies in its potential applicability to various real-world phenomena that can be modeled as interacting particle systems, such as epidemic modeling and wildfire prediction. The paper makes a valuable contribution by addressing the challenge of learning from partial and noisy observations in complex systems. If further validated through experiments on real-world datasets, this methodology could have far-reaching implications in fields ranging from epidemiology to environmental science and multi-agent systems. However, without real-world validation, the immediate significance of the results remains speculative, though the theoretical contributions are undoubtedly significant in advancing inference techniques for continuous-time stochastic systems.


List of Pros and Cons:


Pros:


- Novel Approach: Introduces an innovative two-stage learning procedure using the multi-marginal Schrödinger bridge for trajectory reconstruction and prediction.
- Scalability: The variational approximation is scalable, making it potentially suitable for high-dimensional spatiotemporal processes.
- Flexibility: The method allows for the incorporation of domain knowledge through informative priors, making it adaptable to different applications.
- Application Potential: The framework can be applied to diverse phenomena, such as epidemic spread and wildfire dynamics, offering  broad applicability.

Cons:



- Lack of Real-World Validation: The method has only been tested on simulated datasets, and its performance on real-world data remains to be demonstrated.
- Complexity: The use of neural networks for parameterizing the CTMC generator increases computational complexity, which might limit the method's scalability in very large systems.
- Accessibility: Some sections are mathematically dense, which might limit accessibility for non-expert readers or those not intimately familiar with the underlying concepts.
- Limited Experimental Evaluation: While the theoretical contributions are substantial, the paper could benefit from more extensive empirical evaluation, including comparisons with state-of-the-art methods.

Conclusion:


Overall, this paper is a high-quality, original contribution that proposes a novel approach for variational inference in interacting particle systems. While the paper is mathematically dense, it presents a significant advancement in the field of spatiotemporal modeling, particularly for systems governed by discrete latent states. The lack of real-world validation and the computational complexity introduced by the neural networks are drawbacks, but the paper's theoretical contributions and potential applications make it a valuable addition to the literature.

----
title: title: Need more details
main_review: main_review: Pro:

1. The paper introduces a novel Bayesian learning framework for interacting particle systems with discrete latent states, which addresses the challenge of inferring dynamics from partial, noisy observations. This is a valuable contribution to the field.
2. The paper formulates the problem as a multi-marginal Schrödinger bridge with aligned samples, which is a well-established and effective approach for solving similar problems.
3. The proposed method incorporates an emission distribution for decoding latent states, which is a practical and important component in real-world applications.

Cons:

1. Can you discuss the computational complexity and scalability of the proposed approach? What are the limitations and potential challenges in applying the method to high-dimensional spatio-temporal processes?
2. Can you provide more insights into the variational approximation used in the approach? What assumptions are made and how do these assumptions affect the accuracy of the inference?
3. Can you provide more information on the learning procedure? How did you optimize the evidence lower bound? Did you use any specific optimization algorithms or techniques?
4. How does the computational complexity of your approach scale with the dimensionality of the spatio-temporal processes? Have you performed any experiments to evaluate the scalability of your method?
5. How does the proposed method compare to existing methods for inferring dynamics from partial, noisy observations? Can you provide a more detailed discussion of the advantages and limitations of the proposed approach compared to previous approaches?