SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We evaluate A-NESI on the Multi-Digit MNISTAdd task (Section 2). For the perception model, we use the same CNN as in DeepProbLog [35]. The prediction model has  $N + 1$  factors  $q_{\phi}(y_i | \mathbf{y}_{1:i - 1}, \mathbf{P})$ , while the explanation model has  $2N$  factors  $q_{\psi}(w_i | \mathbf{y}, \mathbf{w}_{1,i - 1}, \mathbf{P})$ . We model each factor with a separate MLP.  $y_i$  and  $w_i$  are one-hot encoded digits, except for the first output digit  $y_1$ : it can only be 0 or 1. We used a shared set of hyperparameters for all  $N$ . For more details and a description of the baselines, see Appendix I.1.
Table 1 reports the accuracy of predicting the sum. For all  $N$ , A-NESI is close to the reference accuracy, meaning there is no significant drop in the accuracy of the perception model as  $N$  increases. For small  $N$ , it is slightly outperformed by DeepStochLog [56], which can not scale to  $N = 15$ . A-NESI also performs slightly better than DeepProbLog, showing approximate inference does not hurt the accuracy. With neural prediction, we get the same accuracy for low  $N$ , but there is a significant drop for  $N = 15$ , meaning the prediction network did not perfectly learn the problem. However, compared to training a neural network without background knowledge (Embed2Sym with
N=1	N=2	N=4	N=15
Symbolic prediction
LTN	80.54 ± 23.33	77.54 ± 35.55	T/O	T/O
DeepProbLog	97.20 ± 0.50	95.20 ± 1.70	T/O	T/O
DPLA*	88.90 ± 14.80	83.60 ± 23.70	T/O	T/O
DeepStochLog	97.90 ± 0.10	96.40 ± 0.10	92.70 ± 0.60	T/O
Embed2Sym	97.62 ± 0.29	93.81 ± 1.37	91.65 ± 0.57	60.46 ± 20.36
A-NESI (predict)	97.66 ± 0.21	95.96 ± 0.38	92.56 ± 0.79	75.90 ± 2.21
A-NESI (explain)	97.37 ± 0.32	96.04 ± 0.46	92.11 ± 1.06	76.84 ± 2.82
A-NESI (pruning)	97.57 ± 0.27	95.82 ± 0.33	92.40 ± 0.68	76.46 ± 1.39
A-NESI (no prior)	76.54 ± 27.38	95.67 ± 0.53	44.58 ± 38.34	0.03 ± 0.09
Neural prediction
Embed2Sym	97.34 ± 0.19	84.35 ± 6.16	0.81 ± 0.12	0.00
A-NESI (predict)	97.66 ± 0.21	95.95 ± 0.38	92.48 ± 0.76	54.66 ± 1.87
A-NESI (explain)	97.37 ± 0.32	96.05 ± 0.47	92.14 ± 1.05	61.77 ± 2.37
A-NESI (pruning)	97.57 ± 0.27	95.82 ± 0.33	92.38 ± 0.64	59.88 ± 2.95
A-NESI (no prior)	76.54 ± 27.01	95.28 ± 0.62	40.76 ± 34.29	0.00 ± 0.00
Reference	98.01	96.06	92.27	73.97
Table 1: Test accuracy of predicting the correct sum on the Multi-digit MNISTAdd task. "T/O" (timeout) represent computational timeouts. Reference accuracy approximates the accuracy of an MNIST predictor with 0.99 accuracy using  $0.99^{2N}$ . Bold numbers signify the highest average accuracy for some  $N$  within the prediction categories. predict is the prediction-only variant, explain is the explainable variant, pruning adds symbolic pruning (see Appendix H), and no prior is the prediction-only variant without the prior  $p(\mathbf{P})$ .
Figure 4: Inference time for a single input  $\mathbf{x}$  for different nr. of digits. The left plot uses a log scale, and the right plot a linear scale. DSL stands for DeepStochLog. We use the GM variant of DPLA*.
neural prediction), it is much more accurate already for  $N = 2$ . Therefore, A-NESI's training loop allows training a prediction network with high accuracy on large-scale problems.
Comparing the different A-NeSI variants, we see the prediction-only, explainable and pruned variants perform quite similarly, with significant differences only appearing at  $N = 15$  where the explainable and pruning variants outperform the predict-only model, especially on neural prediction. However, when removing the prior  $p(\mathbf{P})$ , the performance degrades quickly. The prediction model sees much fewer beliefs  $\mathbf{P}$  than when sampling from a (high-entropy) prior  $p(\mathbf{P})$ . A second and more subtle reason is that at the beginning of training, all the beliefs  $\mathbf{P} = f_{\theta}(\mathbf{x})$  will be uniform because the perception model is not yet trained. Then, the prediction model learns to ignore the input belief  $\mathbf{P}$ .
Figure 4 shows the runtime for inference on a single input  $\mathbf{x}$ . Inference in DeepProbLog (corresponding to exact inference) increases with a factor 100 as  $N$  increases, and DPLA* (another approximation) is not far behind. Inference in DeepStochLog, which uses different semantics, is efficient due to caching but requires a grounding step that is exponential both in time and memory. We could not ground beyond  $N = 4$  because of memory issues. A-NESI avoids having to perform grounding altogether: it scales slightly slower than linear. Furthermore, it is much faster in practice as parallelizing the computation of multiple queries on GPUs is trivial.
N=4	N=9
CNN	51.50 ± 3.34	51.20 ± 2.20
Exact inference	86.70 ± 0.50	T/O
NeuPSL	89.7 ± 2.20	51.50 ± 1.37
A-NESI (symbolic prediction)	89.70 ± 2.08	62.15 ± 2.08
A-NESI (neural prediction)	89.80 ± 2.11	62.25 ± 2.20
4.2 Visual Sudoku Puzzle Classification
In this task, the goal is to recognize whether an  $N \times N$  grid of MNIST digits is a Sudoku. We have 100 examples of correct and incorrect grids from [4]. We use the same MNIST classifier as in the previous section. We treat  $\mathbf{w}$  as a grid of digits in  $\{1, \dots, N\}$  and designed an easy-to-learn representation of the label. For sudokus, all pairs of digits  $\mathbf{w}_i, \mathbf{w}_j$  in a row, column, and block must differ. For each such pair, a dimension in  $\mathbf{y}$  is 1 if different and 0 otherwise, and the symbolic function  $c$  computes these interactions from  $\mathbf{w}$ . The prediction model is a single MLP that takes the probabilities for the digits  $\mathbf{P}_i$  and  $\mathbf{P}_j$  and predicts the probability that those represent different digits. For additional details, see Appendix I.2.
Table 2 shows the accuracy of classifying Sodomu puzzles. A-NESI is the only method to perform better than random on  $9 \times 9$  sudokus. Exact inference cannot scale to  $9 \times 9$  sudokus, while we were able to run A-NESI for 3000 epochs in 38 minutes on a single NVIDIA RTX A4000.
4.3 Wcrafcraft Visual Path Planning
Figure 5: Example of the Wcrafcraft task.
Table 2: Test accuracy of predicting whether a grid of numbers is a Selenium. For A-NeSI, we used the prediction-only variant.
12 × 12	30 × 30
ResNet18 [43]	23.0 ± 0.3	0.0 ± 0.0
SPL [1]	78.2	T/O
I-MLE [41]	97.2 ± 0.5	93.7 ± 0.6
RLOO [29]	43.75 ± 12.35	12.59 ± 16.38
A-NESI	94.57 ± 2.27	17.13 ± 16.32
A-NESI + RLOO	98.96 ± 1.33	67.57 ± 36.76
Table 3: Test accuracy of predicting the lowest-cost path on the Wcraft Visual Path Planning task. For A-NeSI, we used the prediction-only variant.
The Wcraft Visual Path Planning task [43] is to predict a shortest path from the top left corner of an  $N \times N$  grid to the bottom right, given an image from the Wcraft game (Figure 5). The perception model has to learn the cost of each tile of the Wcraft image. We use a small CNN that takes the tile  $i,j$  (a  $3 \times 8 \times 8$  image) and outputs a distribution over 5 possible costs. The symbolic function  $c(\mathbf{w})$  is Dijkstra's algorithm and returns a shortest path  $\mathbf{y}$ , which we encode as a sequence of the 8 (inter)-cardinal directions starting from the top-left corner. We use a ResNet18 [22] as the prediction model, which we train to predict the next direction given beliefs of the tile costs  $\mathbf{P}$  and the current location. We pretrain the prediction model on a fixed prior and train the perception model with a frozen prediction model. See Appendix I.3 for additional details.
Table 3 presents the accuracy of predicting a shortest path. A-NESI is competitive on  $12 \times 12$  grids but struggles on  $30 \times 30$  grids. We believe this is because the prediction model is not accurate enough, resulting in gradients with too high bias. Still, A-NESI finds short paths, is far better than a pure neural network, and can scale to  $30 \times 30$  grids, unlike SPL [1], which uses exact inference. Since both SPL and I-MLE [41] have significantly different setups (see Appendix I.3.2), we added experiments using REINFORCE with the leave-one-out baseline (RLOO, [29]) that we implemented with the Storchastic PyTorch library [53]. We find that RLOO has high variance in its performance. Since RLOO is unbiased with high variance and A-NESI is biased with no variance, we also tried running A-NESI and RLOO simultaneously. Interestingly, this is the best-performing method on the  $12 \times 12$  grid, and has competitive performance on  $30 \times 30$  grids, albeit with high variance (6 out 10 runs get to an accuracy between  $93.3\%$  and  $98.5\%$ , while the other runs are stuck around  $25\%$  accuracy).
A-NESI can approximate multiple PNL methods [16]. DeepProbLog [35] performs symbolic reasoning by representing  $\mathbf{w}$  as ground facts in a Prolog program. It enumerates all possible proofs of a query  $\mathbf{y}$  and weights each proof by  $p(\mathbf{w}|\mathbf{P})$ . NeurASP [58] is a PNL framework closely related to DeepProbLog, but is based on Answer Set Programming [12]. Some methods consider constrained structured output prediction [19]. In Appendix B, we discuss extending A-NESI to this setting. Semantic Loss [57] improves learning with a loss function but does not guarantee that formulas are satisfied at test time. Like A-NESI, Semantic Probabilistic Layers [1] solves this with a layer that performs constrained prediction. These approaches perform exact inference using probabilistic circuits (PCs) [60]. Other methods perform approximate inference by only considering the top-k proofs in PCs [36, 23]. However, finding those proofs is hard, especially when beliefs have high entropy, and limiting to top-k significantly reduces performance. Other work considers MCMC approximations [31]. Using neural networks for approximate inference ensures computation time is constant and independent of the entropy of  $\mathbf{P}$  or long MCMC chains.
Other neurosymbolic methods use fuzzy logics [5, 17, 15, 18], which are faster than PNL with exact inference. Although traversing ground formulas is linear time, the grounding is itself often exponential [37], so the scalability of fuzzy logics often fails to deliver. A-NESI is polynomial in the number of ground atoms and does not traverse the ground formula. Furthermore, background knowledge is often not fuzzy [54, 20], and fuzzy semantics does not preserve classical equivalence.
A-NESI performs gradient estimation of the WMC problem. We can extend our method to biased but zero-variance gradient estimation by learning a distribution over function outputs (see Appendix C). Many recent works consider continuous relaxations of discrete computation to make them differentiable [42, 24] but require many tricks to be computationally feasible. Other methods compute MAP states to compute the gradients [41, 11, 47] but are restricted to integer linear programs. The score function (or 'REINFORCE') gives unbiased yet high-variance gradient estimates [40, 53]. Variance reduction techniques, such as memory augmentation [32] and leave-one-out baselines [29], exist to reduce this variance.
6 Conclusion, Discussion and Limitations
We introduced A-NESI, a scalable approximate method for probabilistic neurosymbolic learning. We demonstrated that A-NESI scales to combinatorially challenging tasks without losing accuracy. A-NESI can be extended to include explanations and hard constraints without loss of performance.
However, there is no 'free lunch': when is A-NESI a good approximation? We discuss three aspects of learning tasks that could make it difficult to learn a strong and efficient inference model.
- Dependencies of variables. When variables in world  $\mathbf{w}$  are highly dependent, finding an informative prior is hard. We suggest then using a prior that can incorporate dependencies such as a normalizing flow [45, 14] or deep generative models [51] over a Dirichlet distribution.
- Structure in symbolic reasoning function. We studied reasoning tasks with a relatively simple structure. Learning the inference model will be more difficult when the symbolic function  $c$  is less structured. Studying the relation between structure and learnability is interesting future work.
- Problem size. A-NESI did not perfectly train the prediction model in more challenging problems, which is evident from the divergence in performance between symbolic and neural prediction for 15 digits in Table 1. We expect its required parameter size and training time to increase with the problem size.
Promising future avenues are studying if the explanation model produces helpful explanations [50], extensions to continuous random variables [21] (see Appendix C for an example), and extensions to unnormalized distributions such as Markov Logic Networks[46], as well as (semi-) automated A-NESI solutions for neurosymbolic programming languages like DeepProbLog [35].
We want to thank Robin Manhaeve, Giusseppe Marra, Thomas Winters, Yaniv Aspis, Robin Smet, and Anna Kuzina for helpful discussions. For our experiments, we used the DAS-6 compute cluster [6]. We also thank SURF (www.surf.nl) for its support in using the Lisa Compute Cluster.
[1] Kareem Ahmed, Stefano Teso, Kai-Wei Chang, Guy Van den Broeck, and Antonio Vergari. Semantic probabilistic layers for neuro-symbolic learning. Advances in Neural Information Processing Systems, 35:29944-29959, 2022.
[2] Kareem Ahmed, Eric Wang, Kai-Wei Chang, and Guy Van den Broeck. Neuro-Symbolic Entropy Regularization. In The 38th Conference on Uncertainty in Artificial Intelligence, June 2022.
[3] Yaniv Aspis, Krysia Broda, Jorge Lobo, and Alessandra Russo. Embed2Sym - Scalable Neuro-Symbolic Reasoning via Clustered Embeddings. In Proceedings of the Nineteenth International Conference on Principles of Knowledge Representation and Reasoning, pages 421-431, Haifa, Israel, July 2022. International Joint Conferences on Artificial Intelligence Organization. ISBN 978-1-956792-01-0. doi: 10.24963/kr.2022/44.
[4] Eriq Augustine, Connor Pryor, Charles Dickens, Jay Pujara, William Yang Wang, and Lise Getoor. Visual神经系统: A suite of collective neuro-symbolic tasks. In International Workshop on Neural-Symbolic Learning and Reasoning (NeSy), Windsor, United Kingdom, 2022.
[5] Samy Badreddine, Artur d'Avila Garcez, Luciano Serafini, and Michael Spranger. Logic Tensor Networks. Artificial Intelligence, 303:103649, February 2022. ISSN 0004-3702. doi: 10.1016/j.artint.2021.103649.
[6] Henri Bal, Dick Epema, Cees de Laat, Rob van Nieuwpoort, John Romein, Frank Seinstra, Cees Snoek, and Harry Wijshoff. A medium-scale distributed system for computer science research: Infrastructure for the long term. Computer, 49(5):54-63, 2016.
[7] Tomas Balyo, Marijn J. H. Heule, Markus Iser, Matti Järvisalo, Martin Suda, Helsinki Institute for Information Technology, Constraint Reasoning and Optimization research group / Matti Järvisalo, and Department of Computer Science. Proceedings of SAT Competition 2022: Solver and Benchmark Descriptions. 2022.
[8] Marc G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforcement learning. In International Conference on Machine Learning, pages 449-458. PMLR, 2017.
[9] Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Advances in Neural Information Processing Systems, 34:27381-27394, 2021.
[10] Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J. Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet Foundations, August 2022.
[11] Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. Advances in Neural Information Processing Systems, 35:5230–5242, 2022.
[12] Gerhard Brewka, Thomas Eiter, and Miroław Truszczyński. Answer set programming at a glance. Communications of the ACM, 54(12):92-103, 2011.
[13] Mark Chavira and Adnan Darwiche. On probabilistic inference by weighted model counting. Artificial Intelligence, 172(6):772-799, April 2008. ISSN 0004-3702. doi: 10.1016/j.artint.2007.11.002.
A Derivation of prediction-only variant loss
$$
\mathbb {E} _ {p (\mathbf {P})} \left[ - \mathbb {E} _ {p (\mathbf {y} | \mathbf {P})} [ \log q _ {\phi} (\mathbf {y} | \mathbf {P}) ] \right] \tag {10}
$$
$$
= - \mathbb {E} _ {p (\mathbf {P})} \left[ \mathbb {E} _ {p (\mathbf {w}, \mathbf {y} | \mathbf {P})} [ \log q _ {\phi} (\mathbf {y} | \mathbf {P}) ] \right] \tag {11}
$$
$$
= - \mathbb {E} _ {p (\mathbf {P}, \mathbf {w})} [ \log q _ {\phi} (c (\mathbf {w}) | \mathbf {P}) ] \tag {12}
$$
$$
\mathcal {L} _ {\text {P r e d}} (\phi) = - \log q _ {\phi} (c (\mathbf {w}) | \mathbf {P}), \quad \mathbf {P}, \mathbf {w} \sim p (\mathbf {P}, \mathbf {w}) \tag {13}
$$
In line 11, we marginalize out  $\mathbf{w}$ , and use the fact that  $\mathbf{y}$  is deterministic given  $\mathbf{w}$ .
B Constrained structured output prediction
We consider how to implement the constrained structured output prediction task considered in (for example) [2, 1, 57] in the A-NESI framework. Here, the goal is to learn a mapping of some  $X$  to a structured output space  $W$ , where we have some constraint  $c(\mathbf{w})$  that returns 1 if the background knowledge holds, and 0 otherwise. We can model the constraints using  $Y = \{0,1\}$ ; that is, the 'output' in our problem setup is whether  $\mathbf{w}$  satisfies the background knowledge  $c$  or not. We give an example of this setting in Figure 6.
Then, we design the inference model as follows. 1)  $q_{\phi}(y|\mathbf{P})$  is tasked with predicting the probability that randomly sampled outputs  $\mathbf{w} \sim p(\mathbf{w}|\mathbf{P})$  will satisfy the background knowledge. 2)  $q_{\psi}(\mathbf{w}|y = 1,\mathbf{P})$  is an approximate posterior over structured outputs  $\mathbf{w}$  that satisfy the background knowledge  $c$ .
This setting changes the interpretation of the set  $W$  from unobserved worlds to observed outputs. We will train our perception module using a "strongly" supervised learning loss where  $\mathbf{x}, \mathbf{w} \sim \mathcal{D}_{\mathcal{L}}$ :
$$
\mathcal {L} _ {\text {P e r c}} (\boldsymbol {\theta}) = - \log q _ {\psi} (\mathbf {w} | y = 1, \mathbf {P} = f _ {\boldsymbol {\theta}} (\mathbf {x})). \tag {14}
$$
If we also have unlabeled data  $\mathcal{D}_U$ , we can use the prediction model to ensure the perception model gives high probabilities for worlds that satisfy the background knowledge. This approximates Semantic Loss [57]: Given  $\mathbf{x} \sim \mathcal{D}_U$ ,
$$
\mathcal {L} _ {S L} (\boldsymbol {\theta}) = - \log q _ {\phi} (y = 1 | \mathbf {P} = f _ {\boldsymbol {\theta}} (\mathbf {x})). \tag {15}
$$
That is, we have some input  $\mathbf{x}$  for which we have no labelled output. Then, we increase the probability that the belief  $\mathbf{P}$  the perception module  $f_{\theta}$  predicts for  $\mathbf{x}$  would sample structured outputs  $\mathbf{w}$  that satisfy the background knowledge.
Training the inference model in this setting can be challenging if the problem is very constrained. Then, random samples  $\mathbf{P},\mathbf{w}\sim p(\mathbf{P},\mathbf{w})$  will usually not satisfy the background knowledge. Since we are only in the case that  $y = 1$ , we can choose to sample from the inference model  $q_{\phi}$  and exploit the symbolic pruner to obtain samples that are guaranteed to satisfy the background knowledge. Therefore, we modify equation 8 to the on-policy joint matching loss
$$
\mathcal {L} _ {E x p l} (\mathbf {P}, \phi , \psi) = \mathbb {E} _ {q _ {\psi} (\mathbf {w} | y = 1, \mathbf {P})} \left[ \left(\log \frac {q _ {\phi , \psi} (\mathbf {w} , y = 1 | \mathbf {P})}{p (\mathbf {w} | \mathbf {P})}\right) ^ {2} \right] \tag {16}
$$
Here, we incur some sampling bias by not sampling structured outputs from the true posterior, but this bias will reduce as  $q_{\phi}$  becomes more accurate. We can also choose to combine the on- and off-policy losses. Another option to make learning easier is using the suggestions of Section 3.2.2: factorize  $y$  to make it more fine-grained.
C A-NESI as a Gradient Estimation method
In this appendix, we discuss using the method A-NESI introduced for general gradient estimation [40]. We first define the gradient estimation problem. Consider some neural network  $f_{\theta}$  that predicts the parameters  $\mathbf{P}$  of a distribution over unobserved variable  $\mathbf{z} \in Z$ :  $p(\mathbf{z}|\mathbf{P} = f_{\theta}(\mathbf{x}))$ . This distribution corresponds to the distribution over worlds  $p(\mathbf{w}|\mathbf{P})$  in A-NESI. Additionally, assume we have some deterministic function  $g(\mathbf{z})$  that we want to maximize in expectation. This maximization requires estimating the gradient
$$
\nabla_ {\boldsymbol {\theta}} \mathbb {E} _ {p (\mathbf {z} | \mathbf {P} = f _ {\boldsymbol {\theta}} (\mathbf {x}))} [ g (\mathbf {z}) ]. \tag {17}
$$
Common methods for estimating this gradient are reparameterization [26], which only applies to continuous random variables and differentiable  $r$ , and the score function [40, 49] which has notoriously high variance.
Instead, our gradient estimation method learns an inference model  $q_{\phi}(r|\mathbf{P})$  to approximate the distribution of outcomes  $r = g(\mathbf{z})$  for a given  $\mathbf{P}$ . In A-NESI, this is the prediction network  $q_{\phi}(\mathbf{y}|\mathbf{P})$  that estimates the WMC problem of Equation 1. Approximating a distribution over outcomes is similar to the idea of distributional reinforcement learning [8]. Our approach is general: Unlike reparameterization, we can use inference models in settings with discrete random variables  $\mathbf{z}$  and non-differentiable downstream functions  $g$ .
We derive the training loss for our inference model similar to that in Section 3.2.1. First, we define the joint on latents  $\mathbf{z}$  and outcomes  $r$  like the joint process in 6 as  $p(r,\mathbf{z}|\mathbf{P}) = p(\mathbf{z}|\mathbf{P})\cdot \delta_{g(z)}(r)$ , where  $\delta_{g(z)}(r)$  is the dirac-delta distribution that checks if the output of  $g$  on  $\mathbf{z}$  is equal to  $r$ . Then we introduce a prior over distribution parameters  $p(\mathbf{P})$ , much like the prior over beliefs in A-NESI. An obvious choice is to use a prior conjugate to  $p(\mathbf{z}|\mathbf{P})$ . We minimize the expected KL-divergence between  $p(r|\mathbf{P})$  and  $q_{\phi}(r|\mathbf{P})$ :
$$
\mathbb {E} _ {p (\mathbf {P})} \left[ D _ {K L} (p | | q _ {\phi}) \right] \tag {18}
$$
$$
= \mathbb {E} _ {p (\mathbf {P})} \left[ \mathbb {E} _ {p (r | \mathbf {P})} [ \log q _ {\phi} (r | \mathbf {P}) ] \right] + C \tag {19}
$$
$$
= \mathbb {E} _ {p (\mathbf {P})} \left[ \int_ {\mathbb {R}} p (r | \mathbf {P}) \log q _ {\phi} (r | \mathbf {P}) ] d r \right] + C \tag {20}
$$
Next, we marginalize over  $\mathbf{z}$ , dropping the constant:
$$
\mathbb {E} _ {p (\mathbf {P})} \left[ \int_ {Z} \int_ {\mathbb {R}} p (r, \mathbf {z} | \mathbf {P}) \log q _ {\phi} (r | \mathbf {P}) ] d r d \mathbf {z} \right] \tag {21}
$$
$$
= \mathbb {E} _ {p (\mathbf {P})} \left[ \int_ {Z} p (\mathbf {z} | \mathbf {P}) \int_ {\mathbb {R}} \delta_ {g (\mathbf {z})} (r) \log q _ {\phi} (r | \mathbf {P}) ] d r d \mathbf {z} \right] \tag {22}
$$
$$
= \mathbb {E} _ {p (\mathbf {P})} \left[ \int_ {Z} p (\mathbf {z} | \mathbf {P}) \log q _ {\phi} (g (\mathbf {z}) | \mathbf {P}) ] d \mathbf {z} \right] \tag {23}
$$
$$
= \mathbb {E} _ {p (\mathbf {P}, \mathbf {z})} [ \log q _ {\phi} (g (\mathbf {z}) | \mathbf {P}) ] \tag {24}
$$
This gives us a negative-log-likelihood loss function similar to Equation 7.
$$
\mathcal {L} _ {I n f} (\phi) = - \log q _ {\phi} (g (\mathbf {z}) | \mathbf {P}), \quad \mathbf {P}, \mathbf {z} \sim p (\mathbf {z}, \mathbf {P}) \tag {25}
$$
where we sample from the joint  $p(\mathbf{z},\mathbf{P}) = p(\mathbf{P})p(\mathbf{z}|\mathbf{P})$
We use a trained inference model to get gradient estimates:
$$
\nabla_ {\mathbf {P}} \mathbb {E} _ {p (\mathbf {z} | \mathbf {P})} [ g (\mathbf {z}) ] \approx \nabla_ {\mathbf {P}} \mathbb {E} _ {q _ {\phi} (r | \mathbf {P})} [ r ] \tag {26}
$$