Research map / Architectures and efficiency / RNNs and normalization
KLay: Accelerating Arithmetic Circuits for Neurosymbolic AI
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces KLay, a layered tensor representation of arithmetic circuits that reduces circuit evaluation to indexing and scatter-reduce operations, enabling fast GPU/CPU evaluation for neurosymbolic learning.
Problem
Arithmetic circuits compiled from logic formulas are highly irregularly sparse and hard to evaluate efficiently on modern tensor accelerators, bottlenecking neurosymbolic AI.
Approach
Layerize d-DNNF/arithmetic circuits by node height (inserting unary nodes so children lie in the previous layer), deduplicate nodes via Merkle hashes, and evaluate each layer as an index-then-scatter-reduce (sum/product or logsumexp) using standard tensor library primitives, with no custom kernels.
Key results
- On Warcraft (12x12), KLay (cuda) forward+backward takes 1.59 ms vs 18930.67 ms for naive CUDA, roughly four orders of magnitude faster.
- On MNIST-addition with 1 digit, KLay (cuda) takes 1.26 s per epoch vs 241.63 s for DeepProbLog (cpu), about two orders of magnitude faster.
- On large synthetic circuits, KLay on GPU outperforms all baselines by over one order of magnitude, with JAX giving the best results.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace naive PyTorch circuit evaluation with KLay on CUDA | Neurosymbolic circuit evaluation · Sudoku (4 × 4), batch size 128 | runtime (ms, forward and backward pass): $42.70 \pm 2.76$ -> $0.46 \pm 0.17$ | Helped |
| Replace naive PyTorch circuit evaluation with KLay on CUDA | Neurosymbolic circuit evaluation · Grid (4 × 4), batch size 128 | runtime (ms, forward and backward pass): $107.78 \pm 4.67$ -> $0.49 \pm 0.19$ | Helped |
| Replace naive PyTorch circuit evaluation with KLay on CUDA | Neurosymbolic circuit evaluation · HMLC, batch size 128 | runtime (ms, forward and backward pass): $175.58 \pm \phantom{00}1.26$ -> $1.03 \pm \phantom{00}0.35$ | Helped |
| Replace naive PyTorch circuit evaluation with KLay on CUDA | Neurosymbolic circuit evaluation · Warcraft (12 × 12), batch size 128 | runtime (ms, forward and backward pass): $18930.67 \pm 496.16$ -> $1.59 \pm \phantom{00}0.25$ | Helped |
| Use KLay instead of DeepProbLog | MNIST-addition training · 1 digit, batch size 128 | training time per epoch (s): $241.63 \pm 4.38$ -> $1.26 \pm 0.03$ | Helped |
| Use KLay instead of Scallop (k=3) | MNIST-addition training · 1 digit, batch size 128 | training time per epoch (s): $52.30 \pm 0.37$ -> $1.26 \pm 0.03$ | Helped |
| Use KLay instead of DeepProbLog | MNIST-addition training · 2 digits, batch size 128 | training time per epoch (s): $323.26 \pm 22.20$ -> $1.25 \pm \phantom{0}0.01$ | Helped |
| Use KLay instead of Scallop (k=3) | MNIST-addition training · 2 digits, batch size 128 | training time per epoch (s): $601.20 \pm 14.65$ -> $1.25 \pm \phantom{0}0.01$ | Helped |
| Use KLay instead of Scallop (k=3) | MNIST-addition training · 3 digits, batch size 128 | training time per epoch (s): $5345.78 \pm 171.47$ -> $8.85 \pm \phantom{00}0.02$ | Helped |
Surprising findings
- Juice's GPU implementation is slower than its own CPU implementation, indicating missed parallelization opportunities.
- Contrary to prior claims, arithmetic circuits can be run efficiently on GPUs despite extreme sparsity (fewer than one in a million edges present for the largest circuits).
Limitations
- KLay node counts can exceed the original SDD for some compilers (e.g., D4) due to unary node chains overhead.
- Exact inference only; the method relies on d-DNNF/SDD compilation, which is #P-hard in general.
- Evaluated mainly on PyTorch and JAX; no hardware beyond one RTX 4090 and one CPU.
Experimental setup
Datasets: MNIST addition Sudoku Grid Warcraft HMLC random 3-CNF formulas
Benchmarks: Synthetic 3-CNF SDD circuits Sudoku Grid Warcraft HMLC MNIST-addition
Metrics: runtime in milliseconds runtime in seconds per epoch number of nodes and layers sparsity
Baselines: Naive PyTorch node-by-node evaluation Juice PySDD post-order traversal DeepProbLog Scallop
Builds on: Juice DeepProbLog SDD knowledge compilation Merkle hashing
Closest papers by meaning
- A Fast Convoluted Story: Scaling Probabilistic Inference for Integer Arithmetics (NeurIPS 2024 · similarity 0.84)
- Arithmetic Without Algorithms: Language Models Solve Math with a Bag of Heuristics (ICLR 2025 · similarity 0.84)
- Learning to Add, Multiply, and Execute Algorithmic Instructions Exactly with Neural Networks (NeurIPS 2025 · similarity 0.84)
- Improving Soft Unification with Knowledge Graph Embedding Methods (ICML 2025 · similarity 0.83)
- DOLPHIN: A Programmable Framework for Scalable Neurosymbolic Learning (ICML 2025 · similarity 0.83)
- Graph Neural Networks and Arithmetic Circuits (NeurIPS 2024 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.