【選考チューニング】あるプロンプトをLLMに与えたときに、ある出力が得られる確率を求める

選考チューニングのDPOを勉強していて、理論は理解できたのですが、どうしても$\pi_\phi(y^1|x)$をどうやってLLMから算出するかがわからないままでした。
そんなこと知らなくても、ライブラリを使えば簡単に学習できるわけですが...ライブラリがあるから理論を知らなくていいということは決してないのでちゃんと調べてみました。

調べたことをメモがてら書いておきます（間違ってる可能性もありますが）。

定義

LLM: $\pi_\phi$
入力プロンプト: $x$
任意の出力: $y = w_1, w_2, \dots, w_N$
狙った出力: $y^1 = w_1^1, w_2^1, \dots, w_N^1$

語彙: $V$
語彙数: $|V|$
LLMモデルの出力: logitsと呼ぶ(入力のトークン数$\times|V|$ のテンソル)
logitをsoftmaxしたもの: Probabilitiesと呼ぶ(入力のトークン数$\times|V|$ のテンソル)

本題

プロンプト$x$をLLMに与えたときに出力$y$を得る確率

パラメータ$\phi$を持つLLM$\pi_\phi$にプロンプト$x$を渡したときに、出力$y = w_1, w_2, \dots, w_N$を得られる確率は下記の式で表されます。

$$P(y|x, \phi) = \pi_\phi(y|x) = \prod_{i=1}^N \pi_\phi (w_i | x, w_{<i})$$

これを求めるだけなら、行けそうですね。

まずは$\pi_\phi (w_1 | x)$の確率を求める。
これは、LLMに入力$x$を渡したときの出力(logit)にsoftmax変換をさせたProbabilitiesの最後の要素から確率を求められますね。Python風に書くとProbabilities[-1][出力トークンID]でしょうか。

$\pi_\phi (w_2 | x, w_q)$の確率を求める。
これは、LLMに入力$x, w_1$を結合して渡したときの、Probabilities[-1][出力トークンID]です。

と言うように、順にトークン$w_i$を出力させて、その確率を使えば$P(y|x, \phi)$を計算することができますね。

プロンプト$x$をLLMに与えたときに、狙った出力$y^1$を得る確率

上記のように任意の出力$y$を得る確率$P(y|x, \phi)$は簡単に求められることがわかりました。
これはLLMを自然に実行したときに得られるProbabilitiesを使えば簡単に求められました。
では、狙った出力$y^1$を出力する確率$P(y^1|x, \phi)$を得るにはどうしたらよいでしょう。

結論、このように算出しているみたいです。

まず、LLMに$x, y$を結合させたものを入力し、Probabilitiesを得ます。
言葉でうまく説明できないので、ところどころPythonコード風になりますが、
Probabilities[len(x)-1][y^1の1トークン目のトークンID]
Probabilities[len(x)][y^1の2トークン目のトークンID]
Probabilities[len(x)+1][y^1の3トークン目のトークンID]
...
Probabilities[len(x)+len(y)-2][y^1のlen(y)トークン目のトークンID]
がそれぞれ、$\pi_\phi (w_1^1 | x), \pi_\phi (w_2^1 | x, w_1^1), \pi_\phi (w_3^1 | x, w_1^1, w_2^1), \dots, \pi_\phi (w_N | x, w_{<N})$に対応します。

さて、なぜ$x, y^1$を結合させたものを入力とした出力が、$P(y^1|x, \phi)$となるのでしょう？
未来の情報が入ってしまって大丈夫でしょうか？

じつは、大丈夫なのです。
LLMにはCausal Attentionという、未来の情報をマスクするAttentionが使われてます。このことから、上記のように$x, y^1$を入力として確率を計算しても問題のです。

調査終わり。