path: root/experiments/bp_support_sparsity.py

"""
BP Support Sparsity Analysis.
A1: threshold sweep, log-gradient stats, active-subset Gamma, energy-weighted Gamma
A2: same for CIFAR
All from checkpoints — no retraining.
"""
import os, sys, csv, json, argparse, numpy as np, torch, torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from models.residual_mlp import ResidualMLP

class StudentBlock(torch.nn.Module):
    def __init__(self, d, alpha=1.0):
        super().__init__()
        self.ln=torch.nn.LayerNorm(d);self.w=torch.nn.Linear(d,d,bias=False)
        torch.nn.init.normal_(self.w.weight,std=0.01);self.alpha=alpha
    def forward(self, h):
        return self.w(((1-self.alpha)*self.ln(h)+self.alpha*torch.tanh(self.ln(h))))

class StudentNet(torch.nn.Module):
    def __init__(self, d, C, L, alpha=1.0):
        super().__init__()
        self.blocks=torch.nn.ModuleList([StudentBlock(d,alpha) for _ in range(L)])
        self.out_head=torch.nn.Linear(d,C);self.num_blocks=L;self.d_hidden=d
    def forward(self, x, return_hidden=False):
        h=x;hi=[h] if return_hidden else None
        for b in self.blocks:
            h=h+b(h)
            if return_hidden:hi.append(h)
        lo=self.out_head(h)
        return (lo,hi) if return_hidden else lo

class TeacherNet(torch.nn.Module):
    def __init__(self, d, C, L, alpha=1.0, seed=0):
        super().__init__()
        self.alpha=alpha;rng=torch.Generator().manual_seed(seed)
        self.Ws=torch.nn.ParameterList()
        for _ in range(L):
            W=torch.randn(d,d,generator=rng)*0.3/(d**0.5)
            U,S,Vh=torch.linalg.svd(W,full_matrices=False)
            self.Ws.append(torch.nn.Parameter(U@torch.diag(S.clamp(max=0.3))@Vh,requires_grad=False))
        self.U=torch.nn.Parameter(torch.randn(C,d,generator=rng)/(d**0.5),requires_grad=False)
    def forward(self, x):
        h=x
        for W in self.Ws:h=h+((1-self.alpha)*h+self.alpha*torch.tanh(h))@W.T
        return h@self.U.T


def get_bp_grads(model, x, y, device, is_cifar=False):
    """Get per-layer BP gradients via manual forward."""
    model.eval()
    L = model.num_blocks
    if is_cifar:
        h0 = model.embed(x.detach())
    else:
        h0 = x.detach()
    h_list = [h0.clone().requires_grad_(True)]
    for b in model.blocks:
        h_list.append(h_list[-1] + b(h_list[-1]))
    if is_cifar:
        lo = model.out_head(model.out_ln(h_list[-1]))
    else:
        lo = model.out_head(h_list[-1])
    loss = F.cross_entropy(lo, y)
    grads = torch.autograd.grad(loss, h_list)
    return {l: grads[l].detach() for l in range(L)}


def analyze_model(model, x, y, device, is_cifar=False):
    """Full sparsity analysis for one model."""
    L = model.num_blocks
    bp = get_bp_grads(model, x, y, device, is_cifar)

    thresholds = [1e-8, 1e-7, 1e-6, 1e-5, 1e-4]
    batch = x.size(0)

    results = {'thresholds': {}, 'log_grad_norms': [], 'per_layer': []}

    # Per-layer analysis
    all_norms = []
    for l in range(L):
        g = bp[l]
        norms = g.norm(dim=-1)  # (batch,)
        all_norms.append(norms)
        log_norms = torch.log10(norms.clamp(min=1e-20)).cpu().numpy()
        results['log_grad_norms'].append(log_norms.tolist())

        layer_res = {'layer': l}
        for tau in thresholds:
            s = (norms > tau).float().mean().item()
            layer_res[f's_tau_{tau}'] = s
        results['per_layer'].append(layer_res)

    # Threshold sweep (averaged over layers)
    for tau in thresholds:
        mean_s = np.mean([res[f's_tau_{tau}'] for res in results['per_layer']])
        results['thresholds'][str(tau)] = mean_s

    # Active-subset Gamma and energy-weighted Gamma for each threshold
    # (self-cosine for now — comparing BP with BP; real cross-method needs credit source)
    # We store raw norms for post-processing
    results['mean_grad_norm'] = np.mean([n.mean().item() for n in all_norms])
    results['median_grad_norm'] = np.mean([n.median().item() for n in all_norms])
    results['grad_norm_percentiles'] = {}
    stacked = torch.cat(all_norms)
    for p in [10, 25, 50, 75, 90, 95, 99]:
        results['grad_norm_percentiles'][str(p)] = np.percentile(stacked.cpu().numpy(), p)

    return results


def run_analysis(args):
    device = torch.device(f'cuda:{args.gpu}')
    os.makedirs(args.output_dir, exist_ok=True)
    seeds = [42, 123, 456, 789, 1024, 2048, 3000, 4000, 5000, 6000]
    thresholds = [1e-8, 1e-7, 1e-6, 1e-5, 1e-4]

    # ===== A1 Synthetic =====
    print("=== A1 Synthetic ===", flush=True)
    alphas = [0.0, 0.5, 1.0]; depths = [4, 8]; d = 128; C = 10
    methods = ['bp', 'dfa', 'state_bridge', 'credit_bridge']

    a1_threshold_rows = []
    a1_histogram_data = {}

    for alpha in alphas:
        for L in depths:
            # Use seed=42 for histogram (representative)
            teacher = TeacherNet(d, C, L, alpha, seed=0).to(device)
            torch.manual_seed(42 + 10000)
            X_test = torch.randn(512, d, device=device)
            with torch.no_grad():
                Y_test = teacher(X_test).argmax(-1)

            for method in methods:
                # Aggregate over seeds for threshold table
                seed_results = []
                for seed in seeds:
                    ckpt = f'results/confirmatory/checkpoints_A1/a{alpha}_L{L}_{method}_s{seed}.pt'
                    if not os.path.exists(ckpt):
                        continue
                    torch.manual_seed(seed)
                    model = StudentNet(d, C, L, alpha).to(device)
                    model.load_state_dict(torch.load(ckpt, map_location=device))
                    res = analyze_model(model, X_test, Y_test, device, is_cifar=False)
                    seed_results.append(res)

                    # Threshold rows
                    for tau in thresholds:
                        a1_threshold_rows.append({
                            'alpha': alpha, 'depth': L, 'method': method, 'seed': seed,
                            'threshold': tau, 'support_fraction': res['thresholds'][str(tau)]
                        })

                # Histogram data for seed=42 only
                if seed_results:
                    key = f"a{alpha}_L{L}_{method}"
                    a1_histogram_data[key] = {
                        'log_grad_norms': seed_results[0]['log_grad_norms'],
                        'percentiles': seed_results[0]['grad_norm_percentiles'],
                        'mean_norm': seed_results[0]['mean_grad_norm'],
                        'median_norm': seed_results[0]['median_grad_norm'],
                    }

                if seed_results:
                    mean_s = np.mean([r['thresholds']['1e-06'] for r in seed_results])
                    print(f"  a={alpha} L={L} {method}: s(1e-6)={mean_s:.4f}, "
                          f"mean_norm={np.mean([r['mean_grad_norm'] for r in seed_results]):.2e}", flush=True)

    # Save A1 threshold CSV
    out1 = os.path.join(args.output_dir, 'A1_threshold_sweep.csv')
    with open(out1, 'w', newline='') as f:
        w = csv.DictWriter(f, fieldnames=['alpha','depth','method','seed','threshold','support_fraction'])
        w.writeheader(); w.writerows(a1_threshold_rows)
    print(f"A1 threshold: {len(a1_threshold_rows)} rows -> {out1}", flush=True)

    # Save A1 histogram JSON
    out1h = os.path.join(args.output_dir, 'A1_histogram_data.json')
    def to_serializable(obj):
        if isinstance(obj, (np.floating, np.integer)): return float(obj)
        if isinstance(obj, np.ndarray): return obj.tolist()
        return obj
    with open(out1h, 'w') as f:
        json.dump(a1_histogram_data, f, indent=2, default=to_serializable)
    print(f"A1 histogram data -> {out1h}", flush=True)

    # ===== A2 CIFAR =====
    print("\n=== A2 CIFAR ===", flush=True)
    import torchvision, torchvision.transforms as transforms
    tv = transforms.Compose([transforms.ToTensor(),
                              transforms.Normalize((0.4914,0.4822,0.4465),(0.2470,0.2435,0.2616))])
    tel = DataLoader(torchvision.datasets.CIFAR10('./data', False, download=True, transform=tv),
                     256, False, num_workers=4)
    for x, y in tel:
        x = x.view(x.size(0), -1).to(device); y = y.to(device); break

    L_c, d_c = 4, 256
    a2_threshold_rows = []
    a2_histogram_data = {}

    for method in methods:
        seed_results = []
        for seed in seeds:
            ckpt = f'results/confirmatory/checkpoints_A2/{method}_s{seed}.pt'
            if not os.path.exists(ckpt):
                continue
            torch.manual_seed(seed)
            model = ResidualMLP(3072, d_c, 10, L_c).to(device)
            model.load_state_dict(torch.load(ckpt, map_location=device))
            res = analyze_model(model, x, y, device, is_cifar=True)
            seed_results.append(res)

            for tau in thresholds:
                a2_threshold_rows.append({
                    'method': method, 'seed': seed,
                    'threshold': tau, 'support_fraction': res['thresholds'][str(tau)]
                })

        if seed_results:
            key = method
            a2_histogram_data[key] = {
                'log_grad_norms': seed_results[0]['log_grad_norms'],
                'percentiles': seed_results[0]['grad_norm_percentiles'],
                'mean_norm': seed_results[0]['mean_grad_norm'],
                'median_norm': seed_results[0]['median_grad_norm'],
            }
            mean_s = np.mean([r['thresholds']['1e-06'] for r in seed_results])
            print(f"  {method}: s(1e-6)={mean_s:.4f}, "
                  f"mean_norm={np.mean([r['mean_grad_norm'] for r in seed_results]):.2e}, "
                  f"median_norm={np.mean([r['median_grad_norm'] for r in seed_results]):.2e}", flush=True)

    out2 = os.path.join(args.output_dir, 'A2_threshold_sweep.csv')
    with open(out2, 'w', newline='') as f:
        w = csv.DictWriter(f, fieldnames=['method','seed','threshold','support_fraction'])
        w.writeheader(); w.writerows(a2_threshold_rows)
    print(f"A2 threshold: {len(a2_threshold_rows)} rows -> {out2}", flush=True)

    out2h = os.path.join(args.output_dir, 'A2_histogram_data.json')
    with open(out2h, 'w') as f:
        json.dump(a2_histogram_data, f, indent=2, default=to_serializable)
    print(f"A2 histogram data -> {out2h}", flush=True)

    # ===== Summary =====
    print(f"\n{'='*70}", flush=True)
    print("SUMMARY: Support fraction s(τ) at τ=1e-6 (mean over 10 seeds)", flush=True)
    print(f"{'='*70}", flush=True)

    print("\nA1 Synthetic:")
    for alpha in alphas:
        for L in depths:
            print(f"  alpha={alpha}, L={L}:")
            for method in methods:
                vals = [r['support_fraction'] for r in a1_threshold_rows
                        if r['alpha']==alpha and r['depth']==L and r['method']==method and r['threshold']==1e-6]
                if vals:
                    print(f"    {method}: s(1e-6) = {np.mean(vals):.4f} ± {np.std(vals):.4f}")

    print("\nA2 CIFAR:")
    for method in methods:
        vals = [r['support_fraction'] for r in a2_threshold_rows
                if r['method']==method and r['threshold']==1e-6]
        if vals:
            print(f"  {method}: s(1e-6) = {np.mean(vals):.4f} ± {np.std(vals):.4f}")

    print("\nDONE", flush=True)


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--gpu', type=int, default=0)
    p.add_argument('--output_dir', type=str, default='results/confirmatory')
    args = p.parse_args()
    run_analysis(args)

if __name__ == '__main__':
    main()