Direct Preference Optimizationを読む（その2）

　その1

でDPOの損失関数

$\mathcal{L} _ \mathrm{DPO} (\pi _ \theta; \pi _ \mathrm{ref})= - \mathbb{E} _ {(x, y_w, y_l) \sim \mathcal{D}} \left\lbrack \log \sigma \left( \beta \log \frac{\pi _ \theta (y _ w|x)}{\pi _ \mathrm{ref}(y _ w|x)} - \beta \log \frac{\pi _ \theta (y _ l|x)}{\pi _ \mathrm{ref}(y _ l|x)} \right) \right\rbrack \quad(7)$

が導出できたので、この関数の性質を分析してみます。

勾配がどうなっているか

　まず微分してみます。整理するために

$\hat{r} _ \theta (x, y) = \beta \log \frac{\pi _ \theta (y _ w|x)}{\pi _ \mathrm{ref}(y _ w|x)}$

と置きます。後にこれは暗黙の報酬モデルであることが明らかになりますが、それはさておいて、まずは勾配を求めます。 $f = \hat{r} _ \theta(x, y _ w) - \hat{r} _ \theta(x, y _ l)$ とすると、損失関数の中身は

$h = \log \sigma(f)$

という形式になっているので、連鎖律で

$\frac{\partial h}{\partial \theta} = f' \sigma(f) (1 - \sigma(f)) \frac{1}{\sigma(f)}$

となり、ロジスティック関数の性質から $1 - \sigma(f) = \sigma(-f)$ なので、結局勾配は $f' \sigma(-f)$ となります。この1つ目は

$f' _ \theta = \nabla _ \theta \log \pi _ \theta (y _ w | x) - \nabla _ \theta \log \pi _ \theta (y _ l | x)$

であり、単純に良いものを増加させ、悪いものを低下させるように働きます。2つ目の

$\sigma(-f) = \sigma(-\hat{r} _ \theta(x, y _ w) + \hat{r} _ \theta(x, y _ l))$

がその係数であり、つまり報酬によって上手いこと重み付けがなされると解釈できます。

DPOのさらなる理論解析

　まず次のことを定めます。

定義1. 2つの報酬関数 $r(x, y)$ と $r'(x, y)$ が、ある関数 $f$ に対して $r(x, y) - r'(x, y) = f(x)$ となる場合、その2つの報酬関数は等価であるという。

　これは報酬関数の集合をクラスに分割する同値関係を定めていることになります。

　この同値関係について、次の補題2つがあります。

補題1. 同じクラスにある2つの報酬関数は、Plackett-Luceモデル（特にBradley-Terryモデル）の下で同じ嗜好分布を導く
補題2. 同じクラスにある2つの報酬関数は、制約付きRL問題の下で同じ最適方策を導く

　嗜好分布とは回答の良し悪し $y _ w, y _ l$ の組についての分布で、補題1はunder-specification問題としてPlackett-Luceモデル族については知られている話のようです。報酬が一意に上手く定まらないわけですが、逆に補題2から、クラスさえ決まってしまえばどの報酬関数でも問題ないことがわかります。

【補題1の証明】

　一般的なPlackett-Luceモデル（Bradley-Terryモデルはこの内 $K = 2$ の特殊なケース）を考えます。特定の報酬関数 $r'(x, y) = r(x, y) + f(x)$ によって誘導されるランキング上の確率分布を $p _ r$ とします。手順はほぼ自明で、expなので和が掛け算としてくくり出せて約分できるので同じということになります。

$p _ r' (\tau | y _ 1, \dots, y _ K, x) = \Pi _ {k = 1} ^ {K} \frac{ \exp(r'(x, y _ {\tau(k)})) }{ \sum _ {j = k} ^ {K} \exp (r'(x, y _ {\tau(j)} ))} \\ = \Pi _ {k = 1} ^ {K} \frac{ \exp(r(x, y _ {\tau(k)}) + f(x)) }{ \sum _ {j = k} ^ {K} \exp (r(x, y _ {\tau(j)} ) + f(x))} \\ = \Pi _ {k = 1} ^ {K} \frac{ \exp(r(x, y _ {\tau(k)})) }{ \sum _ {j = k} ^ {K} \exp (r(x, y _ {\tau(j)} ))} \\ = p _ r (\tau | y _ 1, \dots, y _ K, x)$

　補題2の証明もほぼ同様なので省略します。最適方策が $\pi _ \mathrm{ref}(y | x) \exp \left( \frac{1}{\beta} r(x, y) \right)$ に比例する形でかけるのでexpが打ち消し合います。

　そして、以下の定理が成り立ちます。

定理1. 適当な仮定の下で、Plackett-Luce（特にBradley-Terry）モデルと一致するすべての報酬クラスは、再パラメータ化によって表現できる。この再パラメータ化は、ある方策 $\pi (y | x)$ と与えられた参照方策 $\pi _ \mathrm{ref}(y | x)$ を用いて、 $r(x, y) = \beta \log\left(\frac{\pi(y|x)}{\pi _ \mathrm{ref}(y|x)}\right)$ の形をする。

（ちょっと書き方がわかりにくいのですが、要するに報酬をクラスとしてしか指定できていなかったところから、そのうちの一つに定まるという意味合いなのだと思います）

　証明の概要として、最適方策 $\pi _ r (x, y)$ を導く報酬のクラスに属する任意の関数 $r(x, y)$ に対して、以下のような写像 $f$ を考えます。

$f(r; \pi _ \mathrm{ref}, \beta) = r(x, y) - \beta \log \sum _ y \pi _ \mathrm{ref} (y|x) \exp \left( \frac{1}{\beta} r(x, y) \right)$

これは要するに $r(x, y)$ を $\pi _ r$ の分配関数で正規化するような操作になっています。また、第二項は $x$ についてだけの関数なので、定義1より報酬関数としてのクラスは変わらないことがわかります。ここで、その1で求めた報酬関数についての等式 $(5)$ （これはどの報酬関数についても成り立ちます）