KLay: Accelerating Arithmetic Circuits for Neurosymbolic AI

Abstract

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.

Figures (9)

• Figure 1: Example of (a) a Boolean circuit, (b) its d-DNNF version, and (c) the corresponding arithmetic circuit. The arithmetic circuit replaces Boolean variables with their probabilities. Consequently, the arithmetic circuit computes the probability of the underlying Boolean circuit being satisfied. In a neurosymbolic setting, the input probabilities are predicted by a neural network.
• Figure 2: A layered version of the d-DNNF in Figure \ref{['fig:circuit_example:ddnnf']}. Dashed nodes indicate nodes that are not present in the original d-DNNF and which are introduced during the layerization of the circuit.
• Figure 3: (\ref{['fig:lbcircuitroot2']}) contains a layered d-DNNF circuit over the variables $a$, $c$, $d$, and $e$. The dashed circles indicate again nodes introduced during the layerization process. (\ref{['fig:lbcircuitmerged']}) illustrates a multi-rooted circuit where the upper left node computes the circuit from Figure \ref{['fig:lbcircuit']} and where the upper right node computes the circuit from Figure (\ref{['fig:lbcircuitroot2']}). Note how nodes at the same height in the multi-rooted circuit are of the same type and how the sub-circuit colored in orange is used to compute both root nodes.
• Figure 4: The evaluation of the first layer in Figure \ref{['fig:circuit_example:ac']}, as an indexing and aggregation operation. The symbols ${\bm{\mathsfit{N}}}_0$ and ${\bm{\mathsfit{N}}}_1$ denote the node values at layer $L_0$ and $L_1$ respectively. First, we index in ${\bm{\mathsfit{N}}}_0$ using the edge indices ${\bm{\mathsfit{S}}}_1$, effectively creating a vector with the values of all edges in between ${\bm{\mathsfit{N}}}_0$ and ${\bm{\mathsfit{N}}}_1$. Next, the aggregation indices ${\bm{\mathsfit{R}}}_1$ determine what edges are reduced together. As a result we obtain the node values ${\bm{\mathsfit{N}}}_1 = \text{scatter}({\bm{\mathsfit{N}}}_0[{\bm{\mathsfit{S}}}_1], {\bm{\mathsfit{R}}}_1, \text{reduce}=\text{'product'})$.
• Figure 5: Tensorized KLay representation of the arithmetic circuit in Figure \ref{['fig:circuit_example:ac']}. We depict the segments of the edge vectors ${\bm{\mathsfit{E}}}_l$ with alternating colors.

Theorems & Definitions (3)

• Theorem 1
• Theorem 1
• proof