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

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

Generative Adversarial Networks and Image Synthesis参考文献 32被引用数 36

ひとこと要約

本論文はカテゴリカルデータを正規化フローと拡散でモデル化する 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 の確率的逆変換を介して離散データと連続密度モデルを橋渡しする。
• Argmax Flows と Multinomial Diffusion の二つのフレームワークを用いてカテゴリデータの直接学習を可能にする。
• 言語とセグメンテーション課題で dequantization のベースラインより改善された log-likelihood を示す。

提案手法

• Argmax Flows を導入する: 連続密度 p(v) と離散 x への argmax 写像を結合する; argmax 制約を満たす確率的逆 q(v|x) を学習する。
• 厳密には解けない積分を介して P(x) を用いた変分下界（ELBO）を定義し、argmax 制約の下で q(v|x) を用いて最適化する。
• 3つの q(v|x) 構成を提示: 閾値処理ベース、Gumbel ベース、Gumbel 閾値処理、各々で対数密度計算を含む。
• Argmax Flows のデカルト積を許可し、異なる基底表現（K進 x を D 次元にまたがる）を扱えるようにする。
• Multinomial Diffusion を導入する: ワンホットカテゴリベクトル上に直接定義された拡散過程で、固定のマルコフ連鎖 q(x_t|x_{t-1}) とノイズ除去生成モデル p(x_{t-1}|x_t) を持つ。
• KL ベースの変分目的関数と閉形式の後方分布 q(x_{t-1}|x_t,x_0) による訓練を説明する。

実験結果

リサーチクエスチョン

• RQ1 Argmax Flows は離散的カテゴリデータを連続空間に効果的に持ち上げ、密度モデリングを可能にするか。
• RQ2 argmax の確率的逆変換は訓練を扱いやすくし、カテゴリデータにおける dequantization よりも良い log-likelihood を実現するか。
• RQ3 拡散をカテゴリデータ上に直接定義し、離散フローのベースラインと競合する、あるいはそれを上回ることができるか。
• RQ4 言語とセグメンテーション課題において、Argmax Flows と Multinomial Diffusion は自己回帰・非自己回帰のベースラインとどう比較されるか。
• RQ5 表現（デカルト積）が性能と効率に与える影響は何か。

主な発見

• Argmax Flows は text8 と enwik8 において log-likelihood で uniform および variational dequantization を上回る。
• 自己回帰の Argmax Flows は言語データに対して、いくつかの VAE や非 argmax 离散フローのベースラインより良い性能を達成する。
• Multinomial Diffusion はカテゴリテキストデータで強力な結果を示し、セグメンテーションマップでは ELBO/IWBO の競争力のある値を示すが、ある設定で過学習が見られ拡張が必要。
• 非自己回帰の多項拡散は一部の設定で argmax coupling flows を上回ることがあり、拡散軌道とモデルの堅牢性のトレードオフを浮き彫りにする。
• デカルト積表現は、対称性を次元性と交換することで、異なる基底 K 表現に適応する。
• サンプルは妥当なテキストとセグメンテーション出力を示し、綴りチェックは拡散モデルの副産物として実証されている。