Add beta support for jsd

[Tcc0403]

Summary

Resolve #278 .

Details

Forward:

$$\begin{align} JSD(X, Y, \beta) &= JSD_{\beta}(P \Vert Q)\\\ &= \beta\ KL(P \Vert \beta P + (1-\beta)Q) + (1-\beta)\ KL(Q \Vert \beta P + (1-\beta)Q)\\\ &= \sum \beta\ PY + (1-\beta)QX - M\ logM \end{align}$$

where $X=logQ$, $Y=logP$ and $M=\beta P + (1-\beta)Q$.

Gradients:

$$\frac{\partial}{\partial X_i} JSD(X, Y, \beta) = (1-\beta)Q_i(X_i - logM_i)$$

Testing Done

• Hardware Type: H100
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[Tcc0403]

This is the correct order of input and target (student and teacher) respectively. would it be too confusing?

[qingquansong]

yeah, the name is a bit confusing, or we can add some descriptions here to clarify

[Tcc0403]

@qingquansong @yundai424 ready for review!

[qingquansong]

LGTM in general! In case you're interested, I think one good future work is to make those KL or JSD losses similar to the fuse CE loss: feed teacher and student model last projection layer to the kernel and fuse it with the losses. Here teacher weight does not need grad and student will need grad.

[qingquansong]

yeah, the name is a bit confusing, or we can add some descriptions here to clarify

[ByronHsu]

awesome work! waiting for the final nit review

[Tcc0403]

I think one good future work is to make those KL or JSD losses similar to the fuse CE loss: feed teacher and student model last projection layer to the kernel and fuse it with the losses. Here teacher weight does not need grad and student will need grad.

@qingquansong sure, I'm in.

[Tcc0403]

Forgot to add jsd in readme and liger_kernel.transformer