Distillation KL

Problem

Knowledge distillation loss: T^2 * KL(teacher || student), where KL = sum_k p_t[k] * (log p_t[k] - log p_s[k]) and both distributions come from stable log-softmax applied to logits / T.

The signature is distillation_kl(student_logits, teacher_logits, temperature).

Hint

The T^2 factor keeps the gradient magnitude comparable when softening with temperature.

Test cases

case 1

in: [[1,2,3],[1,2,3],1]

expected: 0

case 2

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

expected: 1.5231883119

case 3

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

expected: 0

case 4

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

expected: 0.3555882856

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