Sequence Distillation KL

Problem

Compute the mean token-level distillation KL over a batch.

For each position scale both student_logits and teacher_logits by 1 / temperature, convert them to distributions p (student) and q (teacher) with a stable softmax, and average KL(q || p) = sum q * (log q - log p) across positions. An empty batch returns 0.0.

Hint

Higher temperature softens the teacher distribution so the student sees dark knowledge.

Test cases

case 1

in: [[[1,0]],[[1,0]],1]

expected: 0

case 2

in: [[[1,0]],[[0,0]],1]

expected: 0.12011450695827752

case 3

in: [[[2,0]],[[0,0]],0.5]

expected: 1.325002747357864

case 4

in: [[],[],1]

expected: 0

Starter codePython

import math
def distillation_kl(student_logits, teacher_logits, temperature):
    # Your code here
    pass

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