path: root/experiments/snapshot_evolution_no_outln.py

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