""" 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()