Scaling Probabilistic Circuits via Monarch Matrices
Authors: Honghua Zhang
Meihua Dang
Benjie Wang
Stefano Ermon
Nanyun Peng
Guy Van den Broeck
Venue: ICML 2025 poster (2025)
Keywords: Probabilistic Circuits
Density Estimation
Primary area: probabilistic_methods->graphical_models

Abstract: Probabilistic Circuits (PCs) are tractable representations of probability distributions allowing for exact and efficient computation of likelihoods and marginals. Recent advancements have improved the scalability of PCs either by leveraging their sparse properties or through the use of tensorized operations for better hardware utilization. However, no existing method fully exploits both aspects simultaneously. In this paper, we propose a novel sparse and structured parameterization for the sum blocks in PCs. By replacing dense matrices with sparse Monarch matrices, we significantly reduce the memory and computation costs, enabling unprecedented scaling of PCs. From a theory perspective, our construction arises naturally from circuit multiplication; from a practical perspective, compared to previous efforts on scaling up tractable probabilistic models, our approach not only achieves state-of-the-art generative modeling performance on challenging benchmarks like Text8, LM1B and ImageNet, but also demonstrates superior scaling behavior, achieving the same performance with substantially less compute as measured by the number of floating-point operations (FLOPs) during training.

Scaling Probabilistic Circuits via Monarch Matrices
Honghua Zhang"
Meihua Dang" 22
Benjie Wang"
Stefano Ermon
Nanyun Peng
Guy Van den Broeck
Abstract
2021). The tractability of PCs have been exploited in number of domains, including fair and explainable machine learning (Choi et al., 2021; Wang et al., 2021b), causal infer- ence (Wang et al., 2021a; Zečević et al., 2021; Wang et al., 2022; Wang & Kwiatkowska, 2023), controllable generation (Liu et al., 2024b; Zhang et al., 2024), and neuro-symbolic AI (Ahmed al., 2022; Maene et al., 2024).
Probabilistic Circuits (PCs) are tractable repre- sentations of probability distributions allowing for exact and efficient computation of likelihoods and marginals. Recent advancements have im- proved the scalability of PCs either by leveraging their sparse properties or through the use of ten- sorized operations for better hardware utilization. However, no existing method fully exploits both aspects simultaneously. In this paper, we propose novel sparse and structured parameterization for the sum blocks in PCs. By replacing dense matrices with sparse Monarch matrices, we signif- icantly reduce the memory and computation costs, enabling unprecedented scaling of PCs. From a theory perspective, our construction arises natu- rally from circuit multiplication; from a practi- cal perspective, compared to previous efforts on scaling up tractable probabilistic models, our ap- proach not only achieves state-of-the-ar genera- tive modeling performance on challenging bench- marks like Text8, LMIB and ImageNet, but also demonstrates superior scaling behavior, achieving the same performance with substantially less com- pute measured by the number of floating-point operations (FLOPs) during training.
Recent advancements in PC learning techniques (Liu et al., 2023b; Gala et al., 2024) and efficient tensorized implemen- tations (Peharz et al., 2020; Dang et al., 2021; Liu et al., 2024a) have significantly enhanced the expressiveness and scalability of PCs. However, to further boost the perfor- mance of PCs, simply scaling up the model size is insuffi- cient; we need to better utilize the available capacity. To this end, Dang et al. (2022) found that learned PCs empirically exhibit significant sparsity, and leveraged this observation to iteratively learn PC structures through pruning and grow- ing. However, the sparse connections learned are arbitrary, making them difficult to tensorize and parallelize effectively.
In this paper, we focus on leveraging structured sparse pa- rameterizations for the sum blocks in PCs, which represent linear maps. To the best of our knowledge, all previous cir- cuit architectures utilizing tensorized operations have relied on dense matrices to parameterize these linear maps (Peharz et al., 2020), which incur a quadratic cost in the number of nodes. Inspired by recent advances in low low-rank approxima- tions for transformers (Hu et al., 2022), we propose a novel, more efficient parameterization for the sum blocks.
1. Introduction
We begin by illustrating how our parameterization naturally arises from circuit multiplication. Previous analysis (Shen et al., 2016; Vergari et al., 2021) showed that multiplying two (compatible) circuits could result in a quadratic increase in size in the worst case. However, we observe that the linear maps in the sum blocks generated by multiplication are not dense but rather the Kronecker product of the linear maps from the original blocks, which can be implemented more efficiently. Further, by explicitly materializing this map as interleaving sums and permutations, we identify inter- esting connection between these product circuits and some recently introduced class of structured matrices including Butterfly matrices (Dao et al., 2019) and Monarch matri- ces (Dao et al., 2022). Building on this insight, we propose replacing the dense linear maps in PCs with Monarch layers.
Probabilistic circuits (PCs) are a unifying representation of tractable probability distributions through computation graphs (Choi al., 2020; Darwiche, 2003). The key prop- erty that separates PCs from other deep generative models such flow-based models (Papamakarios al., 2021) and VAEs (Kingma& & Welling, 2013) is their tractability. This property enables PCs to compute various queries. including marginal probabilities, exactly and efficiently (Vergari et al.,
'Equal contribution 'Department of Computer Science. Uni- versity of California, Los Angeles 'Deparmment of Computer Sci- ence, Stanford University. Correspondence to: Honghua Zhang <hzhang I 99s @ gmail.com>.
Proceedings of the 4rnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

Scaling Probabilistic Circuits via Monarch Matrices
In our empirical evaluation', we demonstrate that by re- placing dense matrices in PCs with structured Monarch matrices, we are able to scale PCs to orders-of-magnitue larger hidden sizes and, among a variety of tractable genera- tive models, we are able to achieve state-of-the-art density estimation performance on various benchmarks, including ImageNec32/64 (Deng et al., 2009) for image modeling and Text8 (Mahoney, 2011) and LMIB (Chelba et al., 2013) for language modeling. Furthermore, we show that, compared to circuits with dense layers, the ones with Monarch layers yield significantly better scaling curves: they can achieve the same performance with substantially fewer floating-poni operations (FLOPs) in training and inference.
Definition 2.2 (Smoothness and Decomposability). A sum node is smooth if all of its children have the same scope. A product node is decomposable if its children have disjoint scope. A PC is smooth (resp. decomposable) if all of its sum (resp. product) nodes are smooth (resp. decomposable).
In practice, probabilissic circuit graphs are typically de- signed in a tensorized manner, in which sets of nodes of the same type (sum, product, leaf) and with the same scope grouped together as a block; the computation graph is then specified through connections between the blocks (Peharz et al., 2020; Liu et al., 2024a; Loconte et al., 2024a). We write n to denote a node block and |r|] for the number of nodes in the block.
2. Tensorized Probabilistic Circuits
Definition 2.3 (Sum Block). A sum block has set of child blocks {c(+) )m T== such that each sum node in the block is connected to every node in each of the child blocks. We can write W E R Definition 2.4 (Product Block), A product block n has set of child blocks {c! Jm We define two types of product node block with different connectivity:
Notation We use uppercase to denote variables (e.g. X)
and lowercase to denote values of variables (e.g. xr). We use boldface to denote sets of variables/values (e.g. X,x).
e( for the weight matrix.
Definition 2.1 (Probabilistic Circuit). A PC C = (G, 0) represents a joint probability distribution over random vari- ables X through a directed acyclic (computation) graph (DAG) G parameterized by 0. Specifically, the DAG g consists of three types of nodes- sum, product, and leaf nodes. Each leaf node n is associated with a non-negative function (X,) over some variable X,, called its scope X,, == }. The scope of any sum or product node n is defined to be X, = Ueechin) Xx, where ch(n) denotes the children of in G. Each node represents a probability distribution P., over its scope X,,, defined recursively by:
Hadamard O: If |c01| |r2| for all i = I,... then we define Hadamard product block where n i=1 cl0)
Kronecker Ø: If |n| = II lc! then we can define a Kronecker product node block where n = (
Sum blocks represent parameterized linear maps, while product blocks represent fixed multilinear maps. Typically, for smooth and decomposable PCs, one builds the circuit by alternating between sum and product blocks (i.e., chil- dren of sum blocks are product blocks, children of product blocks are sum/leaf blocks). In this paper, we focus on the parameterization of the sum blocks, which is independent from the choice of Hadamard or Kronecker product blocks (we use Hadamard product blocks for our experiments).
(n))
if nt is leaf node if n is a product node
Pr (x(),,)) = IIEEch(n) Pe(XX) Ecech(nn Peln Pe PX() if is a sum node
where for each leaf node, function fn (X,,) represents a nor- malized univariate probability mass/density function (e.g. Categorical, Gaussian); and for every sum node n. Ocl. is non-negative weight associated with the edge (n, c) in the DAG. If Ecech(n) 4ej. = 1, then the PC computes normal- ized joint probability mass/density function. The function represented by a PC, denoted (x()). is the function repre- sented by its root node; and the size of a PC, denoted [C], is the number of edges in its graph.
Measuring the sizes of PCs The size of a probabilistic circuit (number of edges) determines the number of FLOPS needed for forward pass. This can be roughly characterized as a function of two parameters: the number of sum blocks n, and the input/output dimension of the largest sum block, which we call the hidden size h. The size of the circuit is then bounded by (nh?)).
It is immediate from the definition that one can evaluate a PC's function with a single traversal through its computation graph. The distinguishing feature of PCs compared to other computation graphs such as neural networks is that one can also efficiently compute marginals under the following restrictions on the node scopes:
3. From Circuit Multiplication to Generalized Monarch Matrices
In this section, we describe from first principles our con- struction of PCs parameterized by (generalized) Monarch matrices. Firstly, in Section 3.1, we focus on the funda- mental operation of circuit multiplicatiin, and show how
'Code available https: // ht///gtttttt wtt.bb.comwwaannt cow.bb//////wwaaaaa/// Mona rehCi rcuits
2

---

Scaling Probabilistic Circuits via Monarch Matrices
B
A Ree
p
(a)
T R A N O S
Kronecker Product X B
o
(b)
(c)
Figure 1: Probabilistic circuits architecture illustration with Monarch matrices. (a) Two sum blocks in PCs with weight matrices A, of arbitrary dimensions. (b) A constructed sum block with weight matrix as the Kronecker product B, representing the eircuit product of two sum blocks. (c) Efficient eircuit representation for the linear transformation A @ B.
Kronecker product A @ B (Vergari et al., 2021). Figure 1b shows a circuit materialization of A 69 B, consisting of O(rspg) edges. However, the linear transformation given by A & B can actually be executed in a significantly more efficient way: letting æ E R2x4 be an input tensor, we compute the linear transformation (A & B)x as
TRANSPOSE
Arxrpp
B B':
((A 2 B)x)
Figure 2: Generalized two-layer Monarch matrix.
Σ (A B)..Aja. En ByyB
Σ, Bis E, Auiti = -E, (1x(Ax)
the resulting circuit exhibits structured sparsity, leading to PCs parameterized by Monarch matrices. In Section 3.2, we extend this analysis to the multiplication of more than two PCs, and derive a new multi-layer Monarch matrix parame- terization. Finally, in Section 3.3, for a PC with hidden size h, we show how this can be used to control the compute and memory used between (h log(h)) and Ohh)), which ef- fectively interpolates between the Monarch and the sparser Butterfly matrices.
BIs (B(Ax) (B(Ax) yT:
hence, we have
(A & B)x (B(Ax)
(I)
We attempt to materialize Equation as a circuit, as shown in Figure Ic. with æ well R2x4 viewed as a flattened 1-d tensor of dimension pq. The total number of edges bounded by O + srq). For rough comparison against the naive construction of A & B, if A and B are both of dimension m X m, then the naive circuit construction contains ((m²) edges while the construction based on Equation 1 contains only (mm3) edges.
3.1. Materializing Circuit Multiplication
Given two circuits A and B, the goal of circuit multiplication is to construct tractable (i.e. smooth and decomposable) circuit C such that (x) (=( (xr)). If A and B are structured-decompossse decomposable with respect to the same vtree, i.e., the product nodes in A and B always factor the same way, then C can be constructed with at most a quadratic increment in size (Shen et al., 2016; Vergari et al., 2021). The corresponding algorithm essentially constructs a new node for each pair of nodes with the same scope.
A Although the circuit shown in Figure Ic constructed by multiplying two sum blocks, it can also be interpreted as
"The code for circuit multiplication the official implementa- tions of Wang & Kwiatkowska (2023) and Loconte et al. (2024b) also achieve this complexity through einsum operations; though this is not explicitly described in either of the papers, which re- port the looser bound of O(m²), i.c., square of the circuit size. Crucially, we show that the product can be explicitly material- ized as a circuit, which was missed by these prior works. This allows us to directly apply standard PC inference and learning algorithms to the product eircuit, such as parameter learning with expectation-maxx maximizaaii.n.....
Here we focus on the local operation of multiplying two sum blocks. Given two sum blocks with weight matrices A E R' and B E Rex (Figure la), their product can be represented as a sum block with input dimension pq and output dimension rs,, and weight matrix given by the
3

---

Scaling Probabilistic Circuits via Monarch Matrices
sparse representation for some linear transformation M from Re4 to R Furthermore, in this interpretation, each of the A blocks (and each of the B blocks) do not need to have the same parameters for M to be valid. Hence, we "untie" the parameters of A and B to obtain the (generalized) Monarch transformation M.
Th X m4+1 x ng my ... 14-1- We define the
generalized Monarch matrix M E R(m1x--ma) Xx-m)) (m1 x--na)
(Mæ) Janmnnd
Σ
II
A Ji
Firmis
Definition 3.1 (Generalized Monarch Matrices 1). Given A RTxxx. and B Raxexr, we define M be the linear transformation from Rey to Ih's such that
i1,ea.nd
Istsd
(A* * (AA-1
(...* (A² (A' * 2 2)))s....))
Here x² = (A' * ..... * (At * x) ...) ) is of dimension
(Mz) = Σ Akyy By = (B (A* (A+ E
Rd m X-.. .. my and A xx denotes a batched
Re+1 X--.
linear transformation such that
where Ax * denotes "batched" linear transformation
(A' * x ) -.,, A jcheriate 1 Miieriahrn
= x), = EAusa
with i441-.iia1 ... Jt- enumerating over the batch di- mension. Note that the circuit materialization of this con- struction of Monarch matrix consists of consecutive sum blocks, so we call a d-layer Monarch matrix.
with; j enumerating through the "batch" dimension.
We visualize the generalized Monarch matrix in Figure 2. Here, the 2D matrices A and Bt are actually slices of the 3D tensors A, B in Definition 3.1; specifically (A9) )aa Aki and (B*),
Despite the complexity of this definition, construct a PC A with Monarch layers of hidden size h = IIsts hi, we really just need to multiply d PCs Ar. each with dense layers of hidden size ht. As in the case of two circuits, we can first train the smaller PCs A,, with dense layers and use their parameters as an initialization point for the training of A where (=)))) x II. P.A. (z)). Such parameter initialization significantly improves the result of training PCs with Monarch layers of large hidden sizes, and we refer readers to Section 5.2 for details.
Bijk-
Thus, to enable structured sparsity in PCs, we propose to replace the dense matrices in sum blocks with generalized Monarch matrices, parameterized by the tensors A, B. In the case where we have h input and output nodes (with p = 9 r s vh), then the compute and memory requirements (123/2) as compared to the O(h²) cost for dense PC. The interpretation circuit multiplication gives us a principled means to initialize the parameters of such circuit: namely, train two smaller PCs A. B with dense layers and hidden sizes Vh and then multiply them to obtain a circuit representing the distribution pc(x) x ( (a))))p(() pz(()) We can then untie the parameters of C during training to leverage the fully general Monarch paramterization.
3.3. Interpolating Butterfly and Monarch Matrices
To conclude this section, we draw an interesting connection between the generalized d-layer Monarch matrices and the butterfly matrices (Parker, 1995; Dao et al., 2019; Meng et al., 2022). Butterfly matrices are a class of expressive structured matrices, constructed as the product of sparse matrices known as butterfly factor matrices. For ease of exposition, we will describe here butterfly factor matrices of size D x D, where D is power of 2.
3.2. Multiplying Multiple PCs
natural way to further generalize the construction above would be to consider products of multiple PCs, or more specifically, Kronecker products of multiple matrices. We generalize Equation 1 as:
Definition 3.3 (Butterfly Factor). Given any D 2d and < i < d, butterfly factor matrix B(i, D) is sparse matrix with the following sparsity pattern:
(A' @ 22... 8 1d
= (A (AT...AA(AA
B(1, D) has non-zero elements only along the diago- nals of the four 2 x Z
D submatrices.
with A E R.mexne and x R the superscript S denotes the left shifting operation where, e.g., x² jki Tiyk- Similarly, we can materialize Equation 2 as circuit, untie the shared parameters among the blocks ^ ,!....,4" Ad, and then obtain the (further) generalized construction of Monarch matrices.
B(i, D) is a biock-diagonal matrix with block size - x LP, where each block is a B (1, 2 butterfly factor.
Examples of butterfly factors for D = 16 are shown in Figure 3. A D x D butterfly matrix is given by the ma- trix multiplication B(D) == B(1, B2, D) ... B(d, D).
Definition 3.2 (Generalized Monarch Matrices II). Let {A*) Sisgd be d tensors, where A has dimensions m, x
4

---

Scaling Probabilistic Circuits via Monarch Matrices
be nicely "factored" into squared matrices, we resort to the almost-squared factorizations like 2 = 2 2 In Section 5.1 and 5.2, we conduct an empirical study on the scaling behavior of Monarch matrices with different number of layers and show that even though Monarch matrices with more layers exhibit slightly better scaling behavior com- pared to two-layer Monarch matrices, they induce higher memory consumption and thus are less desirable in practice.
16)
8(2, 16)
1
B(4,16)
Figure 3: Illustration of Butterfly factors (Meng et al., 2022).
Initialization The probabilistic interpretation of Monarch layers as circuit multiplication further provides an intuitive way for initializiig large circuits. For instance, to initialize a 2-layer Monarch with hidden size of h = 2200 we can multiply two smaller models with hidden size of h = 20 (or any possible factorizations). The smaller models are trained separately to fit the same data distribution and then combined to initialize the larger model's parameters.
Even though this is a product of sparse matrices, we show that in fact can be interpreted as a product of dense tensors, specifically, generalized Monarch matrices:
Theorem 3.4. Suppose that mt T4 for all 1 S.S. d. Then the generalized d-layer Monarch matrix M is a butterfly matrix of dimension 2 x 2
More generally, suppose that we want to construct a gen- eralized Monarch matrix of dimension h h. We can in- terpolate between butterfly (consisting of loga (h) different ())
We use stochastic mini-batch version of aopeecatiio-Maximmzaann optimization (Peharz et al., 2016). Empirically, this approach converges faster and is better regularized compared to EM on the whole dataset.
Training
tensors of shape x x 2) and Monarch matrices (consisting of 2 different 3-dimensional tensors of shape h 1/2 x h/22 x h 1/2 as follows. We can first pick a number e such that cd = h for some d, and then the construction in Defintion 3.2 would give us a log, (h)- layer Monarch matrix with each layer (in terms of the circuit materialization) consisting of he edges; i... the # of FLOPs needed for the linear transformation is given by he logo (h). We have the butterfly matrix if c = 2 and the commonly known Monarch matrix (Dao et al., 2022) if e = Vh. If is nice enough, e.g., 2 we can freely choose c to in- terpolate between 2 and vh, which corresponds to d-layer Monarch matrices with d ranging from logz(h) 2, each having different degrees of sparsity.
Research Questions In the following experiments, we seek to verify the following hypotheses:
Scaling law for Monarch PCs. Sparse structures achieve better scaling behavior compared to dense ones in terms of bits-per per-dimension (BPD) versus floating point operations (FI.Pss).
Effect of Circuit Mulriplication. Initializations using circuit multiplication leads to better performance.
Increasing model size induces sparsity. For a fixed hidden size, denser matrices perform better, but the performance gap between sparse and dense matrices diminishes quickly as the hidden size increases.
4. Scaling PCs via Monarch Matrices
Employing structured matrices improves two aspects of probabilistic circuits. First, fixing the hidden size, the PCs require fewer FLOPs so training will be more efficient. Sec- models with
5. Experiments
ond, it reduces the memory consumption, larger hidden sizes able fit in the GPU memory. These two benefits make it possible to scale up probabilistic cir- cuits to much larger hidden sizes. In this section, we discuss some of the practical choices to be made for learning PCs with (generalized) Monarch parameterizations, the effect of which we investigate empirically through ablation studies.
We evaluate our method using generative modeling bench- marks for both text and image data. We use log-likelihoods as a measurement of a model's performance and the number of floating point operations (FLOPs) per dimension a measurement of a model's efficiency. Given hidden size of h, the FL.OPs per token is 1t2 for an HMM and 2/3/2 for a two layer Monarch-HMM. Details are in Appendix B.
Choices of Monarch Structures As described at the end of Section 3, given desired linear transformation from h to RAA there are different choices of Monarch struc- tures. For example, f = 28 can be factorized as 29 x 29, 26 x 26, x 26, (23)6 etc., where 2 x 2 gives the standard two-layer Monarch matrix and 26 x 26 x 26 gives a three- layer Monarch matrix. For hidden sizes like 219 that cannot
5.1. Character-level Language Modeling
Dataset Text88 (Mahoney, 2011) is character-level lan- guage modeling dataset with a vocabulary of 27 tokens: the letters 'a'-'z' and the whitespace token. We follow the stan- dard practice of training and evaluating text8 in chunks of length 256 without preprocessing (Hoogeboom al., 2021).
5

---

Scaling Probabilistic Circuits via Monarch Matrices
Type
Model
BPC (4)
Time (s) (4)
2.3
HMM
Flow
IAFISCF
1.88
0.04
Monarch-2
Flow
Argmax Coup Flow
1.80
0.40
Monarch-3
Diffusion
D3PM Uniform
S 1.61
3.60
Test BPC
2.0
Monarch-4
Diffusion
SEDD Uniform
< 1.47
PC
SparsePC
2.60
1.8
PC
NPC²
3.17
PC
HMM
1.69
0.006
PC
Monarch-HMM 1.57
0.017
1.5
1E+04
1E+06
1E+08
1E+10
Table 1: Averaged test set BPC on text8. Sample times are for generating an example of length 256. Our method outperforms PC baselines, and significantly close the gap between PCs and the other less tractable generative models.
Inference FLOPs/character
Figure 4: Scaling curves comparing HMM and varying Monarch structures. BPC (L) as a function of training FLOPs per character. Monarch-HMM demonstrates greater efficiency than HMM across varying computational budgets. Sparser Monarch further leads to better scaling behaviors.
Baselines We compare our method with probabilistic eircuit structures. including SparsePC (Dang et al., 2022), NPC² (Loconte et al., 2024b), and HMM (Zhang et al., 2023; 2024), as well other types of less tractable gen- erative models as a reference, such as flow-based mod- els (IAF/SCF (Ziegler& & Rush, 2019), Argmax Coupling Flow (Hoogeboom et al., 2021)) and diffusion models (D3PM (Austin et al., 2021), SEDD (Lou et al., 2024)) with uniform transition matrices. The results for SparsePC and NPC² are obtained by rerunning their official imple- mentations. SparsePC constructs arbitrary sparse structures by iteratively pruning and growing an initial model, which prevents it from scaling to very large hidden size due to sys- tem limitations. For HMM, we use dense transition matrices with the largest possible hidden state and run it using our codebase to ensure optimal performance. The results for other generative models are taken from (Austin al., 2021) and (Lou et al., 2024).
Hidden Size
2
213
214
2
Random Initialization
2.25
2.10
1.98
1.88
PC Multiplication
2.14
2.01
1.90
1.81
Table 2: Test set BPC (L) of a 3-layer Monarch-HMM for different hidden sizes comparing random initialization and initialization from PC multiplication with 3 dense HMMs.
sparsity leads to better scaling behavior. However, sparser structures are constrained by memory consumption (Sec- tion 5.2), so we use Monarch-2 as our main result in Table 1.
5.2. Ablation and Analysis
We investigate the relationship between hidden size, FLOPs and memory consumption for varying structures, including dense HMM and differernt layers of Monarch- HM.. From Section 5.1, we conclude that sparser structures achieve bet- ter FLOPs. In this section, we aim to verify the following hypotheses: (1) Initializations from circuit multiplications give better performance. (2) For fixed hidden size, denser matrices outperform sparse ones but the performance gap dimish quickly as hidden size increases; and (3) Monarch- HMMs with more layers requires significantly higher mem- ory consumption.
Benchmark We report log-likelihood results in bits-per- character (BPC) in Table 1. The results show that Monarch- HMM outperforms all PC models and significantly narrows the gap between PC models and other less tractable gen- erative models. We improve upon the baseline HMM by replacing dense matrices with structured sparse (Monarch) matrices; thus, we are able scale up hidden size from 255 to 2n Additionally, Monarch-HMM achieves significantly faster inference, 200x times faster than diffusion models.
Scaling laws of structured matrices We also include a plot in Figure 4 comparing HMM and Monarch-HMM in terms of BPC as a function of inference FLOPs. Monarch- HMM trained with hidden size ranging from 212 to 219 Monarch-HNor demonstrates greater efficiency than HMM across different computational budgets consistently. We additionally include Monarch-3 and Monarch-4, which rep- resent Monarch-HMM with three and four Monarch layers respectively (Section 3). Since models with more layers have more sparse structures, this suggests that increased
Initialization matters The probabilistic interpretation of Monarch matrices as circuit multiplication, as introduced in Section 3, provides an effective approach for initializ- ing the parameters of PCs with Monarch layers. We com- pare the performance of Monarch-HMM trained with two iraimeter-nitiaaiiztton strategies: random initialization vs. initialization from multiplying dense HMMs. As shown in Table 2, circuit multiplication as initialization leads to consistent improvement across all scales.
6

---

Scaling Probabilistic Circuits via Monarch Matrices
2.3
2.0
HMM
2
Monarch-2 Monarch-3
2
Test BPC
2.0
Monarch-4
Test BPC
1.8
1.8
1.5
0.4
0.8
1.6
0
35000
70000
Pruned Hidden States %
Hidden Size
Figure 6: We can effectively prune up to 90% hidden states with no sigificant performance drop for Monarch-HMM.
Figure 5: Test set BPC as a function of hidden size for HMM and varying Monarch structures.
Type
Model
Perplexity (J)
Time (s) (1)
Model
FLOPs (()
Mem Complexity
GPU Mem
Diffusion
D3PM Uniform
< 137.9
1.82
O(-)
GB
Diffusion
SEDD Uniform
X 40.24
PC
HMM
320.78
0.0026
HMM
h?2
288
PC
Monarch-HIMM
190.34
0.0095
Monarch 2
2/3/2
2nhB+ + +2/
65
Monarch 3 3/4/3
vnhB+3h4/3
96
Table 4: Averaged test set perplexity on LMIB. Sample
Monarch 4
4/5/1
AnhB+4h"0
128
times are for generating an example of length 128. Our method outperforms PC baselines, and significantly close the gap between PC and other generative models.
Table 3: FLOPS and memory consumption for varying structures. GPU memory is computed when batch size 128, sequence length 252 and hidden size h 2
Sparse hidden representations Dang al. (2022) show that learned dense layers in PCs are inherently sparse, mo- tivating our use of Monarch layers. Using Monarch layers, we scale circuits to larger hidden sizes, where memory con- sumption from caching hidden states (node values) for back- ward propagation becomes the new bottleneck (O(nn)B) in Table 3). We investigate whether these representations are also sparse by pruning hidden states during the forward pass of a Monarch-HMM.. Figure 6 shows up to 90% of hidden states across all layers can be pruned without significant drop in log-likelihood.
Scaling up hidden size As shown in Figure 5, bits-per- character improves as hidden size increases. For a fixed hidden size, denser structures always perform better as a dense matrix can always represent a sparse one by incorpo- rating zero values, while having greater capacity in theory. However, this comes at the cost of increased FLOPs shown in Figure 4. Very interestingly, as the hidden size in- creases, the performance gap between dense and sparse ones diminishes quickly. This suggests that larger PC models are inherently more sparse, making structured sparse lay- ers increasingly effective in capturing the underlying data distribution as we scale up.
5.3. Token-level Language Modeling
Language modeling for large-scale datasets with large vo- cabularies using probabilistic circuits has not been previ- ously demonstrated. We present results on the One Billion Word dataset (Chell et al., 2013) as a proof of concept, demonstrating the scalability of Monarch-HMN Following D3PM (ustinn et al., 2021), all models are trained and evaluated packed sequences of length 128 using SentencePiece 3 vocabulary of size 8192. We use HMM as our baseline, as other probabilistic circuit models are unable to scale to this dataset. Additionally, we include results from diffusion models from Austin et al. (2021) and Lou et al. (2024) for reference. Monarch- same setup as in the text8 experiments for two epochs.
Sparse structures and memory consumption For prob- abilistic circuit with either dense or sparse matrices, let n be the sequence length, h the hidden size, d the number of Monarch layers, and B the batch size. The training memory consumption consists of: (1) parameter caching, which is also linear with FLOPs per character, requiring O(12) for dense HMM and (ddh-1/dd for a layer Monarch); (2) gradient caching for backward propagation, which is pro- portional to the number of nodes (hidden states), requiring (nhBB for a dense HMM and O(dnh B) for a d layer Monarch. This connection is summarized in Table 3. The last column provides an example on GPU memory usage: though Monarch-3/4 3/4 have better FLOPs efficiency, training them is impractical due to their memory consumption.
trained using the
intops://githin http:/tthhb..coooooooeeg wwt.:o////oooooooooo sentencet ece

---

Scaling Probabilistic Circuits via Monarch Matrices
5.0
ImageNet x2x32 Lossless
ImageNet 64
HCLT
Lossy
Lossy
Lossless
Monarch-HCLL
LVD
4.39
4.12
Test BPD
4.9
LVD-PG
4.06
3.80
QPC
4.46
5.08
4.42
5.05
4.7
Monarch
4.01
4.62
3.74
4.33
RealNVP
4.28
3.98
4.6
1E+03
1E-04
1E+05
1E+06
1E+07
Glow
4.09
3.81
Inference FLOPs/dim
VDM
3.72
3.40
Figure 7: Results for Training on ImageNet32 Dataset. Test bits-per-dimension (L) a function of training FLOPs per pixel. Monarch-HCLI demonstrates greater efficiency than HCLT as the computation budget increases.
Table 5: Density estimation on image datasets. Test set log-likelsooods are in bits-pes det--per-iimenioon (lower is better). Our method performs favorably relative to all PC baselines.
We report test set perplexity and sampling tims in Table 4. While other PC models cannot scale to this range, Monarch- HMM significantly bridges the gap between probabilistic circuits and diffusion models, while retaining the tractability advantage of PCs over diffusion models.
is due to the composition). We show the scaling plot for both PC variants ImageNet64 (YCoCg-R) in Figure 7. We find that dense PCs plateau quickly in terms of their performance, with the more compute and memory-efficiet Monarch layers often having similar performance to the corresponding dense layers with the same number of hidden states. On the other hand, scaling up the number of hidden states further via Monarch parameterizations remains effec- tive; this is because we can use much larger hidden size with the same compute/memory.
5.4. Image Modeling
In this section, we conduct experiments on the ImageNet32 and ImageNet64 datasets, which are downscaled 32 x 32 and 64 x 64 versions of ImageNet (Deng al., 2009). To improve modeling performance, and following prior work (Liu et al., 2023a; 2024b;a; Gala et al., 2024), we apply color transform to the original RGB data, and fit PC models on this transformed data. In particular, we employ a loss- less YCoCg-R R transform (Malvar& & Sullivan, 2003) (see Appendix B for details); as such likelihoods on this trans- formed data are comparable those on the original RGB dataset. As some prior works on PC learning (Liu et al., b023a;b) employ lossy quantized YCoCg transform; we also report results using this transform for fair comparison.
We further show comparison between our models and the state-of-the-art PCs on these datasets. For reference, we also show some neural baselines, namely, RealNVP (Dinh et al., 2017), Glow (Kingma & Dhariwal, 2018), and VDM (Kingma et al., 2021). It can be seen that Monarch-HCLL show improved scaling beyond what is achievable with stan- dard HCLTs. Our largest model, with states, achieves state-of-the-ar performance for PCs (Ta- ble 5), beating even LVD-PG (Liu et al., 2023b), which is an image-specialized PC-based model that uses a high-level PC over 4 x 4 patch PCs, together with a complex optimiza- tion procedure involving latent variable distillation (Liu et al., 2023a) for initializing the parameters and progres- sive growing technique. In contrast, simply use random initializations, the generic HCLT variable decomposition, and train end-to-end using only the well-established EM algorithm for PCs.
16384 hidden
To improve training efficiency and enable greater scaling up of hidden size, we split each image into 8 x patches and train and evaluate
PCs over these 8 x 8 images (8 x 8 x 3 192 dimensions, accounting for the color channels) rather than the full image. We evaluate models using test-set bits-per- pens-e-dddmesionn (bpd): as this is normalized for dimension, our bpds are directly comparable with bpds on the entire image.
Despite the improvements achieved by employing Monarch matrices to scale up the hidden state size, the modeling per- formance still currently lags behind neural image models. Relative to language modeling, this might be because of the lack of an effective homogeneous architecture like HMMs, which limits scaling up the hidden size further. In particu- lar QPCs, which are based on the quad-graph architecture (Mari et al., 2023), have significantly worse performance compared to HCLT's
Our Monarch-HCLL HCLTs replace the dense sum blocks with Monarch matrices. In particular, we use composition of two Monarch-2 layers (as described in Dao et al. (2022)) as we found this to improve performance for the image datasets. Analogously with HMMs for language modeling, we use the number of floating point operations per pixel as a measure of model efficiency. This is h² for a HCLT and 3/2/2 for a Monarch-HCLT" (the factor of instead of
8

---

Scaling Probabilistic Circuits via Monarch Matrices
6. Related Work
controllable language generation.
Our work is connected to recent efforts in the probabilistic circuits community to find more efficient parameterizations of tensorized circuits. While early implementations of ten- sorized circuits used Kronecker product blocks (Peharz et al., 2020), the recent trend has been to prefer Hadamard product blocks, Loconte et al. (2024a) noted that the composition of Hronecker/Hddamard product blocks with sum blocks can be interpreted as Tucker (Tucker, 1964) / canonical polyaiccl-polyaiic (CP) (Carroll & Chang, 1970) tensor decompositions re- spectively. Our work tackles an orthogonal aspect that it focuses on the sum-to con-to-prodde to--duudctttt connection, which has thus far always been implemented
A Acknowledgements
We thank Anji Liu for assistance with the PyJuice PC library and discussions on image dataset evaluation. The Authors acknowledge the National Artificial Intelligence Research Resource (NAIRR) Pilot and TAMU ACES for contributing to this research result. This work was funded in part by the DARPA ANSR, CODORD, and SAFRON programs under awards FA8750-23-2-0004, HR001 12590089, and HR001 12530141, NSF grant IIS194364,, and gifts from Adobe Research, Cisco Research, and Amazon. Approved for public release; distribution unlimited.
dense matrix.
Circuit multiplication is a fundamental operation which has been studied extensively in the probabilistic circuits literature, with work on deriving theoretical conditions for tractability (Shen et al., 2016; Vergari et al., 2021; Wang et al., 2024; Zhang et al., 2025), use as a building block in various applications of PCs (Choi et al., 2015; Khosravi et al., 2019; Wang& & Kwiatkowska, 2023; Zhang et al., 2024), and improving the expressivity of PCs (Loconte et al., 2024b; 2025; Wang & Van den Broeck, 2025). Our work shows how to explicitly materialize and implement these products efficiently for tensorized circuits.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
References
Ahmed, K., Teso, S., Chang, K.-W., Van den Broeck, G., and Vergari, A. Semantic probabilistic layers for neuro- symbolic learning. Advances Neural Information Pro- cessing Systems, 25:59944-99959, 2022.
There is extensive research structured matrices for im- proving neural network efficiency and scalability. Qiu al. (2024) study scaling laws across structures and introduce Block Tensor- Train (BTT), which generalizes tensor-train (TT) (Oseledets, 2011) and Monarch matrices, achieving competitive performance with lower computational costs. Potapczynski et al. (2024) extend this to BTT-MoE, a Mixture-of-Experts model that sparsifies BTT computation. While our generalized Monarch matrices align with the TT-t-BTTT generalization generaBiTnnnerlizzaatonnnn (rank-1 case), our approach pro- vides practical probabilistic semantics, enabling PCs with Monarch matrices to be formed by multiplying smaller PCs with dense matrices. Additionally, Schanobish et al. (2024) propose structured unrestricted- matrices for efficient Transformer fine-tuning.
Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and Van Den Berg, R. Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems, 2021.
Carroll, J. D. and Chang, J.JJ. Analysis of individual differ- ences in multidimensional scaling via an n-way general- ization of "eckart-young decomposition. Psychometrika, 35(3):283-319) 1970.
Chelba, C., Mikolov, T., Schuster, M., Ge, Q., Brants, T., Kochn, P., and Robinson, T. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint urXiv: 2013.
7. Conclusion
Choi, A., Van den Broeck, G., and Darwiche, A. Tractable
learning for structured probability spaces: case study in learning preference distributions. In Proceedings of 24th International Joint Conference on Artificial Intelligence (IJCAI), volume 2015, pp. 2861-2868. IJCAI-INT JOINT CONF ARTIF INTELL, 2015.
We propose scaling up probabilistic circuits by replacing dense sum blocks with structured sparse matrices. Our ap- proach not only provides theoretical insight by establishing the connection between circuit multiplications and struc- tured matrices but also demonstrates significant empirical improvements, including in modeling performance normal- ized for compute. Future work could examine sparsifying the hidden state representation to overcome the memory bot- tleneck, as well as utilizing the improved scaling to unlock improvements in downstream applications of PCs such as
Choi, Y., Vergari, A., and Van den Broeck, G. Probabilistic circuits: A unifying framework for tractable probabilistic models. oct 2020.
Choi, Y., Dang, M., and Van den Broeck, G. Group fairness by probabilistic modeling with latent fair decisions. In

---

Scaling Probabilistic Circuits via Monarch Matrices
Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 35, pp. 12051 2021.
Kingma, D., Salimans, T., Poole, B., and Ho, Varia-
tional diffusion models. Advances in neural information processing systems, 34:21696-2111 2021.
Dang, M., Khosravi, P., Liang, Y., Vergari, A., and Van den Broeck, G. Juice: julia package for logic and probabilis- tic circuits. In Proceedings of the 35th AAAI Conference on Artificial Intelligence (Demo Track), Feb 2021.
Kingma, D. P. and Dhariwal, P. Glow: Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems, 31, 2018.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint 1rXX.v:31226./ 2013.
Dang, M., Liu, A., and Van den Broeck, G. Sparse prob- abilistic circuits via pruning and growing. In Advances in Neural Information Processing Systems 35 (NeurIPS), dec 2022.
Liu, and Van den Broeck, G. Tractable regularization of probabilistic eircuits. In Advances in Neural Information Processing Systems 34 (NeurIPS), dec 2021.
Dao, T., Gu, A., Eichhorn, M., Rudra, A., and Ré, C. Learn- ing fast algorithms for linear transforms using butterfly factorizations. In International conference on machine learning, pp. 1517-1527. PMLR, 2019.
Liu, A., Zhang, H., and Van den Broeck, G. Scaling up probabilistic circuits by latent variable distillation. In Proceedings of the International Conference Learning Representations (ICLR), may 2023a.
Dao. T., Chen, B., Sohoni, N. S., Desai, A., Poli, M., Gro- gan, J., Liu, A., Rao, A., Rudra, A., and Ré, Monarch: Expressive structured matrices for efficient and accurate training. In International Conference Machine Learn- ing. pp. 419446694999 4721. PMLR, 2022.
Liu, A., Ahmed, K., and Van den Broeck, G. Scaling tractable probabilistic circuits: A systems perspective. In Proceedings of the 41th International Conference on Machine Learning (ICML), jul 2024a
Liu, A., Niepert, M., and Van den Broeck, G. Image in- painting via tractable steering of diffusion models. In Proceedings of the Twelfth International Conference on Learning Representations (ICLR), may 2024b.
Darwiche, A. A differential approach to inference in bayesian networks. Journal of the ACM (JACM), 50 (3):280-305, 2003.
Deng, J., Dong, W., Socher, R., Li, L.-J,, J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Liu, X., Liu, A., Van den Broeck, G., and Liang, Y. Un- derstanding the distillation process from deep generative models to tractable probabilistic circuits. In Proceed- ings of the 40th International Conference on Machine Learning (ICMI.). jul 2023b.
Dinh, L., Sohl-Dickstein, J., and Bengio, S. Density esti- mation using real nvp. In International Conference on Learning Representations, 2017.
Loconte, L., Mari, A., Gala, G., Peharz, R., de Campos, C., Quaeghebeur, E., Vessio, G., and Vergari, A. What is the relationship between tensor factorizations and cir- cuits (and how can we exploit it)? arXiv preprint urXiv:2409.07953 2024a.
Gala, G., de Campos, C., Vergari, A., and Quaeghebeur, E. Scaling continuous latent variable models as probabilistic integral circuits. The Thirty-eighuh Annual Conference Neural Information Processing Systems, 2024.
Loconte, L., Sladek, A., Mengel, S., Trapp, M., Solin, A., Gillis, N., and Vergari, A. Subtractive mixture models via squaring: Representation and learning. In International Conference on Learning Representations (ICLR), 2024b.
Hoogeboom, E., Nielsen, D.. Jaini, P., Forré, P., and Welling, M. Argmax flows and multinomial diffusion: Towards non-autoregresssve language models. arXiv preprint arXiv:2102.053379 3(4):5, 2021.
Loconte, L., Mengel, S., and Vergari, A. Sum of squares cir- cuits. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 19077-19085, 2025.
Hu. E. J. Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022.
Lou, A., Meng, C., and Ermon, S. Discrete diffusion mod- eling by estimating the ratios of the data distribution. In Forty-first International Conference on Machine Learn- ing, 2024.
Khosravi, P., Choi, Y., Liang, Y., Vergari, A., and Van den Broeck, G. On tractable computation of expected pre- dictions. In Advances in Neural Information Processing Systems 32 (NeurIPS), dec 2019.
Maene, J., Derkinderen, V., and Martires, P. Z. D. Klay: Accelerating neurosymbolic ai. arXiv preprint arXiv:2410. /415, 2024.
10

---

Scaling Probabilistic Circuits via Monarch Matrices
Mahoney, M. Text8 dataset, 2011. URL http:// Accessed:
Shen, Y., Choi, A., and Darwiche, A. Tractable operations for arithmetic circuits of probabilistic models. Advances in Neural Information Processing Systems, 29, 2016.
ma montmahoney..o 2024-12-1.
Malvar, H. and Sullivan, G. Ycoeg-r: A color space with rgb reversibility and low dynamic range. ISO/IEC JTCI/S22/WGII and ITU-T SG16 Q, 6, 2003.
Tucker, L. R. The extension of factor analysis to three- dimensional matrices. In Gulliksen, H. and Frederiksen, N. (eds.), Contributions mathematical psychology. pp. 110-127. Holt, Rinchart and Winston, New York, 1964.
Mari, A., Vessio, G., and Vergari, A. Unifying and under- standing overparameterized circuit representations via low-rank tensor decompositions. In The 6th Workshop on Tractable Probabilistic Modeling, 2023.
Vergari, A., Choi, Y., Liu, A., Teso, S., and Van den Broeck, G. compositional atlas of tractable circuit operations for probabilistic inference. In Advances in Neural Infor- mation Processing Systems 34 (NeurIPS), dec 2021.
Meng, C., Zhou, L., Choi, K., Dao, T., and Ermon, S. Butter- flyflow: Building invertible layers with butterfly matrices. In International Conference on Machine Learning, pp- 15360-15375. PMLR, 2022.
Wang. B. and Kwiatkowska, M. Compositional probabilistic and causal inference using tractable circuit models. In International Conference on Artificial Intelligence and Statistics, pp. 9488-9498, PMLR, 2023.
Oseledets, I. V. Tensor-train decomposition. SIAM Journal on Scientific Computing, 33(5):2295-2311 2011.
Papamakarios, G., Nalisnick, E., Rezende, D. J., Mohamed, S., and Lakshminarayanann B. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 2021.
Wang, B. and Van den Broeck, G. On the relationship be- tween monotone and squared probabilistic circuits. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 39, pp. 21026-21034, 2025.
Parker, D. Random Butterfly Transformations with Appli- cations in Computational Linear Algebra. CSD (Series). UCLA Computer Science Department, 1995.
Wang, B., Lyle, C., and Kwiatkowska, M. Provable guar- antees on the robustness of decision rules to causal inter- ventions. In Proceedings of the International Joint Con- ference on Artificial Intelligence. International Joint Con- ferences on Artificial Intelligence Organization, 2021a.
Peharz, R., Gens, R., Pernkopf, F., and Domingos, P. On the latent variable interpretation in sum- cum-product networks. IEEE transactions on pattern analysis and machine intel- ligence, 39(10):203--2444 2016.
Wang, B., Wicker, M. R., and Kwiatkowska, M. Tractable uncertainty for structure learning. In International Con- ference on Machine Learning, pp. 23131-23150 PMLR, 2022.
Peharz, R., Lang, S., Vergari, A., Stelzner, K., Molina, A., Trapp. M., Van den Broeck, G., Kersting, K., and Ghahra- mani, Z. Einsum networks: Fast and scalable learning of tractable probabilistic circuits. In Proceedings of the 37th International Conference on Machine Learning (ICML), jul 2020.
Wang, B., Mauá, D., Van den Broeck, G., and Choi, Y. A compositional atlas for algebraic circuits. Advances in Neural Information Processing Systems, 37:141318- 141355, 2024.
Potapczynski, A., Qiu, S., Finzi, M. A., Ferri, C., Chen, Z., Goldblum, M., Bruss, C. B., Sa, C. D., and Wilson, A. G. Searching for efficient linear layers over a continuous space of structured matrices. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Wang. E., Khosravi, P., and Van den Broeck, G. Probabilis- tic sufficient explanations. In Proceedings of the 30th International Joint Conference on Artificial Intelligence (IJCAI), aug 2021b.
Zečević, M., Dhami, D., Karanam, A., Natarajan, S., and Kersting, K. Interventional sum-product networks: Causal inference with tractable probabilistic models. Ad- vances in neural information processing systems, 34: 15019-15031, 2021.
Qiu, S., Potapczynski, A., Finzi, M. A., Goldblum, M., and Wilson, A. G. Compute better spent: Replacing dense lay- ers with structured matrices. In Forty-first International Conference on Machine Learning, 2024.
Sehanobish, A., Dubey, K. A., Choromanski, K. M., Chowd- hury, S. B. R., Jain, D., Sindhwani, V., and Chaturvedi, S. Structured unrestricted -rank matrices for parameter effi- cient finetuning. In The Thirty-eighth Anmual Conference on Neural Information Processing Systems, 2024.
Zhang, H., Dang, M., Peng, N., and Van den Broeck, G. Tractable control for autoregressive language generation. Proceedings of the 40th International Conference on Machine Learning (ICML). jul 2023.

---

Scaling Probabilistic Circuits via Monarch Matrices
Zhang, H., Kung, P.-N., Yoshida, M., den Broeck, G. V., and Peng, N. Adaptable logical control for large lan- guage models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Zhang, H., Wang, B., Arenas, M., and Van den Broeck, G. Restructuring tractable probabilistic circuits. In Proceed- ings of the 28th International Conference Artificial Intelligence and Statistics (AISTATS), may 2025. URL https:/ /tt///aaiii wwtp///ggg////oooooop 12256.
Ziegler, Z. and Rush, A. Latent normalizing flows for dis- crete sequences. In Proceedings of the 36th International Conference on Machine Learning, 2019.
12

---

Scaling Probabilistic Circuits via Monarch Matrices
A. Interpolating Butterfly and Monarch Matrices
Butterfly matrices (Dao et al., 2019; Meng et al., 2022) are a class of expressive structured matrices. They are constructed as the product of sparse matrices known as butterfly factor matrices (Parker, 1995). this section, we will show how Butterfly matrices (products of sparse Butterfly factor matrices) can be viewed generalized Monarch matrices (products of dense high-dimensional tensors), stated in Theorem 3.4. For ease of exposition, we will describe here Butterfly factor matrices of size D D, where D is power of 2.
Definition A.1 (Butterfly Factor). Given any D = 2d and 1 S S d, butterfly factor matrix B(i, D) is sparse matrix with the following sparsity patter::
(1,, D) has non-zero elements only along the diagonals of the four
D x D submatrices. 2 2
+B(i, D) is a block doico-daagonal matrix with block size
2 where each block B (1, 2 butterfly factor.
B61
B2, 16)
B3, 16)
(4,,66
Figure 8: Illustration of Butterfly factors (Meng et al., 2022).
Examples of butterfly factors for D 16 are shown in Figure 8. Notice that each butterfly factor B(i, D) for i E [d] has 2D nonzero elements, but different sparsity patterns. We now present new interpretation of butterfly factors as a reshaping of a dense (d + di-dmeensional tensor of shape 2 x 2 x ... x 2. In particular, consider the following reshaping:
Definition A.2 (Butterfly Unfurling). Let Ak be a (d + 1)-dimensional 2 x 2 x x 2 tensor. We define the 7th butterfly unfurling of to be the 2d 2d matrix defined by:
B3 = A A d in Ja
(3)
cgfi
where we interpret e.g. jij2-..Jaa the integer in {0, 20-1] represented by the binary string, and of is equal to if je j' and 0 otherwise.
Intuitively, this takes the first and last dimensions of A and inserts them into the 7th bit of the row and column of B respectively; while all other bits of B between the rows and columns are tied. This produces the sparse structure of a Butterfly matrix. For instance, the th butterfly unfurling creates a diagonal matrix with block size 2.
Proposition A.3. The hh butterfly unfurling of a (d + dimeessioo tensor is a butterfly factor matrix Bii, 2d).
Proof. We begin by considering the 1st butterfly unfurling, which is:
B Bo Ajjiamadd Jach
c#1
Notice that the j1.ji are the most significant bits of the row/column respectively, and as each value of (i1,í) corresponds to one of the D/2 x D/ shaped submatrices. The entry B Bj.j....o../i....1 is unless je = j. for all c # 1, L.c. if the entry lies along the diagonal of the submatrix.
Now let us consider the thh butterfly unfurling for i + 1:
Bj = A Aju Ajuii.mmii in
13

---

Scaling Probabilistic Circuits via Monarch Matrices
Notice that, the most significant bits, J1, -..: Ji-1 and Jí Ji..... -i. index the submatrices of size 2 Note that the entry -p is zero unless je = Je for all c 1,-,, 8 nonzero elements. Now let us consider any such submatrix; w.l.o.g. let choose ji ji-1 jí =j'-1 0. Then we have:
thus only the block submatrices along the diagonal can have
Boo.. 00... == Ao,,0..000 4immddad
II
c>1
This submatrix satisfies the sparsity pattern of the 1st butterfly unfolding i dimensional tensor, i.c. is butterfly factor matrix (11, (2)) required.
Notice that, since a butterfly unfurling is just reshaping, given any butterfly factor B(i, D)) we can invert the unfurling operation and get dense tensor Aii,D)
Butterfly matrices Butterfly matrices are constructed as the product of butterfly factors of a given dimension, i.c. B(24) := ((1,24)B(2 2"2, (d, 24). If we instead parameterize these matrices using a butterfly unfurling, we get a product dense tensors with some repeated indices. For example, when d = 2, we get the following expression:
B(4)
EB(1, ()(1,4) 4) B(, ((244) f)
B1,
13.44
-
A(1, j. 0/ A(, 4) jo. 8 38.34
A(1, 4) J) 1(2, 4) j2.js.jo
We can generalize this idea to higher dimensions:
Theorem A.4. Any butterfly matrix B(24) matrix of size 2d x 2d can be written as:
Bi Mim-maaaaaa 12.
A(1,2 ) Je.....ee A(2, 2d)
2( Jem... ... A(d, 2d) 134
where each Ali, 2 is a ... x 2 dense tensor of dimension d + 1.
Proof. Essentially this is replacing each butterfly factor with its furled version, and by inspection of the unfurling formula in Equation 3. For each of the butterfly factors, let us write:
Bii, 1o 5i4t 20th = A(i, 2"i ,/ "h 1i Si 1f ((
II
GID
5
where we have attached superscript to each index indicate it is associated with the 7th butterfly factor; in particular, j! je and je = jd4o-
(i+ whenever e + i. By iterative application of this rule, one can see that the indices of A(i, 2d) correspond to ji+d (up to some permutation). For example, the first j Ji. In general, je = ja+e whenever e < i, and ((()) = jo whenever 2 i.
Notice that the delta functions require (for non-zero entry) that (0)
index Ji J; (1-1)
Interpreting jaji as a matrix Kjssia we get sequence of tensor contractions matching the structure of the generalized Monarch matrices. By generalizing this idea to larger d, we can encode and execute any butterfly matrix-vector multiplication series of dense tensor contractions.
B. Additional Experimental Results and Details
B.I. Additional experiment setup for text8
We train Monarch-HMM using two-layer Monarch matrices and a 219 hidden state for 20 epochs. Following prior works (Zhang et al., 2024), optimization performed using the stochastic EM algorithm with a mini batch of 4096 and a linearly decaying learning rate from to
14

---

Scaling Probabilistic Circuits via Monarch Matrices
B.2. Additional results for text8
Hidden Size
Monarch-2
Monarch-3
HMM
Monarch-4
FLOPs
BPC
FLOPs
BPC
FLOPs
BPC
FLOPs
BPC
2 4096
524288
2.041
196608
2.135
16777216
1.908
131072
2.188
213 8192
1572864
1.926
524288
2.013
67108864
1.786
327680
2.065
214 16384
4194304
1.824
1310720
1.895
268435456
1.726
786432
1.948
215 32768
12582912
1.747
3145728
1.814
1073741824
1.687
1835008
1.853
2 65536
33554432
1.690
8388608
1.741
2 = 131072
100663296
1.637
21 262144
268435456
1.609
2 524288
805306368
1.578
Table 6: We report FLOPs and test set BPC of HMM and variying Monarch structures (Monarch-2, Monarch-3, and Monarch-4) for hidden size ranging from 2 to 2 on dataset text8.
B.3. Additional experimental setup for image modeling
We use hidden Chow-Liu trees (HCLT) (Liu & Van den Broeck, 2021) to define the variable decomposition (viree) the PC. We train all models using stochastic EM. In particular, we use cosine learning rate decay scheduler over the course of the entire training. The mini-batch size M 20000 and number of epochs E = 20 used in all experiments were chosen based on an initial hyperparameter search in the range M E [1000, 5000, 20000, 60000] and E E [5, 10, 20).
Color Transforms The original ImageNet data has three color channels R, G, B, each taking integer values in [0, 255). We apply a lossless YCoCg-R transform (Malvar & Sullivan, 2003) to this data to obtain three transformed color channels Y, Co, Cg, given as follows in integer arithmetic:
Co = R - B
Imp = B + [Co/2]
Cg = G tmp
Y tmp + [Cg/2]
where Y and Cg take integer values in 255] and Co takes integer values in 2-255, 255J. This transform is exactly invertible (lossless) and as such likelihoods on this transformed data are comparable to likelihoods on the RGB dataset.
To compare against prior PC results, we also report results using lossy YCoCg transform defined as follows. Firstly, the data is normalized to take values the range [0, 1]:
r R R25::; 9 = G/255; B/255
Then, the following linear transformation is applied:
co
tmp b + co/2
cg 9 lmp
tmp * 2 + cg
The variables y. co,cg eg take values in [-1,1]. To convert back to categorical data, quantize the interval [-1,1] uniformly into 256 bins and convert y, co, eg into their quantized versions Y, Co, Cg which take integer values in [0, 255). This transformation from R. G, B to Y, Co, Cg is not invertible and so likelihoods on this dataset cannot be compared to the original RGB dataset.
Image Patches We train PCs with 192 variables, modeling 8 x aligned image patches (with 3 color channels). This data obtained by splitting each image from the the ImageNet32 (resp. ImageNet64) dataset into (resp. 64) of these patches.
15


Reviews:
[{"review_id": "UE3wFvfAoL", "reviewer": "Reviewer_Gwtt", "summary": "This paper introduces a novel method to scaling Probabilistic Circuits (PCs) by replacing dense matrices in sum nodes with Monarch matrices which is a type of structured sparse matrices constructed by Kronecker products. The key idea is to leverage the sparsity and structure of Monarch matrices to reduce memory and computation costs which is demonstrated on generative modeling benchmarks. Furthermore, this paper provides a theoretical foundation on the connection between Monarch matrices and circuit multiplication.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback.\n\n```\nThe paper does not compare Monarch matrices to alternative structured representations like Block Tensor-Train (BTT) decomposition [1] or Toeplitz-like structured layers. \n```\n\nOur study is not only limited to the Monarch matrices defined in [1]. Our construction of Monarch matrices as product of PCs with dense layers naturally generalizes the definition of Monarch matrices: the construction in [2] corresponds to the special case of multiplying **2** dense PCs, and we consider Monarch structures obtained by multiplying **$k$** dense PCs, showing empirical results for k = 2, 3, 4.\n\nSecondly, in terms of BTTs, we would like to note that BTTs of rank 1 correspond exactly to the Monarch matrices proposed by [2], and [1] made an important observation that BTTs of rank higher than 2 do not lead to better scaling behaviors while increasing higher memory cost. \n\nAs a quick corroboration for PCs, we studied PCs with BTT layers on the ImageNet32 (lossless) dataset. As shown in the Figure at https://anonymous.4open.science/r/MonarchRebuttal-D0D3/a.pdf, we compare Monarch PCs with BTT layers of rank 2, 4, and 8: BTT-2 has a similar scaling curve with Monarch, and the scaling curves of BTT-4 and BTT-8 get worse as rank increases, echoing the findings of [1].\n\nFinally, we would like to note that prior work on PCs has not considered using non-dense matrices, and our work opens up a new way of scaling PCs with great promise as illustrated by empirical achievements. We believe that investigating Toeplitz-like structured layers or other efficient representations would be excellent topics for future work. \n\n```\nThe paper highlights performance improvements but does not discuss potential failure cases.\n```\n\nWe agree that if we fix the hidden size of PCs, replacing the dense linear layers with Monarch matrices indeed reduces the expressive power, which is exactly shown in Figure 5, where we measure the performance of PCs with varying hidden sizes and different structures of linear layers. We can see that when fixing the hidden sizes, dense matrices always perform better, but at the same time, as the hidden size grows: (1) the performance gap between the dense PCs and the Monarch PCs diminishes rapidly while (2) the FLOPs gap between dense PCs ($O(h^2)$) and Monarch PCs($O(2h^{3/2}$) grows rapidly; and these two factors together lead to the large gap between the scaling curves of the dense PCs and Monarch PCs shown in Figure 4. \n\nFrom a theory perspective, one could ask whether there exists any distribution represented as compact dense PCs which would require exponential size to represent as a Monarch PC. The answer is no, as we have shown that a Monarch circuit can be interpreted as a relaxed version of the product of two circuits. Thus, any dense PC can be simulated by multiplying it with a PC representing the uniform distribution. \n\nWe will incorporate this discussion into our revision and highlight that our study on the expressive power of Monarch matrices in PCs is focused on the empirical side. \n\n```\nThe experiments focus on generative modeling benchmarks, but PCs are also used in other applications … Monarch-based PCs in non-generative tasks (e.g., probabilistic reasoning) would further strengthen the paper's impact.\n```\n\nPrior work has shown that the better the PCs model the desired distributions, the better they perform on downstream applications. For example, one line of work [3, 4] on applying PCs for controllable text generation from LLMs shows that the better the PCs approximate the LLMs, the higher the text generation quality, as shown in Figure 3 of [3] and Table 1 of [4]. Similarly, in the application of PCs to group fairness, Figure 3 and Table 1 of [5] show that the PC model achieving the best likelihood also leads to better classification accuracy and lower discrimination scores. Based on these findings, we believe that Monarch PCs, by achieving significantly better generative modeling performance, should naturally give rise to better downstream performance.\n\nHence, instead of testing our state-of-the-art PCs on downstream applications, we devote our effort, as well as the limited computation resources, to ablation studies that may help people better understand the behavior of PCs with Monarch matrices, such as: what is a good number of dense PCs to multiply to form Monarch structures, how does initialization of PC parameters via circuit products benefit training, is the hidden states of PCs also sparse and etc.\n\n[1] Qiu et al. “Compute Better Spent: Replacing Dense Layers with Structured Matrices” ICML 2024\n\n[2] Dao et al. “Monarch: Expressive Structured Matrices for Efficient and Accurate Training” ICML 2022\n\n[3] Zhang et al. “Tractable Control for Autoregressive Language Generation” ICML 2023\n\n[4] Zhang et al. “Adaptable Logical Control for Large Language Models” NeurIPS 2024\n\n[5] Choi et al. “Group Fairness by Probabilistic Modeling with Latent Fair Decisions” AAAI 2021", "claims_and_evidence": "The majority of the claims in the paper are supported by clear evidence, such as those related to computational efficiency, performance improvements.", "methods_and_evaluation_criteria": "The methods and evaluation criteria make sense for the problem and application at hand, but evaluating on larger datasets and broader tasks could further strengthen the claims.", "theoretical_claims": "The theoretical framework is sound, but further formalization would strengthen the paper's contributions.", "experimental_designs_or_analyses": "Yes, the experiments support the claims made in the paper, but additional analyses could further strengthen the results.", "relation_to_broader_scientific_literature": "This paper builds upon and extends multiple streams of prior research in PCs, structured matrices, and efficient model scaling.", "essential_references_not_discussed": "I think this paper has discussed enough relevant work in this area and provided comprehensive references to previous studies to illustrate its main contributions.", "other_strengths_and_weaknesses": "strength:\n1. The paper introduces Monarch matrices as a structured sparse parameterization for sum blocks in PCs, bridging the gap between tractability and scalability.\n2. The paper provides solid theoretical justification for using Monarch matrices, linking them to circuit multiplication and structured sparsity.\n3. Comprehensive experiments on generative modeling benchmarks (Text8, LM1B, ImageNet) demonstrate state-of-the-art performance with reduced computational cost (FLOPs).\n\nweakness:\n1. The paper does not compare Monarch matrices to alternative structured representations like Block Tensor-Train (BTT) decomposition [1] or Toeplitz-like structured layers. How do Monarch matrices perform relative to these alternatives in terms of efficiency and expressiveness?\n2. The paper highlights performance improvements but does not discuss potential failure cases. For example, does the structured sparsity of Monarch layers introduce expressiveness limitations compared to dense PCs?\n3. The experiments focus on generative modeling benchmarks, but PCs are also used in other applications, such as causal inference, fairness, and tractable reasoning. Evaluating Monarch-based PCs in non-generative tasks (e.g., probabilistic reasoning) would further strengthen the paper's impact.\n\n[1] Qiu et al. Compute better spent: Replacing dense layers with structured matrices. ICML 2024.", "other_comments_or_suggestions": "I believe the first two weaknesses highlight the main issues. As for the third weakness, since it involves additional experiments, if there isn’t enough time to conduct them, a detailed verbal explanation of the method’s generalization would be helpful.", "questions_for_authors": "Please check my questions in the weakness section.\n\n==Post Rebuttal==\n\nI think the empirical evaluation of this paper can be more sufficient. I am ok if AC and the other reviewers decide to accept this paper.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fRFELVgfJR", "reviewer": "Reviewer_NJT4", "summary": "Despite many advantages of probabilistic circuits (PC), their implementations are often difficult due to computational burden, even with block structures. In this paper, the authors proposed an alternative method that replaces dense sum blocks with Monarch matrices, and the method significantly reduce the memory and computation costs. Ultimately, the authors claimed that this method  significantly bridges the gap between highly tractable models (probabilistic circuits) and less tractable models (diffusion models).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your encouraging feedback. Please feel free to follow up if you have any questions.", "claims_and_evidence": "Claim 1. Monarch matrices are efficient.\nThis claim is well-supported via theoretical arguments on Page 3 by showing the improvement from $O(m^4)$ to $O(m^3)$ edges.\n\nClaim 2. The Monarch models have outstanding performance.\nThe performance is supported via empirical tests, shown in Tables 1 and 5.\n\nClaim 3. The Monarch models have better scaling behavior.\nThe performance is supported via empirical tests, shown in Figures 4 and 7.", "methods_and_evaluation_criteria": "This paper has extensive justifications in theory. I am not knowledgable in experiments, but the benchmark datasets are standard for generative models. Using FLOPs to measure efficiency is a commonly accepted choice.", "theoretical_claims": "I have checked all definitions and linear algebraic results: Discussions on page 3, and the theorems and proofs in Appendix A.", "experimental_designs_or_analyses": "The results (in tables and figures) support the paper's claims, but I am not qualified to comment the empirical part.", "relation_to_broader_scientific_literature": "This paper extensively discusses its connections with relevant fields, including previous works on probabilistic circuits (in particular block parameterizations) and relevant background in linear algebra (butterfly matrices and Monarch matrices). They also discussed other types of models such as the diffusion model and flow-based models, and gave a high level picture on the status of PCs.", "essential_references_not_discussed": "I am not qualified to comment on this section.", "other_strengths_and_weaknesses": "Strengths:\n1. The paper is well-motivated, studies an important problem, and has rigorous theoretical justifications.\n2. The claim on bridging performance gaps is justified and promising to future research in this direction.\n3. The empirical comparisons take care of a wide range of metrics.", "other_comments_or_suggestions": "I do not have any other comments or suggestions.", "questions_for_authors": "I do not have any other questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kBTETAsyMA", "reviewer": "Reviewer_AzZ2", "summary": "This paper proposes a novel parameterization for probabilistic circuits (PCs) to improve their scalability, using structured sparse matrices called Monarch matrices. By replacing dense matrices in sum blocks of PCs with Monarch matrices, the proposed methods can reduce computational costs and allow larger scale PC training. The authors conducted experiments on text and image datasets. To my knowledge, the proposed idea is novel and the experimental results look compelling (at least compared to other PC methods).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful feedback.\n\n```\nI find the introduction of Butterfly matrices in line 205 a bit redundant and irrelevant. Why are they mentioned in the methods section rather than the related work? Have they been applied in PCs?\n```\n\nOne of the main contributions of this work is identifying the connection between structured Monarch matrices and the multiplication of PCs: the product of two dense PCs gives exactly the Monarch matrix while the product of $d$ dense PCs of hidden size 2 gives the Butterfly matrix.\n\nButterfly matrices are a well-known type of structured matrix which require only $O(h\\log h)$ compute/memory (compared with $O(h^{3/2})$ for Monarch); however, despite their theoretical efficiency, Butterfly matrices are known to be less efficient on modern GPUs (due to their sparsity). \n\nOur construction of Monarch matrices as circuit product immediately implies that (i) one can interpolate between Butterfly and Monarch matrices through the generalized Monarch matrix; and (ii) this corresponds precisely to the multiplication of k PCs, for integers k >= 2. This then motivated us to perform an empirical study for $k = 2, 3, 4$ to see if one could achieve better scaling using the generalized Monarch matrix.\n\nWe will incorporate this explanation to motivate this section more clearly.\n\n```\nThe memory saveup shown in Table 3 is not significant, which limits the extension of the method.\n```\n\nWe would like to note that Table 3 shows the **training** memory consumption with a reasonably large batch size $B = 128$ and the # of variables $n = 256$. At **inference** time (i.e. no gradient computation), Monarch PCs can be significantly more memory-efficient: with $h = 2^{18}$, a dense PC would still use $\\approx 256$ GB while a Monarch PC (2-layer) would only use $\\approx 1$ GB. That is, even when training a large Monarch PC is relatively memory-consuming, at inference time, researchers can still apply large powerful Monarch PCs with very little cost.\n\nIn addition, the seemingly inefficient memory consumption of Monarch PCs during training is due to the need of gradient computation, which can be optimized by various well-known techniques such as gradient checkpointing, mixed-precision training, etc. Further, as shown at the end of Sec. 5.2, the hidden states of large Monarch PCs are actually sparse, motivating future research to exploit such sparsity to reduce the memory consumption of training large Monarch PCs.\n\n```\nThere is still a large performance gap between the proposed method and other SOTA methods (e.g. diffusion models).\n```\n\nWe agree that there is still a large gap between state-of-the-art PCs and the other deep generative models, but at the same time we would like to highlight that we have already scaled PCs to a degree far beyond existing tractable probabilistic models, while the SotA generative models are substantially less tractable compared to PCs (e.g. none of transformers, normalizing flows or diffusion models allow for tractable computation of marginal probabilities). Closing the performance gap between PCs and other deep generative models is precisely the motivation of our research.\n\n```\nI don't think the sampling time comparison throughout the paper is convincing. \n```\nThe sampling time comparison in Table 4 is not meant to argue that PCs are superior to discrete diffusions in terms of efficiency but more of a sanity check: for readers who are not particularly familiar with PCs, we just want to show that our models can be efficiently implemented. More specifically, for example, the D3PM Uniform model achieves $\\leq 1.61$ bpc with 1000 diffusion steps (3.6s) and $\\leq 1.79$ bpc with 20 diffusion steps (0.077s), where the latter number is more comparable to that of Monarch HMMs. In our revision, we will add both D3PM results to the table and carefully rephrase our statement about the runtime of our models. Thank you for your suggestion.\n\n```\nWhy is the YCoCg transform used in the image modeling experiments?\n```\nTable 5 shows the experiment results on both images transformed via YCoCg and transformed via YCoCg-R (the Lossless column). Since the YCoCg-R transform is lossless, the BPDs are directly comparable to the original RGB dataset. The two strongest PC baselines [3, 4] only reported their results on YCoCg transformed images (which is not reported in those papers, but is confirmed in [5] and the software implementations); we choose to report results for both settings for a fair comparison to [3, 4] and to models trained on RGB datasets, following the practice of [5]. \n\n[3] Liu et al. \"Scaling Up Probabilistic Circuits by Latent Variable Distillation.\" ICLR 2023\n\n[4] Liu et al. \"Understanding the distillation process from deep generative models to tractable probabilistic circuits.\" ICML 2023\n\n[5] Gala et al. \"Scaling Continuous Latent Variable Models as Probabilistic Integral Circuits.\" NeurIPS 2024", "claims_and_evidence": "The claims made in the paper are well supported by the conducted experiments. The authors are honest about the remaining gap between the proposed method and state-of-the-art generative modeling methods.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria (efficiency in terms of FLOPs and performance in terms of test log-likelihood) are reasonable to the problem studied in this paper.", "theoretical_claims": "I checked the derivations for equations (1) and (2) and they look correct to me. But the flow from the Monarch matrix operations to the construction of a PC is not clear to me by reading sections 3 and 4.", "experimental_designs_or_analyses": "The experiments are reasonably designed and analyzed for the studied problem of the paper.", "relation_to_broader_scientific_literature": "The key contributions of the paper build upon [1] and the implementations and training of the proposed method rely on [2] and [3].\n\n\n[1] Dao et al. Monarch: Expressive structured matrices for efficient and accurate training. ICML 2022.\n\n[2] Peharz et al. Einsum Networks: Fast and Scalable Learning of Tractable Probabilistic Circuits. ICML 2020.\n\n[3] Liu et al. Tractable regularization of probabilistic circuits. NeurIPS 2021.", "essential_references_not_discussed": "To my knowledge, no.", "other_strengths_and_weaknesses": "Strengths:\n- This paper advances scalable PC training and shows empirically compelling results, which is critical to the PC community.\n\nWeaknesses:\n- The memory saveup shown in Table 3 is not significant, which limits the extension of the method.\n- There is still a large performance gap between the proposed method and other SOTA methods (e.g. diffusion models). In image modeling experiments, the method is weaker than flow-based models which have been proposed years ago.\n- I don't think the sampling time comparison throughout the paper is convincing. According to the paper, the baseline results are often taken from the original papers. The comparison in such case is not fair under varied implementation details and hardware setup.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "- I find the introduction of Butterfly matrices in line 205 a bit redundant and irrelevant. Why are they mentioned in the methods section rather than the related work? Have they been applied in PCs?\n- Why is the YCoCg transform used in the image modeling experiments? Is it necessary in the setup of this paper? If not, how does the proposed method perform using the original RGB images?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6Znp1rB2LY", "reviewer": "Reviewer_zQvJ", "summary": "This paper replaces dense matrices with sparse Monarch matrices, reducing the computation cost and maintaining accuracy.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. \n\n```\nTo be honest, I do not know much about hybrid models since I have only read several papers such as Mamba. However, I think the main issue of the hybrid model is that we cannot scale up the size of the hybrid model. That is my understanding. Consequently, almost no leading companies choose to train and deploy hybrid models.\n```\n\nTo clarify, in this work we investigate scaling of probabilistic circuit models. As opposed to hybrid models (and indeed transformers), probabilistic circuits have the key property of tractability, which enables efficient computation of quantities relating to the probability distribution such as marginals. Tractability has been used to achieve state-of-the-art results in applications such as controllable generation [1] and image inpainting [2], among others, beating transformer and diffusion-based models. Thus, we respectfully believe that research on alternative directions for generative modeling is valuable and this should not be viewed as a weakness of the work.\n\n\n```\nAs for the paper, could you show me more results if you increase the length of the example from 256 to 1024?\n```\n\nWe appreciate the reviewer’s request for further experiments, but we do not have the computational resources to run the requested experiments during the rebuttal period. Please note that we chose the sequence lengths and ImageNet downscaled resolutions to match with those considered by our baselines for more direct comparison; in particular, our method shows significant and consistent advantage over the PC baselines. Given the existing range of experiments we do not believe that there would be any change in the qualitative conclusions. \n\n\n```\nAlso, could you show me the model size in the Table?\n```\n\nThe PCs for image modeling have around 1.5B parameters. The Monarch HMM on text8 with hidden size $h=2^{19}$ has 0.75B parameters and the Monarch HMM on lm1b has 4.75B parameters. We will update the Tables with model sizes.\n\n\n[1] Zhang et al. “Adaptable Logical Control for Large Language Models” NeurIPS 2024 \n\n[2] Liu et al. “Image Inpainting via Tractable Steering of Diffusion Models” ICLR 2024", "claims_and_evidence": "Yes, the claims made in the submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The proposed methods may be suitable for the problem or application at hand.", "theoretical_claims": "Yes, I have checked the proofs for correctness. Any minor issues present should not impede understanding of the overall paper.", "experimental_designs_or_analyses": "I have checked the soundness/validity of any experimental designs or analyses.", "relation_to_broader_scientific_literature": "This paper presents several interesting results.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n1. The paper is exceptionally well-written and easy to follow.\n2. This paper presents some theoretical analysis.\n3. This paper replaces dense matrices with sparse Monarch matrices to reduce memory and computation cost.\n\nWeaknesses:\n1. To be honest, I do not know much about hybrid models since I have only read several papers such as Mamba. However, I think the main issue of the hybrid model is that we cannot scale up the size of the hybrid model. That is my understanding. Consequently, almost no leading companies choose to train and deploy hybrid models.\n2. As for the paper, could you show me more results if you increase the length of the example from 256 to 1024?\n3. In addition, could you show me the results evaluated over imagenet-256?", "other_comments_or_suggestions": "1. Enhancing Tables 4 and 5 with additional results would further improve the completeness of the paper.\n2. Also, could you show me the model size in the Table?", "questions_for_authors": "See strengths and weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces a novel parameterisation for the sum blocks in probabilistic circuits using Monarch matrices.  Monarch matrices allow for sparsity and tensorized operations, and thus their use significantly improves the scalability of probabilistic circuits.  The reviews were rather bi-modal with two accepts and two weak rejects (4, 4, 2, 2).  One of the weak reject reviews was rather short and signaled low confidence or even misunderstanding about the subject area.  Reviewers found that the paper was "exceptionally" well written, easy to follow, well-motivated, theoretically sound and empirically strong.  One reviewer would have liked to see comparisons to other parameterizations such as Toeplitz matrices or Block Tensor-Train matrices and evaluation of probabilistic circuits on other applications beyond generative modeling.  Another reviewer pointed out that while the results seem to be good for probabilistic circuits, they still lag behind other common generative models. 
Overall it seems that the work is novel, well-motivated and convincingly enables better scaling of probabilistic circuits.   Therefore, the recommendation is to accept.
Decision: Accept (poster)