從“RL for LLM”視角重新理解KL近似：關於“近似KL散度”的筆記

PPO和GRPO中使用的KL散度估計方法有什麼區別？

John Schulman的部落格文章“近似KL散度”討論瞭如何透過取樣（蒙特卡羅）近似KL散度，並介紹了三種估計器（\(k_1\)、$k_{2}k\_2$、$k_{3}k\_3$）及其偏差-方差行為。但原始文章是在一般機率分佈的背景下提出的，並未涉及大型語言模型（LLM）的強化學習訓練設定。本文記錄了我在閱讀時遇到的問題、將內容對映到RL for LLM後形成的思考，以及一些我認為原始解釋可以進一步闡述的地方。

“近似KL散度”說了什麼（用我自己的話）

在本節中，我假定讀者尚未閱讀原始文章，因此我們快速瀏覽最重要的部分。簡單來說，這篇文章是關於當我們無法直接計算KL散度時，如何構建合理的蒙特卡羅式估計器。

$$\mathrm{K}\mathrm{L}(q,p)=\sum _{x}q(x)\log \frac{q(x)}{p(x)}={\mathbb{E}}_{x\sim q}\text{\hspace{0.17em}⁣}[\log \frac{q(x)}{p(x)}]\mathrm{.}\backslash mathrm\{KL\}(q,\; p)\; =\; \backslash sum\_x\; q(x)\backslash ,\backslash log\backslash frac\{q(x)\}\{p(x)\}\; =\; \backslash mathbb\{E\}\_\{x\backslash sim\; q\}\backslash !\backslash left[\backslash log\backslash frac\{q(x)\}\{p(x)\}\backslash right].$$

如公式所示：當估計兩個（複雜）分佈之間的KL散度時，人們常用一種編碼技巧：即透過從$qq$中抽樣，用$\log \text{\hspace{0.17em}⁣}\left(\frac{q(x)}{p(x)}\right)\backslash log\backslash !\backslash big(\backslash frac\{q(x)\}\{p(x)\}\backslash big)$的樣本均值來近似KL（而不是試圖精確評估完整的期望）。文章接著指出另一種方法：使用$\frac{1}{2}(\log r{)}^2\backslash tfrac\{1\}\{2\}(\backslash log\; r)^2$的樣本均值來替代更“標準”的$\log r\backslash log\; r$形式，其中$r=\frac{q(x)}{p(x)}r=\backslash frac\{q(x)\}\{p(x)\}$。本文解釋了為什麼這個表示式可以成為KL的良好（儘管有偏）估計器，以及如何在保持低方差的同時使其無偏。

我們計算KL的方式取決於我們如何訪問$pp$和$qq$。這裡我們假設我們可以評估任何$xx$的$p(x)p(x)$和$q(x)q(x)$（機率或密度），但我們**無法**對$xx$進行解析求和/積分。我們為什麼不能進行解析求和/積分呢？可能是因為精確計算在計算或記憶體方面過於昂貴，可能沒有閉合形式，或者我們為了簡化程式碼，只儲存對數機率而不是完整的分佈，尤其是在KL僅用於診斷時（強化學習中常出現這種情況）。近似求和或積分最常見的策略是**蒙特卡羅**。給定從$qq$中抽取的樣本$x_1,x_2,\dots ,x_n\sim qx\_1,\; x\_2,\; \backslash dots,\; x\_n\; \backslash sim\; q$，我們如何構建一個好的估計器？

一個好的估計器應該**無偏**（平均值正確）且**方差低**。我們知道一個無偏估計器

$$k_1=\log \frac{q(x)}{p(x)}\mathrm{.}k\_1\; =\; \backslash log\backslash frac\{q(x)\}\{p(x)\}.$$

但它的方差很高：根據定義，KL是一個非負量，然而對於上述估計器，大約“一半”的樣本值可能是負的（如果我們不對$pp$和$qq$做任何先驗假設），這使得平均值波動很大，因此方差很高。為了符號方便，設$r=\frac{q(x)}{p(x)}r\; =\; \backslash frac\{q(x)\}\{p(x)\}$。那麼原始的KL可以寫成

$$\mathrm{K}\mathrm{L}[q,p]={\mathbb{E}}_{x\sim q}[\log r]\mathrm{.}\backslash mathrm\{KL\}[q,\; p]\; \backslash ;=\backslash ;\; \backslash mathbb\{E\}\_\{x\backslash sim\; q\}\backslash ,[\backslash log\; r].$$

為了減少方差，我們可以設計一個替代估計器：$$k_2=\frac{1}{2}(\log r{)}^2\mathrm{.}k\_2\; =\; \backslash frac\{1\}\{2\}(\backslash log\; r)^2.$$它的方差較低，但有偏。直觀上，$k_{2}k\_2$感覺更好，因為每個樣本都給出了$pp$和$qq$之間的非負“距離”，因此它保持正值。經驗上，$k_{2}k\_2$的方差確實比$k_{1}k\_1$低得多，而且偏差可以很小。至於為什麼$k_{2}k\_2$相比$k_{1}k\_1$能夠大幅降低方差，原始文章使用了f-散度檢視給出了分析解釋，這裡我不再贅述。

現在，我們能否得到一個既**無偏**又**低方差**的估計器呢？一個通用的技巧是使用**控制變數**：從無偏的$k_{1}k\_1$開始，並新增一個期望值為零且與它負相關的量以降低方差。這裡一個非常方便的零均值量是$r-1r-1$。因此，對於任意$\lambda \backslash lambda$，$$k=-\log r+\lambda (r-1)k\; \backslash ;=\backslash ;\; -\backslash log\; r\; +\; \backslash lambda\backslash ,(r-1)$$仍然是一個無偏的KL估計器。理論上，我們可以在$\lambda \backslash lambda$上最小化方差，但其閉合形式取決於$pp$和$qq$，不容易得到。然而請注意，由於$\log (x)\backslash log(x)$是凹函式，$$\log (x)\le x-1,\backslash log(x)\; \backslash ;\backslash le\backslash ;\; x-1,$$所以如果我們選擇$\lambda =1\backslash lambda=1$，該表示式保證非負。在這裡，$r-1r-1$是$\log r\backslash log\; r$在$r=1r=1$處的切線。因此，當$\lambda =1\backslash lambda=1$時，我們實際上測量的是$\log (x)\backslash log(x)$與其切線之間的垂直距離。這導致了估計器$$k_3=(r-1)-\log r,k\_3\; \backslash ;=\backslash ;\; (r\; -\; 1)\; \backslash ;-\backslash ;\; \backslash log\; r,$$它總是非負的。而$k_{3}k\_3$正是實際中GRPO與PPO在KL估計方式上有所不同的地方（PPO使用$k_{1}k\_1$）。

從“RL for LLM”的角度討論KL估計

在強化學習（例如PPO、GRPO等）中，我們通常會在損失函式中加入一個KL散度項，以防止新策略偏離舊策略太遠。這裡，$qq$是舊策略分佈（${\pi }_{\text{old}}\backslash pi\_\{\backslash text\{old\}\}$），$pp$是新策略分佈（${\pi }_{\text{new}}\backslash pi\_\{\backslash text\{new\}\}$），而$xx$是一個完整的動作樣本（在LLM中，這表示一個token或一個token序列）。我們通常用$ss$表示狀態（在LLM中，這是提示或上下文），$xx$是在該上下文中生成的特定token。當我們計算KL時，我們實際上是在**給定狀態下的動作分佈**上計算KL，然後對狀態進行平均：

在取樣時，我們通常會固定一個提示（狀態），然後為該提示估計此KL散度。

那麼**為什麼我們不能直接精確計算KL散度，而非要估計它呢？**原因與原始部落格文章中列出的完全相同；在LLM的強化學習中，主要癥結在於**原因1**：*動作空間（token空間）太大，無法對所有可能的$xx$進行求和/積分*。例如，如果一個分詞器有50,000個詞彙條目，即使計算單個token的KL散度也意味著對50,000個動作求和；而在強化學習中，我們通常進行多步（序列）生成，因此空間呈指數級增長，這完全不切實際。還有一個實用原因：在訓練過程中，我們通常不儲存完整的分佈（所有token的機率）；我們只保留沿軌跡實際生成的token的對數機率，以節省GPU記憶體和I/O。因此，我們必須使用**蒙特卡羅取樣**：從某個分佈（通常是$qq$，即舊策略）中抽取$xx$，並使用這些樣本來近似KL散度。這就把我們直接帶入了部落格文章所討論的領域。

在該文章中，我們一直談論的**估計器**實際上只是樣本的一個函式：它接收某個取樣$xx$的$p(x)p(x)$和$q(x)q(x)$（或它們的比率$r=\frac{q(x)}{p(x)}r\; =\; \backslash frac\{q(x)\}\{p(x)\}$)，並輸出一個數字。然後，我們對這些數字在樣本上求平均，以近似KL散度。例如：

• $k_1(x)=-\log rk\_1(x)\; =\; -\backslash log\; r$
• $k_3(x)=(r-1)-\log rk\_3(x)\; =\; (r\; -\; 1)\; -\; \backslash log\; r$

這些$k_{i}k\_i$只是不同的KL估計器公式。它們都透過**對樣本求平均**來近似KL散度，但在偏差和方差上有所不同。一旦我們選擇了一個估計器，我們實際上就承諾使用一個特定的公式來近似KL散度。這個過程看起來像這樣：

1. 取樣
   從舊策略$qq$中取樣一批token（或序列）$x_1,x_2,\dots ,x_{N}x\_1,\; x\_2,\; \backslash dots,\; x\_N$。
2. 計算對數機率
   對於每個樣本，計算新舊策略下的對數機率

$$\log p(x_i),\log q(x_i)\backslash log\; p(x\_i),\backslash \; \backslash log\; q(x\_i)$$

並得到$r_i=q(x_i)p(x_i)r\_i\; =\; \backslash frac\{q(x\_i)\}\{p(x\_i)\}$或$\log r_i\backslash log\; r\_i$。3. **代入估計器公式**
例如，如果我們選擇$k_{3}k\_3$

$$k_3(x_i)=(r_i-1)-\log r_{i}k\_3(x\_i)\; =\; (r\_i\; -\; 1)\; -\; \backslash log\; r\_i$$

4. 平均分

$$\widehat{\mathrm{K}\mathrm{L}}\approx \frac{1}{N}\sum _{i=1}^N{k}_3(x_i)\backslash widehat\{\backslash mathrm\{KL\}\}\; \backslash approx\; \backslash frac1N\; \backslash sum\_\{i=1\}^N\; k\_3(x\_i)$$

這是近似的 KL 值，代表了真實的 KL。

如果我們將這與離散機率分佈（LLM 單令牌步長）的真實 KL 計算（無估計）進行比較：我們需要遍歷每個可能的令牌 $xx$： $$\mathrm{K}\mathrm{L}(p\mathrm{\parallel }q)=\sum _{x}p(x)\log \frac{p(x)}{q(x)}\backslash mathrm\{KL\}(p\backslash |q)\; =\; \backslash sum\_x\; p(x)\; \backslash log\; \backslash frac\{p(x)\}\{q(x)\}$$ 您可以立即看到，使用估算器，計算量比進行完整求和小得多，尤其是在高維動作空間中。

談論不同 KL 估計器的方差

重要提示：我們這裡討論的“方差”是估計器在樣本上輸出值的方差： $${\mathrm{V}\mathrm{a}\mathrm{r}}_{x\sim q}[k(x)]\backslash mathrm\{Var\}\_\{x\; \backslash sim\; q\}[k(x)]$$ 也就是說， $k(x)k(x)$ 在樣本空間中的波動程度。一個**無偏**估計器意味著在無限多的樣本下，其均值等於真實 KL。但高方差估計器意味著即使均值正確（無偏），在少量樣本下，平均值也可能偏差很大。在 LLM 的強化學習中，KL 項通常是損失中的正則化因子（例如， $\beta \cdot \mathrm{K}\mathrm{L}\backslash beta\; \backslash cdot\; \backslash mathrm\{KL\}$）。如果 KL 估計器的方差很大，會使損失變得嘈雜，進而使梯度嘈雜並導致訓練不穩定。

在原帖中，為了讓讀者直觀理解為什麼 $k_{1}k\_1$ 不是低方差的，作者寫道：

然而，它（$k_{1}k\_1$）具有高方差，因為它對一半的樣本是負的，而 KL 始終是正的。

作者指出，儘管 $k_{1}k\_1$ 是無偏的，但如果沒有對 $pp$ 和 $qq$ 的先驗約束，一半的樣本會一個比另一個大，所以一半的 $k_{1}k\_1$ 值是正的，一半是負的。到目前為止，我都同意。但隨後作者說：因為 KL 總是大於 0（一個基本不等式），所以 $k_{1}k\_1$ 因此必須具有高方差。而在這裡，我認為因果關係並不成立：你不能用期望的符號來決定單個樣本的符號。一個簡單的反例：在計算期望時， $p(x)\log \frac{p(x)}{q(x)}p(x)\; \backslash log\; \backslash frac\{p(x)\}\{q(x)\}$ 也時而為正，時而為負；這個事實本身並不能說明方差。實際上，單樣本的**對數比率**（無論是 $\log \frac{q(x)}{p(x)}\backslash log\; \backslash frac\{q(x)\}\{p(x)\}$ 或 $\log \frac{p(x)}{q(x)}\backslash log\; \backslash frac\{p(x)\}\{q(x)\}$）都可以是正的或負的，就像 $k_{1}k\_1$ 一樣，所以**單獨的符號翻轉並不是高方差的唯一原因**。

根據 KL 定義： $$\mathrm{K}\mathrm{L}(q\mathrm{\parallel }p)={\mathbb{E}}_{x\sim q}[\log \frac{q(x)}{p(x)}]\backslash mathrm\{KL\}(q\; \backslash |\; p)\; =\; \backslash mathbb\{E\}\_\{x\backslash sim\; q\}\backslash left[\; \backslash log\; \backslash frac\{q(x)\}\{p(x)\}\; \backslash right]$$ 期望值**保證非負**，但被積函式 $\log \frac{q(x)}{p(x)}\backslash log\backslash frac\{q(x)\}\{p(x)\}$ 可以對單個樣本是正的或負的。而 $k_{1}k\_1$ 正是這個被積函式： $$k_1(x)=\log \frac{q(x)}{p(x)}k\_1(x)\; =\; \backslash log\; \backslash frac\{q(x)\}\{p(x)\}$$ 所以每個樣本值確實可以是正的或負的，與 KL 定義中的被積函式相同。

那麼為什麼 $k_{1}k\_1$ 會有高方差？

這不僅僅是“符號翻轉”。真正的原因是 $k_{1}k\_1$ 的值分佈通常很寬（重尾）。例如，如果 $p(x)p(x)$ 對於某些樣本來說很小，那麼 $\log \frac{q}{p}\backslash log\backslash frac\{q\}\{p\}$ 可能會非常大（正或負）。這些極端值主導有限樣本平均值，推高了方差。換句話說，它是**極端值 + 正負抵消**的組合：抵消意味著你需要更多的樣本才能收斂到真實平均值，而極端值會使樣本方差本身更大。因此，部落格中“一半為負”的評論更多的是一種直覺提示，而不是完整的解釋。

從這個角度來看，如果我們看其他估計器 $k_{2}k\_2$ 和 $k_{3}k\_3$，我們發現： $k_2=\frac{1}{2}(\log r{)}^{2}k\_2\; =\; \backslash frac12\; (\backslash log\; r)^2$ 總是正的，所以沒有抵消，但這引入了偏差；平方也平滑了幅度，降低了方差。$k_{3}k\_3$ 使用控制變數來消除部分波動源，在保持無偏性的同時降低方差（詳細資訊見下文）。

在 PPO/GRPO 中，如果您使用 $k_{1}k\_1$ 並且批次很小或分佈相距很遠，KL 估計值將跳來跳去（因為少數極端樣本會使平均值劇烈波動）。這使得 KL 懲罰係數不穩定：它可能突然變得過強或過弱。切換到低方差估計器（ $k_{2}k\_2$ 或 $k_{3}k\_3$ ）使每個樣本的 KL 貢獻更穩定，更不容易被少數極端樣本主導。

為什麼 $k_{3}k\_3$ 既能無偏又能低方差？

乍一看， $k_{3}k\_3$ 總是正的，所以你可能會認為它的平均值必須大於 $k_{1}k\_1$ 的平均值。
但請記住： $k_{3}k\_3$ 是透過**控制變數**從 $k_{1}k\_1$ 匯出的。部落格的推理如下： $$\tilde{k}(x)=k_1(x)+\lambda \cdot h(x)<!--</mo-->\backslash tilde\{k\}(x)\; =\; k\_1(x)\; +\; \backslash lambda\; \backslash cdot\; h(x)$$ 其中 $h(x)=r-1h(x)\; =\; r\; -\; 1$，並且在 $x\sim qx\backslash sim\; q$ 下，其期望值為： $${\mathbb{E}}_{x\sim q}[h(x)]={\mathbb{E}}_q[\frac{p(x)}{q(x)}-1]=\sum _{x}p(x)-1=1-1=0.\backslash mathbb\{E\}\_\{x\backslash sim\; q\}[h(x)]\; =\; \backslash mathbb\{E\}\_q\backslash left[\backslash frac\{p(x)\}\{q(x)\}\; -\; 1\backslash right]\; =\; \backslash sum\_x\; p(x)\; -\; 1\; =\; 1\; -\; 1\; =\; 0.$$ 因此，新增任何 $h(x)h(x)$ 的倍數都不會改變期望值。當 $\lambda =1\backslash lambda\; =\; 1$ 時： $$\tilde{k}(x)=-\log r+(r-1)=(r-1)-\log r=k_3(x)\mathrm{.}\backslash tilde\{k\}(x)\; =\; -\backslash log\; r\; +\; (r\; -\; 1)\; =\; (r\; -\; 1)\; -\; \backslash log\; r\; =\; k\_3(x).$$ 這解釋了為什麼 $k_{3}k\_3$ 的期望值等於 $k_{1}k\_1$ 的期望值，並等於 KL，使其成為一個無偏估計器。

$k_{3}k\_3$ 比 $k_{1}k\_1$ 具有更低方差的原因是： $k_{1}k\_1$ 只有 $-\log r-\backslash log\; r$，其值可能劇烈波動（既有正有負，偶爾出現巨大值）。但是 $r-1r\; -\; 1$ 和 $-\log r-\backslash log\; r$ 在數值上高度相關（一個增長，另一個也增長/收縮），並且這種相關性是**負的**。新增 $(r-1)(r\; -\; 1)$ 就像引入一個**負相關項來抵消波動**。抵消後， $k_{3}k\_3$ 中剩下的值範圍更緊密，始終為正，因此樣本方差更低。