[論文レビュー] Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions

Emiel Hoogeboom, Didrik Nielsen+3 more|arXiv (Cornell University)|Feb 10, 2021

Music and Audio Processing被引用数 79

ひとこと要約

この論文は Argmax Flows と Multinomial Diffusion を導入し、カテゴリ分布を学習させることで正規化フローと拡散モデルがテキストやセグメンテーションマップのような離散データを扱えるようにし、標準のデクォンタイズ法より対数尤度が向上することを示す。

ABSTRACT

Generative flows and diffusion models have been predominantly trained on ordinal data, for example natural images. This paper introduces two extensions of flows and diffusion for categorical data such as language or image segmentation: Argmax Flows and Multinomial Diffusion. Argmax Flows are defined by a composition of a continuous distribution (such as a normalizing flow), and an argmax function. To optimize this model, we learn a probabilistic inverse for the argmax that lifts the categorical data to a continuous space. Multinomial Diffusion gradually adds categorical noise in a diffusion process, for which the generative denoising process is learned. We demonstrate that our method outperforms existing dequantization approaches on text modelling and modelling on image segmentation maps in log-likelihood.

研究の動機と目的

• 生成フローと拡散モデルで高次元カテゴリデータの学習を動機づける。
• Argmax ベースのリフティングを介して直接カテゴリ分布へ正規化フローを拡張。
• カテゴリ変数上で拡散モデルを提案（Multinomial Diffusion）
• 言語モデリングとセグメンテーションデータセットで、デクォンタイズ基準と比較して対数尤度が改善されることを示す。

提案手法

• Argmax Flows は連続密度 p(v) と argmax v -> x マッピング、および argmax 制約を満たす確率的逆関数 q(v|x) を組み合わせる。
• q(v|x) の3つの確率적逆関数を検討：閾値処理、Gumbel ベース、Gumbel 閾値処理、それぞれ訓練のための対数決定項を持つ。
• 最適化は変分下限 ELBO を用い、log P(x) >= E_{v~q(v|x)}[log p(v) - log q(v|x)]。
• Argmax Flows のカルテジアン積は base-K 空間へ拡張され、対称性と次元性をトレードオフ。
• Multinomial Diffusion は離散変数上の拡散過程を固定のデノイズ経路と閉形式の事後 q(x_{t-1}|x_t,x_0) を持って直接定義。
• モデルはflows による p(v)（Argmax Flows）を予測するか、カテゴリ空間で denoising 経路 p(x_{t-1}|x_t) を学習する。

実験結果

リサーチクエスチョン

• RQ1 Argmax Flows は離散的カテゴリデータを正規化フローで密度モデリングするために連続空間へ効果的にリフトできるか？
• RQ2 argmax 制約を満たす確率的逆関数 q(v|x) をどのように構築・訓練するか？
• RQ3 カテゴリ変数上の拡散過程（Multinomial Diffusion）は離散データのデクォンタイズ基準より対数尤度を改善するか？
• RQ4 テキストとセグメンテーションマップをモデル化する際、自己回帰型と非自己回帰型（結合型）フローアーキテクチャのトレードオフは？
• RQ5 Argmax Flows のカルテジアン積は大規模カテゴリ数に対して性能とモデル容量にどのような影響を与えるか？

主な発見

• Argmax Flows は text8 と enwik8 において、対数尤度（文字/バイトあたりのビット）で一様デクォンタイズおよび変分デクォンタイズを一貫して上回る。
• 自己回帰型 Argmax Flows は言語データ上で複数のVAEおよび離散フロー基準より良好な性能を達成。
• 非自己回帰型 Argmax Flows（結合型）はいくつかの離散フロー基準より改善するが、テキストでは自己回帰型に遅れをとる。一方、Multinomial Diffusion は特定の設定でいくつかの非自己回帰型 Argmax 変体より上回ることがある。
• Cityscapes の画像セグメンテーションマップでは、Argmax Flows と Multinomial Diffusion は競争力のある ELBO および IWBO スコアを達成し、従来のデクォンタイズ法を上回る。
• 標準の自己回帰モデル、連続自己回帰モデル、および非自己回帰連続モデル間には性能ギャップが存在し、離散から連続へのリフティングと自己回帰コンポーネントの将来の研究方向を示唆している。