SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
h _ {i} ^ {l + 1} = \psi^ {l} \left(h _ {i} ^ {l}, \bigoplus_ {j \in \mathcal {N} (i)} \phi^ {l} \left(h _ {j} ^ {l}\right)\right), \tag {9}
$$
where  $\bigoplus$  is a permutation-invariant aggregation function, performed over neighboring nodes. Usually,  $\psi^l$  and  $\phi^l$  would be parameterized with MLPs, but since  $h_i^l$  lives on a grid domain, we choose both functions to be convolutional neural nets. We refer to this architecture as SpatialConv-GNN.
Model design. We again start by laying out the overall model structure, also shown in Figure 3b:
- Forward model  $f_{\theta}$ : SpatialConv-GNN( $x^t$ )
- Conditional prior:  $p_{\theta}(z|h^t)\colon \mathcal{N}\left(\mu^{\mathrm{prior}}(h^t),\sigma^{\mathrm{prior}}(h^t)\right)$
- Decoder:  $p_{\theta}(x^{t + 1}|z,h^t)$ :  $\mathcal{C}\left(\pi^{\mathrm{decoder}}(z,h^t)\right)$
Since  $f_{\theta}$  is parameterized by a SpatialConv-GNN, it is equivariant to translations and permutations. At any point in time, each cell is only affected by itself and the directly adjacent cells, so we design  $p_{\theta}(z|h^{t})$  to be equivariant with respect to permutations, but invariant with respect to translations; concretely, we perform alternating convolutions and spatial pooling operations for each node  $h_{i}$  to infer distributions over a latent variable  $z$  for each cell. Finally, the decoder is parameterized by  $\pi^{\mathrm{decoder}}(z,h^t)$ , a SpatialConv-GNN which equivariantly maps to the parameters of a pixel-wise categorical distribution over the set of cells  $C$ . More details can be found in Appendix C.2.
5.1 Experiment setup
Objectives. Our experiments aim to a) assess the effectiveness of EPNS in terms of simulation faithfulness, calibration, and stability; b) empirically verify the equivariance property of EPNS, and investigate the benefits of equivariance in the context of probabilistic simulation; and c) compare EPNS to existing methods for probabilistic simulation of dynamical systems. We report the results conveying the main insights in this section; further results and ablations are reported in Appendix E.
Data. For both problem settings, we generate training sets with 800 trajectories and validation and test sets each consisting of 100 trajectories. The length of each trajectory is 1000 for the celestial dynamics case, saved at a temporal resolution of  $dt = 0.01$  units per step. We train the models to predict a time step of  $\Delta t = 0.1$  ahead, resulting in an effective rollout length of 100 steps for a full trajectory. Similar to earlier works [38, 44, 59], we choose  $n = 5$  bodies for our experiments. For the cellular dynamics case, we use trajectories with  $|C| = 64$  cells. The length of each trajectory is 59 steps, where we train the models to predict a single step ahead. Additionally, we generate test sets of 100 trajectories starting from the same initial condition, which are used to assess performance regarding uncertainty quantification.
Implementation and baselines. The EPNS implementations are based on the designs explained in Section 4. The maximum rollout lengths used for multi-step training are 16 steps (celestial dynamics) and 14 steps (cellular dynamics). We use a linear KL-annealing schedule [8], as well as the 'free bits' modification of the ELBO as proposed in [22] for the cellular dynamics data. For celestial dynamics, we compare to neural SDEs (NSDE) [27] and interacting Gaussian Process ODEs (iGPODE) [59], as these methods can probabilistically model dynamics of (interacting) particles. For cellular dynamics, we compare to ODE $^{2}$ VAE [58], as it probabilistically models dynamics on a grid domain. For each baseline, we search over the most relevant hyperparameters, considering at least seven configurations per baseline. We also compare to a non-equivariant counterpart of EPNS, which we will refer to as PNS. All models are implemented in PyTorch [35] and trained on a single NVIDIA A100 GPU. Further details on the baseline models can be found in Appendix D. Our code is available at https://github.com/kminartz/EPNS.
Qualitative results. Figure 4 shows a sample from the test set, as well as simulations produced by neural simulators starting from the same initial condition  $x^0$  for both datasets. More samples are provided in Appendix E. For celestial dynamics, both EPNS and iGPODE produce a simulation that looks plausible, as all objects follow orbiting trajectories with reasonable velocities. As expected, the simulations start to deviate further from the ground truth as time proceeds, due to the accumulation
Figure 4: Qualitative results for celestial dynamics (left) and cellular dynamics (right).
Table 1: Comparison of EPNS to related work. For celestial dynamics, the kinetic energy of the system is used to calculate  $D_{\mathrm{KS}}$ , while for cellular dynamics, we use the number of cell clusters of the same type. We report the mean±standard error over three independent training runs.
Celestial dynamics	Cellular dynamics
Model	LL ↑·103	DKS(KE) ↓t=50	t=100	Model	LL ↑·104	DKS(#clusters) ↓t=30	t=45
NSDE [27]	-5.0±0.0	0.80±0.01	0.65±0.03	ODE²VAE [58]	-30.1±0.0	0.98±0.01	0.96±0.04
iGPODE [59]	-8.5±0.1	0.42±0.07	0.57±0.02
PNS (ours)	10.9±0.4	0.61±0.18	0.20±0.06	PNS (ours)	-16.4±0.3	0.70±0.10	0.77±0.05
EPNS (ours)	10.8±0.1	0.14±0.03	0.14±0.04	EPNS (ours)	-5.9±0.1	0.58±0.09	0.58±0.05
of randomness. In the cellular dynamics case, the dynamics produced by EPNS look similar to the ground truth, as cells of the same type tend to cluster together. In contrast, ODE $^2$ VAE fails to capture this behavior. Although cells are located at roughly the right location for the initial part of the trajectory, the clustering and highly dynamic nature of the cells' movement and shape is absent.
Comparison to related work. We consider two metrics to compare to baselines: the first is the ELBO on the log-likelihood on the test set, denoted as LL. LL is a widely used metric in the generative modeling community for evaluating probabilistic models, expressing goodness-of-fit on unseen data, see for example [15, 22, 23, 59]. Still, the calculation of this bound differs per model type, and it does not measure performance in terms of new sample quality. This is why the second metric, the Kolmogorov-Smirnov (KS) test statistic  $D_{\mathrm{KS}}$ , compares empirical distributions of newly sampled simulations to ground truth samples. Specifically,  $D_{\mathrm{KS}}(y) = \max_y|F^* (y) - F_\theta (y)|$
the largest absolute difference between the ground truth and model empirical cumulative distribution functions  $F^{*}$  and  $F_{\theta}$ .  $D_{\mathrm{KS}} = 0$  corresponds to identical samples, whereas  $D_{\mathrm{KS}} = 1$  corresponds to no distribution overlap at all. We opt for this metric since we have explicit access to the ground-truth simulator, and as such we can calculate  $D_{\mathrm{KS}}$  over 100 simulations starting from a fixed  $x^0$  in order to assess how well EPNS matches the ground-truth distribution over possible trajectories. The results are shown in Table 1. EPNS performs comparative to or better than the other methods in terms of LL, indicating a distribution that has high peaks around the ground truth samples. In terms of  $D_{\mathrm{KS}}$ , EPNS performs best, meaning that simulations sampled from EPNS also match better to the ground truth in terms of the examined properties of Table 1.
Uncertainty quantification. In addition to Tables 1 and 3, we further investigate how distributions produced by EPNS, PNS, and EPNS trained with one step training (EPNS-one step) align with the ground truth. Figure 5 shows how properties of 100 simulations, all starting from the same  $x^0$ , evolve over time. In the case of celestial dynamics, the distribution over potential energy produced by EPNS closely matches the ground truth, as shown by the overlapping shaded areas. In contrast, EPNS-one step effectively collapses to a deterministic model. Although PNS produces a distribution that overlaps with the ground truth, it doesn't match as closely as EPNS. For the cellular dynamics
Figure 5: Distributions of the potential energy in a system (top row, celestial dynamics) and the number of cell clusters of the same type (bottom row, cellular dynamics) over time. The shaded area indicates the  $(0.1, 0.9)$ -quantile interval of the observations; the solid line indicates the median value.
case, the distribution produced by EPNS is slightly off compared to the ground truth, but still matches considerably better than EPNS-one step and PNS. These results indicate that both equivariance and multi-step training improve uncertainty quantification for these applications.
Verification of equivariance. To empirically verify the equivariance property of EPNS, we sample ten simulations for each  $x^0$  in the test set, leading to 1000 simulations in total, and apply a transformation  $\rho (g)$  to the final outputs  $x^{T}$ .  $\rho (g)$  is a random rotation and random permutation for celestial
dynamics and cellular dynamics respectively. Let us refer to the transformed outcomes as distribution 1. We then repeat the procedure, but now transform the inputs  $x^0$  by the same transformations before applying EPNS, yielding outcome distribution 2. For an equivariant model, distributions 1 and 2 should be statistically indistinguishable. We test this by comparing attributes of a single randomly chosen body or cell using a two-sample KS-test. We also show results of PNS to investigate if equivariance is learned by a non-equivariant model. The results are shown in Table 2. In all cases, distributions 1 and 2 are statistically indistinguishable for EPNS ( $p \gg 0.05$ ), but not for PNS ( $p \ll 0.05$ ). As such, EPNS indeed models equivariant distributions over trajectories, while PNS does not learn to be equivariant.
Table 2: KS-test results for equivariance verification.
Celestial Dynamics	Cellular Dynamics
Model	x-coordinate D_KS	p-value	y-coordinate D_KS	p-value	z-coordinate D_KS	p-value	distance D_KS	traveled p-value
PNS	0.09	0.00	0.22	0.00	0.09	0.00	0.503	0.00
EPNS	0.04	0.37	0.02	0.99	0.02	0.98	0.036	0.536
Data efficiency. We now examine the effect of equivariance on data efficiency. Table 3 shows the ELBO of both EPNS and PNS when trained on varying training set sizes. In the celestial dynamics case, the ELBO of EPNS degrades more gracefully when shrinking the training set compared to PNS. For cellular dynamics, the ELBO of EPNS does not worsen much at all when reducing the dataset to only 80 trajectories,
while the ELBO of PNS deteriorates rapidly. These results demonstrate that EPNS is more data efficient due to its richer inductive biases.
Table 3: Test set ELBO values  $(\cdot 10^{3})$  for varying amounts of training samples.
# Training samples	Celestial Dynamics	Cellular Dynamics
EPNS	PNS	EPNS	PNS
80	9.9	7.7	-59.6	-326.5
400	10.2	10.0	-61.0	-173.9
800	10.8	10.9	-59.4	-163.9
Rollout stability. As we cannot measure stability by tracking an error metric over time for stochastic dynamics, we resort to domain-specific stability criteria. For celestial dynamics, we define a run to be unstable when the energy in the system jumps by more than 20 units, which we also used to ensure
(a) Celestial dynamics.
(b) Cellular dynamics.
Figure 6: Fraction of simulations that remain stable over long rollouts.
the quality of the training samples. For cellular dynamics, we define a run to be unstable as soon as  $20\%$  of the cells have a volume outside the range of volumes in the training set. Although these are not the only possible criteria, they provide a reasonable proxy to the overall notion of stability for these applications. The fraction of stable simulations over time, starting from the initial conditions of the test set, is depicted in Figure 6. For both datasets, EPNS generally produces simulations that remain stable for longer than other models, suggesting that equivariance improves rollout stability. Although stability still decreases over time, most of EPNS' simulations remain stable for substantially longer than the rollout lengths seen during training.
Conclusions. In this work, we propose EPNS, a generally applicable framework for equivariant probabilistic spatiotemporal simulation. We evaluate EPNS in the domains of stochastic celestial dynamics and stochastic cellular dynamics. Our results demonstrate that EPNS outperforms existing methods for these applications. Furthermore, we observe that embedding equivariance constraints in EPNS improves data efficiency, stability, and uncertainty estimates. In conclusion, we demonstrate the value of incorporating symmetries in probabilistic neural simulators, and show that EPNS provides the means to achieve this for diverse applications.
Limitations. The autoregressive nature of EPNS results in part of the simulations becoming unstable as the rollout length increases. This is a known limitation of autoregressive neural simulation models, and we believe that further advances in this field can transfer to our method [9, 48, 49]. Second, the SpatialConv-GNN has relatively large memory requirements, and it would be interesting to improve the architecture to be more memory efficient. Third, we noticed that training can be unstable, either due to exploding loss values or due to getting stuck in bad local minima. Our efforts to mitigate this issue resulted in approximately two out of three training runs converging to a good local optimum, and we believe this could be improved by further tweaking of the training procedure. Finally, our experimental evaluation has shown the benefits of equivariant probabilistic simulation using relatively small scale systems with stylized dynamics. An interesting avenue for future work is to apply EPNS to large-scale systems with more complex interactions that are of significant interest to domain experts in various scientific fields. Examples of such applications could be the simulation of Langevin dynamics or cancer cell migration.
The authors would like to thank Jakub Tomczak and Tim d'Hondt for the insightful discussions and valuable feedback. This work used the Dutch national e-infrastructure with the support of the SURF Cooperative using grant no. EINF-3935 and grant no. EINF-3557.
[49] Jingtong Su, Julia Kempe, Drummond Fielding, Nikolaos Tsilivis, Miles Cranmer, and Shirley Ho. Adversarial noise injection for learned turbulence simulations. In NeurIPS 2022 Workshop on Machine Learning and the Physical Sciences, 2022.
[50] Sophia Huiwen Sun, Robin Walters, Jinxi Li, and Rose Yu. Probabilistic symmetry for multiagent dynamics. In Nikolai Matni, Manfred Morari, and George J. Pappas, editors, Proceedings of The 5th Annual Learning for Dynamics and Control Conference, volume 211 of Proceedings of Machine Learning Research, pages 1231-1244. PMLR, 15-16 Jun 2023.
[51] Thomas Thenard, Anita Catapano, Michel Mesnard, and Rachele Allena. A cellular potts energy-based approach to analyse the influence of the surface topography on single cell motility. Journal of Theoretical Biology, 509:110487, September 2020.
[52] Mark Vogelsberger, Shy Genel, Volker Springel, Paul Torrey, Debora Sijacki, Dandan Xu, Greg Snyder, Dylan Nelson, and Lars Hernquist. Introducing the illustris project: simulating the coevolution of dark and visible matter in the universe. Monthly Notices of the Royal Astronomical Society, 444(2):1518-1547, August 2014.
[53] Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. In International Conference on Learning Representations, 2021.
[54] Rui Wang, Robin Walters, and Rose Yu. Approximately equivariant networks for imperfectly symmetric dynamics. In Proceedings of the 39th International Conference on Machine Learning, 2022.
[55] S. Wiesen, D. Reiter, V. Kotov, M. Baelmans, W. Dekeyser, A.S. Kukushkin, S.W. Lisgo, R.A. Pitts, V. Rozhansky, G. Saibene, I. Veselova, and S. Voskoboynikov. The new SOLPS-ITER code package. Journal of Nuclear Materials, 463:480–484, August 2015.
[56] Tailin Wu, Takashi Maruyama, Qingqing Zhao, Gordon Wetzstein, and Jure Leskovec. Learning controllable adaptive simulation for multi-resolution physics. In International Conference on Learning Representations, 2023.
[57] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
[58] Cagatay Yildiz, Markus Heinonen, and Harri Lahdesmaki. Ode2vae: Deep generative second order odes with bayesian neural networks. In Advances in Neural Information Processing Systems, 2019.
[59] Cagatay Yildiz, Melih Kandemir, and Barbara Rakitsch. Learning interacting dynamical systems with latent gaussian process ODEs. In Advances in Neural Information Processing Systems, 2022.
A Proof of equivariance
Here, we prove Lemma 1, which we repeat below:
Lemma 1. Assume we model  $p_{\theta}(x^{1:T}|x^0)$  as described in Section 3.1, that is,  $p_{\theta}(x^{1:T}|x^0) = \prod_{t=0}^{T-1} p_{\theta}(x^{t+1}|x^t)$  and  $p_{\theta}(x^{t+1}|x^t) = \int_z p_{\theta}(x^{t+1}|z,h^t)p_{\theta}(z|h^t)dz$ , where  $h^t = f_{\theta}(x^t)$ . Then,  $p_{\theta}(x^{1:T}|x^0)$  is equivariant to linear transformations  $\rho(g)$  of a symmetry group  $G$  in the sense of definition 2 if:
(a) The forward model is  $G$ -equivariant:  $f_{\theta}(\rho(g)x) = \rho(g)f_{\theta}(x)$ ;
(b) The conditional prior is  $G$ -invariant or  $G$ -equivariant:  $p_{\theta}(z|\rho(g)h^t) = p_{\theta}(z|h^t)$ , or  $p_{\theta}(\rho(g)z|\rho(g)h^t) = p_{\theta}(z|h^t)$ ;
(c) The decoder is  $G$ -equivariant with respect to  $h^t$  (for an invariant conditional prior), or  $G$ -equivariant with respect to both  $h^t$  and  $z$  (for an equivariant conditional prior):  $p_{\theta}(x^{t + 1}|z,h^{t}) = p_{\theta}(\rho (g)x^{t + 1}|z,\rho (g)h^{t})$ , or  $p_{\theta}(x^{t + 1}|z,h^{t}) = p_{\theta}(\rho (g)x^{t + 1}|\rho (g)z,\rho (g)h^{t})$ .
To support the proof, we first show that, if a  $G$ -equivariant distribution  $p(y|x)$  exists, i.e.  $p(y|x) = p(\rho(g)y|\rho(g)x)$ , then the absolute value of the determinant of  $\rho(g)$  must equal 1. To this end, let  $y_g = \rho(g)y$  and  $x_g = \rho(g)x$
$$
\begin{array}{l} 1 = \int_ {y _ {g}} p \left(y _ {g} \mid x _ {g}\right) d y _ {g} \\ = \int_ {y _ {g}} p (y | x) d y _ {g} \quad \left\{p \text {i s e q u i v a r i a n t} \right\} \\ = \int_ {y} p (y | x) | \det  \rho (g) | d y \quad \left\{\text {s u b s t i t u t e} y _ {g} = \rho (g) y \right\} \\ = | \det  \rho (g) | \int_ {y} p (y | x) d y \\ = | \det  \rho (g) | \\ \end{array}
$$
For the lemma to hold for a first-order Markovian data-generating process, it suffices to have an autoregressive model of which the one-step transition distribution  $p_{\theta}(x^{t + 1}|x^t)$  is equivariant. To show this, let  $x_{g}^{t} = \rho (g)x^{t}$ :
$$
\begin{array}{l} p _ {\theta} (\rho (g) x ^ {1: T} | \rho (g) x ^ {0}) = p _ {\theta} (x _ {g} ^ {1: T} | x _ {g} ^ {0}) \\ = \prod_ {t = 1} ^ {T} p _ {\theta} \left(x _ {g} ^ {t} \mid x _ {g} ^ {t - 1}\right) \quad \{\text {M a r k o v p r o p e r t y} \} \\ = \prod_ {t = 1} ^ {T} p _ {\theta} \left(x ^ {t} \mid x ^ {t - 1}\right) \quad \text {E q u i v a r i a n t o n e - s t e p} p _ {\theta} \} \\ = p _ {\theta} \left(x ^ {1: T} \mid x ^ {0}\right) \quad \left\{\text {I n v e r s e M a r k o v p r o p e r t y} \right\} \\ \end{array}
$$
As such, it remains to show that  $p_{\theta}(x^{t + 1}|x^t)$  is equivariant when conditions (a)-(c) of the lemma hold. We show here the case where  $p_{\theta}(z|h^t)$  is equivariant, the invariant case is analogous, where  $g$  acts on  $z$  as  $\rho(g)z = Iz$ . Again, let  $x_g^t = \rho(g)x^t$ , and let  $z_g = \rho(g)z$ :
$$
\begin{array}{l} p _ {\theta} (\rho (g) x ^ {t + 1} | \rho (g) x ^ {t}) = p _ {\theta} \left(x _ {g} ^ {t + 1} \mid x _ {g} ^ {t}\right) \\ = p _ {\theta} \left(x _ {g} ^ {t + 1} \mid h _ {g} ^ {t}\right) \quad \left\{\text {e q u i v a r i a n t} f _ {\theta} \right\} \\ = \int_ {z _ {g}} p _ {\theta} (x _ {g} ^ {t + 1} | z _ {g}, h _ {g} ^ {t}) p _ {\theta} (z _ {g} | h _ {g} ^ {t}) d z _ {g} \\ = \int_ {z _ {g}} p _ {\theta} \left(x ^ {t + 1} \mid z, h ^ {t}\right) p _ {\theta} \left(z \mid h ^ {t}\right) d z _ {g} \quad \text {e q u i v a r i a n t} p _ {\theta} \} \\ = \int_ {z} p _ {\theta} \left(x ^ {t + 1} \mid z, h ^ {t}\right) p _ {\theta} \left(z \mid h ^ {t}\right) | \det  \rho (g) | d z \quad \{z _ {g} = \rho (g) z \} \\ = \int_ {z} p _ {\theta} \left(x ^ {t + 1} \mid z, h ^ {t}\right) p _ {\theta} \left(z \mid h ^ {t}\right) d z \quad \left\{\left| \det  \rho (g) \right| = 1 \right\} \\ = p _ {\theta} \left(x ^ {t + 1} \mid x ^ {t}\right) \\ \end{array}
$$
B Parameters of the Cellular Potts model
Recall that the Hamiltonian for the cell sorting simulation is defined as follows:
$$
H = \sum_ {l _ {i}, l _ {j} \in \mathcal {N} (L)} J \left(x \left(l _ {i}\right), x \left(l _ {j}\right)\right) \left(1 - \delta_ {x \left(l _ {i}\right), x \left(l _ {j}\right)}\right) + \sum_ {c \in C} \lambda_ {V} \left(V (c) - V ^ {*} (c)\right) ^ {2}, \tag {10}
$$
For an explanation on the interpretation of this equation, please consult Section 4.2. The parameter values used in Equation 10 can be found in Tables 4 and 5 below:
Table 4: Parameter values of the Cellular Potts simulation.
Parameter	Value
J(x(li),x(lj))	See Table 5
λV	2
V*(c)	25
T	8
Table 5: Contact energy  $J\left( {x\left( {l}_{i}\right) ,x\left( {l}_{j}\right) }\right)$  lookup table.
Cell type	Medium	Type 1	Type 2	Type 3	Type 4
Medium	0	16	16	16	16
Type 1	16	8	11	11	11
Type 2	16	11	8	11	11
Type 3	16	11	11	8	11
Type 4	16	11	11	11	8
These parameters were not chosen to be biologically plausible, but rather to result in behavior in which there is a clear tendency of clustering of same-type cells as time proceeds, and to strike a balance between the stochastic fluctuations and clustering tendencies in the system.
C EPNS model details
C.1 Celestial dynamics
Before going into the specifics, we first discuss some general remarks on the EPNS model implementation. Recall that the data is modeled as a geometric graph consisting of  $n$  nodes. The celestial dynamics EPNS model distinguishes between E(n)-invariant features (mass  $m_{i} \in \mathbb{R}^{1}$ ) and equivariant features (position  $\mathbf{p}_{\mathbf{i}}^t \in \mathbb{R}^3$  and velocity  $\mathbf{v}_{\mathbf{i}}^t \in \mathbb{R}^3$ ) for each node  $i$ . We provide the Euclidean distance  $(||\mathbf{p}_{\mathbf{i}}^t - \mathbf{p}_{\mathbf{j}}^t||_2)$ , as well as the magnitude of the difference between the velocity vectors  $(||\mathbf{v}_{\mathbf{i}}^t - \mathbf{v}_{\mathbf{j}}^t||_2)$  between all nodes as edge features  $e_{ij}^t \in \mathbb{R}^2$  to the model. ReLU is used throughout the model as activation function. Our implementation of the FA-GNN differs slightly from [38], as we do not construct and invert the frame at each message passing layer. Instead, we construct the frame once, then apply all message passing layers, and then invert the frame again to arrive at invariant or equivariant output. Overall, the model has about 1.8 million trainable parameters.
Forward model:
$$
\left\{m \in \mathbb {R} ^ {n \times 1}, e ^ {t} \in \mathbb {R} ^ {n ^ {2} \times 2}, \mathbf {p} ^ {\mathbf {t}} \in \mathbb {R} ^ {n \times 3}, \mathbf {v} ^ {\mathbf {t}} \in \mathbb {R} ^ {n \times 3} \right\}\rightarrow f _ {\theta} \rightarrow
$$
$$
\left\{h ^ {t} \in \mathbb {R} ^ {n \times 1 2 8}, e ^ {t} \in \mathbb {R} ^ {n ^ {2} \times 2}, \mathbf {p} ^ {t} \in \mathbb {R} ^ {n \times 3}, \mathbf {v} ^ {t} \in \mathbb {R} ^ {n \times 3} \right\}
$$