Probabilistic Neural Circuits (2403.06235v1)

Abstract: Probabilistic circuits (PCs) have gained prominence in recent years as a versatile framework for discussing probabilistic models that support tractable queries and are yet expressive enough to model complex probability distributions. Nevertheless, tractability comes at a cost: PCs are less expressive than neural networks. In this paper we introduce probabilistic neural circuits (PNCs), which strike a balance between PCs and neural nets in terms of tractability and expressive power. Theoretically, we show that PNCs can be interpreted as deep mixtures of Bayesian networks. Experimentally, we demonstrate that PNCs constitute powerful function approximators.

• The paper introduces PNCs that merge the tractability of probabilistic circuits with the expressiveness of neural networks by relaxing decomposability constraints.
• It presents a layered architecture using conditional probabilistic circuits and convolutional neural sum layers to efficiently approximate complex dependencies.
• Experimental results demonstrate superior density estimation on MNIST and competitive classification performance, underscoring their practical potential.

This paper introduces Probabilistic Neural Circuits (PNCs), a class of models designed to bridge the gap between the tractability of Probabilistic Circuits (PCs) and the expressive power of neural networks. PCs, while allowing for efficient probabilistic inference like marginalization, suffer from limited expressiveness due to structural constraints (smoothness and decomposability). PNCs aim to relax these constraints selectively to gain expressiveness while retaining some tractability.

Core Concepts and Implementation:

1. Conditional Probabilistic Circuits (CPCs): The foundation for PNCs. CPCs generalize PCs by introducing a partial order ($\sqsubset$) over the random variables, inspired by Bayesian Networks. Computations within a CPC unit are conditioned on the values of parent variables according to this order.
   • A CPC unit $k$ computes $p_k(\mathcal{X}_n | \mathcal{X}_{pa(n)})$, where $\mathcal{X}_n$ is the scope of the unit and $\mathcal{X}_{pa(n)}$ are the parents according to the partial order.
   • CPCs maintain conditional versions of smoothness and decomposability.
   • Theoretically, valid CPCs (conditionally smooth, decomposable, normalized sum weights) represent deep mixtures of Bayesian networks (Eq. \ref{eq:f_tree_decomp_PCP}).
2. Probabilistic Neural Circuits (PNCs): PNCs are introduced as a tractable approximation of CPCs. The key challenge in CPCs is that sum units require evaluating potentially exponentially many conditional distributions based on parent instantiations. PNCs approximate this:

   • The sum unit computation is modified (Definition \ref{def:pnc}, Eq. \ref{eq:pnc}):

     $${p}_k(\mathcal{X}_n \mid \mathcal{X}_{pa(n)}) = \sum_{j\in\mathcal{I}(k)} \phi_{kj} (\mathcal{X}_{an(n)}) p_j(\mathcal{X}_n)$$

     Here, $\phi_{kj}$ is a neural network that takes the values of the ancestors $\mathcal{X}_{an(n)}$ as input and outputs normalized weights ($\sum_j \phi_{kj} = 1$). $p_j(\mathcal{X}_n)$ are the outputs of the child units, without the explicit conditioning seen in CPCs.

   • This neural network $\phi_{kj}$ replaces the complex conditional probability ratio derived from Bayes' rule, making the computation depend only on the ancestor values and the unconditioned child distributions.

3. Layered PNC Construction: A practical method for building PNCs is proposed, based on layered PC structures like those from Shih et al. [shih2021hyperspns] but with added neural dependencies.
   • The structure alternates sum and product operations layer-wise, operating on partitions of variables.

   • Neural Sum Layer: The core implementation detail. This layer computes:

     $${\kappa}_{l,p,2,c} = \sum_{c'=1}^{N_C} \phi_{l,p,c,c'}(\mathcal{K}_{ancestors}) \times {\kappa}_{l,p,1,c'}$$

     where $\kappa$ represents the value of a computational component (indexed by layer $l$, partition $p$, input/output $i$, component $c$). The neural network $\phi_{l,p,c,c'}$ computes the weights for component $c$ in partition $p$ based on the outputs ($\mathcal{K}_{ancestors}$) of components in preceding partitions ($p-\nu_{l,p}$ to $p-1$) within the same layer $l$. The hyperparameter $\nu$ determines the range of this dependency.

   • These intra-layer dependencies via $\phi$ induce the partial order ($\sqsubset$) on the variables.

4. Convolutional Implementation: The NeuralSumLayer is efficiently implemented using Convolutional Neural Networks (CNNs).
   • Partitions are treated as the spatial dimension(s).
   • Components ($N_C$) are treated as channels.
   • "Half kernels" (masked convolutions, see Figure \ref{fig:kernel}) are used to ensure that the computation for a partition only depends on preceding partitions, respecting the induced variable order. This is crucial for maintaining tractability.
   • A final softmax activation ensures the neural network outputs ($\phi$) sum to 1 for each partition's summation.

Tractability:

• Density Evaluation: Computing $p(\mathbf{x})$ for a full evidence vector $\mathbf{x}$ is linear in the size of the circuit, just like standard PCs.
• Ordered Marginals/Conditionals: Computing $p(\mathcal{X}_e)$ or $p(\mathcal{X}_o | \mathcal{X}_e)$ is tractable (polynomial time) if the variables being marginalized out ($\mathcal{X}_m$) come after the evidence/query variables in the partial order induced by the neural dependencies (i.e., $\mathcal{X}_e \sqsubset \mathcal{X}_m$, $\mathcal{X}_o \sqsubset \mathcal{X}_m$).
  • The marginalization algorithm (Algorithm \ref{alg:layer}) works by setting leaf nodes corresponding to marginalized variables to 1 and performing a forward pass. The conditional decomposability and the specific structure of the neural sum units (where integration results in 1 due to normalized weights) allow integrals/summations to be pushed down correctly, similar to standard PCs, provided the order is respected. Arbitrary marginalization is generally not tractable.

Experimental Results and Applications:

• Density Estimation: PNCs achieve state-of-the-art or competitive performance (measured in bits per dimension) on MNIST variants, outperforming standard PCs (PSCs), their implementation of SPQNs (PQCs), and other PC methods like HCLT and RAT-SPN. This demonstrates the practical benefit of increased expressiveness from the neural components.
• Classification: PNCs were adapted for discriminative tasks by training one circuit per class and using a cross-entropy loss. They outperformed PQCs but lagged behind Logistic Circuits and RAT-SPNs. The paper suggests this might be due to suboptimal regularization, as PNCs achieved very high training accuracy, indicating potential overfitting.
• Implementation Details: Experiments used PyTorch Lightning on Nvidia V100 GPUs. Models had roughly 2.6-2.8M parameters for the PNC/PQC/PSC comparison. A specific layered architecture alternating row/column merges with 12 components per partition was used for image tasks.

Limitations and Future Work:

• Tractability: Only ordered marginalization is guaranteed to be tractable.
• Structure Learning: The paper uses a predefined structure; learning optimal PNC structures remains an open question.
• Regularization: Effective regularization techniques for discriminative training with PNCs need further investigation.
• Future Directions: Exploring sampling, applying PNCs to tabular data, lossless compression, and potentially linking to any-order autoregressive models.

Comparison to Related Work:

• SPQNs: PNCs are presented as a generalization and simplification. They achieve the goal of relaxing decomposability using neural networks within sum units, rather than introducing explicit quotient units. The paper shows CMO-SPQNs are a restricted form of PNCs.
• Conditional SPNs (Shao et al.): These condition sum weights on external variables, modelling a single conditional distribution. PNCs condition internally based on ancestor variables within the circuit, modelling mixtures of Bayesian networks.

In summary, PNCs offer a practical way to enhance the expressiveness of probabilistic circuits using neural networks embedded within their sum units. This is achieved by approximating conditional dependencies, leading to models that perform well on density estimation tasks while retaining tractability for ordered inference queries. The convolutional implementation provides an efficient pathway for applying PNCs to structured data like images.