Scaling Tractable Probabilistic Circuits: A Systems Perspective

Probabilistic Circuits（PCs）是一种通用的框架，用于可处理的深度生成模型，支持对其学习分布进行准确和高效的概率推理。最近的建模和训练进展使其能够应用于复杂的现实任务。然而，现有PC实现的时间和内存效率阻碍了进一步的扩展。本文提出了PyJuice，一种适用于PC的通用GPU实现设计，它在多个方面改进了之前的技术。具体而言，PyJuice比现有系统（包括最近的系统）快1-2个数量级，可训练大规模PC。此外，PyJuice消耗的GPU内存比现有系统少2-5倍，这使我们能够训练更大的模型。我们系统的核心是编译过程，将PC转换为紧凑表示，以便进行高效的块状并行处理，从而显着减少了IO并使其能够利用现代GPU中可用的张量核心。实证上，PyJuice可用于改进在图像（例如ImageNet32）和语言（例如WikiText，CommonGen）数据集上训练的最先进的PC。我们还通过对现有PC结构进行基准测试，但使用更大的规模和更多的训练时期，建立了自然图像和语言数据集的新一组基线，希望激励未来的研究。代码可在https://github.com/Tractables/pyjuice上获得。

本文旨在解决现有PC实现中时间和内存效率低下的问题，提出了一种GPU实现设计PyJuice，以支持更大规模的PC模型训练。

PyJuice通过编译过程将PC转换为紧凑表示形式，利用块状并行化显著减少IO，从而实现了1-2个数量级的训练速度提升和2-5倍的GPU内存消耗降低。

实验结果表明，PyJuice可用于改进在图像和语言数据集上训练的PC，并在自然图像和语言数据集上建立了新的基准。代码已开源。

与本文相关的研究包括PC的先前实现，如TensorFlow Probability和Probabilistic Torch，以及最近的PC应用研究，如在图像生成和自然语言处理中的应用。