diff options
Diffstat (limited to 'experiments/vector_credit_audit.py')
| -rw-r--r-- | experiments/vector_credit_audit.py | 844 |
1 files changed, 844 insertions, 0 deletions
diff --git a/experiments/vector_credit_audit.py b/experiments/vector_credit_audit.py new file mode 100644 index 0000000..048efb7 --- /dev/null +++ b/experiments/vector_credit_audit.py @@ -0,0 +1,844 @@ +""" +Phase 5A: Vector Credit Field Audit. + +Verify that the vector field's gains are real, not implementation artifacts. + +4 mandatory sanity checks: +A. Train/eval direction split (independent random directions) +B. Shuffled-target control (permute g_j within batch) +C. No-terminal ablation (L_term = 0) +D. One-sided vs symmetric finite difference +""" +import os +import sys +import json +import argparse +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +import copy + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from models.value_net import ValueNet, SinusoidalTimeEmbed, create_ema_model, update_ema +from metrics.credit_metrics import ( + cosine_similarity_batch, perturbation_correlation, nudging_test +) + + +# ============================================================================= +# Synthetic teacher-student +# ============================================================================= +class TeacherNet(nn.Module): + def __init__(self, d_hidden, num_classes, num_blocks, alpha=1.0, seed=0): + super().__init__() + self.d_hidden = d_hidden + self.num_blocks = num_blocks + self.alpha = alpha + rng = torch.Generator().manual_seed(seed) + self.Ws = nn.ParameterList() + for _ in range(num_blocks): + W = torch.randn(d_hidden, d_hidden, generator=rng) * 0.3 / (d_hidden ** 0.5) + U, S, Vh = torch.linalg.svd(W, full_matrices=False) + S_clamped = S.clamp(max=0.3) + W = U @ torch.diag(S_clamped) @ Vh + self.Ws.append(nn.Parameter(W, requires_grad=False)) + self.U = nn.Parameter( + torch.randn(num_classes, d_hidden, generator=rng) / (d_hidden ** 0.5), + requires_grad=False) + + def phi(self, z): + return (1 - self.alpha) * z + self.alpha * torch.tanh(z) + + def forward(self, x): + h = x + for W in self.Ws: + h = h + self.phi(h @ W.T) + return h @ self.U.T + + +class StudentBlock(nn.Module): + def __init__(self, d_hidden, alpha=1.0): + super().__init__() + self.ln = nn.LayerNorm(d_hidden) + self.w = nn.Linear(d_hidden, d_hidden, bias=False) + nn.init.normal_(self.w.weight, std=0.01) + self.alpha = alpha + + def phi(self, z): + return (1 - self.alpha) * z + self.alpha * torch.tanh(z) + + def forward(self, h): + return self.w(self.phi(self.ln(h))) + + +class StudentNet(nn.Module): + def __init__(self, d_hidden, num_classes, num_blocks, alpha=1.0): + super().__init__() + self.blocks = nn.ModuleList([StudentBlock(d_hidden, alpha) for _ in range(num_blocks)]) + self.out_head = nn.Linear(d_hidden, num_classes) + self.d_hidden = d_hidden + self.num_blocks = num_blocks + + def forward(self, x, return_hidden=False): + h = x + hiddens = [h] if return_hidden else None + for block in self.blocks: + f = block(h) + h = h + f + if return_hidden: + hiddens.append(h) + logits = self.out_head(h) + if return_hidden: + return logits, hiddens + return logits + + def forward_from_layer(self, h, start_layer): + for i in range(start_layer, self.num_blocks): + f = self.blocks[i](h) + h = h + f + return self.out_head(h) + + +class VectorCreditNet(nn.Module): + """Direct vector credit field: a_phi(h_l, t_l, s) -> R^d.""" + def __init__(self, d_hidden, s_dim, time_embed_dim=32, hidden_dim=256, num_layers=3): + super().__init__() + self.ln = nn.LayerNorm(d_hidden) + self.time_embed = SinusoidalTimeEmbed(time_embed_dim) + input_dim = d_hidden + time_embed_dim + s_dim + layers = [] + for i in range(num_layers): + in_d = input_dim if i == 0 else hidden_dim + layers.append(nn.Linear(in_d, hidden_dim)) + layers.append(nn.GELU()) + layers.append(nn.Linear(hidden_dim, d_hidden)) + self.net = nn.Sequential(*layers) + + def forward(self, h, t, s): + h_normed = self.ln(h) + t_emb = self.time_embed(t) + inp = torch.cat([h_normed, t_emb, s], dim=-1) + return self.net(inp) + + +def generate_batch(teacher, d_hidden, num_classes, batch_size, device): + x = torch.randn(batch_size, d_hidden, device=device) + with torch.no_grad(): + teacher_logits = teacher(x) + y = teacher_logits.argmax(dim=-1) + return x, y + + +# ============================================================================= +# Training: vector field with audit controls +# ============================================================================= +def train_vector_field_audit(model, teacher, device, args, M=4, + use_terminal=True, + shuffle_targets=False, + use_central_diff=True, + tag='vec'): + """ + Train vector credit field with configurable audit controls. + + Args: + use_terminal: if False, L_term = 0 (no-terminal ablation) + shuffle_targets: if True, permute g_j within batch (leak check) + use_central_diff: if True, central difference; if False, one-sided + tag: label for printing + """ + d = model.d_hidden + L = model.num_blocks + num_classes = args.num_classes + + vector_net = VectorCreditNet(d_hidden=d, s_dim=num_classes, time_embed_dim=32, + hidden_dim=256, num_layers=3).to(device) + + Bs = [torch.randn(d, num_classes, device=device) / np.sqrt(num_classes) for _ in range(L)] + + block_opts = [optim.AdamW(b.parameters(), lr=args.lr, weight_decay=0.01) for b in model.blocks] + head_opt = optim.AdamW(model.out_head.parameters(), lr=args.lr, weight_decay=0.01) + vec_opt = optim.Adam(vector_net.parameters(), lr=args.lr_fb) + + warmup_epochs = max(1, int(args.epochs * args.warmup_ratio)) + eps = args.pert_eps + beta = args.pert_beta + + for epoch in range(1, args.epochs + 1): + model.train() + vector_net.train() + + if epoch <= warmup_epochs: + credit_blend = 0.0 + else: + credit_blend = min(1.0, (epoch - warmup_epochs) / max(1, warmup_epochs)) + + total_loss, correct, total = 0, 0, 0 + total_vloss = 0 + + for _ in range(args.steps_per_epoch): + x, y = generate_batch(teacher, d, num_classes, args.batch_size, device) + batch = x.size(0) + + with torch.no_grad(): + logits, hiddens = model(x, return_hidden=True) + loss_val = F.cross_entropy(logits, y) + e_T = logits.softmax(dim=-1) + e_T[torch.arange(batch), y] -= 1 + s = e_T.detach() + + hL_det = hiddens[-1].detach() + + # --- Terminal matching --- + loss_term = torch.tensor(0.0, device=device) + if use_terminal: + t_L = torch.ones(batch, device=device) + a_terminal = vector_net(hL_det, t_L, s) + hL_req = hL_det.clone().requires_grad_(True) + logits_tgt = model.out_head(hL_req) + ce = F.cross_entropy(logits_tgt, y, reduction='sum') + delta_L = torch.autograd.grad(ce, hL_req, create_graph=False)[0].detach() + loss_term = ((a_terminal - delta_L) ** 2).sum(dim=-1).mean() + + # --- Perturbation directional targets --- + # IMPORTANT: training directions are sampled fresh each step. + # Evaluation uses independently sampled directions (see compute_diagnostics). + loss_proj = torch.tensor(0.0, device=device) + for l in range(L): + h_l_det = hiddens[l].detach() + t_l = torch.full((batch,), l / L, device=device) + a_l = vector_net(h_l_det, t_l, s) + + layer_proj_loss = 0.0 + for _ in range(M): + v = torch.randn_like(h_l_det) + v = v / (v.norm(dim=-1, keepdim=True) + 1e-8) + + with torch.no_grad(): + if use_central_diff: + # Central difference: [loss(h+eps*v) - loss(h-eps*v)] / (2*eps) + logits_plus = model.forward_from_layer(h_l_det + eps * v, l) + loss_plus = F.cross_entropy(logits_plus, y, reduction='none') + logits_minus = model.forward_from_layer(h_l_det - eps * v, l) + loss_minus = F.cross_entropy(logits_minus, y, reduction='none') + g_j = (loss_plus - loss_minus) / (2 * eps) + else: + # One-sided difference: [loss(h+eps*v) - loss(h)] / eps + logits_base = model.forward_from_layer(h_l_det, l) + loss_base = F.cross_entropy(logits_base, y, reduction='none') + logits_plus = model.forward_from_layer(h_l_det + eps * v, l) + loss_plus = F.cross_entropy(logits_plus, y, reduction='none') + g_j = (loss_plus - loss_base) / eps + + # Shuffled-target control: permute g_j within batch + if shuffle_targets: + perm = torch.randperm(batch, device=device) + g_j = g_j[perm] + + pred_j = (a_l * v).sum(dim=-1) + layer_proj_loss = layer_proj_loss + ((pred_j - g_j.detach()) ** 2).mean() + + loss_proj = loss_proj + layer_proj_loss / M + loss_proj = loss_proj / L + + vec_loss = loss_term + beta * loss_proj + vec_opt.zero_grad() + vec_loss.backward() + torch.nn.utils.clip_grad_norm_(vector_net.parameters(), 1.0) + vec_opt.step() + total_vloss += vec_loss.item() * batch + + # --- Block updates --- + with torch.no_grad(): + vec_credits = [] + for l in range(L): + h_l_det = hiddens[l].detach() + t_l = torch.full((batch,), l / L, device=device) + a_l = vector_net(h_l_det, t_l, s) + vec_credits.append(a_l.detach()) + + dfa_credits = [(e_T @ Bs[l].T).detach() for l in range(L)] + + credits = [] + for l in range(L): + if credit_blend >= 1.0: + credits.append(vec_credits[l]) + elif credit_blend <= 0.0: + credits.append(dfa_credits[l]) + else: + vc_rms = (vec_credits[l] ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6 + dfa_rms = (dfa_credits[l] ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6 + credits.append(credit_blend * vec_credits[l] / vc_rms + + (1 - credit_blend) * dfa_credits[l] / dfa_rms) + + logits_out = model.out_head(hL_det) + loss_out = F.cross_entropy(logits_out, y) + head_opt.zero_grad() + loss_out.backward() + head_opt.step() + + for l in range(L): + h_l = hiddens[l].detach() + a = credits[l] + rms = (a ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6 + f_l = model.blocks[l](h_l) + local_loss = (f_l * (a / rms)).sum(dim=-1).mean() + block_opts[l].zero_grad() + local_loss.backward() + torch.nn.utils.clip_grad_norm_(model.blocks[l].parameters(), 1.0) + block_opts[l].step() + + total_loss += loss_val.item() * batch + correct += (logits.argmax(1) == y).sum().item() + total += batch + + if epoch % 20 == 0 or epoch == 1: + acc = correct / total + print(f" [{tag}] Ep {epoch}: loss={total_loss/total:.4f}, acc={acc:.4f}, " + f"vloss={total_vloss/total:.6f}") + + return vector_net + + +def train_scalar_cb(model, teacher, device, args): + """Scalar credit bridge baseline.""" + d = model.d_hidden + L = model.num_blocks + num_classes = args.num_classes + + value_net = ValueNet(d_hidden=d, s_dim=num_classes, time_embed_dim=32, + hidden_dim=256, num_layers=3).to(device) + value_net_ema = create_ema_model(value_net) + + Bs = [torch.randn(d, num_classes, device=device) / np.sqrt(num_classes) for _ in range(L)] + + block_opts = [optim.AdamW(b.parameters(), lr=args.lr, weight_decay=0.01) for b in model.blocks] + head_opt = optim.AdamW(model.out_head.parameters(), lr=args.lr, weight_decay=0.01) + value_opt = optim.Adam(value_net.parameters(), lr=args.lr_fb) + + warmup_epochs = max(1, int(args.epochs * args.warmup_ratio)) + + for epoch in range(1, args.epochs + 1): + model.train() + value_net.train() + + if epoch <= warmup_epochs: + credit_blend = 0.0 + else: + credit_blend = min(1.0, (epoch - warmup_epochs) / max(1, warmup_epochs)) + + total_loss, correct, total = 0, 0, 0 + for _ in range(args.steps_per_epoch): + x, y = generate_batch(teacher, d, num_classes, args.batch_size, device) + batch = x.size(0) + + with torch.no_grad(): + logits, hiddens = model(x, return_hidden=True) + loss_val = F.cross_entropy(logits, y) + e_T = logits.softmax(dim=-1) + e_T[torch.arange(batch), y] -= 1 + s = e_T.detach() + true_loss = F.cross_entropy(logits, y, reduction='none').detach() + + hL_det = hiddens[-1].detach() + t_L = torch.ones(batch, device=device) + V_term = value_net(hL_det, t_L, s) + loss_term = ((V_term - true_loss) ** 2).mean() + + hL_req = hL_det.clone().requires_grad_(True) + V_at_L = value_net(hL_req, t_L, s) + grad_V_L = torch.autograd.grad(V_at_L.sum(), hL_req, create_graph=True)[0] + hL_req2 = hL_det.clone().requires_grad_(True) + logits_tgt = model.out_head(hL_req2) + ce = F.cross_entropy(logits_tgt, y, reduction='sum') + a_L_exact = torch.autograd.grad(ce, hL_req2, create_graph=False)[0].detach() + loss_tgrad = ((grad_V_L - a_L_exact) ** 2).sum(dim=-1).mean() + + loss_bridge = 0.0 + for l in range(L): + h_l_det = hiddens[l].detach() + t_l = torch.full((batch,), l / L, device=device) + t_next = torch.full((batch,), (l + 1) / L, device=device) + V_l = value_net(h_l_det, t_l, s) + with torch.no_grad(): + h_next = hiddens[l + 1].detach() + log_terms = [] + for k in range(args.K): + noise = args.sigma_bridge * torch.randn_like(h_next) + V_next = value_net_ema(h_next + noise, t_next, s) + log_terms.append(-V_next / args.lam) + log_stack = torch.stack(log_terms, dim=-1) + V_target = -args.lam * (torch.logsumexp(log_stack, dim=-1) - np.log(args.K)) + loss_bridge += ((V_l - V_target.detach()) ** 2).mean() + loss_bridge /= L + + vloss = loss_term + loss_bridge + args.term_grad_weight * loss_tgrad + value_opt.zero_grad() + vloss.backward() + torch.nn.utils.clip_grad_norm_(value_net.parameters(), 1.0) + value_opt.step() + update_ema(value_net, value_net_ema, args.ema_momentum) + + cb_credits = [] + for l in range(L): + h_l_det = hiddens[l].detach().requires_grad_(True) + t_l = torch.full((batch,), l / L, device=device) + V_l = value_net(h_l_det, t_l, s) + a_l = torch.autograd.grad(V_l.sum(), h_l_det, create_graph=False)[0] + cb_credits.append(a_l.detach()) + + dfa_credits = [(e_T @ Bs[l].T).detach() for l in range(L)] + + credits = [] + for l in range(L): + if credit_blend >= 1.0: + credits.append(cb_credits[l]) + elif credit_blend <= 0.0: + credits.append(dfa_credits[l]) + else: + cb_rms = (cb_credits[l] ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6 + dfa_rms = (dfa_credits[l] ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6 + credits.append(credit_blend * cb_credits[l] / cb_rms + + (1 - credit_blend) * dfa_credits[l] / dfa_rms) + + logits_out = model.out_head(hL_det) + loss_out = F.cross_entropy(logits_out, y) + head_opt.zero_grad() + loss_out.backward() + head_opt.step() + + for l in range(L): + h_l = hiddens[l].detach() + a = credits[l] + rms = (a ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6 + f_l = model.blocks[l](h_l) + local_loss = (f_l * (a / rms)).sum(dim=-1).mean() + block_opts[l].zero_grad() + local_loss.backward() + torch.nn.utils.clip_grad_norm_(model.blocks[l].parameters(), 1.0) + block_opts[l].step() + + total_loss += loss_val.item() * batch + correct += (logits.argmax(1) == y).sum().item() + total += batch + + if epoch % 20 == 0 or epoch == 1: + print(f" [scalar_cb] Ep {epoch}: loss={total_loss/total:.4f}, acc={correct/total:.4f}") + + return value_net + + +def train_dfa(model, teacher, device, args): + """DFA baseline.""" + d = model.d_hidden + L = model.num_blocks + num_classes = args.num_classes + Bs = [torch.randn(d, num_classes, device=device) / np.sqrt(num_classes) for _ in range(L)] + + block_opts = [optim.AdamW(b.parameters(), lr=args.lr, weight_decay=0.01) for b in model.blocks] + head_opt = optim.AdamW(model.out_head.parameters(), lr=args.lr, weight_decay=0.01) + + for epoch in range(1, args.epochs + 1): + model.train() + total_loss, correct, total = 0, 0, 0 + for _ in range(args.steps_per_epoch): + x, y = generate_batch(teacher, d, num_classes, args.batch_size, device) + batch = x.size(0) + with torch.no_grad(): + logits, hiddens = model(x, return_hidden=True) + loss_val = F.cross_entropy(logits, y) + e_T = logits.softmax(dim=-1) + e_T[torch.arange(batch), y] -= 1 + hL_det = hiddens[-1].detach() + logits_out = model.out_head(hL_det) + loss_out = F.cross_entropy(logits_out, y) + head_opt.zero_grad() + loss_out.backward() + head_opt.step() + for l in range(L): + h_l = hiddens[l].detach() + a = (e_T @ Bs[l].T).detach() + rms = (a ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6 + f_l = model.blocks[l](h_l) + local_loss = (f_l * (a / rms)).sum(dim=-1).mean() + block_opts[l].zero_grad() + local_loss.backward() + torch.nn.utils.clip_grad_norm_(model.blocks[l].parameters(), 1.0) + block_opts[l].step() + total_loss += loss_val.item() * batch + correct += (logits.argmax(1) == y).sum().item() + total += batch + if epoch % 20 == 0 or epoch == 1: + print(f" [DFA] Ep {epoch}: loss={total_loss/total:.4f}, acc={correct/total:.4f}") + return Bs + + +# ============================================================================= +# Diagnostics — uses INDEPENDENT eval directions (check A) +# ============================================================================= +def compute_diagnostics(model, teacher, device, method_name, args, + value_net=None, vector_net=None, dfa_Bs=None): + """ + Compute Gamma, rho, nudging per layer. + IMPORTANT: perturbation_correlation uses its own freshly-sampled directions, + completely independent of any training directions. This ensures check A. + """ + model.eval() + if value_net is not None: + value_net.eval() + if vector_net is not None: + vector_net.eval() + + d = model.d_hidden + L = model.num_blocks + num_classes = args.num_classes + + # Use a fixed eval seed different from training + eval_rng = torch.Generator(device=device) + eval_rng.manual_seed(99999) + + x = torch.randn(512, d, device=device, generator=eval_rng) + with torch.no_grad(): + teacher_logits = teacher(x) + y = teacher_logits.argmax(dim=-1) + batch = x.size(0) + + # BP gradients (evaluation only — never used for training) + h = x.detach().requires_grad_(True) + hiddens_bp = [h] + for block in model.blocks: + f = block(hiddens_bp[-1]) + h_next = hiddens_bp[-1] + f + hiddens_bp.append(h_next) + logits_bp = model.out_head(hiddens_bp[-1]) + loss_bp = F.cross_entropy(logits_bp, y) + grads = torch.autograd.grad(loss_bp, hiddens_bp, retain_graph=False) + bp_grads = {l: grads[l].detach().clone() for l in range(L + 1)} + + with torch.no_grad(): + logits, hiddens = model(x, return_hidden=True) + e_T = logits.softmax(dim=-1) + e_T[torch.arange(batch), y] -= 1 + s = e_T.detach() + + results = {'bp_cosine': [], 'perturbation_rho': [], 'nudging': []} + + for l in range(L): + h_l = hiddens[l].detach() + t_l = torch.full((batch,), l / L, device=device) + + if method_name == 'dfa': + a_l = (s @ dfa_Bs[l].T).detach() + elif method_name == 'scalar_cb': + h_l_req = h_l.clone().requires_grad_(True) + V_l = value_net(h_l_req, t_l, s) + a_l = torch.autograd.grad(V_l.sum(), h_l_req, create_graph=False)[0].detach() + elif method_name.startswith('vec'): + a_l = vector_net(h_l, t_l, s).detach() + else: + raise ValueError(f"Unknown: {method_name}") + + bp_cos = cosine_similarity_batch(a_l, bp_grads[l]) + results['bp_cosine'].append(float(bp_cos)) + + # perturbation_correlation uses its own random directions internally + # (from metrics/credit_metrics.py — independent of training directions) + def make_fwd_fn(start_l): + def fwd_fn(h): + with torch.no_grad(): + logits = model.forward_from_layer(h, start_l) + return F.cross_entropy(logits, y, reduction='none') + return fwd_fn + + fwd_fn = make_fwd_fn(l) + rho = perturbation_correlation(h_l, a_l, fwd_fn, epsilon=1e-3, M=32) + results['perturbation_rho'].append(float(rho)) + + nud = nudging_test(h_l, a_l, fwd_fn, eta=0.003) + results['nudging'].append(float(nud)) + + return results + + +# ============================================================================= +# Main +# ============================================================================= +def run_experiment(args): + device = torch.device(f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu') + print(f"Using device: {device}") + os.makedirs(args.output_dir, exist_ok=True) + + all_results = [] + + for L in args.depths: + for seed in args.seeds: + print(f"\n{'='*60}") + print(f"L={L}, seed={seed}") + print(f"{'='*60}") + + teacher = TeacherNet(args.d_hidden, args.num_classes, L, + alpha=args.alpha, seed=seed * 1000).to(device) + + # --- DFA --- + print("\n --- DFA ---") + torch.manual_seed(seed) + np.random.seed(seed) + torch.cuda.manual_seed_all(seed) + model_dfa = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device) + Bs = train_dfa(model_dfa, teacher, device, args) + diag = compute_diagnostics(model_dfa, teacher, device, 'dfa', args, dfa_Bs=Bs) + r = {'method': 'dfa', 'L': L, 'seed': seed, + 'mean_gamma': float(np.mean(diag['bp_cosine'])), + 'mean_rho': float(np.mean(diag['perturbation_rho'])), + 'mean_nudge': float(np.mean(diag['nudging'])), + 'per_layer': diag} + print(f" Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}") + all_results.append(r) + + # --- Scalar CB --- + print("\n --- Scalar CB ---") + torch.manual_seed(seed) + np.random.seed(seed) + torch.cuda.manual_seed_all(seed) + model_cb = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device) + vnet = train_scalar_cb(model_cb, teacher, device, args) + diag = compute_diagnostics(model_cb, teacher, device, 'scalar_cb', args, value_net=vnet) + r = {'method': 'scalar_cb', 'L': L, 'seed': seed, + 'mean_gamma': float(np.mean(diag['bp_cosine'])), + 'mean_rho': float(np.mean(diag['perturbation_rho'])), + 'mean_nudge': float(np.mean(diag['nudging'])), + 'per_layer': diag} + print(f" Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}") + all_results.append(r) + + # --- Vector Field M4 (central diff, with terminal) --- + print("\n --- vec_eT_M4 (central, +term) ---") + torch.manual_seed(seed) + np.random.seed(seed) + torch.cuda.manual_seed_all(seed) + model_v4 = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device) + vnet4 = train_vector_field_audit(model_v4, teacher, device, args, M=4, + use_terminal=True, shuffle_targets=False, + use_central_diff=True, tag='vec_eT_M4') + diag = compute_diagnostics(model_v4, teacher, device, 'vec_eT_M4', args, vector_net=vnet4) + r = {'method': 'vec_eT_M4', 'L': L, 'seed': seed, + 'mean_gamma': float(np.mean(diag['bp_cosine'])), + 'mean_rho': float(np.mean(diag['perturbation_rho'])), + 'mean_nudge': float(np.mean(diag['nudging'])), + 'per_layer': diag} + print(f" Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}") + all_results.append(r) + + # --- Vector Field M8 (central diff, with terminal) --- + if 8 in args.M_values: + print("\n --- vec_eT_M8 (central, +term) ---") + torch.manual_seed(seed) + np.random.seed(seed) + torch.cuda.manual_seed_all(seed) + model_v8 = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device) + vnet8 = train_vector_field_audit(model_v8, teacher, device, args, M=8, + use_terminal=True, shuffle_targets=False, + use_central_diff=True, tag='vec_eT_M8') + diag = compute_diagnostics(model_v8, teacher, device, 'vec_eT_M8', args, vector_net=vnet8) + r = {'method': 'vec_eT_M8', 'L': L, 'seed': seed, + 'mean_gamma': float(np.mean(diag['bp_cosine'])), + 'mean_rho': float(np.mean(diag['perturbation_rho'])), + 'mean_nudge': float(np.mean(diag['nudging'])), + 'per_layer': diag} + print(f" Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}") + all_results.append(r) + + # ================================================================= + # SANITY CHECKS (only for first seed to save time, unless full mode) + # ================================================================= + if seed == args.seeds[0] or args.full_audit: + # --- Check B: Shuffled-target control --- + print("\n --- vec_eT_M4_shuffleCtrl ---") + torch.manual_seed(seed) + np.random.seed(seed) + torch.cuda.manual_seed_all(seed) + model_shuf = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device) + vnet_shuf = train_vector_field_audit(model_shuf, teacher, device, args, M=4, + use_terminal=True, shuffle_targets=True, + use_central_diff=True, tag='vec_shuffleCtrl') + diag = compute_diagnostics(model_shuf, teacher, device, 'vec_shuffleCtrl', args, vector_net=vnet_shuf) + r = {'method': 'vec_eT_M4_shuffleCtrl', 'L': L, 'seed': seed, + 'mean_gamma': float(np.mean(diag['bp_cosine'])), + 'mean_rho': float(np.mean(diag['perturbation_rho'])), + 'mean_nudge': float(np.mean(diag['nudging'])), + 'per_layer': diag} + print(f" Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}") + all_results.append(r) + + # --- Check C: No-terminal ablation --- + print("\n --- vec_eT_M4_noTerm ---") + torch.manual_seed(seed) + np.random.seed(seed) + torch.cuda.manual_seed_all(seed) + model_nt = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device) + vnet_nt = train_vector_field_audit(model_nt, teacher, device, args, M=4, + use_terminal=False, shuffle_targets=False, + use_central_diff=True, tag='vec_noTerm') + diag = compute_diagnostics(model_nt, teacher, device, 'vec_noTerm', args, vector_net=vnet_nt) + r = {'method': 'vec_eT_M4_noTerm', 'L': L, 'seed': seed, + 'mean_gamma': float(np.mean(diag['bp_cosine'])), + 'mean_rho': float(np.mean(diag['perturbation_rho'])), + 'mean_nudge': float(np.mean(diag['nudging'])), + 'per_layer': diag} + print(f" Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}") + all_results.append(r) + + # --- Check D: One-sided difference --- + print("\n --- vec_eT_M4_onesided ---") + torch.manual_seed(seed) + np.random.seed(seed) + torch.cuda.manual_seed_all(seed) + model_os = StudentNet(args.d_hidden, args.num_classes, L, alpha=args.alpha).to(device) + vnet_os = train_vector_field_audit(model_os, teacher, device, args, M=4, + use_terminal=True, shuffle_targets=False, + use_central_diff=False, tag='vec_onesided') + diag = compute_diagnostics(model_os, teacher, device, 'vec_onesided', args, vector_net=vnet_os) + r = {'method': 'vec_eT_M4_onesided', 'L': L, 'seed': seed, + 'mean_gamma': float(np.mean(diag['bp_cosine'])), + 'mean_rho': float(np.mean(diag['perturbation_rho'])), + 'mean_nudge': float(np.mean(diag['nudging'])), + 'per_layer': diag} + print(f" Result: Gamma={r['mean_gamma']:.4f}, rho={r['mean_rho']:.4f}, nudge={r['mean_nudge']:.6f}") + all_results.append(r) + + # ================================================================= + # Summary + # ================================================================= + print(f"\n{'='*80}") + print("AUDIT SUMMARY") + print(f"{'='*80}") + print(f"{'Method':<30} {'L':>3} {'seed':>5} {'Gamma':>8} {'rho':>8} {'nudge':>10}") + print("-" * 70) + for r in all_results: + print(f"{r['method']:<30} {r['L']:>3} {r['seed']:>5} " + f"{r['mean_gamma']:>8.4f} {r['mean_rho']:>8.4f} {r['mean_nudge']:>10.6f}") + + # Check verdicts + print(f"\n{'='*60}") + print("SANITY CHECK VERDICTS") + print(f"{'='*60}") + + for L in args.depths: + seed0 = args.seeds[0] + vec_main = [r for r in all_results if r['method'] == 'vec_eT_M4' and r['L'] == L and r['seed'] == seed0] + scalar_cb = [r for r in all_results if r['method'] == 'scalar_cb' and r['L'] == L and r['seed'] == seed0] + shuf = [r for r in all_results if r['method'] == 'vec_eT_M4_shuffleCtrl' and r['L'] == L and r['seed'] == seed0] + noterm = [r for r in all_results if r['method'] == 'vec_eT_M4_noTerm' and r['L'] == L and r['seed'] == seed0] + onesided = [r for r in all_results if r['method'] == 'vec_eT_M4_onesided' and r['L'] == L and r['seed'] == seed0] + + if not vec_main or not scalar_cb: + continue + v = vec_main[0] + cb = scalar_cb[0] + + print(f"\n L={L}:") + delta_gamma = v['mean_gamma'] - cb['mean_gamma'] + delta_rho = v['mean_rho'] - cb['mean_rho'] + print(f" vec_M4 vs scalar_cb: delta_Gamma={delta_gamma:+.4f}, delta_rho={delta_rho:+.4f}") + + if shuf: + s = shuf[0] + print(f" Check B (shuffle): Gamma={s['mean_gamma']:.4f}, rho={s['mean_rho']:.4f}") + if s['mean_gamma'] < v['mean_gamma'] * 0.5 and s['mean_rho'] < v['mean_rho'] * 0.5: + print(f" -> PASS: shuffled control collapses (Gamma dropped by {v['mean_gamma']-s['mean_gamma']:.3f})") + else: + print(f" -> FAIL: shuffled control too close to main result!") + + if noterm: + n = noterm[0] + print(f" Check C (noTerm): Gamma={n['mean_gamma']:.4f}, rho={n['mean_rho']:.4f}") + if n['mean_gamma'] < v['mean_gamma'] * 0.8: + print(f" -> PASS: terminal matching contributes (Gamma dropped by {v['mean_gamma']-n['mean_gamma']:.3f})") + else: + print(f" -> NOTE: terminal removal didn't collapse result. Perturbation target alone is sufficient.") + + if onesided: + o = onesided[0] + print(f" Check D (onesided): Gamma={o['mean_gamma']:.4f}, rho={o['mean_rho']:.4f}") + if abs(o['mean_gamma'] - v['mean_gamma']) < 0.15: + print(f" -> PASS: one-sided ≈ central (difference = {abs(o['mean_gamma']-v['mean_gamma']):.3f})") + else: + print(f" -> NOTE: one-sided differs from central by {abs(o['mean_gamma']-v['mean_gamma']):.3f}") + + # Final verdict + print(f"\n{'='*60}") + print("OVERALL AUDIT VERDICT") + print(f"{'='*60}") + all_pass = True + for L in args.depths: + for seed in args.seeds: + v = [r for r in all_results if r['method'] == 'vec_eT_M4' and r['L'] == L and r['seed'] == seed] + cb = [r for r in all_results if r['method'] == 'scalar_cb' and r['L'] == L and r['seed'] == seed] + if v and cb: + dg = v[0]['mean_gamma'] - cb[0]['mean_gamma'] + dr = v[0]['mean_rho'] - cb[0]['mean_rho'] + if dg < 0.2 or dr < 0.2: + print(f" L={L} seed={seed}: delta_Gamma={dg:.3f}, delta_rho={dr:.3f} - BELOW THRESHOLD") + all_pass = False + else: + print(f" L={L} seed={seed}: delta_Gamma={dg:.3f}, delta_rho={dr:.3f} - PASS") + + shuf_results = [r for r in all_results if 'shuffleCtrl' in r['method']] + for s in shuf_results: + if s['mean_rho'] > 0.3: + print(f" SHUFFLE CONTROL WARNING: L={s['L']} rho={s['mean_rho']:.3f} too high!") + all_pass = False + + if all_pass: + print("\n AUDIT PASSED. Vector field gains are real.") + else: + print("\n AUDIT FAILED or INCOMPLETE. Investigate before proceeding.") + + # Save + save_data = [] + for r in all_results: + save_r = {k: v for k, v in r.items() if k != 'per_layer'} + save_r['per_layer_gamma'] = r['per_layer']['bp_cosine'] + save_r['per_layer_rho'] = r['per_layer']['perturbation_rho'] + save_r['per_layer_nudge'] = r['per_layer']['nudging'] + save_data.append(save_r) + + out_path = os.path.join(args.output_dir, 'audit_results.json') + with open(out_path, 'w') as f: + json.dump(save_data, f, indent=2) + print(f"\nResults saved to {out_path}") + + +def main(): + parser = argparse.ArgumentParser(description='Phase 5A: Vector Credit Field Audit') + parser.add_argument('--d_hidden', type=int, default=128) + parser.add_argument('--num_classes', type=int, default=10) + parser.add_argument('--alpha', type=float, default=1.0) + parser.add_argument('--depths', type=int, nargs='+', default=[4]) + parser.add_argument('--M_values', type=int, nargs='+', default=[4, 8]) + parser.add_argument('--epochs', type=int, default=80) + parser.add_argument('--steps_per_epoch', type=int, default=50) + parser.add_argument('--batch_size', type=int, default=256) + parser.add_argument('--lr', type=float, default=1e-3) + parser.add_argument('--lr_fb', type=float, default=1e-3) + parser.add_argument('--warmup_ratio', type=float, default=0.05) + parser.add_argument('--term_grad_weight', type=float, default=1.0) + parser.add_argument('--lam', type=float, default=0.1) + parser.add_argument('--K', type=int, default=4) + parser.add_argument('--sigma_bridge', type=float, default=0.05) + parser.add_argument('--ema_momentum', type=float, default=0.995) + parser.add_argument('--pert_eps', type=float, default=1e-3) + parser.add_argument('--pert_beta', type=float, default=1.0) + parser.add_argument('--seeds', type=int, nargs='+', default=[42]) + parser.add_argument('--gpu', type=int, default=2) + parser.add_argument('--output_dir', type=str, default='results/vector_audit') + parser.add_argument('--full_audit', action='store_true', + help='Run sanity checks for all seeds (default: first seed only)') + args = parser.parse_args() + run_experiment(args) + + +if __name__ == '__main__': + main() |