Tractable Sharpness-Aware Learning of Probabilistic Circuits

Probabilistic Circuits (PCs) are a class of generative models that allow exact and tractable inference for diverse queries. While recent developments have enabled the learning of deep and expressive PCs, this increased capacity can often lead to overfitting, especially when data is limited. In this work, we analyze PC overfitting from a log-likelihood-landscape perspective and show that it is often caused by convergence to $\textit{sharp optima}$ that generalize poorly. Inspired by sharpness aware minimization in neural networks, we propose a Hessian-based regularizer for training PCs. As a key contribution, we show that the trace of the Hessian of the log-likelihood—a sharpness proxy that is typically intractable in deep neural networks- can be computed efficiently for PCs. Minimizing this Hessian trace induces a natural gradient-norm-based regularizer that yields simple closed-form parameter updates for EM, and integrates seamlessly with gradient based methods. Experiments on synthetic and real-world PC benchmarks demonstrate that our method consistently guides PCs toward flatter minima, improving generalization over standard training.