Clarification on TVD loss definition in distillation

[shqii1j]

Hi, thanks for releasing the code and paper.

I have a question regarding the TVD loss used in the distillation objective. From the code, both preds and target appear to be log-probabilities. However, in sdtt/src/sdtt/core/distill/multi_round_sdtt.py (line 349), the TVD loss is implemented as:

(preds - target).abs().sum(-1).mean()

This seems to correspond to an L1 distance in log-probability space, rather than the standard total variation distance defined in the Appendix B.

Could you confirm whether this is an intentional surrogate loss (log-probability L1), and if so, what the motivation is compared to computing TVD directly in probability space?

Thank you!

[jdeschena]

Hi,

We are penalizing the log probabilities on purpose. If I remember correctly, penalizing the probability did not work well, because then both preds and targets contain very small values, hence the student did not really learn.

Did you try training minimizing the distance in probability space and observe a different behaviour? Going through the paper, I see that we could make it more explicit. I will update it accordingly, thanks for the question.

Let me know if there is anything else, or if I can close the issue!

Cheers,
Justin