KLay: Accelerating Arithmetic Circuits for Neurosymbolic AI
This work addresses a performance bottleneck in neurosymbolic AI, enabling scaling to larger real-world applications.
The paper tackles the challenge of efficiently running arithmetic circuits on modern AI accelerators by introducing knowledge layers (KLay), a new data structure that enables GPU parallelization, achieving speedups of multiple orders of magnitude over the state of the art.
A popular approach to neurosymbolic AI involves mapping logic formulas to arithmetic circuits (computation graphs consisting of sums and products) and passing the outputs of a neural network through these circuits. This approach enforces symbolic constraints onto a neural network in a principled and end-to-end differentiable way. Unfortunately, arithmetic circuits are challenging to run on modern AI accelerators as they exhibit a high degree of irregular sparsity. To address this limitation, we introduce knowledge layers (KLay), a new data structure to represent arithmetic circuits that can be efficiently parallelized on GPUs. Moreover, we contribute two algorithms used in the translation of traditional circuit representations to KLay and a further algorithm that exploits parallelization opportunities during circuit evaluations. We empirically show that KLay achieves speedups of multiple orders of magnitude over the state of the art, thereby paving the way towards scaling neurosymbolic AI to larger real-world applications.