A Fast Convoluted Story: Scaling Probabilistic Inference for Integer Arithmetics

Authors: lennert  de
smet  ku
leuven  belgium
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 18:20:29

Abstract: As illustrated by the success of integer linear programming, linear integer arithmetics is a powerful tool for modelling combinatorial problems. Furthermore, the probabilistic extension of linear programming has been used to formulate problems in neurosymbolic AI. However, two key problems persist that prevent the adoption of neurosymbolic techniques beyond toy problems. First, probabilistic inference is inherently hard, #P-hard to be precise. Second, the discrete nature of integers renders the construction of meaningful gradients challenging, which is problematic for learning. In order to mitigate these issues, we formulate linear arithmetics over integer-valued random variables as tensor manipulations that can be implemented in a straightforward fashion using modern deep learning libraries. At the core of our formulation lies the observation that the addition of two integer-valued random variables can be performed by adapting the fast Fourier transform to probabilities in the log-domain. By relying on tensor operations we obtain a differentiable data structure, which unlocks, virtually for free, gradient-based learning. In our experimental validation we show that tensorising probabilistic integer linear arithmetics and leveraging the fast Fourier transform allows us to push the state of the art by several orders of magnitude in terms of inference and learning times.

Main paper: 
 introduction: 
 Integer linear programming (ILP) [15, 32] uses linear arithmetic over integer variables to model intricate combinatorial problems and has successfully been applied to domains such as scheduling [31] , telecommunications [34] and energy grid optimisation [26] . If one replaces deterministic integers with integer-valued random variables, the resulting probabilistic arithmetic expressions can be used to model probabilistic combinatorial problems. In particular, many problems studied in the field of neurosymbolic AI can be described using probabilistic linear integer arithmetic.Unfortunately, exact probabilistic inference for integer arithmetic is a #P-hard problem in general. Consequently, even state-of-the-art probabilistic programming languages with dedicated inference algorithms for discrete random variables, such as ProbLog [8] and Dice [13] , fail to scale. The reason being that they resort to exact enumeration algorithms, as exemplified in Figure 1 . Note that while approximate inference algorithms such as Monte Carlo methods and variational inference can be applied to probabilistic combinatorial problems, they come with their own set of limitations, as discussed by Cao et al. [5] . For instance, conditional inference with low-probability evidence.In order to mitigate the computational hardness of probabilistic inference over integer-valued random variables, we make the simple yet powerful observation that the probability mass function (PMF) of the sum of two random variables is equal to the convolution of the PMFs of the summands. The key advantage of this perspective is that the exact convolution for finite domains can be implemented efficiently using the fast Fourier transform (FFT) in O(N log N ), which avoids the traditionally quadratic behaviour of computing the PMF of a sum of two random variables (Figure 1 ). Moreover, 38th Conference on Neural Information Processing Systems (NeurIPS 2024). pX (5) pX (6) Figure 1 : On the left and in the middle we have two histograms representing the probability distributions of the random variables X 1 and X 2 , respectively. The grid on the right represents the joint probability of the two distributions, with more intense colors indicating events with higher probability. The distribution of the random variable X = X 1 + X 2 can be obtained by summing up the diagonals of the grid as indicated in the figure. While this method of obtaining the distribution for X is valid and used by state-of-the-art neurosymbolic techniques [14, 21] , the explicit construction of the joint is unnecessary and hampers inference and learning times (cf. Section 5) .efficient implementations of the FFT are readily available in modern deep learning libraries such as TensorFlow [1] and PyTorch [28] , making our approach to probabilistic inference end-to-end differentiable by construction. In turn, differentiability allows us to apply our approach to prototypical problems in neurosymbolic AI.Our main contributions are the following. 1) We propose a tensor representation of the distributions of bounded integer-valued random variables that allows for the computation of the distribution of a sum of two such variables in O(N log N ) instead of O(N 2 ) by exploiting the fast Fourier transform (Section 2). 2) We formulate common operations in linear integer arithmetic, such as multiplication by constants and the modulo operation, as tensor manipulations (Section 3). These tensorised operations give rise to PLIA t , a scalable and differentiable framework for Probabilistic Linear Integer Arithmetic. 1 PLIA t supports two exact probabilistic inference primitives; taking expected values and performing probabilistic branching (Section 4). 3) We provide experimental evidence that PLIA t outperforms the state of the art in exact probabilistic inference for integer arithmetic [5] in terms of inference time by multiple orders of magnitude (Section 5.1). Moreover, we deploy PLIA t in the context of challenging neurosymbolic combinatorial problems, where it is again orders of magnitude more efficient when compared to state-of-the-art exact and approximate methods (Section 5.2). 
efficient addition of integer-valued random variables: 
 In what follows, we denote random variables by uppercase letters, while a specific realisation of a random variable is written in lowercase. That is, the value x is an element of the sample space Ω(X) of X. We will also refer to Ω(X) as the domain of the random variable X. Furthermore, Ω(X) is assumed to be integer-valued, i.e. it is a finite subset of the integers Z with lower and upper bounds L(X) and U (X), respectively. In particular, we have that the cardinality |Ω(X)| = U (X)-L(X)+1.We will call these integer-valued random variables probabilistic integers from here on.The distribution of a probabilistic integer X is represented using its probability mass function (PMF) p X : Ω(X) → [0, 1] with the conventional restrictions ∀x ∈ Ω(X) : p X (x) ≥ 0 and x∈Ω(X) p X (x) = 1.(1)At the core of PLIA t and linear arithmetic in general is the addition of two probabilistic integers X 1 and X 2 . Let us assume for now that X 1 and X 2 satisfy L(X 1 ) = L(X 2 ) = 0 and have upper bounds U (X 1 ) = N 1 and U (X 2 ) = N 2 . Just as in Figure 1 , we would now like to find the PMF of the random variable X such that X = X 1 + X 2 . However, contrary to Figure 1 , we wish to avoid the explicit quadratic construction of all possible outcomes. To this end, we exploit that the PMF of the sum of two random variables is equal to the convolution of their respective PMFs [12] pEQUATIONNext, we apply the Fourier transform F to both sides of the equation and use the convolution theorem (CT) [23] that states that the Fourier transform of two convoluted functions is equal to the product of their transformsEQUATIONwhere we also introduce the hat notation F(p X ) = p X for a Fourier transformed PMF p X . As X = X 1 + X 2 , we know that Ω(X) = {0, . . . , N 1 + N 2 }. Consequently, the PMF p X is nonzero for just these M = N 1 + N 2 + 1 domain elements and can be represented using a vector of probabilitiesEQUATIONAll vectors of probabilities will be written using a boldface π and their elements accessed using square brackets. Looking at Equation 3, we would now like to express the Fourier transformed probability vector π X as the point-wise product of the transformed vectorsEQUATIONwhere F M ∈ C M ×M is the M -point discrete Fourier transform (DFT) matrix and the symbol ⊙ denotes the Hadamard product. In order for Equation 5 to hold we need to have thatEQUATIONAt first sight these equalities seem to cause a problem; each probabilistic integer X i has a domain of size N i + 1 while its PMF should be represented with a probability vector π X1 ∈ R M for the multiplication with F M to make sense. Fortunately, this problem is easily resolved by observing that we can extend the domain Ω(X i ) of X i by simply assigning a probability of zero to newly added elements. In practice, we simply pad the probability vectors π X1 and π X2 with N 2 and N 1 zeros at the end to obtain vectors of dimension M . With this issue resolved, we can finally obtain the probability vector π X that represents the PMF p X by using Equation 5viaEQUATIONThe attentive reader might have noticed that, even though we avoid the explicit construction of the joint probability distribution p X1X2 (x 1 , x 2 ), we have not gained much. The matrix-vector products in Equation 8 still take O(M 2 ) to compute. Fortunately, matrix-vector products where the matrix is the DFT matrix or its inverse can be computed in time O(M log M ) by using the fast Fourier transform (FFT), with M being the size of the vector. As a result, we can express Equation 8asEQUATIONComputing the values of the vector π X can now be done in time O(M log M ). First we apply the FFT on the vectors π X1 and π X2 . Then we multiply the transformed vectors pointwise and apply the inverse FFT on the result of this Hadamard product. We note that Equation 9 is a well known result from the signal processing literature, where convolutions are always computed in this fashion [23] .However, applying Equation 9 naively to the problem of probabilistic inference quickly results in numerical stability issues. The problem is that multiplying together small probabilities eventually results in numerical underflow. A well-known and widely used remedy to this problem is the log-sumexp trick, which allows one to avoid underflow by performing computations in the log-domain instead of the linear domain. Inspired by the log-einsum-exp trick [29] , we introduce the fast log-conv-exp trick, which allows us to perform the FFT on probabilities in the log-domain.We first characterise a probability distribution p X not by the vector of probabilities π X but by the vector of log-probabilities λ X = log π X . In terms of log-probabilities, Equation 9 can be written asEQUATIONNow define µ i := max x∈Ω(Xi) λ Xi [x] as the maximum value present in the vector λ Xi , which lets us write Equation 10asexp λ X = IFFT FFT(exp(λ X1 -µ X1 + µ X1 )) ⊙ FFT(exp(λ X2 -µ X2 + µ X2 )) (11) = IFFT FFT(exp(λ X1 -µ X1 )) ⊙ FFT(exp(λ X2 -µ X2 )) exp(µ X1 ) exp(µ X2 ).Crucially, we were able to pull out the scalars exp(µ X1 ) and exp(µ X2 ) due to the linearity of the FFT transform and its inverse. Taking the logarithm of both sides results in the fast log-conv-exp trickEQUATIONwhich expresses the log-probabilities λ X in function of λ X1 and λ X2 . It can still be computed in time O(M log M ) and avoids, at the same time, numerical stability issues by exponentiating λ Xi -µ i instead of λ Xi directly.While using the fast log-conv-exp trick is necessary to scale computations in a numerically stable manner, describing operations on probability mass functions in the log-domain is rather cumbersome. Hence, we will, for the sake of clarity, describe PLIA t using probability vectors π (cf.. Sections 2.3, 3 and 4). We refer the reader to our implementation for the log-domain versions.Another solution to the numerical instability of applying the FFT on probabilities was given in an application of the FFT to open-population [6] N -mixture models [27] . However, it has a major drawback when compared to the fast log-conv-exp trick: it relies on repeated applications of the traditional log-sum-exp trick within each of the N log N iterations of the FFT. This drawback prevents the use of optimised, off-the-shelf FFT algorithms and adds computational overhead. In contrast, we utilise the linearity of the FFT transform to provide an implementation-agnostic solution that works with tensorised representations.In the previous sections we assumed that the first non-zero probability event of all probabilistic integers X was the event X = 0, i.e. L(X) = 0. However, we can remove this assumption by characterizing a PMF p X not only by a vector of probabilities π X , but also by an integer β X = L(X) encoding the non-zero lower bound of the domain Ω(X). Indeed, we can write any PMF p X over an integer domain as the translation of a PMF p 0 X whose first non-zero probability event isEQUATIONSince the PMF p 0 X can be represented by a probability vector π X as in the previous sections, it follows that the PMF of any probabilistic integer can be characterised by such a vector and an integer β X that shifts the domain. The upper bound of the domain is not important for the characterisation of p X as it can be obtained via U (X) = β X + dim(π X ) -1, where dim(π X ) is the dimension of the probability vector π X .The inclusion of translations in our representation of PMFs is compatible with using convolutions to compute the PMF of the sum of two probabilistic integers because of the translational invariance of the convolutionEQUATIONwhere τ k denotes the translations by a scalar k. In general, k can be real-valued, but for PLIA t we limit k to integers. Using Equation 13, we can write the PMF of the sum of two probabilistic integers X 1 and X 2 with non-zero lower bounds β X1 and β X2 asEQUATIONThis final equality shows that we can characterise the PMF p X for X = X 1 + X 2 by the following lower bound and probability vectorEQUATION) 0 1 2 3 4 x 0 1 2 3 4x-1 0 1 2 x -2 -1 0 1 x 0 1 2 3 4 5 6x 0 1 2 3 4 5 6xFigure 2 : (Left) Adding a constant to a probabilistic integer simply means that we have to shift the corresponding histogram, shown here for X ′ = X + 1. (Middle) For the negation X ′ = -X, the bins of the histogram reverse their order and the negation of the upper bound becomes the new lower bound. (Right) For multiplication, here show the case X ′ = 3X by inserting zero probability bins.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 as tensors whose entries are the log-probabilities of the individual events in the sample space of the random variable that is being represented. We define this formally as follows. Definition 2.1 (Probabilistic linear arithmetic expression). Let {X 1 , . . . , X N } be a set of N independent probabilistic integers with bounded domains. A probabilistic linear integer arithmetic expression X is itself a bounded probabilistic integer of the formEQUATIONwhere each f i denotes an operation performed on the specified random variables that can be either one of the operations specified in Section 3 as well as compositions thereof.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, PLIA t will also be provided with probabilistic inference primitives that allow it to compute certain expected values efficiently as well as to perform probabilistic branching.Assuming all f i (X i ) are computable in polytime, we can also compute PLIA t expressions (Equation 18) in polytime in N . However, when computing PLIA t expressions recursively, the domain size of the random variables might grow super-polynomially -manifesting the #P-hard character of probabilistic inference.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We explicitly mention the problem of probabilistic inference is #P-hard in general and that PLIA t , because of its exact inference, can still suffer from this limitation. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
arithmetic on integer-valued random variables: 
 The previous section introduced how PLIA t deals with the addition of two probabilistic integers. We discuss now five further operations: 1) addition of a constant, 2) negation, 3) multiplications by a constant, 4) integer division by a constant and 5) the modulo.Constant Addition. The addition of a probabilistic integer X and constant scalar integer k forms a new probabilistic integer X ′ = X + k. Adding a scalar integer is equivalent to a translation of the distribution of X (Figure 2 , left). In other words, the lower bound and probability vector of X ′ are given by We show the histogram transformation for the integer division X ′ = X /3. The probability mass of three subsequent bins is accumulated in the bins for which x mod 3 = 0 and x /3 ∈ Ω(X). (Right) For the modulo X ′ = X mod 3, the only non-zero elements of Ω(X ′ ) are elements of the set {0, 1, 2}. The bins corresponding to these values then accumulate the probability masses of all other bins as indicated by the colors.EQUATIONNegation. The negation X ′ = -X of a probabilistic integer X is equally straightforward to characterise. Taking a negation mirrors the probability distribution of X around zero (Figure 2 , middle). In terms of lower bound and probability vector, we getβ X ′ = -(β X + dim(π X ) -1)andEQUATIONrespectively. That is, the lower bound of X ′ is equal to the negated upper bound of X while the probability vector is flipped, taking into account that probability vectors have to start at 0.Constant Multiplication. For the multiplication X ′ = X • k of a probabilistic integer X with a scalar integer k, we assume, without loss of generality, that k ≥ 0. Multiplication by a scalar is then characterised asEQUATIONIntuitively, only multiples of k get a non-zero probability equal to the probability of that multiple in π X . The lower bound of X ′ is also immediately given by multiplying the lower bound of X by k. In other words, we obtain π X ′ by inserting k -1 zeros between every two subsequent entries of π X (Figure 2 , right). The case k < 0 is obtained by first negating X.Integer Division and Modulo. For the case of integer division X ′ = X /k and the modulo operation X ′ = X mod k, the probability distribution of X ′ can be obtained by adequately accumulating probability mass from events in Ω(X). We demonstrate these operations by example in Figure 3 and refer the reader to Appendix A for the formal description.4 Probabilistic Inference PrimitivesQuestion: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] Justification: In our more theoretical sections (Sections 2, 3 and 4), our discussion always clearly states the assumptions made, and when and how these can be lifted. Our proposed tensorised representations and the use of the FFT are all theoretically sound. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
computing expected values: 
 PLIA t supports the exact computation of two different forms of expected values. The first is a straightforward expectation of a probabilistic integer X, given by weighing each element of Ω(X) with its probabilityEQUATION)The second is computing the expectation of a linear comparative expression of probabilistic integers. Such a comparison can be an equality, inequality or negated equality. We only consider the equality and strictly larger inequality as the other cases follow from them. The strict inequality can itself always be reduced to an inequality with respect to zero and hence comprises a sum over all domain elements below zeroEQUATIONSimilarly, the computation of the expected value of an equality comparison can always be reduced to a comparison to zero. Hence, the expected value is computable by simple indexingEQUATIONX C X ⊤ X ⊥ X ′ True False g ⊤ g ⊥Figure 4 : Control flow diagram for probabilistic branching. The branching condition is probabilistically true and induces a binary random variable C. In each of the two branches we then have two conditionally independent random variables X ⊤ and X ⊥ to which the functions g ⊤ and g ⊥ are applied in their respective branches. The probabilities of X ′ are then given by the weighted sums of the probabilities of g ⊤ (X ⊤ ) and g ⊥ (X ⊥ ) (Equation 29).As computing expected values is no harder than computing the sum of the elements of a vector, we can conclude that we can compute these expected values in O(dim(π X )). By using prefix sums [18] and harnessing parallel compute on GPUs, the complexity can further be reduced to O(log dim(π X ))Consider an if-then-else statement with condition c(x) = (f (x) ▷◁ 0), where f is a composition of the functions introduced in Section 3 and ▷◁ ∈{<, ≤, =, >, ≥, ̸ =}. Furthermore, x belongs to the domain Ω(X) of a probabilistic integer X. In the case of c(x) being true, a function g ⊤ is executed. If c(x) is false, another function g ⊥ is executed instead. We assume that both g ⊤ and g ⊥ are again linear arithmetic functions expressible in PLIA t (Section 2.4).The if-then-else statement defines a new probabilistic integer X ′ by combining both of its branches (Figure 4 ). These branches depend on X which itself influences a binary random variable C that represents the probabilistic condition of the if-then-else statement. To be precise, the PMF p X ′ is given by the decompositionEQUATIONwhere p X ′ |C is the conditional PMF of X ′ given C. The true branch gives rise to a probabilistic integer X ⊤ with probability distributionEQUATIONIf C = ⊤, then X ′ is given by an application of g ⊤ on the instances x ∈ Ω(X) that satisfy c(x). Consequently, by applying g ⊤ , we find thatEQUATIONWith the right-hand side of Equation 26, we now know how to obtain the probability vector for X ⊤ . Using Equation 27and the operations from Section 3, we can then compute the probability vector π g ⊤ (X ⊤ ) , as well as the lower bound β g ⊤ (X ⊤ ) . Also note that p C (⊤) is nothing but an expected value as described by Equation 23or Equation 24. Analogously, we obtain for the false branch thatEQUATIONBy plugging the expressions for pEQUATIONwhich are all quantities computable using either probabilistic linear arithmetic operations or expected values thereof. Figure 5 : We plot the runtime of Dice [5] and PLIA t against the domain size of the problems.From left to right, we have E[X 1 + X 2 ], E [1 X1+X2<0 ], E [1 X1+X2=0] and probabilistic Luhn. All four plots share the same y-axis on the very left, which is in log-scale. Following the experimental protocol of Cao et al. [5] , we report average runtimes for every integer on the x-axis, both bitwidths and identifier lengths. No significant deviations from the mean were found.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: Sections 2, 3 and 4 clearly describe the theory behind our implementation and all details necessary for reproduction. Moreover, our full implementation, following these descriptions, is provided as supplementary material. Finally, the implementation together with the details given in the appendix, such as architectures, encodings and hyperparameters, are sufficient to reproduce all experiments. Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.Justification: All numerical results are provided with the necessary statistics. Figure 5 reports mean runtime with standard errors aggregated over 10 runs, following the experimental protocol of Cao et al. [5] . Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [No] Justification: We do not explicitly mention societal impact as we do not foresee any immediate positive or negative impact. However, it is our hope that this progression on the state of the art allows solving more challenging combinatorial problems under uncertainty with a positive effect on society in the future. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiments: 
 We first compare PLIA t to the state of the art in probabilistic integer arithmetic [5] in terms of inference speed (Section 5.1). These experiments were performed using an Intel Xeon Gold 6230R CPU @ 2.10GHz, 256GB RAM for CPU experiments and an Nvidia TITAN RTX (24GB) for GPU experiments. In Section 5.2, we then illustrate how PLIA t fares against the state of the art in neurosymbolic AI [14, 33] . These experiments were performed using an Nvidia RTX 3080 Ti (12GB). We implemented PLIA t in TensorFlow [1] using the Einops library [30] . This implementation is open-source and available at https://github.com/ML-KULeuven/probabilistic-arithmeticThe work of Cao et al. [5] exploits the structural properties and symmetries of integer arithmetic by proposing general encoding strategies for an arithmetic expression of probabilistic integers as logical circuits. That is, binary decision diagrams [4] obtained via knowledge compilation [7] . This strategy allows them to avoid redundant calculations and repetition, leading to improved scalability over more naive encodings [8, 13] .We compare PLIA t and Cao et al.'s inference algorithm on four of their benchmark problems. In the first three benchmarks, expected values of the sum of two random variables need to be computed. Concretely, the expectations E [X 1 + X 2 ] (cf. Equation 22), E [1 X1+X2<0 ] (cf. Equation 23) and E [1 X1+X2=0 ] (cf. Equation 24).As a fourth benchmark we use a probabilistic version of the Luhn checksum algorithm [20] , which necessitates summation of two probabilistic integers, negation, addition of a constant, multiplication by a constant, the modulo operations, as well as probabilistic branching. We provide further details on the Luhn algorithm in general and the encoding of its probabilistic variant in PLIA t in Appendix C.As the probabilistic Luhn algorithm takes as input an identifier consisting of a sequence of probabilistic integers with domain {0, . . . , 9}, we can increase the problem size by increasing the length of this sequence. For the other three benchmarks we vary the problem size by varying the domain size of the probabilistic integers in terms of their bitwidth. That is, a bitwidth of i ∈ N indicates that we consider probabilistic integers ranging from 0 up until 2 i -1, increasing the problem size exponentially in terms of i. In our experimental evaluation (Figure 5 ), we measure the time it took for PLIA t and Cao et al.'s method to terminate for varying problem sizes. The measured time includes all computational overhead inherent to each method, such as the construction of computational graphs and compilation time. Each method is also profiled in terms of memory, which we discuss further in Appendix B. Note that Cao et al.'s method is denoted by "Dice", the probabilistic programming language in which it was implemented.For the first three benchmarks, we observe that PLIA t easily scales to probabilistic integers with a domain size of 2 24 on both the CPU and GPU as the highest runtime reached is less than 100 seconds on the CPU and approximately 1 second on the GPU. The similarity of the curves is due to the fact that the run time is dominated by computing the probability vector π X1+X2 and not so much by computing the actual expected value.In contrast, Dice, which only runs on CPU, already reaches a runtime of approximately 1000 seconds for integers with domain size 2 15 , where PLIA t only takes around 10 -1 and 10 -2 seconds on the CPU and GPU, respectively. This is a rather considerable improvement in the order of 10 5 . Dice can outperform PLIA t on the GPU (Figure 5 , bitwidth smaller strictly below 5) due to the computational overhead of running on the GPU. However, 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 smaller strictly below 5).On the probabilistic Luhn benchmark (Figure 5 , extreme right) we observe that both methods exhibit similar linear scaling behaviors. However, the use of tensors as representations instead of logical circuits does result in a significant improvement in terms of run time in the order of 10 2 for the longest sequences of length 350.For the comparison of PLIA t to neurosymbolic systems, we use two standard benchmarks from the literature: MNIST addition [21] and visual sudoku [2] . The common idea for both is to train neural networks to classify MNIST digits while only having access to distant supervision (Figure 6 ).As an MNIST classifier outputs a distribution over the integers {0, . . . , 9}, we can readily encode these predictions as probabilistic integers and enforce the constraints given by the two problems using the arithmetic operations developed in Section 2 and Section 3. We refer the reader to Appendix D (MNIST addition) and Appendix E (visual sudoku) for details on the encodings.In the experimental evaluation we compare PLIA t , which uses exact probabilistic inference, to one other exact method, DeepProbLog (DPL) [21, 22] , and two approximate methods, Scallop [14] and A-NeSI [33] . Similar to Dice (Section 5.1), DPL also relies on expensive knowledge compilation in order to obtain the distribution of the sum of two probabilistic integers. Essentially, it performs explicit enumeration as illustrated in Figure 1 . Scallop approximates this explicit enumeration by only considering the top-k most likely solutions of a problem. The approximation of A-NeSI is based on optimising a neural surrogate model for the combinatorial problem. While this model sidesteps the computational complexity encountered by DPL and Scallop to a certain degree, training the surrogate model becomes prohibitively expensive for larger problems.We compare the different methods along two dimensions, being prediction accuracy and training time. Specific details on training, e.g. neural architectures and hyperparameters, can be found in Appendix F. We report the statistics in Table 1 , where we use numbers consisting of N ∈ {2, 4, 15, 50} digits for the MNIST addition benchmark and grid size G ∈ {4, 9} for the visual sudoku benchmark. We see that PLIA t significantly outperforms the other methods, both exact as well as approximate, in terms of training times. The difference is particularly apparent on the MNIST addition benchmark, where no other method was able to scale up to N = 50 without timing out. The reported accuracies also show that this advantage of training time for PLIA t with respect to the other methods does not come at the cost of predictive performance of the learned neural networks.Table 1 : In the upper part part we report median test accuracies over 10 runs for the MNIST addition and the visual sudoku benchmarks for varying problem sizes and different neurosymbolic frameworks. Sub-and superscript indicate the 25 and 75 percent quantiles, respectively. In the lower part we report the training times in minutes, again using medians with 25 and 75 percent quantiles. We set the time-out to 24 hours (1440 minutes).Visual SudokuN = 2 N = 4 N = 15 N = 50 G = 4 G = 9Accuracy Accuracy DPL 94.20  
conclusion and future work: 
 We introduced PLIA t , an efficient, differentiable and hyperparameter-free framework for linear arithmetic over integer-valued random variables. The efficiency of PLIA t is due to representing probabilistic integers as tensors and exploiting the FFT for computing the sum of two probabilistic integers. Compared to state-of-the-art methods for inference [5] and learning [33] the concepts underlying PLIA t are surprisingly simply: a tensorised calculus and the fast Fourier transform.This simple yet elegant approach has led to improvements in inference and learning times in the order of 10 5 and 10 2 , respectively. We attest this advantage to PLIA t 's formulation in terms of fundamental concepts shared across modern machine learning. As such, any algorithmic or hardware improvements will immediately improve PLIA t 's performance. For instance, incorporating recent advancements for computing FFTs on GPUs [11] would directly benefit PLIA t 's efficiency.In future work we envisage to extend the probabilistic inference primitives of PLIA t to a fullfledged neuroprobabilistic programming language, resulting in a user-friendly interface to an efficient neurosymbolic reasoning engine. In this regard, we deem ideas from the satisfiability modulo theory literature [3] as important. Specifically, probabilistic [17, 25] and neural [9] extensions thereof, as well as novel formula representations [10, 24] . 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This paper poses no immediate risk for misuse, as it claims foundational advances on exact inference. However, this exact probabilistic inference could be used as a safeguard itself in the future. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: Our entire paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Integer Division. Performing a scalar integer division of the random variable X by k yields the random variable X ′ = X k with lower bound β X ′ = β X k and probability vectorEQUATIONwhere we assume that L(X) mod k = 0 and U (X) + 1 mod k = dim(π X ) mod k = 0. This is a weak assumption as it can easily be obtained via zero-padding. The intuition is then that the probability vector π X ′ of an integer division of X by k accumulates probability mass of k sequential entries of π X (Figure 3 , left).Modulo Operation. For the modulo of a random variable by a constant, i.e. X ′ = X mod k, we have the lower bound β X ′ = 0 and the probability vector isEQUATIONwhere we again assumed that L(X) mod k = 0 and U (X) + 1 mod k = dim(π X ) mod k = 0.In other words, the probability vector π X ′ of the modulo of a random variable X accumulates the probability mass of π X by multiples of k (Figure 3 , right).It is well-known that the FFT has, apart from the N log N runtime complexity, also a 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 [5] and PLIA t with respect to memory allocation (Figure 7 ). We found that Dice was only more efficient for smaller problem sizes due to the memory overhead of TensorFlow on the GPU. However, from problem sizes of bitwidth higher than 13, Dice overtook PLIA t in memory usage and we confirmed the better scaling behaviour of PLIA t . Hence, PLIA t not only outshines the current state of the art in terms of runtime but also in terms of memory usage. Memory usage of PLIA t on the CPU could not be properly profiled due to Python's limited functionality, but there is no reason to suspect it would be significanlty different from the GPU memory usage. The traditional Luhn algorithm [20] computes a check digit from a given sequence of integers. If the check digit is correct, then it means that there are no single-entry mistakes in the sequence. Such errors often arise, for example, when transcribing credit card information. Hence, Luhn is often used in this context to quickly identify mistyped numbers.We provide an implementation of the probabilistic Luhn algorithm using PLIA t in Algorithm 1.During its execution, the algorithm sequentially constructs the check digit by iterating over the given sequence. At the start of every step, a deterministic if-then-else statement checks whether the current sequence has the same binary polarity as the length of the total sequence to be checked (Line 4). The first argument of the ifthenselse function is the probabilistic integer on which to perform the test. The second argument gives the test. For Luhn we have digit<5. The two lambda functions then tell us which function to perform on the random variable in each of the two branches. The function outputs a probabilistic integer which we add to the check probabilistic integer.Let n 1 and n 2 be two probabilistic integers. The MNIST addition problem is then usually encoded in one of two different ways. The first approach simply directly constructs the joint distribution over all possible integers. The second approach avoids this explicit constructions by computing conditional distributions for intermediate computations. In other words, the first approach constructs a lookup table for the sum of two integers while the second encoding gives us a probabilistic version of the sum of two integers via a carry. Naturally, the second approach is able to scale to larger numbers and we also use it in our experiment for all methods. We give the implementation for both encodings in Algorithmns 2 and Algorithm 3.For the problem of solving a visual sudoku [2] , the input is a G × G sudoku puzzle filled with handwritten digits (Figure 6 , right). The task itself is to predict whether such a given sudoku puzzle is valid, meaning the values of the digits satisfy the various sudoku constraints. For G = 4, the constraints to satisfy are that every row and every column must have different entries. In case G = 9, all nine inner 3 × 3 grids of the puzzle must also have different entries. Again, because the digits are handwritten, a neural network has to classify the digits into probabilistic integers from 0 to 9. Checking whether a sudoku puzzle is correct requires checking if the sudoku constraints are satisfied. While there are many encodings of these constraints, PLIA t needs an encoding in terms of linear arithmetic operations. The encoding we used is as follows. Every grid cell is associated with 10 binary probabilistic integers, leading to G • G • 10 such integers for the whole puzzle. Intuitively, the probability of the k th binary integer of a grid cell being equal to 1 should be interpreted as the probability that the grid cell contains the number k. We represent these binary integers as a 3-tensor B ijk , where i indexes a row of the grid, j indexes a column and k is the k th binary digit of grid cell (i, j). With this representation, expressing that a row i should only contain different digits is equivalent to obtaining a value of 1 after summing out the index j from B ijk . Concretely, that isEQUATIONIndeed, for a fixed value of i, i.e. for a single row i, there are 10 constraints to satisfy, where each one excludes double presences of one of the 10 possible digits. The constraint that each column only contains different digits is completely analogous to Equation 32, but instead summing over the row index i.In case G = 9, there are the additional constraints that the inner 9 blocks of the puzzle also have no repeating entries. The intuition is similar to the other two constraints; we sum over all grid cells of a block for each of the binary integers and should obtain 1. If we index the 9 inner blocks as a 3 × 3 grid with coordinates (l, m), then the equationsEQUATIONis true if and only if each of the inner blocks has different digits in all of its grid cells.Bringing everything together, PLIA t approximates the probability of a sudoku puzzle being correct by computing the product of the probabilities of all the above linear constraints. Note that this approach computes each of the constraints as independent of one another, allowing for the probability of each constraint to be computed in parallel. Because PLIA t is based on tensor manipulations, our implementation can and does exploit this parallelism. Additionally, note that PLIA t works with log-probabilities in practice, which avoids the problem of numerical underflow often arising when taking long products of probabilities.Choice of Encoding for Comparisons. The other methods we compare to do not necessarily need to use the same binary linear arithmetic encoding as they are not restricted to linear arithmetic. Instead, their provided implementations directly encode every grid cell as a categorical random variable with 10 outcomes, one for every possible digit. Expressing that a row, column or block of grid cells only contains different digits is then done by comparing all pairs of categorical variables in the row, column or block. Each of these pairs then has to be different for the constraint to hold. To again provide a fair comparison, we tested both encodings for Scallop and A-NeSI and took the best performing one. In both cases, the categorical encoding gave better runtime results.Since both the MNIST and visual sudoku task involve the classification of MNIST images, we use the same traditional LeNet [19] in both cases. Next we discuss the hyperparameters involved in training.The optimal parameters for PLIA t and Scallop were found by performing a grid search on a heldout validation set. The main hyperparameters important for all methods used for neurosymbolic learning are the learning rate and the number of epochs to train. For our grid search, we considered values {0.001, 0.005, 0.0001} for the learning rate and {5, 10, 15, 20, 50, 100, 200} for the number of epochs. Scallop and PLIA t both consistently performed best with a learning rate of 0.001 in all experiments. Both Scallop and PLIA t use Adam [16] as optimiser. The number of epochs varies per experiment. On the MNIST task, PLIA t needed 10, 15, 100 and 200 epochs for N = 2, N = 4, N = 15 and N = 50, respectively. For Scallop, N = 2 and N = 4 ran optimally when training for 5 and 10 epochs, respectively.For Scallop, with its top-k-based approximate inference, it is additionally important to choose a good value for k. If k is too small, then the approximation quality is poor. If k is too large, then it becomes prohibitively expensive to compute the approximation. Inspired by the implementation of Scallop, the possible values for k were {1, 3, 5, 6, 7}. For the MNIST addition task, in all cases that did not time out, k = 3 was optimal as it gave state-of-the-art performance and the quickest runtimes. The visual sudoku problem benefitted from a higher value of k, having an optimal value of k = 6. No value of k allowed Scallop to run on a 9 × 9 sudoku.In stark contrast to PLIA t A-NeSI has a multitude of hyperparameters to be set. We used the set of parameters provided in the original paper [33] and its implementation that were already used there to perform the experiments on MNIST addition and visual sudoku.NeurIPS Paper Checklist • The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.The subscript "t" in PLIAt stands for "tensorised" and is not pronounced. 


Decision: This paper describes a simple idea for probabilistic inference over linear, integer systems, based on the well-known property that the distribution of a linear combination of random variables can be performed using the Fourier representation. This allows the authors to frame their inference process using tensor operations, with a useful byproduct being that the entire computation becomes differentiable and usable in learning. Empirically, they show very high speedups and show the application of their ideas to neurosymbolic programming. 
The ideas in the paper are straightforward, with the result that some of the text feels obvious — the discussion of zero-padding to match vector lengths, for example, and even the authors’ log-sum-exp trick contribution. However, in their application they manage to improve over very recent work in significant ways.  One reviewer pointed out a closely related publication, which the authors’ response indicated they would incorporate.  Otherwise, reviewers were agreed on the potential impact of the work, with a few concerns that were generally well addressed by the author response.

Reviews: review: paper_summary: paper_summary: This paper presents an approach to compute probability distributions over sums of random variables. To achieve this, the authors replace as slow quadratic computation with a Fourier transformation and a Hadamard product that can be computed in log-linear time.
The authors introduce the log-sum-exp trick to improve the stability of the computations. In this process, the authors note that their approach is differentiable, which allows them to connect with work on neurosymbolic AI.
The authors present metrics that can be computed, and finally the authors present an empirical evaluation which demonstrates the accuracy and speed improvements of their approach.
strength_weakness: strength_weakness: My only concern with this paper, is that their work shares many similarities with the work in [1]. The authors do not cite this paper, which includes the convolutions and Fourier transform motivation which is known, as well as the use of the log-sum-exp trick. Providing a citation and a more thorough differentiation would be extremely valuable and maybe even required.

[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.
questions: questions: Do you do a special treatment for the case where the pmf is zero i.e., where the log is -inf? or do you simply allow it and let the exponential handle that? A small paragraph on that would be interesting.
limitations: limitations: no issues
paper_summary: paper_summary: The paper addresses key challenges in applying neurosymbolic AI techniques, specifically in integer arithmetic. Leveraging the power of tensor operations and the fast Fourier transform (FFT), the authors propose a novel approach to perform probabilistic inference on integer-valued random variables. Central to their method is the tensor representation of the distributions of bounded integer-valued random variables, enabling efficient computation of distribution sums using FFT, thus significantly reducing computational complexity from quadratic to logarithmic time complexity. The study validates the effectiveness of this approach through experiments showing substantial improvements in inference and learning times, pushing the state of the art by several orders of magnitude. This contribution not only enhances computational efficiency but also facilitates gradient-based learning, which was previously challenging due to the discrete nature of integers.
strength_weakness: strength_weakness: 1.While the paper briefly mentions the inherent #P-hardness of probabilistic inference, there is scant elaboration on how this impacts the scalability or applicability of the proposed method in practical, large-scale scenarios. A more detailed discussion of the practical limitations, potential computational bottlenecks, or scenarios where the method might not perform as expected would be invaluable for readers and practitioners looking to apply these techniques.
2. The clarity and depth provided in the theoretical formulation and experimental setup are excellent. However, the paper could be enhanced by including more detailed information on the implementation, such as the specifics of the software environment, libraries used, and parameter settings. This addition would aid in reproducibility and allow other researchers to more easily validate and build upon the work.
3. The paper could further detail the algorithmic complexity of the proposed method beyond the FFT application. While it highlights the reduction in time complexity, there's minimal discussion about space complexity or the computational overhead introduced by tensor manipulations and log-domain calculations.
questions: questions: 1.  Could the authors provide additional details about the computational environment, versions of libraries used, and precise parameter settings?
2.  How do the authors address potential underflow or overflow problems during computations, and what impact might these issues have on the method's accuracy and stability?
3. Are there plans or ongoing work to test the applicability of this approach in other fields that heavily rely on integer arithmetic, such as cryptographic algorithms or complex financial computations? Insights into such applications could significantly enhance the paper's impact.
limitations: limitations: One potential limitation is the paper's reliance on specific assumptions for the efficiency of FFT in the log-domain. While these assumptions are necessary for the mathematical framework, there is limited discussion on how sensitive the results are to deviations from these assumptions. Exploring the robustness of the proposed method under less ideal conditions or discussing potential modifications that could accommodate more general cases would strengthen the paper.
paper_summary: paper_summary: The paper presents a framework, PLIA_t, to solve the generally intractable problem of probabilistic inference using the fast Fourier transform (FFT) for integer-valued random variables. The paper shows how to use the log-sum-exp trick to solve the numerical stability issue in the FFT setting and defines arithmetic operations on integer-valued random variables. It also defines how to compute expected values and how to handle probabilistic branching. Experiments show that the gains can be realised in an implementation for the given examples.
strength_weakness: strength_weakness: - Even though the paper plans to introduce PLIA_t as a new framework, the framework itself remains rather opaque regarding its inputs, outputs, representations, algorithms, or whatever they consider part of PLIA_t. It almost reads as if a section is missing between Section 4 and 5. As such, I am also not so sure about the reproducibility of the experiments. (I have formulated a corresponding question. This is one of the main reasons for my overall score that I am willing to revise given a convincing rebuttal.)
questions: questions: - Can you please characterise / define PLIA_t formally as a framework in the rebuttal?
limitations: limitations: Yes.
paper_summary: paper_summary: The authors study the problem of probabilistic inference's scalability in the integer arithmetic setting. Their method is based on the representation of integer-valued random variables probability mass functions (PMFs) as vectors and the observation that PMFs of different operations applied on said random variables can be expressed as convolutions of the PMFs vectors by leveraging Fast Fourier Transforms.
strength_weakness: strength_weakness: - The experimental part is a bit limited and could be extended to show more clearly potential and limitations of the method
questions: questions: - More test to understand the limits of the approach in terms of scalability can be useful. How far can you push the size of MNIST addition by allowing, e.g. 1 hour computational time?
limitations: limitations: -

Comments: comment: Thank you for your response. I have decided to keep my score.
Thanks for the reviewer for the detailed response. I acknowledge I have read the rebuttal and I have raised my score accordingly.
Thanks for your comments