[论文解读] Scaling Up Probabilistic Circuits by Latent Variable Distillation

Anji Liu, Honghua Zhang+1 more|arXiv (Cornell University)|Oct 10, 2022

Topic Modeling被引用 5

一句话总结

本文提出潜在变量蒸馏（LVD），一种通过利用深度生成模型（如Transformer）的监督信号来指导潜在变量分配，从而增强大型概率电路（PC）的方法。通过优化该分配所诱导的似然下界，LVD使大型PC能够充分释放其容量，在图像和语言建模基准上实现最先进性能，表现媲美基于流的模型和变分自编码器。

ABSTRACT

Probabilistic Circuits (PCs) are a unified framework for tractable probabilistic models that support efficient computation of various probabilistic queries (e.g., marginal probabilities). One key challenge is to scale PCs to model large and high-dimensional real-world datasets: we observe that as the number of parameters in PCs increases, their performance immediately plateaus. This phenomenon suggests that the existing optimizers fail to exploit the full expressive power of large PCs. We propose to overcome such bottleneck by latent variable distillation: we leverage the less tractable but more expressive deep generative models to provide extra supervision over the latent variables of PCs. Specifically, we extract information from Transformer-based generative models to assign values to latent variables of PCs, providing guidance to PC optimizers. Experiments on both image and language modeling benchmarks (e.g., ImageNet and WikiText-2) show that latent variable distillation substantially boosts the performance of large PCs compared to their counterparts without latent variable distillation. In particular, on the image modeling benchmarks, PCs achieve competitive performance against some of the widely-used deep generative models, including variational autoencoders and flow-based models, opening up new avenues for tractable generative modeling. Our code can be found at https://github.com/UCLA-StarAI/LVD.

研究动机与目标

• 解决尽管模型容量增加，大型概率电路仍出现性能瓶颈的问题。
• 通过来自表达能力强的深度生成模型的潜在变量分配提供额外监督，改进大型PC的优化。
• 通过减少对标准EM优化的依赖，实现大型PC的可扩展、高效训练。
• 证明LVD可使PC在性能上媲美不可训练的深度生成模型（如VAE和基于流的模型）

提出的方法

• LVD使用预训练深度生成模型（如BERT或视觉Transformer）的神经嵌入，对训练样本进行聚类，并为PC中的潜在变量分配语义感知的取值。
• 提出一种改进的训练目标，通过使用诱导的潜在变量分配作为监督，优化边缘似然的下界：∑ᵢ log p(x⁽ⁱ⁾, z⁽ⁱ⁾)。
• 该方法先使用此下界目标进行初始训练以引导优化，随后在完整最大似然目标上进行微调。
• 通过在图像块上构建PC并使用块级条件分布，将框架专门用于图像建模。
• 采用参数共享和聚类等技术，提升样本效率并减少训练时间。
• 该方法具有通用性，可适配多种PC架构和深度生成模型。

实验结果

研究问题

• RQ1 潜在变量蒸馏能否有效克服大型概率电路中的性能瓶颈？
• RQ2 利用深度生成模型的监督是否能提升大型PC的优化效率和表达能力？
• RQ3 LVD训练的PC能否在性能上媲美不可训练的深度生成模型（如VAE和基于流的模型）？
• RQ4 LVD如何影响大型PC的训练效率与可扩展性？

主要发现

• 在ImageNet32上，使用LVD训练的2500万参数PC优于4亿参数的HCLT，展现出更优的样本效率。
• LVD使测试集每维比特数（bpd）随模型规模增大而持续下降，表明容量得到有效利用。
• 在单张A5000 GPU上，LVD训练5亿参数PC仅耗时10小时，远快于基线方法所需的一天以上。
• 在ImageNet32和ImageNet64上，LVD的bpd与最先进基于流的模型和VAEs相差仅0.1–0.3，表现极具竞争力。
• CIFAR-10上的性能差距归因于数据稀缺，表明通过参数共享等架构改进可缩小该差距。