KL-To-Reference Penalty

Problem

Compute the per-token KL penalty against a reference policy using the non-negative estimator exp(d) - d - 1, where d = ref_log_prob - log_prob. An empty batch returns [].

Hint

The estimator is always non-negative and zero only when the policies agree.

Test cases

case 1

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

expected: [0,0]

case 2

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

expected: [0.7182818284590451]

case 3

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

expected: [0.36787944117144233]

case 4

in: [[],[]]

expected: []

Starter codePython

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

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