path: root/experiments/snapshot_fa_crossarch.py

"""
FA-only snapshot evolution for ViT-Mini and ResMLP-no-outLN.
Produces per-epoch ||h_L||, ||g_L||, acc for FA training.
"""
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, torchvision.transforms as transforms

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from models.residual_mlp import ResidualMLP
from models.vit_mini import ViTMini


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(loader, device, n=1024):
    xs, ys = [], []
    for x, y in loader:
        xs.append(x); ys.append(y)
        if sum(xb.size(0) for xb in xs) >= n:
            break
    return torch.cat(xs)[:n].to(device), torch.cat(ys)[:n].to(device)


# ─── Diagnose (works for both ViT and ResMLP) ───────────────────────────

def diagnose_resmlp(model, x_eval, y_eval):
    model.eval()
    x_flat = x_eval.view(x_eval.size(0), -1)
    with torch.no_grad():
        _, hiddens = model(x_flat, return_hidden=True)
    hidden_norms = [h.norm(dim=-1).median().item() for h in hiddens]
    # BP grads
    h0 = model.embed(x_flat.detach())
    hs = [h0.clone().requires_grad_(True)]
    for b in model.blocks:
        hs.append(hs[-1] + b(hs[-1]))
    # Handle both with and without out_ln
    if hasattr(model, 'out_ln'):
        logits = model.out_head(model.out_ln(hs[-1]))
    else:
        logits = model.out_head(hs[-1])
    loss = F.cross_entropy(logits, y_eval)
    grads = torch.autograd.grad(loss, hs)
    g_norms = [g.norm(dim=-1).median().item() for g in grads]
    acc = (logits.argmax(-1) == y_eval).float().mean().item()
    model.train()
    return {'hidden_norms': hidden_norms, 'bp_grad_norms_per_sample_med': g_norms, 'acc_eval': acc}


def diagnose_vit(model, x_eval, y_eval):
    model.eval()
    with torch.no_grad():
        _, hiddens = model(x_eval, return_hidden=True)
    h_cls_norms = [h[:, 0].norm(dim=-1).median().item() for h in hiddens]
    # BP grads via manual forward
    h0 = model.embed(x_eval.detach())
    hs = [h0.clone().requires_grad_(True)]
    for b in model.blocks:
        hs.append(hs[-1] + b(hs[-1]))
    h_cls = model.out_ln(hs[-1][:, 0])
    logits = model.out_head(h_cls)
    loss = F.cross_entropy(logits, y_eval)
    grads = torch.autograd.grad(loss, hs)
    g_cls_norms = [g[:, 0].norm(dim=-1).median().item() for g in grads]
    acc = (logits.argmax(-1) == y_eval).float().mean().item()
    model.train()
    return {'hidden_norms_cls': h_cls_norms, 'bp_grad_per_sample_l2_med': g_cls_norms, 'acc_eval': acc}


# ─── FA training ─────────────────────────────────────────────────────────

def train_fa_resmlp(model, train_loader, x_eval, y_eval, device, epochs, lr, wd, no_outln=False):
    d_hidden = model.d_hidden
    L = model.num_blocks
    Bs = [torch.randn(d_hidden, d_hidden, device=device) / np.sqrt(d_hidden) 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_params = list(model.out_head.parameters())
    if hasattr(model, 'out_ln') and model.out_ln is not None:
        head_params += list(model.out_ln.parameters())
    head_opt = optim.AdamW(head_params, 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_resmlp(model, x_eval, y_eval); d0['epoch'] = 0; log.append(d0)
    print(f"  [FA] Ep 0: 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)
            with torch.no_grad():
                logits, hiddens = model(x, return_hidden=True)
            hL_det = hiddens[-1].detach()
            logits_out = model.out_head(model.out_ln(hL_det)) if hasattr(model, 'out_ln') else model.out_head(hL_det)
            loss_out = F.cross_entropy(logits_out, y)
            head_opt.zero_grad(); loss_out.backward(); head_opt.step()
            # FA credits
            hL_req = hiddens[-1].detach().requires_grad_(True)
            logits_fa = model.out_head(model.out_ln(hL_req)) if hasattr(model, 'out_ln') else model.out_head(hL_req)
            loss_fa = F.cross_entropy(logits_fa, y, reduction='sum')
            a_L = torch.autograd.grad(loss_fa, hL_req)[0].detach()
            credits = [None] * L
            credits[L-1] = a_L
            for ll in range(L-2, -1, -1):
                credits[ll] = (credits[ll+1] @ Bs[ll+1]).detach()
            for l in range(L):
                h_l = hiddens[l].detach()
                a_l = credits[l]
                rms = (a_l**2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
                f_l = model.blocks[l](h_l)
                local_loss = (f_l * (a_l / 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()
            a_0 = credits[0]
            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_resmlp(model, x_eval, y_eval); d['epoch'] = ep; log.append(d)
        if ep % 10 == 0 or ep == 1 or ep == epochs:
            print(f"  [FA] Ep {ep}: ||h_L||={d['hidden_norms'][-1]:.3e} "
                  f"||g_L||={d['bp_grad_norms_per_sample_med'][-1]:.3e} "
                  f"acc={d['acc_eval']:.4f}", flush=True)
    return log


def train_fa_vit(model, train_loader, x_eval, y_eval, device, epochs, lr, wd):
    """Canonical FA for ViT: mean reduction, grad before step, no clipping, top-down."""
    d_model = model.d_hidden
    L = model.num_blocks
    Bs = [torch.randn(d_model, d_model, device=device) / np.sqrt(d_model) for _ in range(L)]
    block_opts = [optim.AdamW(b.parameters(), lr=lr, weight_decay=wd) for b in model.blocks]
    embed_opt = optim.AdamW(
        list(model.patch_embed.parameters()) + [model.cls_token, model.pos_embed],
        lr=lr, weight_decay=wd)
    head_opt = optim.AdamW(
        list(model.out_head.parameters()) + list(model.out_ln.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_vit(model, x_eval, y_eval); d0['epoch'] = 0; log.append(d0)
    print(f"  [FA-vit] Ep 0: acc={d0['acc_eval']:.4f}", flush=True)
    for ep in range(1, epochs + 1):
        model.train()
        for x, y in train_loader:
            x = x.to(device); y = y.to(device)
            with torch.no_grad():
                logits, hiddens = model(x, return_hidden=True)
            # Head update — grad BEFORE step (old head)
            hL_det = hiddens[-1].detach().requires_grad_(True)
            h_cls = model.out_ln(hL_det[:, 0])
            logits_out = model.out_head(h_cls)
            loss_out = F.cross_entropy(logits_out, y)  # mean reduction
            head_opt.zero_grad()
            loss_out.backward()
            a_L_full = hL_det.grad.detach()  # (B, n_tokens, d)
            head_opt.step()
            # Use mean over tokens for the backward signal
            a_credit = a_L_full.mean(dim=1)  # (B, d)
            # Top-down block updates, propagate credit after each
            for l in range(L - 1, -1, -1):
                h_l = hiddens[l].detach()
                a_broadcast = a_credit.unsqueeze(1).expand_as(h_l)
                rms = (a_broadcast ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
                f_l = model.blocks[l](h_l)
                local_loss = (f_l * (a_broadcast / rms)).sum(dim=-1).mean()
                block_opts[l].zero_grad()
                local_loss.backward()
                block_opts[l].step()  # no clipping
                a_credit = (a_credit @ Bs[l]).detach()
            # Embed update with final propagated credit
            a_0_broadcast = a_credit.unsqueeze(1)
            rms_0 = (a_credit ** 2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
            h0 = model.embed(x)
            embed_loss = (h0 * (a_0_broadcast / rms_0.unsqueeze(1))).sum(dim=-1).mean()
            embed_opt.zero_grad(); embed_loss.backward(); embed_opt.step()
        for s in all_sch: s.step()
        d = diagnose_vit(model, x_eval, y_eval); d['epoch'] = ep; log.append(d)
        if ep % 5 == 0 or ep == 1 or ep == epochs:
            print(f"  [FA-vit] Ep {ep}: ||h_L||={d['hidden_norms_cls'][-1]:.3e} "
                  f"||g_L||={d['bp_grad_per_sample_l2_med'][-1]:.3e} "
                  f"acc={d['acc_eval']:.4f}", flush=True)
    return log


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--arch', choices=['vit', 'resmlp_noln'], required=True)
    p.add_argument('--output', type=str, required=True)
    p.add_argument('--epochs', type=int, default=100)
    p.add_argument('--seed', type=int, default=42)
    args = p.parse_args()

    device = torch.device('cuda:0')
    train_loader, test_loader = get_cifar10(128)
    x_eval, y_eval = fixed_eval_buffer(test_loader, device, 1024)

    torch.manual_seed(args.seed); np.random.seed(args.seed); torch.cuda.manual_seed_all(args.seed)

    if args.arch == 'vit':
        # Match ViT snapshot params
        model = ViTMini(d_model=128, n_heads=4, num_blocks=4, num_classes=10).to(device)
        fa_log = train_fa_vit(model, train_loader, x_eval, y_eval, device,
                              args.epochs, lr=1e-3, wd=0.05)
    else:
        # ResMLP without terminal LN — use the same class as the original no-outln experiment
        from experiments.snapshot_evolution_no_outln import ResidualMLP_NoOutLN
        model = ResidualMLP_NoOutLN(3072, 256, 10, 4).to(device)
        fa_log = train_fa_resmlp(model, train_loader, x_eval, y_eval, device,
                                 args.epochs, lr=1e-3, wd=0.01, no_outln=True)

    with open(args.output, 'w') as f:
        json.dump({'fa_log': fa_log, 'arch': args.arch, 'seed': args.seed}, f, indent=2)
    print(f"Saved: {args.output}", flush=True)


if __name__ == '__main__':
    main()