Distillation Loss

Problem

Compute knowledge-distillation loss: soften student and teacher logits with softmax(logits / temperature), then return the cross-entropy -sum(teacher_soft * log(student_soft)) with student probabilities clipped to at least 1e-15.

Use a numerically stable softmax.

Hint

Higher temperature softens the teacher's distribution, exposing dark knowledge.

Test cases

case 1

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

expected: 0.3653338550872077

case 2

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

expected: 1.0443202661482278

case 3

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

expected: 0.7853066497810339

case 4

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

expected: 0.6931471805599453

Starter codePython

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

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