手を動かして拡散言語モデルの仕組みを理解する

はじめに

拡散言語モデル（Diffusion Language Models）が最近注目を集めていますが、公開実装の多くは大規模で、個人の学習環境では動かすのが困難です。論文を読んでもピンとこないし、かといって実際に手を動かす機会はない方も多くいると思います。

そこで本記事では、拡散言語モデルの核心部分である離散拡散モデル（吸収拡散）を、最小限の実装で動かしてみます。ニューラルネットワークは一切使わず、Julia言語でシンプルに実装することで、アルゴリズムの本質的な動作を体験できるようにしました。

「理論はよくわからないけど、とりあえず動かしてみたい」という人向けの内容となっています。

拡散言語モデルとは

最近の高性能画像生成AIモデルでよく採用される拡散モデル（Diffusion Models）の手法を言語分野に持ち込んだものです。結論から話してしまえば、マスク言語モデルを改良し、一度に全てのマスクトークンを推論する代わりに、徐々に確定させていくようなものです。

従来のマスク言語モデルでは「[MASK]は[MASK]で[MASK]です」のような複数のマスクを含むトークン列を一度に「今日は晴れで暖かいです」と予測しますが、拡散言語モデルでは「今日は[MASK]で[MASK]です」→「今日は[MASK]で暖かいです」→「今日は晴れで暖かいです」のように段階的に生成を行います。

(連続値) 拡散モデル

拡散モデルとは、一言で言えば、ノイズ除去を段階的に繰り返すことで、所望の特徴を持つデータを生成する手法です。

拡散モデルの生成過程および拡散過程。初期ノイズ (左) から出発し、ノイズ除去を繰り返すことで所望の画像 (右)^[1] を得る

初期値として完全にランダムなノイズを設定し、弱いノイズ除去を繰り返すことで、目的とする画像が徐々に浮かび上がってきます。この繰り返しノイズ除去を行う手続きは逆拡散過程と呼ばれています。

...「逆」？なぜ生成過程が「逆」と呼ばれているのか、初見で疑問に思う方は多いかと思います。これは歴史的経緯により、データを徐々に崩壊させてノイズにする過程を （順方向の）拡散過程 と定義していることが理由です^[2]。この順方向過程は生成段階では直接利用されませんが、学習段階で「どの程度ノイズが混入した中間データか」を作成するために使われます。

なぜこの方法で良い画像を生み出せるのか？といった疑問は色々あると思いますが、本記事では筆者が手抜きをしたいので扱いません。最近は無料でも高品質な入門文献が増えてきているので、興味がある人はググるなりチャピる^[3]なりして、ご自身でお調べいただければと思います。

離散拡散モデル

先程までの拡散モデルの話を総合すると、拡散言語モデルとは「拡散モデルの適用対象をピクセルからトークンに置き換えたもの」と予想することができます。そして実際にそれは正しいです。したがって、後は画像生成拡散モデルの仕組みを理解するだけで、簡単に言語タスクに応用できます😉✨
...なんてうまい話はありません、あゝ無情。

画像ドメインと言語ドメインには重要な違いがあります。それは 言語モデルが本質的に離散的である という点です。例えば、画像ドメインにおけるピクセル値であれば、「ピクセル値が+10された」という操作は「そのピクセルが明るくなった」と自然に解釈することが可能です。一方で言語ドメインの場合、「トークン値が+10」されたというのは「『ラーメン』という単語が『トーテムポール』に置き換わった」などの自然な解釈が困難な変化になったりします。

このような離散性の問題があるため、言語モデルへの応用においては 離散データに適した改良 が施された手法が主に利用されています。これが離散拡散モデルと呼ばれるものです。

離散拡散モデルと通常の（連続値）拡散モデルの重要な違いの一つは 拡散過程の設計 と言えます。具体的には、離散拡散モデルでは 吸収拡散（別名：マスク拡散）と呼ばれる拡散過程^[4]が一般的に採用されます。この過程では、与えられたデータの各要素が徐々にマスク状態（ピクセルやトークンの中身が隠された状態）という特殊な状態に置き換わり、最終的には全ての要素がマスク状態に置き換わります。

吸収拡散の順方向過程のサンプル一例 (GIFアニメーション)。時間が経つにつれて、それぞれのピクセルが徐々にマスク状態（緑色）に侵食される様子が見て取れます

生成過程ではその逆に、全ての要素がマスク状態のデータから出発し、徐々に要素を確定させていきます。

実装

ここからは、吸収拡散モデルの簡単なサンプル実装を作成していきます。順方向過程（拡散過程）と逆方向過程（逆拡散過程）のそれぞれを実装し、実際に動作を確認してみましょう。

教科書的には離散拡散モデルの理論基盤から説明するべきではありますが、本記事では「とりあえず動かしてみる」を目標としているため、理論解説は実装に必要最小限度のものにとどめています。理論背景をもう少し詳しく知りたい方は、怪文書じみた拙著「お嬢様でもわかる拡散言語モデル」をご参照ください^[5]。

https://zenn.dev/v_ajw63600/books/diffusion-language-models

また本記事では、ニューラルネットワークに基づく推論モデルは一切使用しません。これには以下の理由があります：

• 機械学習ライブラリの使用方法やモデル設計、学習スケジューリングなど、離散拡散モデルの本質とは直接関係ない高度な技術が多く含まれる
• コードの改良や実験の柔軟性を保ちたい
• 筆者の実装スキルの問題

実装には Julia 言語を使用します。機械学習フレームワークを使わない場合、数値計算や行列操作の観点で Python より扱いやすいためです。

順方向

ここでは実際に順方向過程、すなわち、トークン列を与えて徐々にマスクトークンに置き換わるようなシミュレーションのプログラムを作成します。実用的なモデルにおいては、順方向過程のシミュレーションは通常利用されませんが、逆方向過程のイメージを掴む準備には最適なので、あえてこの項目を設けています。まあ、逆方向過程は比べ物にならないほど難しいので、雀の涙程度の助けにしかならないとは思いますが。

順方向過程では、時刻が進むにつれて各トークンが確率的にマスク状態に遷移していきます。この設計において重要なのは、時刻 $t \in [0, 1]$ におけるマスクトークンの割合が $\lambda_t = t$ となるように調整することです。さらに、この変化は段階的で連続的 (前の時刻でマスク状態に置き換わったトークンは、次の時刻でもマスク状態) である必要があります。この条件を満たすような確率過程は、次に示すような方法で実現できることが知られています。

時刻 $t$ におけるトークン系列を $\mathbf{x}_t$、その $i$ 番目要素を $x_t^i$ と表記します。このとき、吸収拡散の順方向過程では、次のようにトークンごと独立にサンプルします。

$$x_{t + \tau}^i = \begin{cases} x_t^i & \mathrm{確率} \ 1 - \beta_t \\ m & \mathrm{確率} \ \beta_t \\ \end{cases}$$

ここで $\beta_t = \frac{\tau}{1 - t}$ としています。言葉で表せば、位置 $i$ のトークンは微小時間 $\tau$ の間に確率 $\beta_t$ でマスク状態 $m$ に遷移することを意味します。

これらの手続きを Julia で実装したものが以下となります:

using LinearAlgebra
using Random
using Distributions: Categorical

const L = 100         # トークン系列長
const K = 10          # カテゴリ数
const MASK_INDEX = 10 # マスク状態に対応するインデックス

"""
one-hot ベクトルを作成する関数
"""
function one_hot(n::Integer)
    Int[i == n for i in 1:K]
end

"""
確率質量ベクトルの系列からサンプリングする関数
"""
function sample_from_categorical(prob_matrix::AbstractMatrix)
    mapslices(rand ∘ Categorical, prob_matrix; dims=1) |> vec
end

"""
時刻 ``t`` におけるトークン系列から、次の時刻 ``t + \\tau`` のトークン系列をサンプリングする関数
"""
function forward_step(
    x::AbstractVector{<:Integer}, # 現在のトークン系列
    t::Real,                      # 現在の時刻
    tau::Real,                    # 微小時間
)
    # マスク状態に対応する one-hot ベクトル
    mask_vec = one_hot(MASK_INDEX)

    # マスク状態への遷移確率を計算
    beta = tau / (1 - t)

    # 各座標ごとに確率分布を計算 (各列ベクトルが遷移確率ベクトルに対応)
    next_probs = zeros(K, L)
    for i in 1:L
        # 現在の状態に対応する one-hot ベクトル
        current_vec = one_hot(x[i])

        # 遷移確率 (確率ベクトル) の計算
        next_probs[:, i] = (1 - beta) * current_vec + beta * mask_vec
    end

    # 各列ごとに確率ベクトルから状態をサンプリング
    sample_from_categorical(next_probs)
end

この実装では確率的な条件分岐を直接記述する代わりに、確率ベクトルの遷移として実現していることに注意してください。

適当に設定された初期トークン系列に対して、上述コードの forward_step() をループで適用してあげれば、シミュレーションの実装は完了です。ということで以降はシミュレーション全体のコードの実装例を示しま... っと、話が長くなる悪い癖が出てしまいました。シミュレーションに関する残りの部分は、手間が掛かるわりにはあまり面白みがないので （←プログラマに対して失礼な発言）、筆者が実装した完全版のスクリプトを共有することで代えたいと思います。

実装スクリプト (Pluto notebook^[6]) は以下から入手できます。ダウンロードして手元の環境で実行してみてください。

https://github.com/v-ajw63600/hands-on-diffusion-lm/blob/main/01-simulation-forward.jl

このシミュレーションスクリプトを実行すると、ランダムに設定した初期トークン系列が徐々にマスクトークンへと置き換わる軌道を可視化した結果を取得できます。その一例が以下のヒートマップです:

順方向シミュレーションのヒートマップ。縦軸がトークン位置、横軸が時間を表す。色はトークンの値を表し、黒がマスクトークンに対応。

全てのトークンが非マスク状態から出発し、最終的にマスク状態に収束していることが確認できると思います。また、状態の遷移は「非マスク → マスク」に限り発生し、非マスク状態は他の非マスク状態に遷移しないことと、一度マスク状態に遷移したら他の状態には遷移しないこと （吸収状態） も同時に確認できると思います。

また、先程のスクリプトでは、シミュレーションで得られたトークン系列のマスク率と、理論値の比較も確認することが可能です。先程のヒートマップの実験に対応するマスク率の比較は以下のようになります:

マスク率の理論値-実測値比較。青が理論値。

冒頭で説明したように、理論的なマスク率は青色の直線 $\lambda_t = t$ で表されます。実際のシミュレーション結果も概ね理論値に沿っていることが確認できます。今回のシミュレーションにおいては偶々理論値に近いグラフが得られましたが、確率過程なので理論値から大きく離れたグラフになることもあります。

ただし、繰り返しになりますが、実用的なモデルでは上述のシミュレーションは通常利用されません。かなりラフな言い方をすれば、わざわざシミュレーションなどせずとも、任意時刻 $t$ のトークン系列は 1 ステップでサンプリングできることが理由です。

逆方向

逆方向（生成過程） も同様に逐次サンプリングによって実現することができます。すなわち、時間の流れを逆向きにして、時刻 $t + \tau$ におけるトークン系列 $\mathbf{x}_{t + \tau}$ が与えられたときの過去の系列 $\mathbf{x}_t$ を予測します。

この逆過程の遷移は、連続時間マルコフ連鎖（CTMC）における順方向遷移率行列 $\mathbf{Q}_t$ に対応する逆方向遷移率行列 $\tilde{\mathbf{Q}}_t$ を構成することでコルモゴロフ方程式 $\frac{d \mathbf{q}_t}{dt} = \tilde{\mathbf{Q}}_t \mathbf{q}_t$ の解である行列指数関数として得られ...
...などと小難しい話をしているといつまで経っても本題に入れないので、細かい理屈は脇において、以下のような結果を天下り的に受け入れることとします^[7]。...説明から逃げているわけではないですよ？理屈はちゃんと理解してますよ本当です信じてくださいお願いします😭

気を取り直して、トークン系列 $\mathbf{x}_{t + \tau}$ が与えられたとき、過去の時刻 $t$ における系列 $\mathbf{x}_t$ は以下の式で与えられます。

$$x_t^i = \begin{cases} x_{t + \tau}^i & x_{t + \tau}^i \neq m \\ y \sim \mathrm{Cat}(\cdot; \mathbf{p}^i (\mathbf{x}_{t + \tau})) & x_{t + \tau}^i = m \\ \end{cases}$$

この式は、現在のトークンが非マスク状態ならば値は変わらず、マスク状態ならば確率的にサンプリングすることを意味します。

ここで問題となるのが、マスク状態からの遷移の具体的な内容です。カテゴリカル分布からのサンプリングであることから、状態 $y$ がサンプリングされる確率は $[\mathbf{p}^i (\mathbf{x}_{t + \tau})]_y$ で与えられることになります。この確率ベクトル $\mathbf{p}$ の値は次のように与えられます（Tweedie $\tau$-leaping からの帰結）：

$$[\mathbf{p}^i (\mathbf{x}_{t + \tau})]_y = \begin{cases} ((1 + \beta_t)^{r^i (\mathbf{x}_{t + \tau})} - 1) \frac{[\mathbf{s}^i (\mathbf{x}_{t + \tau}, t + \tau)]_y}{r^i (\mathbf{x}_{t + \tau})} & y \neq m \\ (1 + \beta_t)^{r^i (\mathbf{x}_{t + \tau})} & y = m \\ \end{cases}$$

ここで $r^i (\mathbf{x}) = -\sum_{z \neq m}[\mathbf{s}^i (\mathbf{x}_{t + \tau}, t + \tau)]_z$ と定義しました。
結局、この $\mathbf{s}^i (\mathbf{x}_{t + \tau}, t + \tau)$ とは何なのか？という謎はまだ残っていますが、その一点を除けば、逆過程のシミュレーションは実行可能です。そのため $\mathbf{s}$ については一旦後回しにして、まずは逆方向のステップを実装したいと思います。

"""
時刻 ``t + \\tau`` におけるトークン系列から、過去の時刻 ``t`` のトークン系列をサンプリングする関数
"""
function reverse_step(
    x::AbstractVector{<:Integer},
    t::Real,
    tau::Real,
    s::AbstractMatrix,
)
    # マスク状態に対応する one-hot ベクトル
    mask_vec = one_hot(MASK_INDEX, K)

    # 順方向過程におけるマスク状態への遷移確率
    beta = tau / (1 - t)

    # 各座標ごとに確率分布を計算
    next_probs = zeros(K, L)
    for i in 1:L
        if x[i] == MASK_INDEX
            # 各座標における `s` を取り出す
            s_vec = (1 .- mask_vec) .* s[:, i]
            r = -sum(s)
            s_vec[MASK_INDEX] = r # NOTE: 計算の単純化のため、自己遷移に対応する `s` の要素を `r` で置き換える

            # 遷移確率の計算
            next_probs[:, i] = mask_vec + ((1 + beta)^r - 1) / r * s_vec
        else
            # 非マスク状態は変化しない（吸収状態）
            next_probs[:, i] = one_hot(x[i], K)
        end
    end

    # サンプリング
    sample_from_categorical(next_probs)
end

密度比とクリーンデータ分布

逆過程の実行には関数 $\mathbf{s}$ が必要であることを前節では説明しました。これは密度比^[8]と呼ばれる値であり、生成データの分布を特徴付ける最も重要な要素です。

順方向過程が吸収拡散で定義される場合、この密度比は次のように与えられます:

$$\mathbf{s}^i (\mathbf{x}, t) = \frac{1 - t}{t} \mathbf{c}^i (\mathbf{x}) \quad \mathrm{where} \quad x^i = m$$

ここで $\mathbf{c}$ はクリーンデータ分布と呼ばれ、訓練データ分布 $p_\mathrm{data}$ を用いて次のように定義されます:

$$[\mathbf{c}^i (\mathbf{x})]_{y} = p_\mathrm{data}\left(y^i|\mathbf{x} \right) \quad \mathrm{where} \quad x^i = m$$

言葉で説明するならば、$i$ 番目トークンがマスク状態であるトークン系列 $\mathbf{x}$ が与えられたときに、 $i$ 番目トークンの真の値が $y$ である確率を表します^[9]。
「どこにも密度の比が出てこないじゃないか！」とツッコミが飛んできそうではありますが、これにはちゃんと理由があります。が、その説明を書くには余白が狭すぎるので、興味がある人は RADD の論文 の Theorem 1 あたりを読んでください（投げやり）。

実践的な拡散言語モデルでは、このクリーンデータ分布をニューラルネットでモデル化します。しかし前述の通り、本記事ではニューラルネットは使わないと宣言しました。そのため、実装では代替物として、経験分布から直接計算する方式を採用します^[10]。

訓練データセットのサンプルを $\{\mathbf{x}_{\mathrm{GT}, n}\}_n$ としたとき、その経験クリーンデータ分布は次のように与えられます:

$$p_\mathrm{data} (y^i|\mathbf{x}) = \frac{\hat{\mathrm{x}}\mathrm{ と非マスクトークンが一致する GT の個数}}{\mathrm{x}\mathrm{ と非マスクトークンが一致する GT の個数}}$$

ここで $\hat{\mathbf{x}}$ は、$\mathbf{x}$ の $i$ 番目トークンを $y$ に置き換えたものと定義しています。導出方法については割愛します。

...色々説明不足なことは明らかではありますが、実装するだけであれば「通常はクリーンデータ分布を Transformer で学習する」ぐらいの理解で十分だと筆者は思っているので、ざっくりとした理論の説明は以上で終わりにします。

以下のコードが、経験クリーンデータ分布ならびに密度比の Julia 実装となります。一致する Ground Truth の個数を数える関数 (count_compatible_sequences) がややショートコーディングの傾向がありますが、それ以外の部分は素直な実装なので、特に説明は不要でしょう。

"""
非マスクトークンが一致する Ground Truth の個数を数える補助関数
"""
function count_compatible_sequences(
    ground_truths::AbstractVector,
    x::AbstractVector{<:Integer},
)
    mask_pattern = (x .!= MASK_INDEX)

    compatible_count = mapslices(ground_truths; dims=2) do row
        # マスクされていない位置が全て一致するかチェック
        all((mask_pattern .* row) .== (mask_pattern .* x))
    end |> count

    return compatible_count
end

"""
クリーンデータ分布を計算

NOTE: 通常はニューラルネットで学習する
"""
function compute_clean_distribution(
    ground_truths::AbstractVector,
    x::AbstractVector{<:Integer}
)
    distribution = zeros(K, L)

    current_compatible = count_compatible_sequences(dist, x)

    if current_compatible == 0
        # 互換性のある系列がない場合（全てのトークンがマスク状態から出発した系列の場合、理論的には起こらないはず）
        return distribution
    end

    for i in 1:L
        if x[i] == MASK_INDEX
            # マスク位置の場合、各可能な値に対する確率を計算
            for k in setdiff(1:K, MASK_INDEX)  # マスク以外の値
                # x[i]をkに置換した場合の系列を構築
                x_modified = copy(x)
                x_modified[i] = k

                # この置換後の互換性のある系列数を計算
                modified_compatible = count_compatible_sequences(dist, x_modified)

                # 確率は互換性のある系列数の比率
                distribution[k, i] = modified_compatible / current_compatible
            end
        else
            # 非マスク位置の場合、状態 x[i] に対応する one-hot ベクトル
            # NOTE: この値は逆過程のシミュレーションでは利用されない
            distribution[x[i], i] = 1.0
        end
    end

    return distribution
end

"""
密度比を計算
"""
function compute_density_ratio(
    ground_truths::AbstractVector,
    x::AbstractVector{<:Integer},
    t::Real,
)
    # NOTE: 厳密に言えば、以下の式は `x[i] != MASK_INDEX` となる要素 `x[i]` については正しくないが、逆過程のシミュレーションでは利用されないため実用上問題ない
    (1 - t) / t * compute_clean_distribution(ground_truths, x)
end

シミュレーション結果

密度比並びにクリーンデータ分布というピースが出揃ったことで、ついに逆過程のシミュレーションの実装は完成します......が、例のごとく、逆過程についても筆者の実装スクリプトの共有に代えたいと思います。説明がめんどくさくなったとか、断じてそんな理由ではありませんからね！

ということで、逆過程の実装スクリプトは以下となります。多少記事中とは異なる実装となっている箇所はありますが、ロジックの違いはないはずです。

https://github.com/v-ajw63600/hands-on-diffusion-lm/blob/main/02-simulation-reverse.jl

このスクリプトを実行すると、順方向過程と同様にシミュレーション結果のヒートマップを取得できます。筆者が実行した結果は以下のとおりです:

逆過程 (生成過程シミュレーション)

左側は Ground Truth の可視化、右側は逆過程 (生成過程) のシミュレーション結果ヒートマップです。生成過程では $t = 1.0$ から $t = 0.0$ の方向に時間が流れることに注意してください。

この図における試行では、マスク状態から出発したトークン系列は、ID #1 の Ground Truth と等しい値に収束しました。スクリプトにおけるシミュレーションのシード値をいじってみると確認できると思いますが、生成結果はシード値によって変化しますが、どの場合でも必ず Ground Truth のどれかと完全一致する値に収束します^[11]。これは経験分布を Ground Truth として直接利用したことに起因するもので、期待通りの結果です。
(もう少しちゃんと実験するのであれば、固定された Ground Truth を用意してシミュレーションを複数回実施し、クロスエントロピーを評価したほうが良いです。面倒なので誰かやってください。)

そして、上述のデータに対応するマスク率の比較グラフは次のようになりました:

逆過程におけるマスク率の比較。青が理論値。

順方向過程と同じく、おおよそ理論値に近いグラフが得られています。$t \in [0.0, 0.5]$ の付近では理論値よりやや離れているように見えますが、確率過程としてはそれほど珍しくはない程度のばらつきです。

おわりに

本記事では、拡散言語モデルの基盤技術である離散拡散モデル（吸収拡散）について、最小限の実装を通じて体験しました。ニューラルネットワークを使わず Julia 言語でシンプルに実装することで、アルゴリズムの本質的な動作を体験していただけたのではないでしょうか。まあ、理論の部分はかなりすっ飛ばしてしまったので、これを読んだだけで「完全に理解した」と言い切るのはどだい無理な話ではありますが...

更に本記事は、実用上の重要度が高いはずの学習方法については、一切取り扱いませんでした。こちらについては、他の文献を読んで理解を深めるのが良いでしょう......と一応言っておきます。正直なところ「Transformer アーキテクチャのモデルを組んで、様々なマスクパターンでクロスエントロピー損失を最小化する」といった程度の話なので、言語モデルの経験がある人物ならば特に躓くことなく理解および実装ができるものだと予想しています。

なんか gdgd になってしまいましたが、この記事を読んで実際に手を動かしたことで、拡散モデルのブラックボックス感が少しでも和らいだなら嬉しいと思います。

脚注

1. サンプルに利用した画像は、筆者の AI 生成による「スチームパンクホラガイ」です。「サザエでしょ？」といったような野暮な指摘は受け付けておりません。悪しからず。

2. 「順」と「逆」の定義は歴史的経緯によるものなので、逆方向が「間違っている」などといったような付加的な意味合いは特にありません。実際、後続研究である Flow Matching や Diffusion Bridge Models の一部では、生成過程を順方向として定義しています。そのせいで文献比較の際に、生成過程の時間方向を毎度確認したり揃えたりする手間が追加されたりするので、非常にめんどくさい。解読する側の身にもなってくれ。

3. ChatGPT や Gemini などのチャット AI を利用して情報を調べることを意味する筆者の造語。筆者の言語センスの低さが (悪い意味で) 光り輝く。

4. もはや数理物理学における拡散現象とは関係がありません。そのため本来ならば、メカニズムに基づく名称としての「吸収拡散 (absorbing diffusion)」は不適当です。そもそも「吸収」と「拡散」は直感的には矛盾するイメージなので、撞着語法的な表現になってしまっています。なぜこのような用語が定着しそうになっているのか、筆者にはその理由がわかりません。しかしながら、人生とは、得てしてそういうものなのです。

5. 概要を見ていただければ分かるように、巫山戯た文体の紛うことなき怪文書ではあるのですが、大学院の講義ノートとして利用可能な程度の内容と品質であるとは自負しています。拡散言語モデルの発展を俯瞰して体系的に解説した日本語の技術文書は、本書を除き、現時点では存在しないと思われます。どうしてこうなったの...

6. Julia 言語における Jupyter のオルタナティブです。使い方は 公式ページのインストールの項目 を見ればすぐに理解できると思います。もちろん、Julia 環境をローカルで構築できる程度の実力は最低限必要ですが。

7. 一応、要点を簡潔に説明すると次のようになります:
   順方向過程の遷移確率 $p(\mathbf{x}_{s}|\mathbf{x}_t)$ は遷移率行列 (transition rate matrix) $\mathbf{Q}_t$ から算出できることが知られています。その順方向に対応する逆方向過程の遷移率は、順方向の転置 $\mathbf{Q}^\top_t$ から構成される遷移率行列 $\tilde{\mathbf{Q}}_t$ で与えられます。逆方向の遷移は一般に全てのトークンを考慮して同時にサンプリングしなければならないのですが、今回利用するサンプリング手法 (Tweedie $\tau$-leaping) では各次元独立に適用する近似を採用します。この近似の下では、各次元 $i$ の遷移確率は、対応する遷移率行列 $\tilde{\mathbf{Q}}_t^i$ によって規定されます。更に、微小時間 $\tau$ の間では $\tilde{\mathbf{Q}}_t^i \approx \beta_t \tilde{\mathbf{Q}}_\mathrm{approx}$ で表せるという近似を追加で適用すると、時刻 $t$ から $t - \tau$ への遷移確率は行列指数関数 $\exp(\overline{\beta}_{t \to t - \tau}\tilde{\mathbf{Q}}_\mathrm{approx})$ で与えられるので、その解析解を求めると、先程の遷移確率を導出できます。
   (...全然簡潔ではない。まさにオタク特有の早口である。)

8. 文献によってはコンクリートスコアと呼ぶ場合がありますが、厳密に言えば誤用なので、筆者は「密度比」の用語を推奨しています（実装における変数名は "score" になっていますが...）。雑に説明すると、コンクリートスコアの本来の定義は「密度比マイナス 1」となります。数値計算のときに取り違えてしまうと痛い目を見ること請け合いです。

9. $\mathbf{x}$ の $i$ 番目要素がマスクでないケースがカバーされていないことに気付いた方も多いかと思いますが、これは意図的なものです。逆過程の実行時には参照されないので、わざわざ計算する必要がないことが理由です。あえて定義するのであれば $[\mathbf{c}^i (\mathbf{x})]_{\hat{x}} = \delta_{x^i, \hat{x}^i} \quad (x^i \neq m)$ となります。

10. 当たり前ですが、実用的なスケールのモデルで経験分布を直接採用することはまずありえません。これは、推論時に必要となるデータ量と計算量が少なくとも線形以上のレートでスケールするためです。

11. 実際には、シード値次第で偶にエラーが発生します。これは実装ミスではなく、サンプリングアルゴリズムと経験分布の組み合わせに起因するものです。各次元独立にサンプリング可能という近似を行ったことで、学習データのパターンからは絶対に有り得ないトークンの組み合わせが出てくるケースが出てきます。そのようなケースに出くわした場合、クリーンデータ分布が適切に計算できずに失敗します。この問題は、適切に訓練されたニューラルネットならば通常は発生しません。