Add Sequence-Level KD

[mst272]

What does this PR do?

In the original paper, they compared Sequence-Level KD 、Supervised KD and GKD (On-policy). In trl GKDTrainer, Supervised KD and GKD have been implemented. So i add Sequence-Level KD in GKDTrainer, control whether to perform Sequence-Level KD through a seq_kd parameter.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[kashif]

thanks @mst272 can you also kindly add these options to the docs-strings and the documentation of GKDTrainer?

[mst272]

hi@kashif,I've added these to the docs-strings and the documentation

[qgallouedec]

LGTM, I leave the last word to @kashif

[moussaKam]

Hi there,

I don't really understand how this PR adds seq_kd. To my understanding the seq_kd computes the standard cross_entropy between student logits and the output generated by the teacher.

In this PR we are simply generating the teacher output then computing the same generalized_jsd_loss. Am I missing something ?

In the documentation it says:

seq_kd: controls whether to perform Sequence-Level KD (can be viewed as supervised FT on teacher-generated out). When seq_kd=True and lmbda=0.0, the loss reduces to supervised JSD, where the teacher generates output sequences and the student receives token-specific feedback on these sequences from the teacher.

But this is the definition of the supervised_kd if I understand correctly.

[kashif]

@moussaKam So recall there are 2 other parameters apart from the seq_kd flag, namely the lmbda so set that to zero and then the beta that interpolates between the forward and reverse KL. SeqKD then would be lmbda=0, and beta=0 if i am not mistaken... its an interesting question and might need exploring, what happens with other values for these 2 hyper-parameters... or do you mean that to be exact we would need to replace the KL div by the cross-entropy loss?

[kashif]

@moussaKam also note that in this case, the KL-div is the same as the CE and a constant term, i.e. the entropy of the target which we assume does not change

[moussaKam]

@kashif thanks for you're reply, yes according to the definition from the paper

Sequence-Level KD (Kim & Rush, 2016). SeqKD maximizes the likelihood of high probability sequences generated by the teacher, and can be viewed as supervised FT on teacher-generated output.

I understand that we compute the cross-entropy in the case of seq_kd. This is what we do in standard sft no?

[moussaKam]

@kashif another point, if seq_kd is put to True we are computing the teacher inference twice. here and here. Do we really need that?

[kashif]

@moussaKam so in the first we generate completions and in the 2nd we calculate the logits of the completions... i suppose we could do that once, and then keep track of it with a bunch of if-else but opted for some cleaner logic here that could work for any of the different hyperparams... any ideas on how to make it a bit more dry?

[kashif]

@moussaKam Mind you there is an orthogonal abstraction i have been working on where instead of the logits (which are assumed to come from the same vocab. size for both the student and teacher) we allow for the student-teacher to have different vocabs: see #2263 and I would welcome any thoughts if this should be a separate class?

[moussaKam]

@kashif, we don't need to compute the logits, we generate the output with the teacher which becomes the new labels, then we run the forward of the student and compute the cross entropy using just the teacher output tokens.

I can implement it in the afternoon if it sounds good for you.