GRPO Masked KL Penalty

Problem

Estimate the GRPO KL penalty on completion tokens only.

For every position with mask[i] = 1 accumulate the non-negative estimator exp(d) - d - 1 where d = ref_log_probs[i] - log_probs[i], then return the masked mean. A zero total mask returns 0.0.

Hint

Only completion tokens should pull the policy toward the reference model.

Test cases

case 1

in: [[-1,-1],[-1,-1],[1,1]]

expected: 0

case 2

in: [[-2],[-1],[1]]

expected: 0.7182818284590451

case 3

in: [[-2],[-1],[0]]

expected: 0

case 4

in: [[-1,-2],[-2,-1],[1,1]]

expected: 0.5430806348152437

Starter codePython

import math
def grpo_kl_penalty(log_probs, ref_log_probs, mask):
    # Your code here
    pass

4 test cases · run your code against them in the editor