The mean_kl implementation is different from that of openai/lm-human-preference

[zhuzilin]

🐛 Describe the bug

Currently, the mean_kl used to update kl_ctl is calculated from:

trlx/trlx/trainer/accelerate_ppo_trainer.py

Lines 437 to 438 in 92b68e4

	log_ratio = (logprobs - ref_logprobs) * attention_mask[:, :-1]
	self.mean_kl = (log_ratio.exp() - 1 - log_ratio).mean().to(device)

which is a mean value over each token.
While in openai/lm-human-preferences, the mean_kl is:

            kl = data['logprobs'] - data['ref_logprobs']
            mean_kl = tf.reduce_mean(tf.reduce_sum(kl, axis=1))

(https://github.com/openai/lm-human-preferences/blob/bd3775f200676e7c9ed438c50727e7452b1a52c1/lm_human_preferences/train_policy.py#L220-L221)
which is not only the mean value upon each response, but also the same form of the kl used in reward.

Also, in anthropic's paper Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback, their kl value could be as large as 25 (as the square root is 5), which is hard to achieve for a token-wise mean kl.

I wonder if there is a specific reason why we use the current form of mean_kl? Thank you!

Gently ping @Dahoas @reciprocated

Which trlX version are you using?

No response

Additional system and package information

No response

[cat-state]

The different formula: this is due to using an unbiased KL estimator, see #88 (comment)

Computing the KL over just the responses rather than the whole prompt + response makes sense to me, although if the goal is to make the policy model stay close to the ref model it would also make sense to constrain the logits of the prompt tokens

[maxreciprocate]

About mean value scope, I agree it should be over a response-mean level instead of a token-mean level. It was deliberately done to make it invariant of response length, although it's perhaps currently too low (looking at more advanced examples) and maybe not in agreement with paper's results.

experiment	\sqrt(D_KL) token level	w&b
ppo_randomwalks	25	https://wandb.ai/sorry/trlx-references/runs/odlulghf
ppo_sentiments	10	https://wandb.ai/sorry/trlx-references/runs/up1mjqal
ppo_hh (1500 steps)	0.2	https://wandb.ai/sorry/trlx-references/runs/y7y4xrql

[lzy37ld]

@reciprocated Hi! Would you mind sharing the code for reproducing such wandb results?

[maxreciprocate]

@lzy37ld Hi, you can reproduce those runs by running this script https://github.com/CarperAI/trlx/blob/9bc08369ca9ec83342c4d7755205dab1a7723006/scripts/benchmark.sh