A Fast Convoluted Story: Scaling Probabilistic Inference for Integer Arithmetic

TL;DR

This paper introduces a tensor-based approach using the fast Fourier transform to enable scalable, differentiable probabilistic inference for linear integer arithmetic, significantly improving efficiency for neurosymbolic AI applications.

Contribution

The authors develop a novel tensor formulation of probabilistic linear integer arithmetic that leverages FFT for fast, differentiable inference, facilitating learning in neurosymbolic AI.

Findings

Achieves several orders of magnitude faster inference and learning times.

Provides a differentiable data structure for integer-valued random variables.

Enables scalable probabilistic reasoning in integer arithmetic using modern deep learning tools.

Abstract

As illustrated by the success of integer linear programming, linear integer arithmetic 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 arithmetic 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 linear integer arithmetic 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.