DINO Self-Distillation Loss

Problem

Compute the DINO cross-entropy between the teacher's probability distribution and the student's logits: -sum_k teacher[k] * log softmax(student)[k], using a stable log-softmax.

The signature is dino_self_distillation(student_logits, teacher_probs).

Hint

The teacher is an EMA of the student, so the loss is a moving-target cross-entropy.

Test cases

case 1

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

expected: 0.6931471806

case 2

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

expected: 1.3132616875

case 3

in: [[-10,10],[0.5,0.5]]

expected: 10.0000000021

case 4

in: [[1,2,3],[0.2,0.3,0.5]]

expected: 1.1076059644

Starter codePython

import math
def dino_self_distillation(student_logits, teacher_probs):
    # Your code here
    pass

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