diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/jlens_v3.py | 54 |
1 files changed, 43 insertions, 11 deletions
diff --git a/src/jlens_v3.py b/src/jlens_v3.py index 69299dc..96af86c 100644 --- a/src/jlens_v3.py +++ b/src/jlens_v3.py @@ -58,7 +58,8 @@ def make_batches(train_data, block_size, batch_size, n_prompts, device, seed=133 return batches -def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32): +def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32, + skip_first=0, source_mean=False): """ Compute the faithful J-lens VECTORS (rows of W_U * J_l) directly. @@ -69,6 +70,14 @@ def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32): probe with the unembedding rows W_U[k] (V=65 probes) — mathematically identical, ~6x cheaper. One batched autograd call per prompt per layer (is_grads_batched=True) computes all V token vectors at once. + + Two estimator modes (mirrors Anthropic's official jacobian-lens repo): + source_mean=False (default): mean over (source, future) PAIRS, all + positions included — the paper-formula estimator. + source_mean=True + skip_first=N: cotangent placed only at target + positions in [N, T-1) and the mean is over VALID SOURCE positions, + excluding the first N (attention sinks) and the last (no next-token + target) — exactly the official repo's valid_position_mask / fit(). """ d = model.config.n_embd V = model.config.vocab_size @@ -77,7 +86,7 @@ def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32): W_U = model.lm_head.weight.detach().float() # (V, d), nanoGPT weight tying accum = torch.zeros(V, d, device='cpu') - n_pairs = 0 + n_positions = 0 for (x, y) in batches: B, T = x.shape @@ -97,8 +106,16 @@ def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32): h_l = resid_captured['h_l'] # (B, T, d) h_final = resid_captured['h_final'] # (B, T, d) - # sum over all future positions t' (causal: t' < t contributes zero grad) - h_final_sum = h_final.sum(dim=1) # (B, d) + if source_mean: + # official-repo valid positions: [skip_first, T-1) + pos = torch.arange(T, device=device) + valid = (pos >= skip_first) & (pos < T - 1) + n_valid = int(valid.sum().item()) + h_final_sum = h_final[:, valid, :].sum(dim=1) # (B, d) + else: + valid = None + n_valid = 0 + h_final_sum = h_final.sum(dim=1) # (B, d); t' < t gives zero grad # grad_outputs[v, b, :] = W_U[v] -> batched VJPs, chunked over vocab # (is_grads_batched with all V=65 probes at once OOMs the 4GB K2200) @@ -110,18 +127,26 @@ def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32): h_final_sum, h_l, grad_outputs=grad_outputs, is_grads_batched=True, retain_graph=(v0 + C < V))[0] # (C, B, T, d) - accum[v0:v0 + C] += grads.detach().cpu().reshape(C, -1, d).sum(dim=1) + if source_mean: + assert valid is not None + grads_v = grads[:, :, valid, :] # (C, B, n_valid, d) + accum[v0:v0 + C] += grads_v.detach().cpu().sum(dim=(1, 2)) + else: + accum[v0:v0 + C] += grads.detach().cpu().reshape(C, -1, d).sum(dim=1) del grads, grad_outputs torch.cuda.empty_cache() - n_pairs += B * T * (T + 1) // 2 # valid (t, t' >= t) pairs - # NOTE: count = T(T+1)/2 per sequence (all source positions t and all - # futures t' >= t). Pairs with t' < t contribute zero gradient by - # causality, so summing over all t' and dividing by this count is exact. + if source_mean: + n_positions += B * n_valid + else: + # NOTE: count = T(T+1)/2 per sequence (all source positions t and all + # futures t' >= t). Pairs with t' < t contribute zero gradient by + # causality, so summing over all t' and dividing by this count is exact. + n_positions += B * T * (T + 1) // 2 del logits, loss, h_l, h_final, h_final_sum torch.cuda.empty_cache() - accum /= n_pairs + accum /= n_positions return accum # (V, d): faithful J-lens vectors, rows of W_U * J_l @@ -176,6 +201,12 @@ def main(): help='directory for per-layer result artifacts') ap.add_argument('--device', default='cuda') ap.add_argument('--chunk', type=int, default=32) + ap.add_argument('--skip_first', type=int, default=0, + help='exclude first N source/target positions (mirrors ' + 'Anthropic official repo skip=16 attention sinks)') + ap.add_argument('--source_mean', action='store_true', + help='mean over valid source positions (official-repo ' + 'estimator) instead of (source, future) pairs') args = ap.parse_args() device = args.device @@ -209,7 +240,8 @@ def main(): print(f"\n--- Layer {layer_idx} ---") print(" Computing faithful J-lens vectors (W_U-probed VJPs)...") faithful_vecs = compute_faithful_jlens(model, layer_idx, batches, device, - args.chunk) # (V, d) + args.chunk, args.skip_first, + args.source_mean) # (V, d) faithful_norms = {k: faithful_vecs[k].norm().item() for k in range(V)} if layer_idx == n_layer - 1: |