diff options
Diffstat (limited to 'experiments/snapshot_evolution_no_outln.py')
| -rw-r--r-- | experiments/snapshot_evolution_no_outln.py | 249 |
1 files changed, 249 insertions, 0 deletions
diff --git a/experiments/snapshot_evolution_no_outln.py b/experiments/snapshot_evolution_no_outln.py new file mode 100644 index 0000000..312a4cb --- /dev/null +++ b/experiments/snapshot_evolution_no_outln.py @@ -0,0 +1,249 @@ +""" +Snapshot evolution on a NO-out_ln variant of the standard ResidualMLP. +Same architecture as ResidualMLP but with the terminal LayerNorm removed +(head reads h_L directly). Trains BP and DFA from scratch on CIFAR-10 and +logs ||h_l||_2 + ||BP grad||_2 per epoch. + +This is the architectural causal control for P4: if removing out_ln from the +SAME architecture rescues the residual-stream pathology, then out_ln is +causally responsible (not just correlated). + +Usage: + CUDA_VISIBLE_DEVICES=2 nohup python experiments/snapshot_evolution_no_outln.py \ + --output_dir results/snapshot_no_outln_v1 --epochs 100 --seed 42 \ + > results/snapshot_no_outln_v1/run_s42.log 2>&1 & +""" +import os, sys, json, argparse +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +import torchvision +import torchvision.transforms as transforms + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from metrics.credit_metrics import cosine_similarity_batch + + +class ResidualBlockPreLN(nn.Module): + """Same as models/residual_mlp.ResidualBlock — pre-LN MLP block.""" + def __init__(self, d_hidden: int): + super().__init__() + self.ln = nn.LayerNorm(d_hidden) + self.w1 = nn.Linear(d_hidden, d_hidden) + self.w2 = nn.Linear(d_hidden, d_hidden) + nn.init.normal_(self.w2.weight, std=0.01) + nn.init.zeros_(self.w2.bias) + def forward(self, h): + z = self.ln(h) + z = self.w1(z) + z = F.gelu(z) + z = self.w2(z) + return z + + +class ResidualMLP_NoOutLN(nn.Module): + """Like ResidualMLP, but WITHOUT out_ln. Head reads h_L directly.""" + def __init__(self, input_dim, d_hidden, num_classes, num_blocks): + super().__init__() + self.embed = nn.Linear(input_dim, d_hidden) + self.blocks = nn.ModuleList([ResidualBlockPreLN(d_hidden) for _ in range(num_blocks)]) + # NO out_ln + self.out_head = nn.Linear(d_hidden, num_classes) + self.num_blocks = num_blocks + self.d_hidden = d_hidden + + def forward(self, x, return_hidden=False): + h = self.embed(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) # NO out_ln + if return_hidden: + return logits, hiddens + return logits + + +def get_cifar10(batch_size=128): + tv_train = transforms.Compose([ + transforms.RandomCrop(32, padding=4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), + ]) + tv = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), + ]) + tr = torchvision.datasets.CIFAR10('./data', True, download=True, transform=tv_train) + te = torchvision.datasets.CIFAR10('./data', False, download=True, transform=tv) + return (DataLoader(tr, batch_size=batch_size, shuffle=True, num_workers=2), + DataLoader(te, batch_size=batch_size, shuffle=False, num_workers=2)) + + +def fixed_eval_buffer(test_loader, device, n_samples=1024): + xs, ys = [], [] + for x, y in test_loader: + xs.append(x.view(x.size(0), -1)); ys.append(y) + if sum(xb.size(0) for xb in xs) >= n_samples: + break + return torch.cat(xs)[:n_samples].to(device), torch.cat(ys)[:n_samples].to(device) + + +def diagnose(model, x_eval, y_eval, dfa_Bs=None): + was_training = model.training + model.eval() + L = model.num_blocks + with torch.no_grad(): + _, hi = model(x_eval, return_hidden=True) + hidden_norms = [h.norm(dim=-1).median().item() for h in hi] + + h0 = model.embed(x_eval.detach()) + hs = [h0.clone().requires_grad_(True)] + for b in model.blocks: + hs.append(hs[-1] + b(hs[-1])) + logits = model.out_head(hs[-1]) # NO out_ln + loss = F.cross_entropy(logits, y_eval) + grads = torch.autograd.grad(loss, hs) + bp_l2 = [g.norm(dim=-1).median().item() for g in grads] + bp_full = [g.detach() for g in grads] + acc = (logits.argmax(-1) == y_eval).float().mean().item() + loss_val = loss.item() + + gamma_dfa = float('nan'); per_layer_gamma = [] + if dfa_Bs is not None: + with torch.no_grad(): + e_T = logits.softmax(-1); e_T[torch.arange(x_eval.size(0)), y_eval] -= 1 + for l in range(L): + a_dfa = (e_T @ dfa_Bs[l].T).detach() + per_layer_gamma.append(cosine_similarity_batch(a_dfa, bp_full[l])) + gamma_dfa = float(np.mean(per_layer_gamma)) + + if was_training: model.train() + return { + 'hidden_norms': hidden_norms, + 'bp_grad_per_sample_l2_med': bp_l2, + 'gamma_dfa': gamma_dfa, + 'gamma_dfa_per_layer': per_layer_gamma, + 'acc_eval': acc, 'loss_eval': loss_val, + } + + +def train_bp(model, train_loader, x_eval, y_eval, device, epochs, lr, wd): + opt = optim.AdamW(model.parameters(), lr=lr, weight_decay=wd) + sch = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs) + log = [] + d0 = diagnose(model, x_eval, y_eval); d0['epoch'] = 0; log.append(d0) + print(f" [BP-noLN] Ep 0: ||h_L||={d0['hidden_norms'][-1]:.3e} ||g||={d0['bp_grad_per_sample_l2_med'][2]:.3e} acc={d0['acc_eval']:.4f}", flush=True) + for ep in range(1, epochs + 1): + model.train() + for x, y in train_loader: + x = x.view(x.size(0), -1).to(device); y = y.to(device) + logits = model(x); loss = F.cross_entropy(logits, y) + opt.zero_grad(); loss.backward(); opt.step() + sch.step() + d = diagnose(model, x_eval, y_eval); d['epoch'] = ep; log.append(d) + if ep % 5 == 0 or ep == 1 or ep == epochs: + print(f" [BP-noLN] Ep {ep}: ||h_L||={d['hidden_norms'][-1]:.3e} ||g||={d['bp_grad_per_sample_l2_med'][2]:.3e} acc={d['acc_eval']:.4f}", flush=True) + return log + + +def train_dfa(model, train_loader, x_eval, y_eval, device, epochs, lr, wd): + d_hidden = model.d_hidden; L = model.num_blocks; C = 10 + Bs = [torch.randn(d_hidden, C, device=device) / np.sqrt(C) for _ in range(L)] + block_opts = [optim.AdamW(b.parameters(), lr=lr, weight_decay=wd) for b in model.blocks] + embed_opt = optim.AdamW(model.embed.parameters(), lr=lr, weight_decay=wd) + head_opt = optim.AdamW(model.out_head.parameters(), lr=lr, weight_decay=wd) + all_sch = [optim.lr_scheduler.CosineAnnealingLR(o, T_max=epochs) for o in block_opts] + \ + [optim.lr_scheduler.CosineAnnealingLR(embed_opt, T_max=epochs), + optim.lr_scheduler.CosineAnnealingLR(head_opt, T_max=epochs)] + log = [] + d0 = diagnose(model, x_eval, y_eval, dfa_Bs=Bs); d0['epoch'] = 0; log.append(d0) + print(f" [DFA-noLN] Ep 0: ||h_L||={d0['hidden_norms'][-1]:.3e} ||g||={d0['bp_grad_per_sample_l2_med'][2]:.3e} acc={d0['acc_eval']:.4f}", flush=True) + for ep in range(1, epochs + 1): + model.train() + for x, y in train_loader: + x = x.view(x.size(0), -1).to(device); y = y.to(device) + batch = x.size(0) + with torch.no_grad(): + logits, hiddens = model(x, return_hidden=True) + e_T = logits.softmax(-1); e_T[torch.arange(batch), y] -= 1 + hL_det = hiddens[-1].detach() + # Head update — NO out_ln + 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() + # Block updates + for l in range(L): + h_l = hiddens[l].detach() + a_dfa = (e_T @ Bs[l].T).detach() + rms = (a_dfa ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6 + a_norm = a_dfa / rms + f_l = model.blocks[l](h_l) + local_loss = (f_l * a_norm).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() + # Embed update + a_0 = (e_T @ Bs[0].T).detach() + rms_0 = (a_0 ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6 + h0 = model.embed(x) + embed_loss = (h0 * (a_0 / rms_0)).sum(dim=-1).mean() + embed_opt.zero_grad(); embed_loss.backward(); embed_opt.step() + for s in all_sch: s.step() + d = diagnose(model, x_eval, y_eval, dfa_Bs=Bs); d['epoch'] = ep; log.append(d) + if ep % 5 == 0 or ep == 1 or ep == epochs: + print(f" [DFA-noLN] Ep {ep}: ||h_L||={d['hidden_norms'][-1]:.3e} ||g||={d['bp_grad_per_sample_l2_med'][2]:.3e} acc={d['acc_eval']:.4f} γ={d['gamma_dfa']:.4f}", flush=True) + return log + + +def main(): + p = argparse.ArgumentParser() + p.add_argument('--output_dir', type=str, default='results/snapshot_no_outln_v1') + p.add_argument('--epochs', type=int, default=100) + p.add_argument('--lr', type=float, default=1e-3) + p.add_argument('--wd', type=float, default=0.01) + p.add_argument('--seed', type=int, default=42) + p.add_argument('--depth', type=int, default=4) + p.add_argument('--d_hidden', type=int, default=256) + args = p.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + device = torch.device('cuda:0') + print(f"NO-OUT_LN VARIANT: depth={args.depth}, d_hidden={args.d_hidden}, " + f"epochs={args.epochs}, seed={args.seed}", flush=True) + + train_loader, test_loader = get_cifar10(batch_size=128) + x_eval, y_eval = fixed_eval_buffer(test_loader, device, n_samples=1024) + + L, d, C = args.depth, args.d_hidden, 10 + + print("\n=== BP training (NO out_ln) ===", flush=True) + torch.manual_seed(args.seed); np.random.seed(args.seed); torch.cuda.manual_seed_all(args.seed) + bp_model = ResidualMLP_NoOutLN(3072, d, C, L).to(device) + bp_log = train_bp(bp_model, train_loader, x_eval, y_eval, device, args.epochs, args.lr, args.wd) + + print("\n=== DFA training (NO out_ln) ===", flush=True) + torch.manual_seed(args.seed); np.random.seed(args.seed); torch.cuda.manual_seed_all(args.seed) + dfa_model = ResidualMLP_NoOutLN(3072, d, C, L).to(device) + dfa_log = train_dfa(dfa_model, train_loader, x_eval, y_eval, device, args.epochs, args.lr, args.wd) + + out = { + 'config': vars(args), 'depth': L, 'd_hidden': d, 'num_classes': C, + 'architecture': 'ResidualMLP_NoOutLN', + 'bp_log': bp_log, 'dfa_log': dfa_log, + } + out_path = os.path.join(args.output_dir, f'snapshot_noLN_s{args.seed}.json') + with open(out_path, 'w') as f: + json.dump(out, f, indent=2) + print(f"\nSaved {out_path}", flush=True) + + +if __name__ == '__main__': + main() |