離散フローマッチング

Discrete Flow Matching

July 22, 2024

著者: Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky T. Q. Chen, Gabriel Synnaeve, Yossi Adi, Yaron Lipman

cs.AI

要旨

Flow Matchingや拡散モデルは、画像や動画などの連続変数に対する強力な生成パラダイムとして登場したものの、言語のような高次元の離散データへの応用はまだ限られています。本研究では、離散データの生成に特化した新しい離散フローパラダイムであるDiscrete Flow Matchingを提案します。Discrete Flow Matchingは以下の主要な貢献を提供します：(i) ソース分布とターゲット分布の間を補間する一般的な確率経路の族と連携する、(ii) 確率デノイザー（x予測）やノイズ予測（ε予測）などの学習済み事後分布を用いて、これらの確率経路からサンプリングするための汎用的な式を可能にする、(iii) 実践的には、異なるスケジューラで定義された特定の確率経路に焦点を当てることで、従来の離散拡散モデルやフローモデルと比較して生成パープレキシティを大幅に改善する、(iv) Discrete Flow Matchingモデルを1.7Bパラメータまでスケールアップすることで、HumanEvalで6.7%のPass@1と13.4%のPass@10、1-shot MBPPコーディングベンチマークで6.7%のPass@1と20.6%のPass@10を達成する。我々のアプローチは、非自己回帰的な方法で高品質な離散データを生成することが可能であり、自己回帰モデルと離散フローモデルの間のギャップを大幅に縮めます。

English

Despite Flow Matching and diffusion models having emerged as powerful generative paradigms for continuous variables such as images and videos, their application to high-dimensional discrete data, such as language, is still limited. In this work, we present Discrete Flow Matching, a novel discrete flow paradigm designed specifically for generating discrete data. Discrete Flow Matching offers several key contributions: (i) it works with a general family of probability paths interpolating between source and target distributions; (ii) it allows for a generic formula for sampling from these probability paths using learned posteriors such as the probability denoiser (x-prediction) and noise-prediction (epsilon-prediction); (iii) practically, focusing on specific probability paths defined with different schedulers considerably improves generative perplexity compared to previous discrete diffusion and flow models; and (iv) by scaling Discrete Flow Matching models up to 1.7B parameters, we reach 6.7% Pass@1 and 13.4% Pass@10 on HumanEval and 6.7% Pass@1 and 20.6% Pass@10 on 1-shot MBPP coding benchmarks. Our approach is capable of generating high-quality discrete data in a non-autoregressive fashion, significantly closing the gap between autoregressive models and discrete flow models.