Tractable Sharpness-Aware Learning of Probabilistic Circuits

TL;DR

This paper introduces a Hessian-based regularizer for Probabilistic Circuits that encourages convergence to flatter minima, thereby reducing overfitting and improving generalization in probabilistic modeling.

Contribution

It proposes a novel regularizer based on the Hessian trace for PCs, enabling efficient computation and integration with existing training methods.

Findings

Regularizer guides PCs toward flatter minima.

Improves generalization on synthetic and real datasets.

Provides closed-form updates for EM algorithm.

Abstract

Probabilistic Circuits (PCs) are a class of generative models that allow exact and tractable inference for a wide range of 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. We analyze PC overfitting from a log-likelihood-landscape perspective and show that it is often caused by convergence to 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 gradient-norm-based regularizer that yields simple closed-form parameter updates for EM, and integrates seamlessly with gradient based learning methods. Experiments on synthetic and real-world datasets demonstrate that our method consistently guides PCs toward flatter minima, improves generalization performance.