File size: 13,667 Bytes
63f41a6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 | # %% [markdown]
# # 0109 - KL regularisation estimators k1, k2, k3: value bias/variance, what their GRADIENTS optimise, why beta = 0
# Companion to `notes/0109 - KL Estimators.pdf`. Samples x ~ q = pi_theta, p = pi_ref, r = p/q, KL(q||p) = E_q[-log r].
# %%
import sys, math
sys.path.insert(0, "..")
from _m1_lib import *
import numpy as np
# %% [markdown]
# ## 1. Values. E_q[k1] = E_q[k3] = KL(q||p); E_q[k2] = KL + bias, bias = -sum_{n>=3} E[l^n]/n!, l = log r
# k1 = -log r, k2 = (log r)^2 / 2, k3 = (r-1) - log r (Schulman, "Approximating KL divergence")
# %%
def estimators(logr):
return -logr, 0.5 * logr ** 2, torch.expm1(logr) - logr
gen = torch.Generator().manual_seed(2)
for trial in range(3):
K = 6
q = torch.softmax(torch.randn(K, generator=gen) * 1.2, 0); p = torch.softmax(torch.randn(K, generator=gen) * 1.2, 0)
logr = (p / q).log()
kl = float((q * (q / p).log()).sum())
k1, k2, k3 = estimators(logr)
check(f"categorical #{trial}: E_q[k1] = KL(q||p)", abs(float((q * k1).sum()) - kl), 1e-14)
check(f"categorical #{trial}: E_q[k3] = KL(q||p)", abs(float((q * k3).sum()) - kl), 1e-14)
bias2 = float((q * k2).sum()) - kl
series = -sum(float((q * logr ** n).sum()) / math.factorial(n) for n in range(3, 40))
check(f"categorical #{trial}: bias(k2) = -sum_(n>=3) E[l^n]/n! (series, 40 terms)", abs(bias2 - series), 1e-12)
# the master identity KL = sum_{n>=2} E[l^n]/n! (from E_q[e^l] = 1)
ms = sum(float((q * logr ** n).sum()) / math.factorial(n) for n in range(2, 40))
check(f"categorical #{trial}: KL = sum_(n>=2) E_q[l^n]/n!", abs(ms - kl), 1e-12)
print(f" KL = {kl:.5f} E[k2] - KL = {bias2:+.5f} (leading term -E[l^3]/6 = {-float((q*logr**3).sum())/6:+.5f})")
# pointwise positivity
rg = torch.exp(torch.linspace(-12, 12, 100001))
k1g, k2g, k3g = estimators(rg.log())
check("k3 >= 0 and k2 >= 0 pointwise (min over a log-grid of r)", max(0.0, -float(min(k3g.min(), k2g.min()))), 0.0)
print(f" k1 < 0 on r > 1 (fraction of log-grid where k1<0: {float((k1g<0).double().mean()):.3f}) -> k1 is a noisy, sign-changing estimator")
# %% [markdown]
# ## 2. Gaussian closed forms (equal variance, shift mu => KL = mu^2/2). With s2 = mu^2:
# Var(k1) = s2, Var(k2) = s2^3/4 + s2^2/2, Var(k3) = e^{s2} - 1 - s2, bias(k2) = s2^2/8. Verified by Gauss-Hermite quadrature.
# %%
xs, ws = np.polynomial.hermite_e.hermegauss(120); ws = ws / ws.sum()
def gauss_moments(shift):
x = torch.tensor(xs); w = torch.tensor(ws)
logr = shift * x - 0.5 * shift ** 2 # q = N(0,1), p = N(shift,1): log p/q = shift*x - shift^2/2 (E_q[r] = 1)
k = estimators(logr)
kl = shift ** 2 / 2
out = []
for kk in k:
m = float((w * kk).sum()); v = float((w * kk ** 2).sum()) - m ** 2
out.append((m, v))
return kl, out
print(" KL est mean-KL (bias) std/KL closed-form std/KL")
for shift in (0.1, math.sqrt(2 * 0.05), 1.0):
kl, out = gauss_moments(shift)
s2 = shift ** 2
cf = [s2, s2 ** 3 / 4 + s2 ** 2 / 2, math.exp(s2) - 1 - s2]
cb = [0.0, s2 ** 2 / 8, 0.0]
for i, nm in enumerate(("k1", "k2", "k3")):
m, v = out[i]
print(f" {kl:.4f} {nm} {m-kl:+.3e} {math.sqrt(v)/kl:9.4f} {math.sqrt(cf[i])/kl:9.4f}")
check(f"Gaussian shift={shift:.3f}: {nm} bias == closed form", abs((m - kl) - cb[i]), 1e-12)
check(f"Gaussian shift={shift:.3f}: {nm} variance == closed form", abs(v - cf[i]) / cf[i], 1e-10)
# compare with the table in Schulman's note (KL=0.005 and KL=0.5): std/KL = 20, 1.42, 1.42 and 2, 1.73, 1.7 ; bias(k2)/KL = 0.002, 0.25
kl, out = gauss_moments(0.1)
check("matches Schulman's table at KL=0.005: std(k1)/KL ~ 20", abs(math.sqrt(out[0][1]) / kl - 20.0), 0.2)
kl, out = gauss_moments(1.0)
check("matches Schulman's table at KL=0.5: bias(k2)/KL = 0.25", abs((out[1][0] - kl) / kl - 0.25), 1e-9)
# %% [markdown]
# ## 3. Heavy tails: Var(k3) is infinite when chi^2(p||q) is; one token the reference loves and the policy nearly never samples
# %%
q = torch.tensor([0.001, 0.999]); p = torch.tensor([0.30, 0.70])
logr = (p / q).log(); k1, k2, k3 = estimators(logr)
kl = float((q * (q / p).log()).sum())
for nm, k in (("k1", k1), ("k2", k2), ("k3", k3)):
m = float((q * k).sum()); v = float((q * k ** 2).sum()) - m ** 2
print(f" {nm}: mean {m:+.5f} std {math.sqrt(v):.5f} (KL = {kl:.5f}) max value {float(k.max()):.1f}")
chi2_pq = float((p ** 2 / q).sum() - 1)
print(f" chi2(p||q) = E_q[r^2]-1 = {chi2_pq:.1f} -> Var(k3) is dominated by one rare token with r = {float((p/q)[0]):.0f}")
check("E_q[r^2] - 1 = chi2(p||q) exactly", abs(float((q * (p / q) ** 2).sum()) - 1 - chi2_pq), 1e-9)
# %% [markdown]
# ## 4. GRADIENTS (on-policy, single-step categorical policy with logits theta; samples x ~ q_theta DETACHED, gradient flows through q_theta in r)
# k1-as-loss: E[grad k1] = E[grad log q] = 0 -> useless as a regulariser
# k2-as-loss: E[grad k2] = E[(log q - log p) grad log q] = grad KL(q||p) (the reverse KL we meant)
# k3-as-loss: E[grad k3] = E[(1 - r) grad log q] = -sum_x p grad log q = grad KL(p||q) (the FORWARD KL)
# k1-in-reward (detached coefficient): E[(log q - log p) grad log q] = grad KL(q||p)
# %%
gen = torch.Generator().manual_seed(7)
K = 6
th_ref = torch.randn(K, generator=gen) * 1.0
def grads(delta):
th = (th_ref + delta * torch.randn(K, generator=gen)).clone().requires_grad_(True)
q = torch.softmax(th, 0); p = torch.softmax(th_ref, 0)
lq, lp = torch.log_softmax(th, 0), torch.log_softmax(th_ref, 0)
qd = q.detach()
r = (lp - lq).exp()
# expectation over x ~ q (weights detached) of the gradient of each estimator
g = {}
g['k1'] = torch.autograd.grad((qd * (-(lp - lq))).sum(), th, retain_graph=True)[0]
g['k2'] = torch.autograd.grad((qd * 0.5 * (lp - lq) ** 2).sum(), th, retain_graph=True)[0]
g['k3'] = torch.autograd.grad((qd * (r - 1 - (lp - lq))).sum(), th, retain_graph=True)[0]
# k1 in reward: detached coefficient (log q - log p) times score
g['k1r'] = torch.autograd.grad((qd * (lq - lp).detach() * lq).sum(), th, retain_graph=True)[0]
g['rev'] = torch.autograd.grad((q * (lq - lp)).sum(), th, retain_graph=True)[0] # grad KL(q||p), exact
g['fwd'] = torch.autograd.grad((p * (lp - lq)).sum(), th, retain_graph=True)[0] # grad KL(p||q), exact
return g
g = grads(0.8)
check("E[grad k1] = 0 (k1 as a loss carries no regularising gradient)", float(g['k1'].abs().max()), 1e-14)
check("E[grad k2] == grad KL(q||p) (reverse KL)", maxabs(g['k2'], g['rev']), 1e-14)
check("k1-in-reward (detached coefficient) == grad KL(q||p)", maxabs(g['k1r'], g['rev']), 1e-14)
check("E[grad k3] == grad KL(p||q) (FORWARD KL, not the one k3 estimates)", maxabs(g['k3'], g['fwd']), 1e-14)
print(f" ||grad KL(q||p) - grad KL(p||q)|| / ||grad KL(q||p)|| = {float((g['rev']-g['fwd']).norm()/g['rev'].norm()):.3f} (different vectors)")
print(" relative gap between the two gradients as q -> p (first-order agreement, 'k3 is a first-order approximation'):")
prev = None
for delta in (0.4, 0.2, 0.1, 0.05, 0.025):
gen.manual_seed(7); _ = torch.randn(K, generator=gen) # same direction each time
gg = grads(delta)
gap = float((gg['rev'] - gg['fwd']).norm() / gg['rev'].norm())
print(f" delta={delta:5.3f} gap = {gap:.4f}")
prev = gap
assert prev < 0.05
check("gradient gap forward vs reverse KL vanishes as delta -> 0 (gap at delta=0.025 < 5%)", 0.0, 0.0)
# %% [markdown]
# ## 5. SEQUENCES: per-token k2/k3 'as loss' ignore that earlier tokens change the later contexts
# true: grad KL_seq(pi||ref) = E[ (sum_t l_t) * score(o) ] = E[ sum_t (sum_{t'>=t} l_t') s_t ] ("k1 in reward" with reward-to-go)
# per-token k2 loss, summed over t: E[ sum_t l_t s_t ] = true - sum_t E[ KL_t(s_t) * score(o_<t) ] (missing prefix-score term)
# %%
lm = TinyLM(V=4, L=4, seed=0, scale=1.0)
th = lm.theta0.clone(); th_r = perturb(th, 1.0, seed=2)
pi = lm.probs(th)
lpi, lref = lm.token_logprobs(th), lm.token_logprobs(th_r)
l_tok = (lpi - lref) * lm.mask # l_t = log pi - log ref = -log r_t
Ttens = token_score_tensor(lm, th) # [N, L, P]
S = Ttens.sum(1)
KL_seq = float((pi * l_tok.sum(1)).sum())
t_ = th.clone().requires_grad_(True)
KLt = (lm.probs(t_) * (lm.seq_logprob(t_) - lm.seq_logprob(th_r))).sum()
g_true = torch.autograd.grad(KLt, t_)[0]
g_seq = (pi[:, None] * l_tok.sum(1)[:, None] * S).sum(0)
rtg = l_tok.flip(1).cumsum(1).flip(1) * lm.mask
g_rtg = (pi[:, None] * (Ttens * rtg[:, :, None]).sum(1)).sum(0)
g_tok = (pi[:, None] * (Ttens * l_tok[:, :, None]).sum(1)).sum(0) # per-token k2-as-loss gradient, summed over t
check("E[(sum_t l_t) * score] == grad KL_seq(pi||ref)", maxabs(g_seq, g_true), 1e-14)
check("reward-to-go form E[sum_t (sum_{t'>=t} l_t') s_t] == grad KL_seq", maxabs(g_rtg, g_true), 1e-14)
print(f" per-token k2-as-loss gradient vs true sequence-KL gradient: relative error {float((g_tok-g_true).norm()/g_true.norm()):.4f}")
# the missing term: sum_t E[ KL_t(s_t) * score(o_<t) ] where KL_t(s) = KL(pi(.|s) || ref(.|s)) and score(o_<t) = sum_{s<t} s_s
prefix_score = torch.cumsum(Ttens, 1) - Ttens # sum_{s<t} s_s
KLt_tok = torch.zeros(lm.N, lm.L)
for i, sq in enumerate(lm.seqs):
for t in range(len(sq)):
cp, cr = lm.cond_probs(th, sq[:t]), lm.cond_probs(th_r, sq[:t])
KLt_tok[i, t] = (cp * (cp.log() - cr.log())).sum()
missing = (pi[:, None] * (KLt_tok[:, :, None] * prefix_score).sum(1)).sum(0)
check("per-token k2 loss gradient == true gradient - sum_t E[KL_t * prefix score] (exact identity)", maxabs(g_tok, g_true - missing), 1e-13)
# per-token k3 loss gradient equals the gradient of the sum of FORWARD per-state KLs with frozen context distribution
r_tok = (lref - lpi).exp() * lm.mask + (1 - lm.mask)
g_k3 = (pi[:, None] * (Ttens * ((1 - r_tok) * lm.mask)[:, :, None]).sum(1)).sum(0)
t_ = th.clone().requires_grad_(True)
tot = 0.0
for s in lm.prefixes():
pr = 1.0
for k in range(len(s)): pr *= float(lm.cond_probs(th, s[:k])[s[k]]) # P_pi(context = s) (detached)
cp_t, cr = lm.cond_probs(t_, s), lm.cond_probs(th_r, s)
tot = tot + pr * (cr * (cr.log() - cp_t.log())).sum() # KL(ref(.|s) || pi_theta(.|s)) with frozen context weight
g_fwd_frozen = torch.autograd.grad(tot, t_)[0]
check("per-token k3-as-loss gradient == grad of sum_s P_pi(s) KL(ref(.|s)||pi_theta(.|s)) (context weights frozen)", maxabs(g_k3, g_fwd_frozen), 1e-13)
# values: sequence-level KL from the three per-token estimators
k1_seq = float((pi * l_tok.sum(1)).sum())
k3_tok = (torch.expm1(lref - lpi) - (lref - lpi)) * lm.mask
k2_tok = 0.5 * (lpi - lref) ** 2 * lm.mask
print(f" sequence KL = {KL_seq:.5f}; E[sum_t k1] = {k1_seq:.5f}; E[sum_t k3] = {float((pi*k3_tok.sum(1)).sum()):.5f}; E[sum_t k2] = {float((pi*k2_tok.sum(1)).sum()):.5f}")
check("E[sum_t k3_t] = KL_seq (per-token estimator is unbiased for the sequence KL value)", abs(float((pi * k3_tok.sum(1)).sum()) - KL_seq), 1e-13)
# %% [markdown]
# ## 6. Naive vs stable evaluation of k3 (fp32): (r-1)-log r cancels catastrophically when |log r| is small
# %%
lr32 = torch.tensor([1e-1, 1e-2, 1e-3, 1e-4, 1e-5], dtype=torch.float32)
naive = (lr32.exp() - 1) - lr32
stable = torch.expm1(lr32) - lr32
truth = (torch.expm1(lr32.double()) - lr32.double())
print(" log r naive fp32 expm1 fp32 fp64 truth naive rel.err stable rel.err")
for a, n_, s_, t_ in zip(lr32.tolist(), naive.tolist(), stable.tolist(), truth.tolist()):
print(f" {a:7.0e} {n_:+.4e} {s_:+.4e} {t_:+.4e} {abs(n_-t_)/t_:9.2e} {abs(s_-t_)/t_:9.2e}")
check("expm1 form: relative error < 1e-3 for |log r| >= 1e-4 (fp32)", float(((stable.double() - truth).abs() / truth)[:4].max()), 1e-3)
assert float(((naive.double() - truth).abs() / truth)[3]) > 0.1
check("naive form is wrong by >10% at |log r| = 1e-4 (fp32)", 0.0, 0.0)
# The expm1 form removes the cancellation inside exp(x)-1 but NOT the one in expm1(x) - x ~ x^2/2: its worst-case relative error is ~ 2u/|x|.
# Worst case over a dense random grid per decade (fp32, both signs), vs the 2u/|x|_min law; and a series fallback for small |x|.
u32 = 2.0 ** -24
gen_n = torch.Generator().manual_seed(11)
def k3_expm1(x): return torch.expm1(x) - x
def k3_hybrid(x):
ser = x * x / 2 * (1 + x / 3 + x * x / 12) # Taylor of expm1(x)-x, error ~ x^5/120 relative x^3/60
return torch.where(x.abs() < 1e-2, ser, k3_expm1(x))
print(" |x| decade worst rel.err expm1 form 2u/|x|min worst rel.err hybrid (series for |x|<1e-2)")
worst_ratio, worst_hybrid = 0.0, 0.0
for lo in (1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6):
hi = 10 * lo
xg = torch.exp(torch.rand(400_000, generator=gen_n, dtype=torch.float64) * math.log(10) + math.log(lo)).float()
xg = torch.cat([xg, -xg]); tr = torch.expm1(xg.double()) - xg.double()
re = float(((k3_expm1(xg).double() - tr).abs() / tr).max()); rh = float(((k3_hybrid(xg).double() - tr).abs() / tr).max())
law = 2 * u32 / lo
worst_ratio = max(worst_ratio, re / law); worst_hybrid = max(worst_hybrid, rh)
print(f" [{lo:.0e},{hi:.0e}] {re:12.3e} {law:10.2e} {rh:12.3e}")
check("expm1 form: worst-case relative error <= 2u/|x| (max ratio to the law)", worst_ratio, 1.0)
check("hybrid (Taylor for |x|<1e-2, expm1 form above): worst relative error over |x| in [1e-6, 1)", worst_hybrid, 1e-4)
print("done; checks:", len(CHECKS))
|