A Fast Convoluted Story: Scaling Probabilistic Inference for Integer Arithmetics

We thank the reviewer for their supportive review. Please allow us to address your experimental concern below.

Our experimental evaluation tackled two of the most prominent neurosymbolic benchmarks, where PLIA$_t$ outperformed both exact and approximate state-of-the-art methods for neurosymbolic learning. It demonstrates that the field of neurosymbolic AI is in need of more challenging benchmarks, as we have shown that it is now possible to solve both MNIST addition and visual sudoku quickly and accurately. The goal of our paper is to scale probabilistic inference on linear integer arithmetic, which we illustrate by outperforming the state-of-the-art in inference (exact method) by 100 000x and the state of the art in learning (approximate method) by 100x in terms of run time. Hence, we fully agree it is high time for more challenging benchmarks, yet proposing new benchmarks is not the aim of this paper.

Further scaling the MNIST addition task to more digits is limited by the size of the dataset, as there are only 60 000 images available for training and 10 000 for testing. Each of these images can only be used once in the construction of the various sums, meaning increasing the number of digits leads to an overall reduction of available sums for training and testing. For example, our case of $N = 50$ requires 100 images to construct a single sum data sample. Consequently, the test dataset only contains 100 samples, which was the lowest number of samples we were still comfortable with to compute test statistics over. We did run learning experiments with $N = 100$ and even $N = 1000$, but we do not deem their test statistics faithful enough to be reported.

We thank the reviewer for their time and effort, and for their suggestion of formalizing PLIA$_t$. For the camera-ready version we propose to add the following subsection at the end of the Section 2:

2.4. Formalizing PLIA$_t$

Definition *(probabilistic linear arithmetic expression) Let $\{X_1, …X_N\}$ be a set of $N$ independent integer-valued random variables with bounded domains. A probabilistic linear integer arithmetic expression X is a bounded integer-valued random variable of the form

In the equation above the functions $f_i$ denote operations performed on the specified random variables that can be either one of the operations specified in Section 3 as well as compositions thereof.*

PLIA$_t$ is concerned with computing the parametric form of the probability distribution of a linear integer arithmetic expression. It does so by representing random variables and linear combinations thereof (cf. Definition of probabilistic linear arithmetic expressions) as tensors whose entries are the log-probabilites of the individual events in the sample space of a the random variable that is being represented. Note that operations within PLIA$_t$ are closed. That is, performing either of the operations delineated in Section 3 will again result in a bounded probabilistic integer representable as a tensor of log-probabilities and an off-set parameter indicating the value of the smallest possible event (cf. Section 2.3). In Section 4 we also equip PLIA$_t$ with probabilistic inference primitives that allow it to compute certain expected values efficiently as well as to perform probabilistic branching.

<END OF SECTION 2.4>

Reproducibility concerns. We would like to stress that the complete implementation of PLIA$_t$, including all operations discussed in Sections 3 and 4, are provided in the supplementary material of our submission. This code also contains our complete experimental setup and automatically installs all necessary dependencies in an easy-to-run script as documented in the provided README file. PLIA$_t$ is a hyperparameter-free framework for inference tasks (cf. Section 5.1). For the neurosymbolic learning tasks (cf. Section 5.2) we provide all our optimal parameters in Appendix E. Upon acceptance, we will also make this implementation publicly available. In this regard we believe we follow best practices concerning reproducibility.

We hope our answers have addressed your concerns and allow for an increase in score. Please let us know if there are any further concerns and we will gladly try to address them.

We thank the reviewer for their thorough review and appreciate the raised concerns. We will start below by answering the main questions:

1. Our precise parameter settings are detailed in Appendix E for the neurosymbolic experiments (cf. Section 5.2), being learning rate, number of training epochs and neural architecture. As PLIA$_t$ itself is a hyperparameter-free framework, there are no other parameters to set when just performing inference in PLIA$_t$ (cf. Section 5.1). All other implementation details are provided in the supplementary code material (submitted together with the paper), where the precise computational environment, including the exact versions of all necessary libraries, are given. Moreover, our complete implementation of PLIA$_t$, its experiments and the precise parameter settings can also be found in the supplementary code material for ease of reproducibility. Finally, we will publicly release PLIA$_t$ upon acceptance.

2. The issue of numerical overflow or underflow is addressed in Section 2.2. There, we introduce a new variant of the well-known log-sum-exp trick for FFT computations in the log domain, called the fast log-conv-exp trick. Using the fast log-conv-exp trick PLIA$_t$ runs without any numerical stability issues.

3. We are currently looking into many different application domains for PLIA$_t$ as the scalability results from our experiments, where we observe orders of magnitude better performance, seem to indicate it is an enabling technology. We did not yet consider applications in cryptography or finance, but these are indeed promising directions! We thank the reviewer for bringing these to our attention.

As for the weaknesses, we would like to address your raised concerns about limitations, complexity and computational overhead.

Limitations. We acknowledge the discussion on the practical applicability of PLIA$_t$ could be extended. We propose to add the following discussion to the camera-ready version of the paper.

The inherent #P-hardness of probabilistic inference does remain the main bottleneck for PLIA$_t$ and can be observed in our experiments. During the comparison between PLIA$_t$ and the state-of-the-art approach of Cao et al. [1] in Figure 5, we can see that PLIA$_t$ also starts to take a significant amount of time when the domain of the probabilistic integers grows. Hence, if the explicit computation of probabilistic integers with a very large domain is required, then PLIA can still struggle. However, PLIA$_t$ does scale orders of magnitude further (5 to be precise) than what was possible so far and considerably pushes the envelope on the state of the art.

Space complexity. It is well-known that the FFT has, apart from the $N \log N$ runtime complexity, also $N \log N$ theoretical space complexity, again in contrast to a $N^2$ space complexity of the naive approach. Moreover, we performed an additional empirical comparison between Dice [1] and PLIA$_t$ with respect to memory allocation (see attached PDF in global rebuttal), where we indeed confirmed the better scaling behaviour of PLIA$_t$. Hence, PLIA$_t$ not only outshines Dice (the current state of the art) in terms of runtime but also in terms of memory usage. We will add this comparison to the appendix and mention it in the main paper in Section 5.1.

Computational overhead. Finally, the computational overhead, including construction of the computational graphs for PLIA$_t$, is already included in the time measurements of all our experiments. For smaller domains, it is indeed possible that PLIA$_t$ on CPU or even Dice can outperform PLIA$_t$ on the GPU (Figure 5, bitwidth strictly below 5). However, it is clear that this computational overhead has a diminishing effect, as PLIA$_t$ clearly scales better than Dice for larger, more practical domain sizes (Figure 5, bitwidth above 5). Moreover, much of this overhead can be avoided by running PLIA$_t$ on the CPU for smaller domain sizes, where it performs on par or better than Dice (Figure 5, bitwidth strictly below 5). We will also make this point clearer in our experimental discussion in Section 5.1.

[1] Cao, W. X., Garg, P., Tjoa, R., Holtzen, S., Millstein, T., & Van den Broeck, G. (2023). Scaling integer arithmetic in probabilistic programs. In Uncertainty in Artificial Intelligence (pp. 260-270). PMLR.

Thank you for your review and for bringing the work of [1] to our attention!

Firstly, we would like to clarify that we do not introduce the log-sum-exp trick as this is indeed a well-known trick to avoid numerical instabilities when summing probabilities in log space. Instead, we introduce a similar trick, dubbed the fast log-conv-exp trick, that avoids numerical instabilities when applying the FFT on probabilities in log space.

Second, we would like to stress the differences between PLIA$_t$ and [1]. Both utilise the FFT to improve computational efficiency of probabilistic inference.

1. PLIA$_t$ and the method of [1] use the FFT for probabilistic inference on different applications. PLIA focuses on scaling probabilistic inference in linear integer arithmetic, while [1] uses the FFT for N-mixture models and applies it to the open-population model of [2].

2. While [1] does discuss numerical instability of computing convolutions in log space, their proposed solution differs from our proposed fast log-conv-exp trick in the following way. As can be seen from algorithm 2 in [1], their approach is to use the traditional log-sum-exp trick inside each iteration of the FFT and inverse FFT, which prevents using any off-the-shelf FFT algorithm and adds computational overhead (note that FFT algorithms are extremely tough to implement efficiently). In contrast, our trick exploits the linearity of the Fourier transform to show it is possible to simply rescale the in- and outgoing log probabilities by their maximum to prevent numerical instabilities. Just like the linearity of the summation does for the log-sum-exp trick. This means we go to and from linear space only once instead of at every iteration in the FFT. Our approach has the advantage of being implementation-agnostic, allowing for off-the-shelf and highly optimised FFT implementations to be used.

3. The proposed tensorised representations and operations of PLIA$_t$ are crucial for the measured increase in performance. Additionally, they also lead to out-of-the-box differentiability. Both of these aspects are not discussed nor examined by [1]. Moreover, [1] only reports an improvement of 6x-30x in computation time for two specific population models, while we observed improvements in the order of 100 000x in our benchmarks compared to the state of the art.

4. The work of [1] only proposes a solution for computing numerically stable convolutions (discussed in points 2 and 3 above). We do provide further operations that can be performed efficiently, e.g. multiplication by scalars and probabilistic branching. This is not discussed at all in [1].

We will add a reference to [1] to the paper and also include the above discussion.

Third, to answer your question, our implementation does allow for -infinity in case the PMF is zero. We utilise the numpy implementation np.inf that adheres to the IEEE 754 [3] industry standard for representing infinity as a floating-point number. Moreover, this implementation is compatible with all operations of our deep learning backend, TensorFlow.

We hope to have clearly differentiated PLIA from the cited work and to have addressed all your concerns. If so, we hope you are open to the idea of raising your score. If not, we will gladly try to discuss any remaining concerns.

[1] Parker, M. R., Cowen, L. L., Cao, J., & Elliott, L. T. (2023). Computational efficiency and precision for replicated-count and batch-marked hidden population models. Journal of Agricultural, Biological and Environmental Statistics, 28(1), 43-58.

[2] Dail, D., Madsen, L. (2011). Models for estimating abundance from repeated counts of an open metapopulation. Biometrics 67(2):577–587.

[3] IEEE Standard for Binary Floating-Point Arithmetic. 1985.

Thanks for the reviewer for the detailed response. I acknowledge I have read the rebuttal and I have raised my score accordingly.