[GKD] `--beta` interpolation is discontinuous at the endpoints: the generalized JSD collapses to 0 as β→0/1, contradicting the documented "degenerates to forward/reverse KL"

Checklist / 检查清单

• I have searched existing issues, and this is a new question or discussion topic. / 我已经搜索过现有的 issues，确认这是一个新的问题与讨论。

Question Description / 问题描述

Summary

The docs (Distillation.md §2.1) describe --beta for GKD as an interpolation between forward KL and reverse KL:

广义 JSD(β): β·KL(P_T‖M) + (1−β)·KL(P_S‖M)，其中 M=βP_T+(1−β)P_S — "在两者之间插值"
β=0 退化为 Forward KL，β=1 退化为 Reverse KL

The endpoint special-cases in jsd_loss (swift/rlhf_trainers/gkd_loss.py) do return the pure KLs, so beta=0 and beta=1 themselves are fine. However, the interior formula has limit 0 at both endpoints, so the loss is jump-discontinuous at β=0 and β=1, and near the endpoints β acts as a signal-strength knob (≈ effective learning-rate scaling), not a direction knob as documented.

The math

For 0 < β < 1 the code computes the "coefficient-matched" generalized JSD (matching the GKD paper, arXiv:2306.13649):

M = (1−β)·S + β·T
D(β) = β·KL(T‖M) + (1−β)·KL(S‖M)

Substituting β = 1−ε:

• KL(S‖M) where M = ε·S + (1−ε)·T — the distance from S to a mixture that is almost T. A first-order expansion gives D(β) ≈ (1−β)·KL(S‖T) + O(ε²) → 0 as β→1⁻.
• Symmetrically D(β) ≈ β·KL(T‖S) → 0 as β→0⁺.

Meanwhile the β==1 branch returns the full KL(S‖T). Hence:

D(0.999) ≈ 0.001·KL(S‖T)      # interior formula
D(1.000) = KL(S‖T)            # endpoint branch

a ~1000× jump for a 0.001 change in β. The documented "β=1 degenerates to reverse KL" holds only via the hardcoded endpoint branch, not as a limit of the interior formula — the interpolation claim is not continuous.

Numerical evidence

Self-contained reproduction (exactly mirrors the three branches of jsd_loss):

import numpy as np
rng = np.random.default_rng(7)
z = rng.normal(0, 1.5, 300); w = rng.normal(0, 1.5, 300) + 0.3; w[:3] += 2.0
softmax = lambda x: np.exp(x - x.max()) / np.exp(x - x.max()).sum()
S, T = softmax(z), softmax(w)
kl = lambda p, q: np.sum(p * np.log(p / q))
R, F = kl(S, T), kl(T, S)          # reverse / forward KL

def D_current(beta):               # == current jsd_loss
    if beta == 0: return F
    if beta == 1: return R
    M = (1 - beta) * S + beta * T
    return beta * kl(T, M) + (1 - beta) * kl(S, M)

def D_swapped(beta):               # suggested fix
    if beta == 0: return F
    if beta == 1: return R
    M = (1 - beta) * S + beta * T
    return (1 - beta) * kl(T, M) + beta * kl(S, M)

for b in [0.0, 0.001, 0.1, 0.5, 0.9, 0.999, 1.0]:
    print(f"beta={b:<6} current={D_current(b):.4f}  swapped={D_swapped(b):.4f}")
# reverse KL = 2.0450, forward KL = 1.9262 on this seed

Output:

β	current (docs=code)	swapped
0.000	1.9262 (endpoint branch)	1.9262
0.001	0.0019	1.9069
0.100	0.1454	1.0605
0.500	0.3492	0.3492 (identical — symmetric JSD)
0.900	0.1474	1.0635
0.999	0.0020	2.0119
1.000	2.0450 (endpoint branch)	2.0450

The swapped form is continuous on [0, 1] with endpoint limits exactly equal to the two pure KLs — i.e. it is the form for which the documented "interpolates / degenerates" semantics is actually true. β=0.5 is bit-identical between the two forms (symmetric JSD), which is why the default configuration never exposes the problem.

Practical impact

We ran a 27B GKD/OPD job with beta=0.9 intending "mostly reverse KL with a mild forward-KL correction". Measured against an otherwise-identical beta=1 run:

	β=1 run	β=0.9 run
loss @ it1	0.478	0.036 (~1/13)
grad_norm @ it1	4.27	0.29 (~1/15)

A finite-difference check on the formula itself: cos(∇D(0.9), ∇D(1)) = 0.91, ‖∇D(0.9)‖/‖∇D(1)‖ ≈ 0.05 — i.e. β=0.9 is effectively reverse KL trained at ~1/15 learning rate, not a directional blend. Any β sweep near the endpoints (e.g. 0.8/0.9/0.95) will mostly measure LR scaling rather than the forward/reverse trade-off, and loss values across β are not comparable.

Suggested fix

Minimal change in jsd_loss (interior branch only; endpoints and β=0.5 untouched):

# current
jsd = beta_t * kl_div_fn(m_log, t_log) + (1 - beta_t) * kl_div_fn(m_log, s_log)
# suggested (endpoint-continuous)
jsd = (1 - beta_t) * kl_div_fn(m_log, t_log) + beta_t * kl_div_fn(m_log, s_log)

• β=0 and β=1 branches unchanged (pure KLs).
• β=0.5 numerically unchanged (symmetric JSD, both forms identical) — default behavior is unaffected.
• The swapped form is still a valid divergence (non-negative, zero iff S=T) and preserves the same mixture M.

Alternatively, if strict fidelity to the GKD paper formula is preferred, the docs should be corrected to state that for 0<β<1 the loss is the paper's generalized JSD whose value does not limit to the KLs at the endpoints (≈ min(β,1−β)·KL near them), and that β near 0/1 mainly scales the signal. But since the docs' current claim ("β=0 退化为 Forward KL，β=1 退化为 Reverse KL") is what users reasonably rely on, making the code match the docs (or clearly documenting the discontinuity) would prevent silent misinterpretation of β sweeps.

Environment

• ms-swift main (checked 2026-09-09): swift/rlhf_trainers/gkd_loss.py, interior branch at L143–144, docstring "beta: JSD interpolation (0=forward KL, 1=reverse KL, 0<beta<1=JSD)".
• Reproduction is pure numpy, no GPU needed.

Start with swift/rlhf_trainers/gkd_loss.py and inspect the jsd_loss branches around the documented beta behavior; compare them with Distillation.md §2.1. Reproduce the endpoint values using the supplied NumPy example, then confirm whether the intended resolution is endpoint-continuous code or corrected documentation, with beta=0, 0.5, and 1 behavior preserved.