Research map / Architectures and efficiency / RNNs and normalization
Scaling Tractable Probabilistic Circuits: A Systems Perspective
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.
Proposes PyJuice, a GPU system for probabilistic circuits that compiles PCs into block-sparse layers for Tensor Core parallelization and uses PC flows for backprop, yielding large speedups and memory savings and new image/language PC baselines.
Problem
Existing probabilistic circuit (PC) implementations are slow and memory-hungry, limiting training of large PCs on large datasets.
Approach
Compile a PC into block-sparse sum layers grouped by child count, so that sum-node computations become block matrix multiplications (using Tensor Cores) with reduced memory reloads; compute parameter gradients via PC flows instead of differentiating the log-space forward pass, and recompute product node outputs in backward to save memory.
Key results
- PyJuice runs 1-2 orders of magnitude faster than SPFlow, EiNet, Juice.jl, and Dynamax on PD, RAT-SPN, HCLT, and HMM structures (e.g., 0.6s vs 6.0s for RAT-SPN with 58K nodes), and uses 2-5x less GPU memory.
- On ImageNet32, HCLT-large trained with PyJuice reaches 4.33 bpd, improving on the 4.82 bpd previously reported for HCLT.
- HMM language model on CommonGen improves perplexity from 9.78 to 8.81 with 4096 states and 8.65 with 8192 states.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use PyJuice instead of EiNet for PD training | PD · PD, 1.38M nodes | runtime per training epoch: 1534.7 ± 0.5 -> 57.1 ± 0.2 | Helped |
| Use PyJuice instead of Juice.jl for RAT-SPN training | RAT-SPN · RAT-SPN, 930K nodes | runtime per training epoch: 375.1 ± 3.4 -> 13.8 ± 0.0 | Helped |
| Use PyJuice instead of Juice.jl for HCLT training | HCLT · HCLT, 1.42M nodes | runtime per training epoch: 143.2 ± 5.1 -> 24.9 ± 0.1 | Helped |
| Use PyJuice instead of Dynamax for HMM training | HMM · HMM, 388K nodes | runtime per training epoch: 4039.8 ± 38.3 -> 39.9 ± 0.1 | Helped |
| Fine-tune the 4096-state HMM for more epochs with PyJuice | HMM language model · CommonGen | perplexity: 9.78 -> 8.81 | Helped |
| Increase HMM hidden states from 4096 to 8192 | PyJuice HMM language model · CommonGen | perplexity: 8.81 -> 8.65 | Helped |
| Fine-tune the entire assembled PC instead of only the top-level model | Sparse image model · ImageNet32 | bits-per-dimension: 4.06 -> 4.04 | Helped |
| Use the larger HCLT benchmark relative to the previous HCLT result | HCLT · ImageNet32 | bits-per-dimension: 4.82 -> 4.33 | Helped |
| Increase block size for block-based parallelization (ablation) | PyJuice sum-layer kernels · PD sum layer | runtime: null | Helped |
| Increase block size for block-based parallelization (ablation) | PyJuice sum-layer kernels · PD sum layer | L1-to-L2 IO: 25-50x fewer IO | Helped |
| Increase the fraction of dropped edge blocks (ablation) | PyJuice block-sparse sum layer · Block-sparse PC layer | runtime: null | Helped |
| Prune edges from HCLTs (ablation) | Sparse HCLT · Sparse HCLTs with 10M and 40M edges | runtime per epoch: null | Helped |
Surprising findings
- Sum layers dominate runtime despite having fewer IO reads than product layers; the gap comes from repeated reloads of child node outputs.
Limitations
- Sparse PCs cannot fully benefit from block-based parallelization, yielding smaller gains over Juice.jl at high pruning rates (over 90%).
- Block-sparse padding adds overhead from pseudo-edges with zero parameters.
Experimental setup
Datasets: ImageNet ImageNet32 CelebA-HQ CommonGen WikiText-103
Benchmarks: ImageNet32 ImageNet CelebA-HQ CommonGen WikiText-103
Metrics: runtime per epoch (s) GPU memory bits-per-dimension perplexity
Baselines: SPFlow EiNet Juice.jl Dynamax
Builds on: EiNets Juice.jl latent variable distillation PC flows
Closest papers by meaning
- Scaling Probabilistic Circuits via Monarch Matrices (ICML 2025 · similarity 0.89)
- The Expressivity Boundary of Probabilistic Circuits: A Comparison with Large Language Models (NeurIPS 2026 · similarity 0.86)
- Sum of Squares Circuits (AAAI 2025 · similarity 0.86)
- Scaling Continuous Latent Variable Models as Probabilistic Integral Circuits (NeurIPS 2024 · similarity 0.85)
- On the Hardness of Approximating Distributions with Tractable Probabilistic Models (NeurIPS 2025 · similarity 0.85)
- Scaling Diffusion Transformers Efficiently via $\mu$P (NeurIPS 2025 · similarity 0.84)
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.