[論文レビュー] Bayesian Flow Networks

Alex Graves, Rupesh Kumar Srivastava+2 more|arXiv (Cornell University)|Aug 14, 2023

Topic Modeling被引用数 5

ひとこと要約

Bayesian Flow Networks (BFNs) は、繰り返しベイズ推論を用いて事前分布を更新する生成モデルフレームワークを導入し、その後ニューラルネットワークを通過させて相関のある出力分布を生成する。この手法は、1文字あたり1.41ビットという最先端の結果をテキスト8の文字レベル言語モデル化タスクで達成しており、すべての既知の離散拡散モデルを上回っている。加えて、離散ドメインにおける微分可能で勾配ベースのサンプルガイドランスを可能としている。

ABSTRACT

This paper introduces Bayesian Flow Networks (BFNs), a new class of generative model in which the parameters of a set of independent distributions are modified with Bayesian inference in the light of noisy data samples, then passed as input to a neural network that outputs a second, interdependent distribution. Starting from a simple prior and iteratively updating the two distributions yields a generative procedure similar to the reverse process of diffusion models; however it is conceptually simpler in that no forward process is required. Discrete and continuous-time loss functions are derived for continuous, discretised and discrete data, along with sample generation procedures. Notably, the network inputs for discrete data lie on the probability simplex, and are therefore natively differentiable, paving the way for gradient-based sample guidance and few-step generation in discrete domains such as language modelling. The loss function directly optimises data compression and places no restrictions on the network architecture. In our experiments BFNs achieve competitive log-likelihoods for image modelling on dynamically binarized MNIST and CIFAR-10, and outperform all known discrete diffusion models on the text8 character-level language modelling task.

研究の動機と目的

• 自然言語のような離散データドメインにおけるエンドツーエンド微分可能訓練と勾配ベースのサンプルガイドランスを可能にする生成モデルフレームワークの開発。
• 離散データによる不連続なノイズ遷移の問題を抱える既存の離散拡散モデルの限界を克服するため、連続的でフローに類似したプロセスを導入すること。
• データ圧縮と生成モデルの原則を統合し、圧縮効率を直接最適化する損失関数を導出すること。
• 確率シンプレックス上の連続的更新を活用して、離散系列における効率的かつ少数ステップでの生成を可能とすること。
• 連続的（画像）および離散的（テキスト）データの両方で競争力のある性能を示し、特に言語モデル化に焦点を当てる。

提案手法

• BFNs は二段階のプロセスを用いる：まず、ノイズの入ったデータサンプルに基づいてベイズ推論により事前分布を更新し、事後分布を生成する。次に、その事後分布のパラメータをニューラルネットワークに通して、新たな相関のある分布を生成する。
• 変分推論の原則に基づいて連続時間および離散時間の損失関数を導出し、エビデンス下限界（ELBO）を用いてデータ圧縮を最適化する。
• 離散データの場合、ネットワークの入力は確率シンプレックス上に位置し、完全な微分可能性と勾配ベースのサンプルガイドランスを可能としている。これは、標準的な離散拡散モデルでは実現不可能である。
• ネットワークアーキテクチャは柔軟であり、制限されていない。論文では言語モデル実験のために24層の深層TransformerとGELU活性化関数を用いている。
• 前向きの拡散プロセスを避けており、代わりに逐次的なベイズ更新に依存し、モデルがデータ分布についての信念を段階的に精緻化する。
• サンプル生成は、単純な事前分布から開始し、繰り返しネットワークを適用して事後分布を更新することで実行される。

Figure 1 : System Overview . The figure represents one step of the modelling process of a Bayesian Flow Network. The data in this example is a ternary symbol sequence, of which the first two variables (‘B’ and ‘A’) are shown. At each step the network emits the parameters of the output distribution b

実験結果

リサーチクエスチョン

• RQ1 離散データに対して微分可能で連続的なフロー的プロセスを設計し、勾配ベースのサンプルガイドランスと少数ステップでの生成を可能にできるか？
• RQ2 ベイズ推論に基づく繰り返し的分布の精緻化が、言語モデル化において既存の離散拡散モデルを上回るか？
• RQ3 同じフレームワークを連続的、量子化済み、離散的データに適用でき、統一された損失関数と訓練手順を用いることができるか？
• RQ4 変分推論の原則に基づいて導出された損失関数が、多様なデータタイプにおける圧縮と生成の両方で有効であるか？
• RQ5 前向きの拡散プロセスを必要とせず、画像とテキストのベンチマークで競争力ある対数尤度を達成できるか？

主な発見

• BFNs は、text8 データセットでテスト時の1文字あたり1.41ビット（BPC）を達成し、文献に報告されたすべての既知の離散拡散モデルを上回った。
• 生成ステップを100ステップに制限した場合でも1.43 BPC を達成しており、少数ステップでの推論に強いロバストネスと潜在能力を示している。
• 1.40 BPC を達成した最良の順序に依存しないモデル（MAC）と同等の性能を示した一方で、他の離散拡散モデル（1.45〜1.72 BPC）を著しく上回った。
• 動的バイナリ化MNISTおよびCIFAR-10においても、競争力ある対数尤度を達成しており、画像データに対する有効性を示している。
• 確率シンプレックス入力を用いることで完全な微分可能性が実現され、離散ドメインにおける勾配ベースのサンプルガイドランスが可能になった。これは、従来の離散拡散モデルには欠落していた機能である。
• 訓練設定では24層のTransformer（1.7億パラメータ）を用い、過学習の兆候が見られたため、正則化によりさらなる性能向上が期待できる。

Figure 2 : Bayesian updates for continuous data . For univariate data $x=0.7$ , the initial input distribution parameters $\theta_{0}=\{\mu_{0}=0,\rho_{0}=1\}$ are updated to $\theta_{1}=\{\mu_{1},\rho_{1}\}$ , $\theta_{2}=\{\mu_{2},\rho_{2}\}$ , $\theta_{3}=\{\mu_{3},\rho_{3}\}$ by iterating Eqs. 4