[flyingdoubleG]

Description

Problem. JSDAdvantage sees only the sampled token's log-probs s = log π_student(y) and t = log π_teacher(y), and y is sampled from the student. The per-token value, used as advantages = -kl_coef * value under the PPO loss, is therefore a REINFORCE cost: training follows -E_{y~student}[value(y) · ∇log π_student(y)].

For 0 < lambda_coef < 1 the value was

β·(t − log m) + (1 − β)·(s − log m),   m = β·teacher + (1 − β)·student

Under student samples, t − log m is not an estimate of KL(teacher‖M), and the expected update does not follow −∇JSD_β. On random categorical student/teacher pairs, computing the expectation over the sampled token exactly, the cosine between the update and −∇JSD_β is:

β	before	after
0.3	−0.22 … −0.58	+1.00
0.5	−0.39 … −0.79	+1.00
0.9	−0.69 … −0.96	+1.00

(3 seeds; the ranges are over seeds.) So with the default lambda_coef: 0.5, used by examples/opd_gsm8k/opd_gsm8k_jsd.yaml, the expected update mostly increases the divergence to the teacher. The logged jsd/mean is also never positive, since β·t + (1−β)·s ≤ log m by concavity.

Fix. Differentiating β·KL(T‖M) + (1−β)·KL(S‖M) with respect to s(y) gives (1−β)·log(s(y)/m(y)), because the β·t/m and (1−β)·s/m terms cancel. The per-token cost is therefore (1 − β)·(s − log m), and its expected REINFORCE gradient equals ∇JSD_β exactly. This PR applies the fix to both the plain and the chunked path. β = 0 (cost s − t, reverse KL) is unchanged.

Not changed: β = 1, which needs your decision. The docstring says β = 1 is KL(teacher‖student), and the code uses cost t − s. Under student samples E[(s − t)∇log s] = ∇KL(student‖teacher), so cost t − s ascends the reverse KL: the cosine with −∇KL(S‖T) is exactly −1, and the cosine with −∇KL(T‖S) is −0.5 to −0.7. Two possible fixes:

1. Keep the documented meaning, KL(teacher‖student). Its gradient is −E_student[e^{t−s} ∇log s], so the cost would be −exp(t − s), an unbounded importance weight that probably needs clipping.
2. Follow the convention of SWIFT/TRL generalized_jsd_loss, where β = 0 is forward KL and β = 1 is reverse KL. That matches the β→0 / β→1 limits of this mixture, but it swaps the current meaning of β = 0.

Happy to add either.

How to test. tests/algorithm/advantage_fn_test.py::TestJSDAdvantageFn::test_jsd_advantage_descends_jsd checks that, for β ∈ {0.3, 0.5, 0.9} on both the plain and the chunked path, the exact expected policy gradient equals −∇JSD_β. It fails on main (8/8 elements mismatched) and passes with this change. The whole file passes (9 tests, CPU).

Checklist

• Code has passed all tests (tests/algorithm/advantage_fn_test.py, CPU; GPU tests not run)
• Docstrings have been added/updated in Google Style
• Documentation has been updated (not needed: the README formula describes the divergence, which is unchanged)
• Code is ready for review

🤖 Generated with Claude Code