SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Extraction and Recovery of Spatio-Temporal Structure in Latent Dynamics Alignment with Diffusion Models
Yule Wang
Atlanta, GA, 30332 USA
yulewang@gatech.edu
Chengrui Li
cnlichengrui@gatech.edu
Zijing Wu
zwu381@gatech.edu
Anqi Wu
anqiwu@gatech.edu
In the field of behavior-related brain computation, it is necessary to align raw neural signals against the drastic domain shift among them. A foundational framework within neuroscience research posits that trial-based neural population activities rely on low-dimensional latent dynamics, thus focusing on the latter greatly facilitates the alignment procedure. Despite this field's progress, existing methods ignore the intrinsic spatio-temporal structure during the alignment phase. Hence, their solutions usually lead to poor quality in latent dynamics structures and overall performance. To tackle this problem, we propose an alignment method ERDiff, which leverages the expressivity of the diffusion model to preserve the spatiotemporal structure of latent dynamics. Specifically, the latent dynamics structures of the source domain are first extracted by a diffusion model. Then, under the guidance of this diffusion model, such structures are well-recovered through a maximum likelihood alignment procedure in the target domain. We first demonstrate the effectiveness of our proposed method on a synthetic dataset. Then, when applied to neural recordings from the non-human primate motor cortex, under both cross-day and inter-subject settings, our method consistently manifests its capability of preserving the spatio-temporal structure of latent dynamics and outperforms existing approaches in alignment goodness-of-fit and neural decoding performance. Codes are available at: https://github.com/alexwangNTL/ERDiff.
A key challenge severely impeding the scalability of behavior-related neural computational applications is their robustness to the distribution shift of neural recordings over time and subjects [1]. Given a behavior model trained on previous neural recordings (e.g., velocity predictor for human with paralysis [2]), it usually suffers performance degradation when applied to new neural recordings due to the neural distribution shift [3, 4]. Thus, for long-term usability and stable performance of the trained neural decoding model, high-quality alignment between the neural recordings used for training (i.e., source domain) and new recordings for testing (i.e., target domain) is of vital importance.
Distribution alignment is an important task at the heart of unsupervised transfer learning [5, 6]. The goal is to align the target domain to the source domain so that the trained model in the source domain can be applied to the target domain after eliminating the distribution shift. However, due to issues such as instabilities and low signal-to-noise ratio [7], raw neural spiking activities are noisy and ambiguous [8, 9], causing difficulties in aligning the distributions of these high-dimensional signals directly.
One promising research direction [10] points out that the trial-based neural activities can always be understood in terms of low-dimensional latent dynamics [11, 12, 13]. Such latent dynamics manifest coordinated patterns of evolution constrained to certain "neural manifolds" [14, 15]. Hence, early studies focusing on the alignment of latent dynamics reach comparably satisfactory results [16, 17]. Generally, most previous methods [16, 18, 19] are based on a pre-defined metric for optimization during latent dynamics alignment, i.e., minimizing the difference evaluated by the metric, between source and target domains within the low-dimensional latent space. However, those metrics are usually non-parametric and handcrafted, which are not guaranteed to suit specific neural recordings or problems well. Methods based on adversarial-learning [20, 21] thus have been introduced since they can implicitly find an adapted metric [22]. However, they suffer from mode collapse and instability issues in practice [23].
Moreover, during the alignment process, we note that the above-mentioned works lack the necessary awareness of the latent dynamics structure, especially when aligning non-linear and lengthy trials. Through an empirical study on the motor cortex of non-human primate (NHP) [8] (shown in Figure 1), we can observe that: a state-of-the-art alignment method JSDM [24] (minimizing the symmetric Jensen-Shannon divergence between distributions) fails to recover the latent dynamics structures of the source domain since JSDM neglects those structures during alignment. From another perspective, in the alignment phase, existing methods fail to effectively model and leverage the information-rich correlations between each time bin and each latent dimension within latent dynamics.
In this paper, we focus on preserving the temporal evolution of each individual latent dimension and the spatial covariation between latent dimensions of the source domain during alignment. The main idea is that we first extract the spatio-temporal structure of latent dynamics from the source domain; and then, we align the target domain by recovering the source domain's underlying structure. However, such a workflow is non-trivial since the underlying spatio-temporal structure is both implicit and complex.
To tackle this problem, we propose a novel alignment method that is capable of Extracting and Recovering the latent dynamics structure with Diffusion model (ERDiff). Firstly, given the source-domain neural observations, we use a diffusion model (DM) [25, 26] to extract the spatiotemporal structure of latent dynamics. Then, in the alignment phase, we propose a maximum likelihood alignment procedure through the guidance of DM, by which the spatio-temporal structure of source-domain latent dynamics can be recovered well in the target domain. The proposed extract-and-recover method nicely encodes and preserves the spatio-temporal structure of latent dynamics, which are significant inductive biases for neural latent dynamics alignment. Furthermore, from the perspective
Figure 1: Empirical study. Latent dynamics (3D visualization) of the source domain and the aligned target domain by JSDM on a primary motor cortex dataset.
of core machine learning, ERDiff introduces an approach of extracting structure knowledge from one distribution and imposing it as the prior to constrain the alignment of another distribution. Note that although we have been emphasizing extraction and recovery of the source-domain structure, ERDiff is not performing a copy-and-paste of the source domain distribution to the target domain. As ERDiff preserves the dynamics structure of the source domain, it also maintains the original characteristics of the target domain. We present experimental results to support this argument. Finally, we conduct extensive experiments to verify the effectiveness of ERDiff on a synthetic dataset and two real-world neural datasets [8, 27]. Visualization of latent dynamics also demonstrates that ERDiff is capable of preserving the spatio-temporal structure consistently during the alignment phase.
Distribution alignment. We denote the source-domain observations of single-trial neural population activities as  $\mathbf{X}^{(s)} = \left[\mathbf{x}_1^{(s)},\dots ,\mathbf{x}_l^{(s)}\right]^\top \in \mathbb{R}^{l\times n}$ , where  $l$  is the trial length (i.e., number of time bins), and  $n$  is the number of observed neurons. We denote its low-dimensional latent dynamics as  $\mathbf{Z}^{(s)} = \left[\mathbf{z}_1^{(s)},\dots ,\mathbf{z}_l^{(s)}\right]^\top \in \mathbb{R}^{l\times d}$ , where  $d$  is the latent dimension size. Generally, we build a variational autoencoder (VAE) to estimate the latent  $\mathbf{Z}^{(s)}$  given the observations  $\mathbf{X}^{(s)}$ . The VAE
consists of a probabilistic encoder  $q(\mathbf{Z}^{(s)} \mid \mathbf{X}^{(s)}; \phi_s)$  and a probabilistic decoder  $p(\mathbf{X}^{(s)}|\mathbf{Z}^{(s)}, \psi_s)$ .  $\phi_s$  and  $\psi_s$  are the parameters of the encoder and decoder. The encoder also serves as an approximated posterior distribution to the intractable true posterior  $p(\mathbf{Z}^{(s)} \mid \mathbf{X}^{(s)})$ . Then in the target domain, given the neural population activities  $\mathbf{X}^{(t)} = [\mathbf{x}_1^{(t)}, \ldots, \mathbf{x}_l^{(t)}]^\top \in \mathbb{R}^{l \times n}$ , we perform distribution alignment by linear probing the probabilistic encoder  $q(\mathbf{Z} \mid \mathbf{X}; \phi)$ . This alignment phase is conducted by minimizing certain probability divergence  $\mathbb{D}(\cdot \mid \cdot)$  between the two posterior distributions:
$$
\min  _ {\phi_ {t}} \mathbb {D} \left(q \left(\mathbf {Z} ^ {(s)} \mid \mathbf {X} ^ {(s)}; \phi_ {s}\right) \| q \left(\mathbf {Z} ^ {(t)} \mid \mathbf {X} ^ {(t)}; \phi_ {t}\right)\right). \tag {1}
$$
Diffusion (probabilistic) model (DM). Given  $l \times d$ -dimensional i.i.d. samples  $\mathbf{Z}$  from an unknown data distribution, a DM [28] aims to approximate such distribution by fitting the parameters of a neural network  $p_{\theta}(\mathbf{Z})$ . DM is composed of a forward process followed by a reverse process. In the forward process, isotropic Gaussian noise is added to diffuse the original data, which can be defined in a linear stochastic differential equation (SDE) form:
$$
\mathrm {d} \mathbf {Z} = \boldsymbol {f} (\mathbf {Z}, t) \mathrm {d} t + g (t) \mathrm {d} \mathbf {w}, \tag {2}
$$
where  $\pmb{f}(\cdot): \mathbb{R}^{l \times d} \times \mathbb{R} \mapsto \mathbb{R}^{l \times d}$  is the drift coefficient,  $g(\cdot): \mathbb{R} \mapsto \mathbb{R}$  is the diffusion coefficient, and  $\mathbf{w}$  is the standard Wiener process. The solution of the SDE is a diffusion process  $\{\mathbf{Z}_t\}_{t \in [0,T]}$ , in which  $[0,T]$  is a fixed time zone. In this paper, we implement them with VP-SDE [25].  $\{\mathbf{Z}_t\}_{t \in [0,T]}$  approaches the standard normal prior distribution  $\pi(\mathbf{Z})$  when  $t = T$ . Under mild conditions on drift and diffusion coefficients [25], the denoising reverse process can be solved in the following closed-form SDE:
$$
\mathrm {d} \mathbf {Z} = \left[ \boldsymbol {f} (\mathbf {Z}, t) - g (t) ^ {2} \nabla_ {\mathbf {Z}} \log p _ {t} (\mathbf {Z}) \right] \mathrm {d} t + g (t) \mathrm {d} \overline {{\mathbf {w}}}, \tag {3}
$$
where  $\nabla_{\mathbf{Z}}\log p_t(\mathbf{Z})$  is the score function, and  $\overline{\mathbf{w}}$  is a reverse-time Wiener process. We train a parameterized network  $s(\mathbf{Z},t;\pmb {\theta})$  to fit the score function  $\nabla_{\mathbf{Z}}\log p_t(\mathbf{Z})$ . However,  $\nabla_{\mathbf{Z}}\log p_t(\mathbf{Z})$  is not directly accessible and we resort to the denoising score matching (DSM) [29] for optimization:
$$
\mathcal {L} _ {\mathrm {D S M}} (\boldsymbol {\theta}) = \mathbb {E} _ {t \sim \mathcal {U} [ 0, T ]} \mathbb {E} _ {\mathbf {Z} _ {0} \sim p, p _ {0 t} \left(\mathbf {Z} _ {t} \mid \mathbf {Z} _ {0}\right)} \left[ \lambda (t) ^ {2} \| \nabla_ {\mathbf {Z} _ {t}} \log p _ {0 t} \left(\mathbf {Z} _ {t} \mid \mathbf {Z} _ {0}\right) - s \left(\mathbf {Z} _ {t}, t; \boldsymbol {\theta}\right) \| _ {2} ^ {2} \right], \tag {4}
$$
where  $\mathcal{U}$  represents the uniform distribution and  $\lambda(t)$  is the weighting function. Under VP-SDE, the transition probability  $p_{0t}(\mathbf{Z}_t \mid \mathbf{Z}_0)$  also follows a Gaussian distribution  $\mathcal{N}(\mu_t \mathbf{Z}_0, \Sigma_t)$ , in which  $\mu_t, \Sigma_t \in \mathbb{R}^{l \times d}$ . On the other hand, according to [28], we can define a noise estimator with the score function as  $\epsilon(\mathbf{Z}_t, t; \boldsymbol{\theta}) = -\mathbf{K}_t^{-T} s(\mathbf{Z}_t, t; \boldsymbol{\theta})$ , in which  $\mathbf{K}_t \mathbf{K}_t^T = \Sigma_t$ . Invoking these expressions, we can thus reformulate the form of DSM loss based on the Fisher Divergence between noise terms:
$$
\mathcal {L} _ {\mathrm {D S M}} (\boldsymbol {\theta}) = \mathbb {E} _ {t \sim \mathcal {U} [ 0, T ]} \mathbb {E} _ {\mathbf {Z} _ {0} \sim p, \boldsymbol {\epsilon} \sim \mathcal {N} (0, I _ {l \times d})} \left[ w (t) ^ {2} \| \boldsymbol {\epsilon} - \boldsymbol {\epsilon} (\mathbf {Z} _ {t}, t; \boldsymbol {\theta}) \| _ {2} ^ {2} \right], \tag {5}
$$
in which  $w(t) = K_t\lambda (t)$  and  $\mathbf{Z}_t = \mu_t\mathbf{Z}_0 + K_t\boldsymbol {\epsilon}$
In this section, we introduce our proposed latent dynamics alignment method ERDiff in detail.
3.1 Maximum likelihood alignment
Given the source-domain neural activities  $\mathbf{X}^{(s)}$ , we infer their latent dynamics  $\mathbf{Z}^{(s)}$  by building a VAE. We use variational inference to find the probabilistic encoder  $q(\mathbf{Z}^{(s)} \mid \mathbf{X}^{(s)}; \boldsymbol{\phi}_s)$  and probabilistic decoder  $p(\mathbf{X}^{(s)} \mid \mathbf{Z}^{(s)}; \boldsymbol{\psi}_s)$  through maximization of the evidence lower bound (ELBO) [30]:
$$
\phi_ {s}, \psi_ {s} = \underset {\phi , \psi} {\operatorname {a r g m a x}} \left[ \mathbb {E} _ {q \left(\mathbf {Z} ^ {(s)} \mid \mathbf {X} ^ {(s)}; \phi\right)} \left[ \log p \left(\mathbf {X} ^ {(s)} \mid \mathbf {Z} ^ {(s)}; \psi\right) \right] - \mathbb {D} _ {\mathrm {K L}} \left(q \left(\mathbf {Z} ^ {(s)} \mid \mathbf {X} ^ {(s)}; \phi\right) \| \bar {q} \left(\mathbf {Z} ^ {(s)}\right)\right) \right], \tag {6}
$$
in which  $\bar{q} (\mathbf{Z}^{(s)})$  is the normal prior. Note that we introduce ERDiff with this basic VAE architecture. But ERDiff can be combined with many variants of latent variable models (LVM) [31, 32]. The essence of ERDiff is to tune the parameter set  $\phi$  of the probabilistic encoder, regardless of the model architecture of the encoder and decoder.
Figure 2: A schematic overview of spatio-temporal structure extraction and recovery in ERDiff. (A) The architecture of DM for spatio-temporal structure extraction. (B) A descriptive diagram of structure recovery schematic. The left presents the extracted spatio-temporal structure of the source-domain latent dynamics; the right illustrates the structure-aware maximum likelihood alignment guidance in ERDiff.
Practically, alignment methods that directly match the discrete samples from the source and target domains in a pair-wise fashion may lead to sub-optimal solutions [33, 34], especially when the collected samples from the target domain are limited. Thus given the target-domain neural activity  $\mathbf{X}^{(t)}$ , we propose to perform neural distribution alignment via maximum likelihood estimation (MLE):
$$
\underset {\phi} {\operatorname {a r g m a x}} \mathbb {E} _ {\mathbf {X} \sim p (\mathbf {X} ^ {(t)})} [ \log p _ {s} (h (\mathbf {X}; \phi)) ] = \underset {\phi} {\operatorname {a r g m a x}} \mathbb {E} _ {\mathbf {Z} \sim q (\mathbf {Z} | \mathbf {X} ^ {(t)}; \phi)} [ \log p _ {s} (\mathbf {Z}) ], \tag {7}
$$
in which  $p_s(\mathbf{Z})$  represents the ground-truth probabilistic density of latent dynamics in the source domain and  $h(\cdot)$  refers to the non-linear transformation from  $\mathbf{X}$  to  $\mathbf{Z}$  underlying the probabilistic encoder  $q(\mathbf{Z} \mid \mathbf{X}; \phi)$ . The objective in Eq. 7 implies that, instead of minimizing a distance metric between source observations and target observations, we aim at maximizing the likelihood where the density comes from the source domain and the data comes from the target domain. The left-hand-side (LHS) is the MLE for observation  $\mathbf{X}$  and the right-hand-side (RHS) is the MLE for latent  $\mathbf{Z}$ . We will focus on the RHS in the following sections. We note that the RHS objective implies that we will optimize the encoder parameter  $\phi$  during alignment so that the latent encoder will map  $\mathbf{X}^{(t)}$  to a proper latent  $\mathbf{Z}^{(t)}$  who fits the source density  $p_s(\mathbf{Z})$  well.
3.2 Spatio-temporal structure extraction and source domain learning
In order to calculate the objective function in Eq. 7, we need to know two density functions:  $q(\mathbf{Z} \mid \mathbf{X}; \phi)$  is defined in the original VAE model with the learnable parameter  $\phi$ ;  $p_s(\mathbf{Z})$  is the density of latent  $\mathbf{Z}$  for the source domain. The latter is inaccessible by building a VAE alone. Therefore, the first step is to learn  $p_s(\mathbf{Z})$  given only  $\mathbf{X}^{(s)}$ . We propose to learn  $p_s(\mathbf{Z})$  through training a DM.
To fully capture  $p_{s}(\mathbf{Z})$ , the DM should consider the overall spatio-temporal structure of latent dynamics. To extract such a structure, the DM can not treat each latent state or time bin within latent dynamics as mutually independent and feed them into the model sequentially. We thus take the entire trial of latent dynamics  $\mathbf{Z}_0^{(s)} \sim q(\cdot \mid \mathbf{X}^{(s)}; \phi_s)$  as input to the DM for training. Specifically, the DM fits  $p_{s}(\mathbf{Z})$  through the training of a denoiser  $\epsilon (\mathbf{Z},t;\pmb {\theta}_s):(\mathbb{R}^{l\times d}\times \mathbb{R})\to \mathbb{R}^{l\times d}$ .
Next, we describe the architecture of  $\epsilon (\mathbf{Z},t;\pmb {\theta}_s)$ , which is refined for extracting the global spatiotemporal structure of latent dynamics. Traditional architecture based on 2D-Convolution Layers [35] focuses on capturing the local features within latent dynamics, which can hardly extract its global spatio-temporal dependency or structure. Thus, we adopt an architecture mainly derived from Diffusion Transformer (DiT) [36, 37]. Specifically, we propose to use Spatio-Temporal Transformer Block (STBlock), shown in Figure 2(A). Each STBlock is composed of a Spatio Transformer layer followed by a Temporal Transformer layer, which are 1-layer encoders based on multi-head selfattention. The Spatio Transformer layer takes latent states of each time bin as inputs to extract spatial structure, whereas the Temporal Transformer layer takes the entire latent trajectory of each latent space dimension as inputs to extract temporal structure. (see Appendix A for details of the architecture of DM).
For the training objective of  $\epsilon (\cdot ;\theta_s)$ , we sample noisy targets  $\mathbf{Z}_t^{(s)}$  and minimize the following DSM loss function:
$$
\boldsymbol {\theta} _ {s} = \underset {\boldsymbol {\theta}} {\operatorname {a r g m i n}} \mathbb {E} _ {t \sim \mathcal {U} [ 0, T ]} \mathbb {E} _ {\mathbf {Z} _ {0} ^ {(s)} \sim q (\cdot | \mathbf {X} ^ {(s)}; \boldsymbol {\phi} _ {s}), \epsilon \sim \mathcal {N} (\mathbf {0}, \mathbf {I} _ {l \times d})} \left[ w (t) ^ {2} \left\| \boldsymbol {\epsilon} - \boldsymbol {\epsilon} \left(\mathbf {Z} _ {t} ^ {(s)}, t; \boldsymbol {\theta}\right) \right\| _ {2} ^ {2} \right]. \tag {8}
$$
We note that  $\mathbf{Z}_0^{(s)}$  here are actually latent dynamics inferred via VAE in Eq. 6. To enrich the input samples and adequately estimate the source density  $p_s(\mathbf{Z})$  as motivated earlier, we propose to learn the VAE objective (Eq. 6) and the diffusion objective (Eq. 8) simultaneously. In each training iteration, conditioning on the current value of  $\phi_s$  and  $\psi_s$ , we obtain a set of  $\mathbf{Z}_0 = h(\mathbf{X}^{(s)};\phi_s)$  and use it as  $\mathbf{Z}_0^{(s)}$  to optimize Eq. 8. We can also optimize VAE first, obtain an optimal  $\phi_s$ , and use it to optimize Eq. 8. Experimental results show that the former approach achieves higher density estimation performance compared to the latter (see Appendix A for details).
3.3 Spatio-temporal structure recovery and distribution alignment
Given the trained denoiser  $\epsilon (\mathbf{Z},t;\pmb{\theta}_s)$ , we go through the reverse process from  $t = T$  to  $t = 0$  in Eq.(3) and obtain the marginal distribution  $p_0(\mathbf{Z};\pmb{\theta}_s)$ . We use  $p_0(\mathbf{Z};\pmb{\theta}_s)$  to approximate  $p_s(\mathbf{Z})$  in Eq. (7). The maximum likelihood estimation can thus be written as
$$
\underset {\phi} {\operatorname {a r g m a x}} \mathbb {E} _ {\mathbf {Z} \sim q (\mathbf {Z} | \mathbf {X} ^ {(t)}; \phi)} \left[ \log p _ {0} (\mathbf {Z}; \boldsymbol {\theta} _ {s}) \right]. \tag {9}
$$
We perform alignment by tuning the parameter set  $\phi$  of the probabilistic encoder while keeping the DM  $p_0(\mathbf{Z};\boldsymbol{\theta}_s)$  fixed. Note that we have already optimized the VAE objective to obtain an optimal  $\phi_s$  using source data. During alignment, we first set  $\phi$  as  $\phi_s$  and then linear probe  $\phi$  (e.g., neural observation read-in layer). Consequently, we not only initialize the model with a good encoder but also make optimization during alignment much faster and more efficient.
In the reverse process, the computation of  $\log p_0(\mathbf{Z};\boldsymbol{\theta}_s)$  is tractable through the probability flow ODE [25] whose marginal distribution at each time step  $t$  matches that of our VP-SDE. However, the direct computation of  $\log p_0(\mathbf{Z};\boldsymbol{\theta}_s)$  will require invoking the ODE solver in each intermediate time step [38, 39]. Such complexity is prohibitively costly for online neural applications. To circumvent this issue, we can reform Eq. (9) as follows:
$$
- \mathbb {E} _ {\mathbf {Z} \sim q (\mathbf {Z} | \mathbf {X} ^ {(t)}; \phi)} [ \log p _ {0} (\mathbf {Z}; \boldsymbol {\theta} _ {s}) ] = \mathbb {D} _ {\mathrm {K L}} \left(q (\mathbf {Z} \mid \mathbf {X} ^ {(t)}; \phi) \| p _ {0} (\mathbf {Z}; \boldsymbol {\theta} _ {s})\right) + \mathbb {H} \left(q (\mathbf {Z} \mid \mathbf {X} ^ {(t)}; \phi)\right), \tag {10}
$$
where the first term is the KL divergence from the DM marginal distribution  $p_0(\mathbf{Z};\boldsymbol{\theta}_s)$  to the probabilistic encoder distribution  $q(\mathbf{Z} \mid \mathbf{X}^{(t)};\phi)$ , and the second term  $\mathbb{H}(\cdot)$  denotes the differential entropy. For the  $\mathbb{D}_{\mathrm{KL}}(\cdot)$  term in Eq. (10), via the Girsanov theorem [40, 41], we have
$$
\mathbb {D} _ {\mathrm {K L}} \left(q (\mathbf {Z} \mid \mathbf {X} ^ {(t)}; \phi) \| p _ {0} (\mathbf {Z}; \boldsymbol {\theta} _ {s})\right) \leqslant \mathcal {L} _ {\mathrm {D S M}} \left(\phi , \boldsymbol {\theta} _ {s}\right) + \mathbb {D} _ {\mathrm {K L}} \left(p _ {T} (\mathbf {Z}; \boldsymbol {\theta} _ {s}) \| \pi (\mathbf {Z})\right), \tag {11}
$$
where  $\mathcal{L}_{\mathrm{DSM}}$  is the denoising score matching loss in Eq. (5), and  $p_T(\cdot)$  is the distribution at final time step  $T$  of Eq. (2). Consequently, we could obtain an upper bound of the maximum likelihood objective, as follows (we provide detailed derivation in Appendix B):
$$
\begin{array}{l} - \mathbb {E} _ {\mathbf {Z} \sim q (\mathbf {Z} | \mathbf {X} ^ {(t)}; \phi)} [ \log p _ {0} (\mathbf {Z}; \boldsymbol {\theta} _ {s}) ] \leqslant \underbrace {\mathbb {D} _ {\mathrm {K L}} (p _ {T} (\mathbf {Z} ; \boldsymbol {\theta} _ {s}) \| \pi (\mathbf {Z}))} _ {\text {C o n s t a n t T e r m}} \\ + \mathbb {E} _ {t \sim \mathcal {U} [ 0, T ]} \mathbb {E} _ {\mathbf {Z} _ {0} \sim q (\mathbf {Z} | \mathbf {X} ^ {(t)}; \phi), \epsilon \sim \mathcal {N} (0, I _ {l \times d})} \left[ \underbrace {w (t) ^ {2} \| \epsilon - \epsilon \left(\mathbf {Z} _ {t} , t ; \boldsymbol {\theta} _ {s}\right) \| _ {2} ^ {2}} _ {\text {W e i g h t e d N o i s e R e s i d u a l}} - \underbrace {2 \nabla_ {\mathbf {Z}} \cdot f (\mathbf {Z} _ {t} , t)} _ {\text {D i v e r g e n c e}} \right]. \tag {12} \\ \end{array}
$$
Since  $\pi (\mathbf{Z})$  is a fixed prior distribution, it does not depend on parameter  $\phi$ . Thus, our optimization objective will include only the latter two terms, which are more computationally tractable. The first objective simplifies to a weighted noise residual for the parameter set  $\phi$  and the second divergence objective can be approximated using the Hutchinson-Skilling trace estimator [42]. We note that the recovery of spatio-temporal structure is primarily conducted by the weighted noise residual part, in which the probabilistic encoder obtains alignment guidance in awareness of the spatio-temporal structure from  $\epsilon (\mathbf{Z},t;\pmb {\theta}_s)$ . This procedure is illustrated in Figure 2(B).