From 28d3c94a772d8cd91b4478619414fda2f752f398 Mon Sep 17 00:00:00 2001 From: flyingdoubleG <53119724+flyingdoubleG@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:07:43 +0800 Subject: [PATCH] fix(jsd): make the JSD advantage follow the gradient of the JSD The JSD advantage only sees the sampled token's log-probs, and the tokens are sampled from the student, so the per-token value is a REINFORCE cost. For 0 < lambda_coef < 1 it used beta*(t - log_m) + (1-beta)*(s - log_m). Under student samples t - log_m does not estimate KL(teacher || M), and the resulting expected update does not follow -grad JSD: on random categorical student/teacher pairs its cosine with -grad JSD is negative for beta in {0.3, 0.5, 0.9}, i.e. training increases the divergence. Differentiating beta*KL(T||M) + (1-beta)*KL(S||M) with respect to s(y) gives (1-beta) * log(s(y)/m(y)), so the cost is (1-beta)*(s - log_m); its expected REINFORCE gradient equals grad JSD exactly. beta = 0 and beta = 1 are unchanged. Co-Authored-By: Claude Opus 5.5 --- tests/algorithm/advantage_fn_test.py | 47 +++++++++++++++++++ .../algorithm/advantage_fn/jsd_advantage.py | 18 +++---- 2 files changed, 57 insertions(+), 8 deletions(-) diff --git a/tests/algorithm/advantage_fn_test.py b/tests/algorithm/advantage_fn_test.py index b9f8dab858a..69dc49ddddc 100644 --- a/tests/algorithm/advantage_fn_test.py +++ b/tests/algorithm/advantage_fn_test.py @@ -394,3 +394,50 @@ def test_step_wise_grpo_with_std_threshold(self): # Verify that all remaining experiences are from task 2 for exp in processed_exps: self.assertEqual(exp.eid.task, 2) + + +class TestJSDAdvantageFn(unittest.TestCase): + """The JSD advantage must follow the gradient of the generalized JSD.""" + + @staticmethod + def _expected_policy_gradient(advantage_fn, theta, teacher_logprobs, chunked): + # Exact expectation over the sampled token: E_{y~student}[A(y) * grad log student(y)]. + student_logprobs = torch.log_softmax(theta, dim=-1) + detached = student_logprobs.detach() + if chunked: + cost = advantage_fn._js_divergence_per_token_chunked( + detached[None], teacher_logprobs[None], torch.ones(1, theta.numel()) + )[0] + else: + cost = advantage_fn._js_divergence_per_token(detached, teacher_logprobs) + surrogate = (detached.exp() * -cost * student_logprobs).sum() + return torch.autograd.grad(surrogate, theta)[0] + + def test_jsd_advantage_descends_jsd(self): + torch.manual_seed(0) + teacher_logprobs = torch.log_softmax(2 * torch.randn(8), dim=-1) + for beta in (0.3, 0.5, 0.9): + for chunked in (False, True): + advantage_fn = ADVANTAGE_FN.get("jsd")(lambda_coef=beta, chunk_size=3) + theta = torch.randn(8, requires_grad=True) + policy_gradient = self._expected_policy_gradient( + advantage_fn, theta, teacher_logprobs, chunked + ) + + student_logprobs = torch.log_softmax(theta, dim=-1) + mixture = torch.logsumexp( + torch.stack( + [ + student_logprobs + np.log(1 - beta), + teacher_logprobs + np.log(beta), + ] + ), + dim=0, + ) + jsd = ( + beta * (teacher_logprobs.exp() * (teacher_logprobs - mixture)).sum() + + (1 - beta) * (student_logprobs.exp() * (student_logprobs - mixture)).sum() + ) + jsd_gradient = torch.autograd.grad(jsd, theta)[0] + + torch.testing.assert_close(policy_gradient, -jsd_gradient) diff --git a/trinity/algorithm/advantage_fn/jsd_advantage.py b/trinity/algorithm/advantage_fn/jsd_advantage.py index 75ffe6894de..90a9eae5df6 100644 --- a/trinity/algorithm/advantage_fn/jsd_advantage.py +++ b/trinity/algorithm/advantage_fn/jsd_advantage.py @@ -22,6 +22,11 @@ class JSDAdvantage(AdvantageFn): - beta=1: JSD = KL(teacher || student) [pure KL] - else: JSD = beta*KL(teacher||M) + (1-beta)*KL(student||M), M = mixture in log-space. + Only the sampled token's log-probs are available and the tokens are sampled from the + student, so the per-token value is a REINFORCE cost whose expected gradient is the + gradient of the divergence. For the mixture this cost is (1-beta) * (log s - log m): + d/ds(y) [beta*KL(T||M) + (1-beta)*KL(S||M)] = (1-beta) * log(s(y) / m(y)). + The teacher_logprobs should be stored in Experience.teacher_logprobs by the workflow during exploration. """ @@ -57,7 +62,7 @@ def _js_divergence_per_token( Compute JSD per token in log-space only (no exp). - beta=0: KL(student || teacher) = student_logprobs - teacher_logprobs - beta=1: KL(teacher || student) = teacher_logprobs - student_logprobs - - else: mixture log_m = logsumexp([s+log(1-beta), t+log(beta)]); JSD = beta*(t-log_m) + (1-beta)*(s-log_m). + - else: mixture log_m = logsumexp([s+log(1-beta), t+log(beta)]); cost = (1-beta)*(s-log_m). """ beta = self.lambda_coef s = student_logprobs @@ -81,10 +86,9 @@ def _js_divergence_per_token( torch.stack([s + log_1_minus_beta, t + log_beta]), dim=0, ) - # KL(teacher || mixture) = t - log_m, KL(student || mixture) = s - log_m (log-space, no exp) - kl_teacher = t - mixture_log - kl_student = s - mixture_log - jsd = beta_t * kl_teacher + (1 - beta_t) * kl_student + # The tokens are sampled from the student, so t - log_m is not an estimate of + # KL(teacher || mixture); the gradient of the whole JSD comes from this term alone. + jsd = (1 - beta_t) * (s - mixture_log) if mask is not None: jsd = jsd * mask @@ -128,9 +132,7 @@ def _js_divergence_per_token_chunked( torch.stack([s_chunk + log_1_minus_beta, t_chunk + log_beta]), dim=0, ) - kl_t = t_chunk - mixture_log - kl_s = s_chunk - mixture_log - jsd_valid[start:end] = beta_t * kl_t + (1 - beta_t) * kl_s + jsd_valid[start:end] = (1 - beta_t) * (s_chunk - mixture_log) out = flat_s.new_zeros(flat_s.shape) out[valid] = jsd_valid