Research map / Architectures and efficiency / RNNs and normalization
Scaling Probabilistic Circuits via Monarch Matrices
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.
Replaces dense sum-block matrices in probabilistic circuits with generalized Monarch (and butterfly-interpolating) matrices derived from circuit multiplication, enabling PCs with hidden sizes up to 2^19 and state-of-the-art PC density estimation on text8, LM1B and ImageNet32/64.
Problem
Scaling probabilistic circuits to large hidden sizes is limited by the quadratic FLOPs and memory cost of dense sum-block weight matrices, and existing methods do not exploit both sparsity and tensorized hardware efficiency.
Approach
Show that multiplying two sum blocks yields a Kronecker-product linear map that can be materialized as sparse circuit layers; untie the parameters to get generalized Monarch matrices, initialized by multiplying smaller dense PCs, and interpolate between Monarch and butterfly structures by number of layers.
Key results
- Monarch-HMM reaches 1.57 BPC on text8, beating the dense HMM (1.69) and other PC baselines and approaching diffusion models (SEDD 1.47).
- Monarch-HMM reaches 190.34 perplexity on LM1B versus 320.78 for dense HMM.
- Monarch-HCLT achieves 4.01 bpd (lossy) on ImageNet32 and 3.74 on ImageNet64, state-of-the-art among PCs and beating LVD-PG (4.06 / 3.80).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace dense HMM transition matrices with Monarch matrices | HMM · text8 test | BPC: 1.69 -> 1.57 | Helped |
| Replace dense HMM transition matrices with Monarch matrices | HMM · LM1B test | perplexity: 320.78 -> 190.34 | Helped |
| Replace dense sum blocks with Monarch matrices | HCLT · ImageNet 32×32 and 64×64, lossy and lossless | bits-per-dimension: 4.06 -> 4.01; 4.46 -> 4.01; 5.08 -> 4.62; 3.80 -> 3.74; 4.42 -> 3.74; 5.05 -> 4.33 | Helped |
| Initialize with PC multiplication using 3 dense HMMs instead of random initialization (ablation) | 3-layer Monarch-HMM · text8 test, hidden sizes 2¹²–2¹⁵ | BPC: 2.25 -> 2.14; 2.10 -> 2.01; 1.98 -> 1.90; 1.88 -> 1.81 | Helped |
| Increase Monarch depth from 2 to 3 or 4 layers (ablation) | Monarch-HMM · text8 test, hidden sizes 2¹²–2¹⁵ | BPC: 2.041 vs 2.135 and 2.188; 1.926 vs 2.013 and 2.065; 1.824 vs 1.895 and 1.948; 1.747 vs 1.814 and 1.853 | Hurt |
| Prune up to 90% of hidden states during the forward pass (ablation) | Monarch-HMM · text8 | log-likelihood | No clear effect |
| Use denser matrices rather than sparse Monarch structures at fixed hidden size (ablation) | Probabilistic circuits · text8 | BPC | Helped |
Surprising findings
- The performance gap between dense and sparse Monarch layers diminishes quickly as hidden size increases, suggesting larger PCs are inherently sparser.
- Up to 90% of hidden states in a Monarch-HMM can be pruned during the forward pass without significant log-likelihood drop.
Limitations
- Image modeling still lags behind neural image models such as VDM.
- Sparser (more-layer) Monarch structures suffer from high memory consumption from caching hidden states, making them impractical to train at scale.
- Scaling is constrained by the lack of an effective homogeneous PC architecture for images, limiting hidden size growth.
Experimental setup
Datasets: text8 LM1B ImageNet32 ImageNet64
Benchmarks: text8 LM1B ImageNet32 ImageNet64
Metrics: bits-per-character perplexity bits-per-dimension FLOPs sampling time
Baselines: SparsePC NPC^2 HMM IAF/SCF Argmax Coupling Flow D3PM SEDD QPC HCLT LVD LVD-PG RealNVP
Builds on: Hidden Markov models Butterfly matrices Monarch matrices Hidden Chow-Liu trees PyJuice
Closest papers by meaning
- Scaling Tractable Probabilistic Circuits: A Systems Perspective (ICML 2024 · similarity 0.89)
- On the Expressive Power of Tree-Structured Probabilistic Circuits (NeurIPS 2024 · similarity 0.88)
- On the Hardness of Approximating Distributions with Tractable Probabilistic Models (NeurIPS 2025 · similarity 0.88)
- The Expressivity Boundary of Probabilistic Circuits: A Comparison with Large Language Models (NeurIPS 2026 · similarity 0.87)
- Sum of Squares Circuits (AAAI 2025 · similarity 0.87)
- On the Relationship Between Monotone and Squared Probabilistic Circuits (AAAI 2025 · similarity 0.86)
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.