[論文レビュー] Scaling Up Probabilistic Circuits by Latent Variable Distillation

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

Topic Modeling被引用数 5

ひとこと要約

本稿では、大規模な確率的回路（PC）の性能を、変換器などの深層生成モデルからの監視を活用して潜在変数の割り当てをガイドする、潜在変数蒸留（LVD）という手法を提案する。これらの蒸留された割り当てを用いて尤度の下界を最適化することで、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が、トレーニング不能な深層生成モデル（VAE や流れベースのモデルなど）と同等の性能を大規模PCが達成できることを示すこと。

提案手法

• 事前学習済みの深層生成モデル（例：BERT やビジョン変換器）のニューラル埋め込みを用いて、トレーニング例をクラスタリングし、PC内の潜在変数に意味的意識のある値を割り当てる。
• 尤度の周辺分布の下界を最適化する修正されたトレーニング目的関数を導入し、誘導された潜在変数の割り当てを監視として使用する：∑ᵢ log p(x⁽ⁱ⁾, z⁽ⁱ⁾)。
• 初期段階でこの下界目的関数を用いて最適化をガイドし、その後、完全な尤度最大化目的関数でのファインチューニングを行う。
• 画像モデリングに特化するため、PCを画像パッチごとに構造化し、パッチ単位の潜在変数条件付き分布を用いる。
• パラメータの共有やクラスタリングなどの技術を用いて、サンプル効率を向上させ、トレーニング時間を短縮する。
• このアプローチは汎用的であり、さまざまなPCアーキテクチャや深層生成モデルに適応可能である。

実験結果

リサーチクエスチョン

• RQ1 潜在変数蒸留は、大規模な確率的回路における性能の頭打ちを効果的に克服できるか？
• RQ2 深層生成モデルからの監視を用いることで、大規模PCの最適化と表現力が向上するか？
• RQ3 LVDでトレーニングされたPCは、トレーニング不能な深層生成モデル（VAE や流れベースのモデルなど）と同等の性能を達成できるか？
• RQ4 LVDは、大規模PCのトレーニング効率とスケーラビリティにどのように影響を与えるか？

主な発見

• ImageNet32では、2500万パラメータのPCをLVDでトレーニングした結果、4億パラメータのHCLTを上回り、優れたサンプル効率を示した。
• LVDにより、モデルサイズが大きくなるにつれてテストセットの次元あたりビット数（bpd）が段階的に低下し、容量の有効活用が示された。
• LVDを用いることで、1枚のA5000 GPUで10時間で5億パラメータのPCをトレーニング可能となり、ベースライン手法が1日以上を要するのと比べて顕著に高速化された。
• ImageNet32およびImageNet64において、LVDは最先端の流れベースのモデルやVAEと比較して、bpdで0.1〜0.3の差で競争力のある結果を達成した。
• CIFAR-10における性能差はデータ不足に起因するとされ、パラメータの共有などのアーキテクチャ的変更によりこのギャップを埋められると示唆された。