KL-Penalized Rewards

The problem, from first principles

When a language model is optimized against a reward model, it quickly finds responses the reward model loves and humans do not: repetitive, exaggerated or simply gibberish that exploits a flaw. The standard defense is to charge the policy for drifting away from the reference model it started from. The charge is a KL penalty, estimated per token from the log-probabilities the two models assign to the sampled tokens. The reward model's score is a single number for the whole response, so it is placed on the last token, while the KL penalty is spread across every token. The result is a per-token reward sequence that a reinforcement learning algorithm can consume.

From theory to code

Implement kl_penalized_rewards and mean_kl.

Constraints

• logp_policy and logp_ref are (B, T) arrays of log-probabilities of the sampled tokens. lengths is an integer array of response lengths, with 1 <= lengths[b] <= T. Tokens at positions >= lengths[b] are padding.
• kl_penalized_rewards(logp_policy, logp_ref, final_reward, lengths, beta) returns a (B, T) array with -beta * (logp_policy - logp_ref) at every valid position, plus final_reward[b] added at position lengths[b] - 1, and exactly 0 at padding.
• mean_kl(logp_policy, logp_ref, lengths) is the mean over all valid tokens of logp_policy - logp_ref, the simple per-token estimator of the KL divergence from sampled tokens.
• Do not modify the inputs.

Hints