【Diffusion勉強会】A continuous time framework for discrete denoising models

DEEP LEARNING JP [DL Papers] “A continuous time framework for discrete denoising models” NeurIPS2022 Shunsuke Sakai, Hasegawa Lab (Univ. Fukui) http://deeplearning.jp/ 1

• http://deeplearning.jp/

拡散言語モデルについて • 拡散言語モデル（Diffusion Language Models; DLMs） – 言語を生成する分布を拡散モデルにより学習． – 自己回帰型（e.g., GPTs）と比較して以下のような利点がある． • 推論が高速（逐次的な単語分布からのサンプリングが不要） • 非自己回帰的な推論タスクに対する能力が高い（e.g., reversal curese [1]への対処） • 分布アニーリング（e.g., nuclear sampling）が不要 – 近年の動向 • Mecury (InceptionLab) • Gemini Diffusion (Google) – 従来の自己回帰型の言語モデルを今後置き換えていくだろうか？ • この発表資料では，拡散言語モデルの初期研究にあたるtauLDR [2]を紹介します． • 論文の中で扱われている内容は今の拡散言語モデルの核となっています． • 拡散言語モデルを調べる際のロードマップは次ページを参考にしてください． 2

拡散言語モデルの主要論文 ※ 著者はまだこの分野には疎いため，あくまで参考程度に． CSM [Meng+, 2022, abs] Diffusion-LLM-Papers tauLDR [Campbell+, 2022, abs] D3PM SDDM [Austin+, 2021, abs] [Sun+, 2022, abs] Multinominal Diffusion [Hoogeboom+, 2021, abs] [GitHub] MDLM [Sahoo+, 2024, abs] RADD [Ou+, 2024, abs] SEDD [Lou+, 2023, abs] MD4 [Shi+, 2024, abs] LLADA [Nie+, 2025, abs] 3

• http://arxiv.org/abs/2205.14987
• https://arxiv.org/abs/2211.16750
• https://arxiv.org/abs/2310.16834
• https://arxiv.org/abs/2406.03736
• https://arxiv.org/abs/2406.04329
• https://arxiv.org/abs/2406.07524
• https://arxiv.org/abs/2502.09992
• https://github.com/ML-GSAI/Diffusion-LLM-Papers
• https://arxiv.org/abs/2102.05379
• https://arxiv.org/abs/2107.03006
• https://arxiv.org/abs/2211.00802

離散状態空間での拡散モデル１ 画像や音声などの拡散モデルと何が異なるのか？ ドメインが連続状態空間から離散状態空間になる． （離散になると困ること） ・対数尤度の勾配（Score）が定義できない． ・非順序尺度データに対する扱い -> 従来のガウシアンノイズは順序尺度を持つデータが前提 （既存手法のアプローチ） ・近傍構造を導入して，代替勾配を導入（Concrete Score Matching） [3] ・離散マルコフ連鎖の遷移行列の設計 [4] 今回紹介するアプローチ 4

離散状態空間での拡散モデル２ • 入力データは𝑆個の離散値のいずれかを取る． • 目標：データ分布をサンプル集合（i.i.d）から近似する • 考え方：連続版と同様に，以下の二つの性質を持つ順過程を定義 – 十分にステップ数を大きくした際に，サンプルが容易な分布に収束 – 任意の時刻tのサンプルを効率的に得られる Ex. 一様カーネル [x] 5

一様カーネルのイメージ ・一様カーネルは現在の拡散言語モデルの主流なアプローチの一つ． From Fig.1 [5] ・もう一つ主流なのは，「吸収状態」を持つ離散拡散モデル． 6

離散状態空間での拡散モデル３ ・順過程の同時分布はマルコフ連鎖として定義 ・逆過程は以下のように表される． 7

離散状態空間での拡散モデル４ ・（Recap）事後分布 ・この事後分布は以下のようにNNにより近似する． 8

離散状態空間での拡散モデル５ ・生成過程 ・学習は負のELBO最小化により行う． 9

離散状態空間での拡散モデル５ ・生成過程 ・学習は負のELBO最小化により行う． 10

連続時間マルコフ連鎖への拡張 ・これまでは離散時間マルコフ連鎖を考えてきた． ➔連続時間へと拡張する． （連続時間へと拡張する利点） ・時間幅が柔軟に取れる ・Samplerによる生成の高速化・高品質化 ・対数尤度の評価…etc 次ページでは，連続時間マルコフ連鎖について説明． 11

連続時間マルコフ連鎖 連続時間マルコフ連鎖（Continuous Time Markov Chain; CTMC） マルコフ性を満たす右連続な確率過程 𝑥𝑡 𝑡 ∈[0,𝑇] ．各状態は離散値を取る． 3状態での連続時間マルコフ連鎖の状態遷移 CTMCは待ち時間（ある状態から異なる状態に遷移するまでの時間）によって特徴付けられる． ある状態𝑥から別の状態に遷移するまでの待ち時間 𝜈(𝑥)は指数分布に従う． ある離散状態空間 𝒳 (𝑆 = 𝒳 )上でのCTMCの遷移率行列 𝑅 ∈ ℝ𝑆 ×𝑆 は以下のような定義． 12

連続時間マルコフ連鎖 逆に，遷移率行列が与えられたもとで，遷移確率は以下のように定義される． （気持ち）非常に短い時間であれば，基本的に同じ状態にとどまり，別の状態に遷移する確率は時間とともに線形に増加する． 遷移率行列は定義（ 前ページ参照）より，以下のような性質を持つ． （性質1）非対角成分は非負 （性質2）対角成分は0以下 （性質3）行和が0 遷移率行列が時間不変な場合の待ち時間分布と遷移確率 13

連続時間マルコフ連鎖 CTMCの遷移確率は以下に示すコルモゴロフの前向き（後ろ向き）方程式を解いて得られる． コルモゴロフの前向き方程式 コルモゴロフの後ろ向き方程式 以上より，遷移率行列 𝑅と初期分布を決めれば，CTMCを定義できる． 次ページ以降では，実際に拡散言語モデルのためのCTMCを定義していく． CTMC v.s. DTMC ・CTMCでは任意の時点で状態遷移が生じうるが，DTMCでは固定間隔 ・CTMCでは遷移率行列により「どこに」「どのくらい後で」遷移するかを表現するが，DTMCでは「どこに」 のみ表現 ・DTMCでは行列の冪乗で複数ステップの遷移を表現したが，CTMCでは微分方程式を解く． 14

連続時間拡散言語モデル 連続時間離散拡散モデルを導入する．従来の離散時間拡散モデルと比較して，離散状態の遷移は 𝑡 ∈ [0, 𝑇]の 任意時点で起こりうる． 15

連続時間マルコフ連鎖の時間可逆性 Proposition1. CTMCの時間可逆性 遷移率行列 𝑅𝑡 ，初期分布 𝑝data(𝑥0 ) ，終端分布 𝑞𝑇 (𝑥𝑇 ) を持つ連続時間マルコフ連鎖（順過程） 𝑥𝑡 𝑡 ∈[0,𝑇] が与えられたとする．この時，初期分布 𝑞𝑇 (𝑥𝑇 )，終端分布 𝑝data(𝑥0 )となるような遷移率 行列 𝑅෠𝑡 によって定まる連続時間マルコフ連鎖 𝑥𝑡 𝑡 ∈[𝑇,0] が存在し，これを逆過程と呼ぶ． また，逆過程の遷移率行列 𝑅෠𝑡 は順過程の遷移率行列 𝑅𝑡 を用いて以下のように表せる． ここで， しかしながら，q0|𝑡 (𝑥0 |𝑥)を解析的に表すことができないので，遷移行列 𝑅෠𝑡 は近似が必要． 16

逆過程の遷移率行列の近似 Recap 解析的に求まらない． そこで，確率モデルp𝜃0|𝑡 (𝑥0 |𝑥)により事後確率 q0|𝑡 (𝑥0 |𝑥)を近似する． ƿ 0 𝑞𝑡∣0 𝑥∣𝑥 𝜃 𝜃 ƶ 𝑅𝑡 (𝑥, 𝑥) ƿ = 𝑅𝑡 (𝑥, ƿ 𝑥)∑𝑥0 𝑝0∣𝑡 𝑥0 ∣ 𝑥 for 𝑥 ≠ 𝑥ƿ 𝑞𝑡∣0 𝑥∣𝑥0 𝑅ƶ 𝑡𝜃 (𝑥 ′ , 𝑥) = − ∑𝑥≠𝑥 ′ 𝑅ƶ 𝑡𝜃 𝑥, 𝑥 ′ for 𝑥 = 𝑥ƿ 同じ状態に留まる遷移率は順過程と同様に行和が0になるように正規化． 17

連続時間マルコフ連鎖のELBO 確率モデルのパラメータ 𝜃 はELBO最大化により学習する．以下に，連続時間マルコフ連鎖のELBOを 示す． Proposition2. CTMCにおけるELBO 遷移率行列R𝜃𝑡 ，終端分布 𝑝0𝜃 (𝑥0 )，初期分布 p𝑟𝑒𝑓 (𝑥𝑇 )を持つ連続時間マルコフ連鎖（逆過程） における負の対数尤度の上界 𝔼𝑝𝑑𝑎𝑡𝑎 𝑥0 −log 𝑝0𝜃 𝑥0 は以下で与えられる． ℒCT (𝜃) = 𝑇𝔼𝑡∼𝒰(0,𝑇)𝑞𝑡(𝑥)𝑟𝑡(𝑥∣𝑥) ∑𝑥 ′≠𝑥 𝑅ƶ 𝑡𝜃 𝑥, 𝑥 ′ ƿ − 𝒵 𝑡 (𝑥)log 𝑅ƶ 𝑡𝜃 (𝑥, ƿ 𝑥) + 𝐶, 𝐶はパラメータ𝜃に依存しない定数であり， 𝒵 𝑡 𝑥 = ∑𝑥′ ≠𝑥 𝑅𝑡 𝑥, 𝑥 ′ ， 𝑟𝑡 𝑥෤ ∣ 𝑥 = 1 − 𝛿𝑥,𝑥 ෤ Τ𝒵 𝑡 𝑥 ෤ 𝑅𝑡 𝑥 𝑥 𝑟𝑡 は時刻𝑡に状態遷移が生じるとわかっている時の遷移先の分布である． これより，CTMCの負のELBO最小化は，以下のような手順で行える． 1. 状態遷移が生じる時刻𝑡をサンプル 2. データ点𝑥0をサンプル 3. データ点を順過程の時刻𝑡まで摂動し， 𝑥をサンプル 4. 𝑟𝑡 に従い， 𝑥の次状態 𝑥を得る ƿ 18

連続時間マルコフ連鎖のELBO Proposition2. CTMCにおけるELBO 遷移率行列R𝜃𝑡 ，終端分布 𝑝0𝜃 (𝑥0 )，初期分布 p𝑟𝑒𝑓 (𝑥𝑇 )を持つ連続時間マルコフ連鎖（逆過程） における負の対数尤度の上界 𝔼𝑝𝑑𝑎𝑡𝑎 𝑥0 −log 𝑝0𝜃 𝑥0 は以下で与えられる． ℒCT (𝜃) = 𝑇𝔼𝑡∼𝒰(0,𝑇)𝑞𝑡(𝑥)𝑟𝑡(𝑥∣𝑥) ∑𝑥 ′ ≠𝑥 𝑅ƶ 𝑡𝜃 𝑥, 𝑥 ′ ƿ − 𝒵 𝑡 (𝑥)log 𝑅ƶ 𝑡𝜃 (𝑥, ƿ 𝑥) + 𝐶, 𝐶はパラメータ𝜃に依存しない定数であり， 𝒵 𝑡 𝑥 = ∑𝑥′ ≠𝑥 𝑅𝑡 𝑥, 𝑥 ′ ， 𝑟𝑡 𝑥෤ ∣ 𝑥 = 1 − 𝛿𝑥,𝑥 ෤ Τ𝒵 𝑡 𝑥 ෤ 𝑅𝑡 𝑥 𝑥 順過程の現在の状態と次状態のペア(𝑥, 𝑥)に対して， ෤ ・期待値中の第一項は，現在の状態 𝑥にとどまり続ける確率を高くする． ・期待値中の第二項は，𝑥෤ → 𝑥への逆過程の遷移率を最大化する． 19

順過程の設計 ここまでは，任意の遷移率行列についての順過程・逆過程の性質について示してきた． 実際には，遷移率行列は以下のような性質を持つ必要がある． (i) 十分速くp𝑟𝑒𝑓 に収束する (ii) 任意の時点tの分布𝑞𝑡|0(𝑥|𝑥0)を解析的に計算できる 時刻𝑡, 𝑡 ′ で遷移率行列𝑅𝑡 , 𝑅𝑡 ′ が可換であればこの条件を満たすことがわかっている．（Appendix E.参照） この研究では，時間非依存の遷移率行列𝑅𝑏 を用いて，異なる時刻の遷移率行列同士が可換であることを保証． 𝑅𝑡 = 𝛽(𝑡)𝑅𝑏 where 𝑅𝑏 ∈ ℝ𝑆×𝑆 時間依存の関数で，ノイズスケジュールのようなもの この時， 𝑞𝑡|0 (𝑥|𝑥0 )は以下のように解析的に求まる． ここで，𝑅𝑏 = 𝑄Λ𝑄−1 とする． Ex. 一様遷移（連続時間版） 𝑅𝑏 = 𝟏𝟏𝑇 − 𝑆Id where Id = diag(1, 1, … , 1) 20

次元間の独立性の仮定 言語や画像などの離散データは，各トークン（画素）において離散的な値を取る． ここでは，各要素が𝑆通りの離散状態のいずれかを取る𝐷次元の離散データ𝒙1:𝐷 ∈ 𝒳 𝐷 , |𝒳| = 𝑆を考える． 最も単純なのは，この離散データを𝑆 𝐷 次元の離散ベクトルとしてみなすことだが，効率的でない． そこでこの研究では，各次元の順過程を独立に考える．今は連続時間を考えているので，同時刻に二つ以 上の次元で遷移が生じる確率はゼロとなる．従って，すべての次元を考慮した場合でも，各遷移は常に一 つの次元だけの変化を意味する． 上記より，実際の遷移率の非ゼロ要素数は元の𝑆 𝐷 ではなく，𝐷 × 𝑆 − 1 + 1となる． （遷移先の候補は現在の状態を除いたS-1通りの状態に次元数を 掛けたものに，自己遷移を足した数） 仮に順過程で次元の独立性を仮定した場合においても，逆過程では次元間は非独立な遷移になりうる．こ れは，順過程の初期分布が次元間に強い依存関係を持つようなデータ分布であるからである． これについて，次ページで詳細に述べる． 21

次元間の独立性の仮定 Proposition3. 次元間の独立性を仮定した場合の逆過程 𝑑 𝑑 順過程が次元間の独立性を仮定しており，𝑞𝑡∣𝑠 𝒙1:𝐷 ∣ 𝒙1:𝐷 = ∏𝐷 𝑠 𝑡 𝑑=1 𝑞𝑡∣𝑠 𝑥𝑡 ∣ 𝑥𝑠 , 𝑡 > 𝑠, のように表されるとする．この時，順過程と逆過程の遷移率行列は以下のように表される． ここで，𝑅𝑡𝑑 ∈ ℝ𝑆×𝑆 であり， 𝛿𝒙1:𝐷∖𝑑 ,𝒙ƿ 1:𝐷∖𝑑 は𝑑番目の次元を除いて要素が等しい場合に1となるクロネッカー のデルタ記号である． 順過程の遷移率行列は次元間で非依存であることがわかる．その一方で逆過程の遷移率行列は全ての次元 の要素によって条件づけられた𝑥0 の事後分布を含んでおり，この事後分布が近似対象となるため生成モデ ルは次元間の依存関係を学習することになる． 22

Tau-Leapingによるサンプリング 学習した遷移率行列𝑅𝑡𝜃 からどのように効率的にサンプリングするか？ 連続マルコフ連鎖では，以下を繰り返せばよい（Gillespie’s Algorithm [6]）． (a) 現在の状態が持続する時間をサンプリング (b) 遷移率行列から，次に遷移する自分以外の状態をサンプリング -> これは高次元になるほど非効率． そこで，τ-leaping [7]と呼ばれる代表的な近似手法を導入． 核となるアイデア ・Gillespie’s Alogorithmは，指数分布に従う待ち時間モデルとしてみれる． ・ある時間間隔τに生じる遷移の回数はこれを拡張し，ポアソン分布に従う． ・一定の時間間隔τで遷移率が変化しないと仮定し，その区間に生じる遷移数をサンプリング． 23

連続時間への拡張による利点 Predictor-Corrector補正 ・τ-leapingによるサンプリングでは真の周辺分布𝑞(𝑥𝑇 )との誤差が大きくなる． ・真の周辺分布 𝑞(𝑥𝑇 )を定常分布として持つ補正遷移率行列を用いてサンプルを補正． Error Boundの導出 ・与えられたデータ分布に対して，τ-leapingによる初期分布の全変動距離を評価． ℒ(𝑦0)・・・ τ-leapingによる初期状態の分布 3𝑀𝑇・・・ 真の逆遷移率行列の近似誤差 ・第二項はτを小さくすることで消せる． -> 精度と速度のトレードオフ ・データ分布を定常分布に混合するまでの時間𝑡 𝑚𝑖𝑥により指数的に誤差が小さくなる 24

実験1- 画像生成 - CIFAR-10における画像生成 - 各画素は0-255の離散値を取る．順過程は離散化ガウシアン． - 提案手法（𝜏LDR−0/10）はD3PM [4]と比較して，生成品質が高い（ELBOでは劣る）． - 連続拡散モデルに匹敵する性能を示す． 25

実験1- 画像生成 - CIFAR-10における画像生成 - 左図：実際のサンプル例 - 右図：τ-leapingの時間間隔を大きくした際の推論速度の変化 26

実験2- 音楽生成 D3PMを上回る性能を示す． 譜面予測も可能． 27

-所感逆過程を辿る際のアプローチについて この研究では，τ-leapingおよびPredictor-Corrector補正による逆過程の高速化を実現している． 一方で，計算化学ではこのような系の高速なシミュレーションに関する他の手法も数多く提案されている． R-leapingやImplict τ-leaping , PP-SAなどでより高速かつ正確なシミュレーションが実現できそう． 遷移率行列の設計について 遷移率行列の制約は，異なる時点の遷移率行列同士が可換であることと，収束の速さ ． この研究では，時間不変な遷移率行列に時間変化する係数を掛けることで収束速度を調整． 上記の制約を満たす行列はこの形しかないのだろうか？ A. 多分そう．ただ，遷移率行列を対角化した際の固有値がどうなっているかに依存する． 固有値が縮退している場合（円環など）では困りそう．もう少しちゃんと考えたい． 28

参考文献 [1] Berglund et al., “The Reversal Curse: LLMs trained on "A is B" fail to learn "B is A"”, abs [2] Campbell et al., “A Continuous Time Framework for Discrete Denoising Models”, abs [3] Meng et al., “Concrete Score Matching: Generalized Score Matching for Discrete Data”, abs [4] Austin et al., “Structured Denoising Diffusion Models in Discrete State-Spaces”, abs [5] Hoogeboom et al., “Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions”, abs [6] Gillespie et al., “Exact stochastic simulation of coupled chemical reactions”, abs [7] Korno et al., “Tau-Leaping”, pdf 29

• https://arxiv.org/abs/2309.12288
• https://arxiv.org/abs/2205.14987
• https://arxiv.org/abs/2211.00802
• https://arxiv.org/abs/2107.03006
• https://arxiv.org/abs/2102.05379
• https://pubs.acs.org/doi/10.1021/j100540a008
• https://lup.lub.lu.se/luur/download?func=downloadFile&recordOId=9149938&fileOId=9149939