Question about the k3 KL estimator implementation

[ZhangShiyue]

I was puzzled by the k3 KL estimator implementation.

It does not seem to implement the k3 estimation from John Schulman's blog. Because in John Schulman's blog, x is just a single variable. But in the context of LM, x is a sequence. If you take x being a sequence into consideration, the current implementation in TRL seems off..

I drafted some detailed derivations in https://zhangshiyue.github.io/#/blog/KL.

I wonder if I missed anything and if my confusion makes sense.

Happy to hear feedback from anyone~

[LeonEricsson]

Hi, thanks for the question. I'm not a math head but let me try provide some context

Why token-level k3 is valid for an LM

For two autoregressive policies

$$ q(y_{1:T})=\prod_{t=1}^{T}q_t,\qquad p(y_{1:T})=\prod_{t=1}^{T}p_t,\qquad r_t = \frac{p_t}{q_t}, $$

the forward KL factorises exactly:

$$ \mathrm{KL}[q\parallel p] = \sum_{t=1}^{T}\mathbb E_q[-\log r_t]. $$

Now note that for each step

Hence

$$ \mathbb E_q\Bigl[\sum_t(r_t-1)\Bigr]=0 \quad\Longrightarrow\quad \mathbb E_q\Bigl[\sum_t\bigl[(r_t-1)-\log r_t\bigr]\Bigr] =\mathrm{KL}[q\parallel p]. $$

That last bracket is exactly Schulman’s k3 term, just evaluated per-token instead of on the whole sequence. So summing those token values (or averaging them) gives an unbiased estimate of the sequence-level KL with the same low-variance property Schulman highlighted.

---

In Eq.(1) of your derivation (in the blog) the product term $\frac{\prod p_t}{\prod q_t}$ is kept intact, which forces a sequence-level ratio. But the crucial control-variate property ($\mathbb E_q[r_t]=1$) still holds per token, allowing us to move $(r_t-1)$ inside the sum without introducing bias. Once that step is taken, Eq.(1) collapses to the TRL form in Eq.(2) and the “bias” disappears.

[ZhangShiyue]

Hey @LeonEricsson, thanks a lot for your feedback! Sorry for the late response.

However, I'm not sure I understand why:

because I think

since the q in E_q is a distribution over sequences. This is the equation below Eq. 2 in my blog.

You seem to suggest:

But I don't understand how we can get the first red =.

Another way to explain my question, you have

which means

but I don't get why they are equivalent?

Hope this makes sense.