""" scripts/run_compression.py -------------------------- Phase 3: Signal verification + compression comparison. Produces all quantitative results from CGC v1 paper Section 5: Signal verification: - Baseline PPL on WikiText-2 test - Wanda importance vs capability density (orthogonality) - Per-head ablation impact vs capability density (correlation) Compression comparison at target_retention: - CGC-L (density-proportional allocation) - Uniform (same budget for all heads) - Inverted (density-inverted — sanity check, should be worst) Outputs (all written to output_dir): signal_check.json -- orthogonality + ablation correlation compression_summary.json -- PPL comparison table wanda_importance.npy -- (N_LAYERS, N_HEADS) Wanda scores ablation_results.npy -- (N_LAYERS, N_HEADS) ΔPPL per head Usage: python scripts/run_compression.py \\ --density_map results/density_map.npz \\ --output_dir results/ \\ [--retention 0.5] [--skip_ablation] Runtime: ~25 min for ablation (384 heads × 32-sequence PPL each). Use --skip_ablation to skip if only PPL comparison is needed. """ import argparse import json import os import random import sys import numpy as np import torch from scipy.stats import pearsonr, spearmanr from tqdm import tqdm sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) # ── Reproducibility ─────────────────────────────────────────────────────────── SEED = 42 random.seed(SEED) np.random.seed(SEED) torch.manual_seed(SEED) torch.cuda.manual_seed_all(SEED) from transformers import GPT2LMHeadModel, GPT2Tokenizer from datasets import load_dataset from cgc.density import ( CapabilityDensityMap, collect_head_activations, N_LAYERS, N_HEADS, HEAD_DIM, ) from cgc.compress import ( compute_wanda_importance_paper, compute_wanda_importance_real, cgc_allocate_budget, uniform_allocate_budget, apply_compression, ) from cgc.evaluate import load_eval_data, compute_perplexity def parse_args(): p = argparse.ArgumentParser(description="CGC v1 compression experiment") p.add_argument("--density_map", type=str, required=True) p.add_argument("--output_dir", type=str, default="results/") p.add_argument("--retention", type=float, default=0.5, help="Global attention weight retention ratio (default: 0.5)") p.add_argument("--n_cal_sequences", type=int, default=128, help="Calibration sequences for Wanda (default: 128)") p.add_argument("--n_eval_sequences",type=int, default=32, help="Evaluation sequences for PPL (default: 32)") p.add_argument("--seq_len", type=int, default=512) p.add_argument("--batch_size", type=int, default=8) p.add_argument("--skip_ablation", action="store_true", help="Skip per-head ablation experiment (~25 min)") p.add_argument("--wanda_mode", type=str, default="paper", choices=["paper", "real"], help="Wanda implementation to use: 'paper' (per-head approximation " "intended to reproduce original experiment conditions) or " "'real' (correct per-weight criterion, Sun et al. 2024). " "Default: paper") return p.parse_args() def load_calibration(tokenizer, n_sequences, seq_len, device): dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split="train") full_text = " ".join([ item["text"].strip() for item in dataset if len(item["text"].strip()) > 50 ]) all_tokens = tokenizer.encode(full_text) chunks = [] for i in range(n_sequences): start = i * seq_len end = start + seq_len if end > len(all_tokens): break chunks.append(torch.tensor(all_tokens[start:end], dtype=torch.long)) return torch.stack(chunks).to(device) def run_head_ablation( model, eval_data: torch.Tensor, baseline_ppl: float, batch_size: int = 1, ) -> np.ndarray: """ Per-head ablation: zero each head's output projection slice, measure ΔPPL, restore weights. This tests individual head importance without cumulative effects. Args: model: GPT2LMHeadModel eval_data: (N_seqs, seq_len) evaluation tokens baseline_ppl: PPL of unmodified model batch_size: sequences per PPL forward pass Returns: ablation_impact: (N_LAYERS, N_HEADS) array of ΔPPL values """ ablation = np.zeros((N_LAYERS, N_HEADS)) n_embd = model.config.n_embd print(f"\nPer-head ablation ({N_LAYERS} layers × {N_HEADS} heads = " f"{N_LAYERS * N_HEADS} runs)...") print(f"Baseline PPL: {baseline_ppl:.4f}\n") for l in tqdm(range(N_LAYERS), desc="Layers"): W = model.transformer.h[l].attn.c_proj.weight.data for h in range(N_HEADS): start = h * HEAD_DIM end = (h + 1) * HEAD_DIM # Zero head h's output projection slice orig = W[:, start:end].clone() W[:, start:end] = 0.0 ablated_ppl = compute_perplexity(model, eval_data, batch_size) ablation[l][h] = ablated_ppl - baseline_ppl # Restore W[:, start:end] = orig print(f"\nAblation complete.") print(f" Most critical: Layer {ablation.max(1).argmax()}, " f"Head {ablation.argmax(1)[ablation.max(1).argmax()]} " f"(ΔPPL={ablation.max():.4f})") print(f" Least critical: ΔPPL={ablation.min():.4f}") print(f" Mean ΔPPL: {ablation.mean():.4f}") return ablation def main(): args = parse_args() os.makedirs(args.output_dir, exist_ok=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}\n") # Load density map density_map = CapabilityDensityMap.load(args.density_map) print(density_map.summary() + "\n") # Load tokenizer and model print("Loading GPT-2 Medium (LM head)...") tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium") tokenizer.pad_token = tokenizer.eos_token model = GPT2LMHeadModel.from_pretrained("gpt2-medium").to(device) model.eval() # Evaluation data (WikiText-2 test, 32 × 512) eval_data = load_eval_data( tokenizer, args.n_eval_sequences, args.seq_len, device ) # Calibration data (WikiText-103 train, 128 × 512) — for Wanda cal_data = load_calibration( tokenizer, args.n_cal_sequences, args.seq_len, device ) # ── Baseline PPL ────────────────────────────────────────────────────────── print("Computing baseline PPL...") baseline_ppl = compute_perplexity(model, eval_data) print(f"Baseline PPL: {baseline_ppl:.4f}") # ── Wanda importance ────────────────────────────────────────────────────── if args.wanda_mode == "paper": print("\nComputing Wanda importance (paper mode: per-head approximation)...") # Collect head activations needed for the paper-mode approximation from cgc.density import collect_head_activations base_model_for_wanda = __import__("transformers").GPT2Model.from_pretrained( "gpt2-medium" ).to(device) head_acts_for_wanda = collect_head_activations( base_model_for_wanda, cal_data, device, args.batch_size ) del base_model_for_wanda wanda_imp = compute_wanda_importance_paper(model, head_acts_for_wanda) else: print("\nComputing Wanda importance (real mode: per-weight criterion)...") wanda_imp = compute_wanda_importance_real(model, cal_data, device, args.batch_size) np.save(os.path.join(args.output_dir, "wanda_importance.npy"), wanda_imp) density_flat = density_map.density.flatten() wanda_flat = wanda_imp.flatten() pearson_r, pearson_p = pearsonr(density_flat, wanda_flat) spearman_rho, spear_p = spearmanr(density_flat, wanda_flat) print(f"\nOrthogonality — density vs Wanda importance (n={len(density_flat)} heads):") print(f" Pearson r = {pearson_r:.4f} (p = {pearson_p:.2e})") print(f" Spearman ρ = {spearman_rho:.4f} (p = {spear_p:.2e})") # ── Per-head ablation ───────────────────────────────────────────────────── ablation_impact = None ab_pr = ab_pp = ab_sr = ab_sp = None if not args.skip_ablation: ablation_impact = run_head_ablation( model, eval_data, baseline_ppl, batch_size=1 ) np.save( os.path.join(args.output_dir, "ablation_results.npy"), ablation_impact, ) # Rescale density to [0,1] for correlation d_min = density_flat.min() d_max = density_flat.max() d_resc = (density_flat - d_min) / (d_max - d_min + 1e-8) ab_flat = ablation_impact.flatten() ab_pr, ab_pp = pearsonr(d_resc, ab_flat) ab_sr, ab_sp = spearmanr(d_resc, ab_flat) print(f"\nDensity vs Ablation Impact (n={len(density_flat)} heads):") print(f" Pearson r = {ab_pr:.4f} (p = {ab_pp:.2e})") print(f" Spearman ρ = {ab_sr:.4f} (p = {ab_sp:.2e})") # Save signal summary signal = { "baseline_ppl": baseline_ppl, "n_heads": int(len(density_flat)), "orthogonality_density_vs_wanda": { "pearson_r": float(pearson_r), "pearson_p": float(pearson_p), "spearman_rho": float(spearman_rho), "spearman_p": float(spear_p), }, "ablation_correlation_density_vs_delta_ppl": { "pearson_r": float(ab_pr) if ab_pr is not None else None, "pearson_p": float(ab_pp) if ab_pp is not None else None, "spearman_rho": float(ab_sr) if ab_sr is not None else None, "spearman_p": float(ab_sp) if ab_sp is not None else None, } if ablation_impact is not None else None, } with open(os.path.join(args.output_dir, "signal_check.json"), "w") as f: json.dump(signal, f, indent=2) print(f"\nSignal check saved: {os.path.join(args.output_dir, 'signal_check.json')}") # ── Budget allocation ───────────────────────────────────────────────────── retention = args.retention cgc_ratios = cgc_allocate_budget(density_map, retention) uniform_ratios = uniform_allocate_budget(retention) # Inverted: mirror of CGC around uniform (sanity check — predicts worst PPL) inverted_ratios = 2 * uniform_ratios - cgc_ratios inverted_ratios = np.clip(inverted_ratios, 0.05, 0.95) inverted_ratios = inverted_ratios * (retention / inverted_ratios.mean()) inverted_ratios = np.clip(inverted_ratios, 0.05, 0.95) print(f"\nBudget allocation at {retention:.0%} global retention:") print(f" CGC-L: mean={cgc_ratios.mean():.4f} " f"min={cgc_ratios.min():.4f} max={cgc_ratios.max():.4f}") print(f" Uniform: mean={uniform_ratios.mean():.4f}") print(f" Inverted: mean={inverted_ratios.mean():.4f} " f"min={inverted_ratios.min():.4f} max={inverted_ratios.max():.4f}") # ── Compression + PPL ───────────────────────────────────────────────────── results = { "baseline_ppl": baseline_ppl, "retention_target": retention, } for name, ratios in [ ("cgc", cgc_ratios), ("uniform", uniform_ratios), ("inverted", inverted_ratios), ]: print(f"\nCompressing: {name.upper()}...") compressed = apply_compression(model, ratios, inplace=False) ppl = compute_perplexity(compressed, eval_data) delta_ppl = ppl - baseline_ppl results[name] = {"ppl": float(ppl), "delta_ppl": float(delta_ppl)} print(f" PPL = {ppl:.4f} (ΔPPL = {delta_ppl:+.4f})") del compressed torch.cuda.empty_cache() # ── Final summary ───────────────────────────────────────────────────────── print("\n" + "=" * 55) print("RESULTS SUMMARY") print("=" * 55) print(f"{'Method':<16} {'PPL':>8} {'ΔPPL':>8}") print("-" * 34) print(f"{'Dense':16} {baseline_ppl:8.4f} {'—':>8}") for name in ["uniform", "cgc", "inverted"]: ppl = results[name]["ppl"] delta = results[name]["delta_ppl"] print(f"{name:<16} {ppl:8.4f} {delta:>+8.4f}") summary_path = os.path.join(args.output_dir, "compression_summary.json") with open(summary_path, "w") as f: json.dump(results, f, indent=2) print(f"\nFull results saved: {summary_path}") print("Next: python scripts/plot_results.py " f"--density_map {args.density_map} " f"--results_dir {args.output_dir}") if __name__ == "__main__": main()