Token-Level KL

Problem

Compute the mean over positions of the per-position KL divergence: mean_t sum_k p[t][k] * log(p[t][k] / q[t][k]), skipping terms where p is zero. p and q are lists of probability distributions, one per position.

The signature is token_level_kl(p, q).

Hint

Per-token KL is used to keep a finetuned model close to its reference at every position.

Test cases

case 1

in: [[[0.5,0.5]],[[0.5,0.5]]]

expected: 0

case 2

in: [[[0.75,0.25]],[[0.5,0.5]]]

expected: 0.1308120359

case 3

in: [[[0.9,0.1],[0.5,0.5]],[[0.8,0.2],[0.5,0.5]]]

expected: 0.018345007

case 4

in: [[[0.25,0.75]],[[0.5,0.5]]]

expected: 0.1308120359

Starter codePython

import math
def token_level_kl(p, q):
    # Your code here
    pass

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