Fomula confusion in distill loss funtion

[yspMing]

The fomula in this this line seems to have a problem

sum_{k=1}^{K} target_prob_k * (logits_k - log sum_{k=1}^K exp(logits_k)) - const
= (sum_{k=1}^{K} target_prob_k * logits_k) - log sum_{k=1}^K exp(logits_k) - const

https://github.com/NVIDIA/NeMo-Aligner/blob/main/nemo_aligner/utils/distributed.py#L622

[odelalleau]

Can you clarify what the problem is?

I'm not familiar with this code/derivation, but at quick glance sum_{k=1}^{K} target_prob_k is equal to 1, which is why it would disappear when multiplying the second term (that doesn't depend on k -- the notations may be a bit confusing in this regard).