diff options
| -rw-r--r-- | ep_run/casc_eq_train.py | 38 |
1 files changed, 18 insertions, 20 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 392d2be..38b687e 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -93,34 +93,33 @@ def tok_sigma(iters=8): v = W.t() @ u; sig = v.norm(); v /= max(sig, 1e-12) return float(sig) -def relax(z0, zs, ins, outs, y, beta, K): - """K fb rounds with GRAPH REUSE: the backward vjps consume the graphs stored by the - previous forward (free pass for round 1, rebuild pass afterwards) — saves one full - graphed chain per round. Rebuild keeps graphs for the next round; oscillation harmless.""" +def relax(z0, zs, ins, outs, y, beta, K, x): + """K fb rounds with GRAPH REUSE + two dedups: (a) the top CE force d_top is refreshed on + even rounds only (states move O(beta) per round -> O(beta^2) error); (b) the LAST rebuild + keeps graphs (layer-0 fed a graphed emb) and returns (ins, outs) so the theta-readout + reuses them instead of re-running a full graphed chain.""" d = [None] * args.L for k in range(K): - zc = zs[args.L - 1].detach().requires_grad_(True) - ce = F.cross_entropy(readout(zc).reshape(-1, vocab), y.reshape(-1)) - d[args.L - 1] = (-beta * NBT * torch.autograd.grad(ce, zc)[0]).detach() + if k % 2 == 0 or d[args.L - 1] is None: + zc = zs[args.L - 1].detach().requires_grad_(True) + ce = F.cross_entropy(readout(zc).reshape(-1, vocab), y.reshape(-1)) + d[args.L - 1] = (-beta * NBT * torch.autograd.grad(ce, zc)[0]).detach() for l in range(args.L - 2, -1, -1): - # d_l = J_{l+1}^T d_{l+1}; outs[l+1] was computed at input ins[l+1] == zs[l] - d[l] = torch.autograd.grad(outs[l + 1], ins[l + 1], grad_outputs=d[l + 1], - retain_graph=(k + 1 < K and False))[0].detach() + d[l] = torch.autograd.grad(outs[l + 1], ins[l + 1], grad_outputs=d[l + 1])[0].detach() last = (k + 1 == K) prev = z0 n_ins, n_outs = [], [] for l in range(args.L): - i = prev.detach().requires_grad_(True) - if last: - with torch.no_grad(): - o = blocks[l](i, mask) + if last and l == 0: + i = tok(x) + pos(torch.arange(args.T, device=dev))[None] # graphed emb for the readout's E-path else: - o = blocks[l](i, mask) + i = prev.detach().requires_grad_(True) + o = blocks[l](i, mask) zs[l] = (o.detach() + d[l]) n_ins.append(i); n_outs.append(o) prev = zs[l] ins, outs = n_ins, n_outs - return zs + return zs, outs def dFdtheta(zs, x, y, beta): """dF/dtheta at fixed relaxed states (z0 rebuilt WITH graph so emb gets its E-path grad).""" @@ -147,7 +146,7 @@ def ep_step(x, y): z0, zs, ins, outs = free_states_graphed(x) zs_free = [z.clone() for z in zs] free_ce = F.cross_entropy(readout(zs_free[-1]).reshape(-1, vocab), y.reshape(-1)).item() - zp = relax(z0, zs, ins, outs, y, +beta_t, GOV['K']) + zp, last_outs = relax(z0, zs, ins, outs, y, +beta_t, GOV['K'], x) with torch.no_grad(): drift = sum(float((a - b).norm()) for a, b in zip(zp, zs_free)) / max( sum(float(b.norm()) for b in zs_free), 1e-9) @@ -155,10 +154,9 @@ def ep_step(x, y): for p in all_params: p.grad = None return free_ce, beta_t, GOV['K'], False GOV['drift'] = drift - prev = tok(x) + pos(torch.arange(args.T, device=dev))[None] E = 0.0 - for z, b in zip(zp, blocks): E = E + 0.5 * ((z - b(prev, mask)) ** 2).sum(); prev = z - obj = E / (NBT * beta_t) + F.cross_entropy(readout(zp[-1]).reshape(-1, vocab), y.reshape(-1)) + for z, o in zip(zp, last_outs): E = E + 0.5 * ((z.detach() - o) ** 2).sum() + obj = E / (NBT * beta_t) + F.cross_entropy(readout(zp[-1].detach()).reshape(-1, vocab), y.reshape(-1)) gs = torch.autograd.grad(obj, all_params, allow_unused=True) gn = 0.0 for g in gs: |