RiteeshGrandhi/sciclaimeval-shared-task huggingface dataset:
Introduction
Probabilistic circuits (PCs) [ 44 , 3 ] are a unifying framework to abstract from a multitude of tractable probabilistic models.
The key property that separates PCs from other deep generative models such as flow-based models [ 31 ] and VAEs [ 19 ] is their tractability . It enables them to compute various queries, including marginal probabilities, exactly and efficiently [ 45 ] .
Therefore, PCs are increasingly used in inference-demanding applications such as enforcing algorithmic fairness [ 2 , 4 ] , making predictions under missing data [ 6 , 18 , 23 ] , data compression [ 26 ] , and anomaly detection [ 13 ] .

Recent advancements in PC learning and regularization [ 40 , 25 ] , and efficient implementations [ 33 , 30 , 8 ] have been pushing the limits of PC’s expressiveness and scalability such that they can even match the performance of less tractable deep generative models, including flow-based models and VAEs.
However, the performance of PCs plateaus as model size increases. This suggests that to further boost the performance of PCs, simply scaling up the model size does not suffice and we need to better utilize the available capacity.

We discover that this might be caused by the fact that the capacity of large PCs is wasted.
As shown in Figure 1 , most parameters in a PC with 2.18M parameters have close-to-zero values, which have little effect on the PC distribution.
Since existing PC structures usually have fully-connected
parameter layers [ 25 , 36 ] , this indicates that the parameter values are only sparsely used.

In this work, we propose to better exploit the sparsity of large PC models by two structure learning primitives — pruning and growing .
Specifically, the goal of the pruning operation is to identify and remove unimportant sub-networks of a PC. This is done by quantifying the importance of PC parameters w.r.t. a dataset using circuit flows , a theoretically-grounded metric that upper bounds the drop of log-likelihood caused by pruning. Compared to L1 regularization, the proposed pruning operator is more informed by the PC semantics, and hence quantifies the global effects of pruning much more effectively.
Empirically, the proposed pruning method achieves a compression rate of 80-98% with at most 1% drop in likelihood on various PCs.

The proposed growing operation increases the model size by copying its existing components and injecting noise. In particular, when applied to PCs compressed by the pruning operation, growing produces larger PCs that can be optimized to achieve better performance. Applying pruning and growing iteratively can greatly refine the structure and parameters of a PC. Empirically, the log-likelihoods metric can improve by 2% to 10% after a few iterations.
Compared to existing PC learners as well as less tractable deep generative models such as VAEs and flow-based models, our proposed method achieves state-of-the-art density estimation results on image datasets including
MNIST, EMNIST, FashionMNIST, and the Penn Tree Bank language modeling task. 1 1 1 Code and experiments are available at https://github.com/UCLA-StarAI/SparsePC .
Probabilistic Circuit Model Compression via Pruning
Figure 1 shows that most parameters in a large PC are very close to zero. Given that these parameters are weights associated with mixture (sum unit) components, the corresponding edges and sub-circuits have little impact on the sum unit output. Hence, by pruning away these unimportant components, it is possible to significantly reduce model size while retaining model expressiveness.
Figure 3(b) illustrates the result of pruning five (red) edges from the PC in Figure 3(a) . Given a PC and a dataset, our goal is to efficiently identify a set of edges to prune, such that the log-likelihood gap between the pruned PC
and the original PC on the given dataset is minimized.
Pruning by parameters.
The parameter value statistics in Figure 1 suggest that a natural criterion is to prune edges by the magnitude of their corresponding parameter. This leads to the eParam (edge parameters) heuristic, which selects the set of edges with the smallest parameters. However, edge parameters themselves are insufficient to quantify the importance of inputs to a sum unit in the entire PC’s distribution. The parameters of a sum unit are normalized to be 1 so they only contain local information about the mixture components.
Specifically, \theta_{c|n} merely defines the relative importance of edge (n,c) in the conditional distribution represented by its corresponding sum unit n , not the joint distribution of the entire PC.
Figure 4(a) illustrates what happens when the edge with the smallest parameter is pruned from the PC in Figure 2 .

However, as shown in Figure 4(b) , pruning another edge delivers better likelihoods as it accounts more for the “global influence” of edges on the PC’s output. This global influence is highly related to the probabilistic “circuit flow” semantics of PCs. We will introduce circuit flows later in this section, along with their corresponding heuristics eFlow . Before that, we first introduce an intermediate concept based on the notion of generative significance of PCs.
Pruning by generative significance.
A more informed pruning strategy needs to consider the global impact of edges on the distribution represented by the output of the PC. To achieve this,
instead of viewing the distribution {p}_{\mathcal{C}} in a feedforward manner following Equation 1 ,
we quantify the significance of a unit or edge by the probability that it will be “activated” when drawing samples from the PC.
Indeed, if the presence of an edge is hardly ever relevant to the generative sampling process, removing it will not significantly affect the PC’s distribution.

Algorithm 1 shows how to draw samples from a PC distribution through a recursive implementation: (1) for an input unit n defined on variable X (line 1 ), the algorithm randomly samples value x according to its input univariate distribution; (2) for a product unit (line 1 ), by decomposability its children have disjoint scope, thus we draw samples from all input units and then concatenate the samples together; (3) for a sum unit n (line 1 ), by smoothness its children have identical scope, thus we first randomly sample one of its input units according to the categorical distribution defined by sum parameters \{\theta_{c|n}:c\in\mathsf{in}(n)\} , and then sample from this input unit recursively. Besides actually drawing samples from the PC, we can also compute the probability that n will be visited during the sampling process. This provides a good measure of the importance of unit n to the PC distribution as a whole, which we define as the top-down probability .

The top-down probability of each unit n in a PC with parameters \bm{\theta} is defined recursively as follows, assuming alternating sum and product layers:

{q}(n;\bm{\theta}):=\begin{cases}1&if $n$ is the root unit,\ \sum_{m\in\mathsf{out}(n)}{q}(m;\bm{\theta})&if $n$ is a sum unit,\ \sum_{m\in\mathsf{out}(n)}\theta_{n|m}\cdot{q}(m;\bm{\theta})&if $n$ is a product unit,\end{cases}

where \mathsf{out}(n) are the units that take n as input in the feedforward computation.
Moreover, the top-down probability of a sum edge (n,c) is defined as {q}(n,c;\bm{\theta})=\theta_{c|n}\cdot{q}(n;\bm{\theta}) .

The top-down probability of the root is always 1; a product unit passes its top-down probability to all its inputs, and a sum unit distributes its top-down probability to its inputs proportional to the corresponding edge weights. Therefore, the top-down probability of a non-root unit is summing over all probabilities it receives from its outputs.

The top-down probability of all PC units and sum edges can be computed in a single backward pass over the PC’s computation graph. Following the intuition that the top-down probability defines the probability that units will be visited during the sampling process, pruning edges with the smallest top-down probability constitutes a reasonable pruning strategy.
Pruning by circuit flows.
The top-down probability {q}(n;\bm{\theta}) represents the probability of reaching unit n in an unconditional random sampling process. Despite its ability to capture global information of PC parameters, the top-down probability is not tailored to a specific dataset.
Therefore, to further utilize the dataset information, we can measure the probability of reaching certain units/edges in the sampling process conditioning on some instance \bm{x} being sampled .
To bridge this gap, we define circuit flow as a sample-dependent version of the top-down probability.

For a given PC with parameters \bm{\theta} and example \bm{x} , the circuit flow of unit n on example \bm{x} is the probability that n will be visited during the sampling procedure conditioned on \bm{x} being sampled. This can be computed recursively as follows, assuming alternating sum and product layers:

\operatorname{F}_{n}(\bm{x})=\begin{cases}1&if $n$ is the root unit,\ \sum_{m\in\mathsf{out}(n)}\operatorname{F}_{m}(\bm{x})&if $n$ is a sum unit,\ \sum_{m\in\mathsf{out}(n)}\frac{\theta_{n|m}\cdot{p}_{n}(\bm{x})}{{p}_{m}(\bm{x})}\cdot\operatorname{F}_{m}(\bm{x})&if $n$ is a product unit.\end{cases}

Similarly, the edge flow \operatorname{F}_{n,c}(\bm{x}) on sample \bm{x} is defined by \operatorname{F}_{n,c}(\bm{x})=\theta_{c|n}\cdot{p}_{c}(\bm{x})/{p}_{n}(\bm{x})\cdot\operatorname{F}_{n}(\bm{x}) . We further define \operatorname{F}_{n,c}(\mathcal{D})=\sum_{\bm{x}\in\mathcal{D}}\operatorname{F}_{n,c}(\bm{x}) as the aggregate edge flow over dataset \mathcal{D} .

Effectively, we can think of \theta_{n|m}^{\bm{x}}:=\theta_{n|m}\cdot{p}_{n}(\bm{x})/{p}_{m}(\bm{x}) as the posterior probability of component n in the mixture of sum unit m conditioned on observing sample \bm{x} . Then, circuit flow is the top-down probability under this \bm{\theta}^{\bm{x}} reparameterization of the circuit: \operatorname{F}_{n}(\bm{x})={q}(n;\bm{\theta}^{\bm{x}}) and \operatorname{F}_{n,c}(\bm{x})={q}(n,c;\bm{\theta}^{\bm{x}}) .

Circuit flow \operatorname{F}_{n}(\bm{x}) defines the probability of reaching unit n in the top-down sampling procedure of Algorithm 1 , given that the sampled instance is \bm{x} . Therefore, edge flow \operatorname{F}_{n,c}(\bm{x}) is a natural metric of the importance of edge (n,c) given \bm{x} .
Intuitively, the aggregate circuit flow measures how many expected samples “flow” through certain edges.
We write eFlow to refer to the heuristic that prunes edges with the smallest aggregate circuit flow.
Empirical Analysis.
Figure 5(a) compares the effect of pruning heuristics eParam , eFlow , as well as an uninformed strategy, prune randomly, which we denote as eRand . It shows that both eParam and eFlow are reasonable pruning strategy, however, as we increase the percentage of pruned parameters, eFlow has less log-likelihoods drop compared with eParam . Using eFlow heuristics we can pruning up to 80% of the parameters without much log-likelihoods drop. As shown in Figure 5(b) , the parameter distribution is more balanced after pruning compared to Figure 1 , indicating a higher significance of each edge. Section 6 will provide more empirical results. Before that, we first theoretically verify the effectiveness of the eFlow heuristic in the next section.
Scalable Structure Learning
The pruning operator improves two aspects of PCs. First, as shown in Figure 5(b) , model parameters are more balanced after pruning. Second, pruning removes sub-circuits with negligible contributions to the model’s distribution.
If we treat PCs as hierarchical mixtures of components, pruning can be regarded as an implicit structure learning step that removes the “unimportant” components for each mixture.
However, since pruning only decreases model capacity, it is impossible to get a more expressive PC than the original one.
To mitigate this problem, we propose a growing operation to increase the capacity of a PC by introducing more components for each mixture.
Pruning and growing together define a scalable structure learning algorithm for PCs.
Growing.
Growing is an operator that increases model size by copying its existing components and injecting noise. As shown in Figure 3 , after applying the growing operation on the original PC in Figure 3(b) , we can get a new grown PC as in Figure 7 . Specifically, the growing operation is applied to units, edges, and parameters respectively: (1) for units, growing operates on every PC unit n and creates another copy n^{\mathtt{new}} ; (2) for edges, the sum edge (n,c) from the original PC (Figure 3(b) ) are copied three times to the grown PC (Figure 7 ): from new parent to new child (n^{\mathtt{new}},c^{\mathtt{new}}) , from old parent to new child (n,c^{\mathtt{new}}) , and from new parent to old child (n^{\mathtt{new}},c) ; product edges are added to connect the copied version of a product unit and its copied inputs;
(3) a new parameter \theta_{c|n}^{\mathtt{new}} is a noisy copy of an old parameter \theta_{c|n} , that is \theta_{c|n}^{\mathtt{new}}\leftarrow\epsilon\cdot\theta_{c|n} where \epsilon\sim\mathcal{N}(1,\sigma^{2}) and \sigma^{2} controls the Gaussian noise variance.
Gaussian noise is added to the copied parameters to ensure that after we apply the growing operation, parameter learning algorithms can find diverse parameters for different copies.
After a growing operation, the PC size is 4 times the original PC size.
Algorithm 3 in appendix shows a feedforward implementation of the growing operation.
Structure Learning through Pruning and Growing.
The proposed pruning and growing algorithms can be applied iteratively to refine the structure and parameters of an initial PC. Specifically, since the growing operator increases the number of PC parameters by a factor of 4, applying growing after pruning 75% of the edges from an initial PC keeps the number of parameters unchanged.
We propose a joint structure and parameter learning algorithm for PCs that uses these two operations.
Specifically, starting from an initial PC, we apply 75% pruning, growing, and parameter learning iteratively until convergence. We utilize HCLTs [ 25 ] as initial PC structure as it has the state-of-the-art likelihood performance. Note that this structure learning pipeline can be applied to any PC structure.
Parameter Estimation.
We use a stochastic mini-batch version of Expectation-Maximization optimization [ 2 ] . Specifically, at each iteration, we draw a mini-batch of samples \mathcal{D}_{B} , compute aggregated circuit flows \operatorname{F}_{n,c}(\mathcal{D}_{B}) and \operatorname{F}_{n}(\mathcal{D}_{B}) of these samples (E-step), and then compute new parameter \theta_{c|n}^{\mathtt{new}}=\operatorname{F}_{n,c}(\mathcal{D}_{B})/\operatorname{F}_{n}(\mathcal{D}_{B}) . The parameters are then updated with learning rate \alpha : \theta^{t+1}\leftarrow\alpha\theta^{\mathtt{new}}+(1-\alpha)\theta^{t} (M-step). Empirically this approach converges faster and is better regularized compared to full-batch EM.
Parallel Computation.
Existing approaches to scaling up learning and inference with PCs, such as Einsum networks [ 33 ] , utilize fully connected parametrized layers (Figure 3(a) ) of PC structures such as HCLT [ 25 ] and RatSPN [ 34 ] . These structures can be easily vectorized to utilize deep learning packages such as PyTorch.
However, the sparse structure learned by pruning and growing is not easily vectorized as a dense matrix operation. We therefore implement customized GPU kernels to parallelize the computation of parameter learning and inference based on Juice.jl [ 8 ] , an open-source Julia package for learning PCs. The kernels segment PC units into layers such that the units in each layer are independent. Thus, the computation can be fully parallelized on the GPU. As a result, we can train PCs with millions of parameters in less than half an hour.
Experiments
We now evaluate our proposed method pruning and growing on two different sets of density estimation benchmarks: (1) the MNIST-family image generation datasets including MNIST [ 22 ] , EMNIST [ 5 ] , and FashionMNIST [ 46 ] ; (2) the character-level Penn Tree Bank language modeling task [ 27 ] .

Section 6.1 first reports the best results we get on image datasets and language modeling tasks via the structure learning procedure proposed in Section 5 .
Section 6.2 then shows the effect of pruning and growing operations via two detailed experimental settings. It studies two different constrained optimization problems: finding the smallest PC for a given likelihood via model compression and finding the best PC of a given size via structure learning.
Settings.
For all experiments, we use hidden Chow-Liu Trees (HCLTs) [ 25 ] with the number of latent states in \{16,32,64,128\} as initial PC structures. We train the parameters of PCs with stochastic mini-batch EM (cf. Section 5 ). We perform early stopping and hyperparameter search using a validation set and report results on the test set. Please refer to Appendix C for more details.
We use mean test set bits-per-dimension (bpd) as the evaluation criteria, where \mathsf{bpd}(\mathcal{D},\mathcal{C})=-\mathcal{LL}(\mathcal{D},\mathcal{C})/(\log(2)\cdot m) and m is the number of features in dataset \mathcal{D} .
Image Datasets.
The MNIST-family datasets contain gray-scale pixel images of size 28\times 28 where each pixel takes values in [0,255] . We split out 5% of training data as a validation set. We compare with two competitive PC learning algorithms: HCLT [ 25 ] and RatSPN [ 34 ] , one flow-based model: IDF [ 17 ] , and three VAE-based methods: BitSwap [ 20 ] , BB-ANS [ 41 ] , and McBits [ 38 ] . For a fair comparison, we implement RatSPN structures ourselves and use the same training pipeline and EM optimizer as our proposed method. Note that EinsumNet [ 33 ] also uses RatSPN structures but with a PyTorch implementation so its comparison is subsumed by comparison with RatSPN. All 7 methods are tested on MNIST, 4 splits of EMNIST and FashionMNIST. As shown in Table 1 , the best results are bold. We see that our proposed method significantly outperforms all other baselines on all datasets, and establishes new state-of-the-art results among PCs, flows, and VAE models.
More experiment details are in Appendix C .
Language Modeling Task.
We use the Penn Tree Bank dataset with standard processing from Mikolov et al. [ 29 ] , which contains around 5M characters and a character-level vocabulary size of 50 . The data is split into sentences with a maximum sequence length of 288 .
We compare with three competitive normalizing-flow-based models: Bipartite flow [ 42 ] and latent flows [ 48 ] including AF/SCF and IAF/SCF, since they are the only comparable work with non-autoregressive language modeling.
As shown in Table 2 , the proposed method outperforms all three baselines.
What is the Smallest PC for the Same Likelihood?
We evaluate the ability of pruning based on circuit flows to do effective model compression by iteratively pruning a k -fraction of the PC parameters and then fine-tuning them until the final training log-likelihood does not decrease by more than 1\% . Specifically, we take pruning percentage k from \{0.05,0.1,0.3\} .
As shown in Figure 8 , we can achieve a compression rate of 80-98% with negligible performance loss on PCs. Besides, by fixing the number of latent parameters (x-axis) and comparing bpp across different numbers of latent states (legend), we discover that compressing a large PC to a get smaller PC yields better likelihoods compared to directly training an HCLT with the same number of parameters from scratch. This can be explained by the sparsity of compressed PC structures, as well as a smarter way of finding good parameters: learning a better PC with larger size and compressing it down to a smaller one.
What is the Best PC for the Same Size?
We evaluate structure learning that combines pruning and growing as proposed in Section 5 . Starting from an initial HCLT, we iteratively prune 75% of the parameters, grow again, and fine-tune until meeting the stopping criteria. As shown in Figure 9 , our method consistently improve the likelihoods of initial PCs for different numbers of latent states among all datasets.
Related Work
Improving the expressiveness of PCs has been a central topic in the literature. Predominant works focus on the structure learning algorithms that iteratively modify PC structures to progressively fit the data [ 7 , 24 , 15 ] . Alternatively, a recent trend is to construct PCs with good initial structures
and only perform parameter learning afterwards [ 36 , 1 , 47 ] . For example, EiNets [ 33 ] , RAT-SPNs [ 34 ] , and XPCs [ 12 ] use randomly generated structures, HCLTs [ 25 ] and ID-SPNs [ 37 ] define PCs dependent on the pairwise correlation on variables, and learn direct and indirect variable correlations among variables.
There are also a few works that boost PC performance with the expressive power of neural networks. CSPNs [ 39 ] harness neural networks to learn expressive conditional density estimators, and HyperSPN [ 40 ] utilizes neural networks for better PC regularization.

Pruning and growing have been introduced in deep neural networks to exploit sparsity [ 16 ] .
Similar strategies can be adopted in probabilistic circuits. Patil et al. [ 32 ] prune decision trees according to validation accuracy reduction. ResSPNs uses the lottery ticket hypothesis for weight pruning to gain more compact PCs [ 43 ] . However, the neural network pruning methods and above PC pruning methods mainly focus on pruning by parameters. In contrast to these, our work develops better pruning strategies based on semantic properties of PCs.
Conclusions
We propose structure learning of probabilistic circuits by combining pruning and growing operations to exploit the sparsity of PC structures. We show significant empirical improvements in the density estimation tasks of PCs compared to existing PC learners and competing flow-based models and VAEs.
All our Sparse-PC learning and inference algorithms are available as GPU-parallel implementations.
Acknowledgements
This work was funded in part by the DARPA Perceptually-enabled Task Guidance (PTG) Program under contract number HR00112220005, NSF grants #IIS-1943641, #IIS-1956441, #CCF-1837129, Samsung, CISCO, a Sloan Fellowship, and a UCLA Samueli Fellowship. We thank Honghua Zhang for proofreading and insightful comments on this paper’s final version.