scripts/compute_mp_statistics.py

#!/usr/bin/env python3
"""Compute Marchenko-Pastur statistics and distributions for all models.

Generates ground truth / null hypothesis reference data for eigenvalue statistics
of random W_QK matrices at various (d_model, d_head) scales.

Outputs:
    mp_statistics.npz - consolidated data file with:
        - Statistics distributions (max, PR, NPR, entropy, cond, stable_rank) for each model
        - Eigenvalue density curves (MP and exact product ensemble)
        - Analytical predictions (MP and exact)
        - Model metadata (d_model, d_head, n_heads, n_layers)
"""

import argparse
import logging
import os
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

MODEL_DIMS = {
    "gpt2":            (12, 12,  64),
    "gpt2-medium":     (24, 16,  64),
    "gpt2-large":      (36, 20,  64),
    "gpt2-xl":         (48, 25,  64),
    "pythia-70m-deduped":   (6,  8,  64),
    "pythia-160m-deduped":  (12, 12, 64),
    "pythia-410m-deduped":  (24, 16, 64),
    "pythia-1b-deduped":    (16, 8, 128),
    "pythia-1.4b-deduped":  (24, 16, 128),
    "pythia-2.8b-deduped":  (32, 32, 80),
    "pythia-6.9b-deduped":  (32, 32, 128),
    "pythia-12b-deduped":   (36, 40, 128),
    "llama-3.1-8b":    (32, 32, 128),
    "mistral-7b-v0.3": (32, 32, 128),
}


def generate_low_rank(d_n, d_m, mu=0, sigma=1):
    W_Q = np.random.normal(loc=mu, scale=sigma, size=d_n * d_m).reshape(d_n, d_m)
    W_K = np.random.normal(loc=mu, scale=sigma, size=d_n * d_m).reshape(d_m, d_n)
    W_QK = W_Q @ W_K / np.sqrt(d_n)
    _, S, _ = np.linalg.svd(W_QK)
    return S


def compute_eig_stat(eigenvalues, stat_type, d_model, d_head):
    lam = eigenvalues
    if stat_type == "max":
        return np.max(lam)
    elif stat_type == "condition_number":
        lam_nz = lam[:d_head]
        if len(lam_nz) > 0 and lam_nz[-1] > 0:
            return lam_nz[0] / lam_nz[-1]
        return 0
    elif stat_type == "participation_ratio":
        return np.sum(lam)**2 / np.sum(lam**2) if np.sum(lam**2) > 0 else 0
    elif stat_type == "normalized_participation_ratio":
        pr = np.sum(lam)**2 / np.sum(lam**2) if np.sum(lam**2) > 0 else 0
        return pr / d_head if d_head > 0 else 0
    elif stat_type == "spectral_entropy":
        s = np.sum(lam)
        if s > 0:
            p = lam / s
            p = p[p > 0]
            return -np.sum(p * np.log(p))
        return 0
    elif stat_type == "stable_rank":
        return np.sum(lam**2) / np.max(lam)**2 if np.max(lam) > 0 else 0
    return 0


def collect_eig_stats(d_model, d_head, n_samples=500, sigma=1.0):
    stat_names = [
        "max", "participation_ratio", "normalized_participation_ratio",
        "spectral_entropy", "condition_number", "stable_rank"
    ]
    stats = {k: np.zeros(n_samples) for k in stat_names}
    singular_values = []
    for i in range(n_samples):
        S = generate_low_rank(d_model, d_head, mu=0, sigma=sigma)
        singular_values.extend(S)
        eigenvalues = S**2
        for k in stat_names:
            stats[k][i] = compute_eig_stat(eigenvalues, k, d_model, d_head)
    return stats, np.array(singular_values)


def marchenko_pastur_density(lam, gamma, sigma_hat_sq, n_points=1000):
    lam_min = sigma_hat_sq * (1 - np.sqrt(gamma))**2
    lam_max = sigma_hat_sq * (1 + np.sqrt(gamma))**2
    lam_grid = np.linspace(lam_min + 1e-10, lam_max - 1e-10, n_points)
    rho = np.zeros_like(lam_grid)
    mask = (lam_grid >= lam_min) & (lam_grid <= lam_max)
    l = lam_grid[mask]
    rho[mask] = np.sqrt((lam_max - l) * (l - lam_min)) / (2 * np.pi * sigma_hat_sq * gamma * l)
    return lam_grid, rho


def product_ensemble_edges(d, d_h, sigma=1.0):
    gamma = d_h / d
    alpha = np.sqrt(1 + 8/gamma)
    z_plus  = (-1 + alpha) / 4
    z_minus = (-1 - alpha) / 4
    def inv_psi(z):
        return (1 + z) * (1 + gamma*z)**2 / z
    scale = d * sigma**4
    return scale * inv_psi(z_minus), scale * inv_psi(z_plus)


def pme_predictions(d, d_h, sigma=1.0):
    """Product Multiplicative Ensemble predictions using S-transform edges."""
    gamma = d_h / d
    s2 = d * sigma**4
    E_lam = s2
    E_lam2_exact = sigma**8 * (d**2 + 2*d*d_h + 2*d + d_h + 3)
    sum_lam = E_lam * d_h
    sum_lam2 = E_lam2_exact * d_h
    pr = sum_lam**2 / sum_lam2 if sum_lam2 > 0 else 0
    npr = pr / d_h if d_h > 0 else 0
    lam_min, lam_max = product_ensemble_edges(d, d_h, sigma)
    stable_rank = sum_lam2 / lam_max**2 if lam_max > 0 else 0
    return {
        "max": lam_max,
        "condition_number": lam_max / lam_min if lam_min > 0 else np.inf,
        "participation_ratio": pr,
        "normalized_participation_ratio": npr,
        "stable_rank": stable_rank,
        "lam_min": lam_min,
        "E_lam": E_lam,
        "E_lam2": E_lam2_exact,
    }


def mp_predictions_direct(d, d_h, sigma=1.0):
    gamma = d_h / d
    s2 = d * sigma**4
    lam_max = s2 * (1 + np.sqrt(gamma))**2
    lam_min = s2 * (1 - np.sqrt(gamma))**2
    cond = lam_max / lam_min if lam_min > 0 else np.inf
    E_lam = s2
    E_lam2 = (1 + gamma) * s2**2
    pr = d_h * E_lam**2 / E_lam2
    npr = pr / d_h
    stable_rank = d_h * E_lam2 / lam_max**2
    return {
        "max": lam_max, "condition_number": cond,
        "participation_ratio": pr, "normalized_participation_ratio": npr,
        "stable_rank": stable_rank,
        "lam_min": lam_min, "E_lam": E_lam, "E_lam2": E_lam2,
    }


def plot_stat_distributions(stats, d_model, d_head, n_bins=40):
    import matplotlib.pyplot as plt

    nice_names = {
        "max": r"$\lambda_{\max}$",
        "participation_ratio": r"PR $= (\sum \lambda)^2 / \sum \lambda^2$",
        "normalized_participation_ratio": r"Normalized PR / $d_h$",
        "spectral_entropy": r"Spectral entropy $-\sum p_i \ln p_i$",
        "condition_number": r"Condition number $\lambda_1 / \lambda_{d_h}$",
        "stable_rank": r"Stable rank $\sum \lambda^2 / \lambda_{\max}^2$",
    }
    fig, axes = plt.subplots(2, 3, figsize=(14, 8))
    axes = axes.flatten()
    for idx, (key, vals) in enumerate(stats.items()):
        ax = axes[idx]
        ax.hist(vals, bins=n_bins, density=True, alpha=0.7, edgecolor='k', linewidth=0.3)
        ax.set_xlabel(nice_names.get(key, key), fontsize=10)
        ax.set_ylabel("Density")
        mean, std = np.mean(vals), np.std(vals)
        ax.axvline(mean, color='red', ls='--', lw=1.2, label=f"mean={mean:.2f}")
        ax.legend(fontsize=8)
        ax.set_title(f"d={d_model}, d_h={d_head}", fontsize=9)
    plt.tight_layout()
    return fig


def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(message)s")

    parser = argparse.ArgumentParser(description="Compute MP statistics for all models")
    parser.add_argument("--out", type=str, default="mp_statistics",
                        help="Output directory")
    parser.add_argument("--n-samples", type=int, default=1000,
                        help="Number of MC samples per model")
    parser.add_argument("--models", nargs="*", default=None,
                        help="Only compute for these models")
    parser.add_argument("--plot", action="store_true",
                        help="Generate plots for each model")
    args = parser.parse_args()

    os.makedirs(args.out, exist_ok=True)

    models_to_run = args.models if args.models else list(MODEL_DIMS.keys())

    all_data = {}

    for model_name in models_to_run:
        if model_name not in MODEL_DIMS:
            logging.warning(f"Unknown model {model_name}, skipping")
            continue

        n_layers, n_heads, d_head = MODEL_DIMS[model_name]
        d_model = n_heads * d_head

        logging.info(f"Processing {model_name}: d_model={d_model}, d_head={d_head}, n_heads={n_heads}")

        stats, singular_values = collect_eig_stats(d_model, d_head, n_samples=args.n_samples)

        mp_pred = mp_predictions_direct(d_model, d_head)
        pme_pred = pme_predictions(d_model, d_head)

        mp_lam, mp_density = marchenko_pastur_density(
            None, d_head / d_model, d_model * 1.0**4, n_points=1000
        )

        all_data[model_name] = {
            "d_model": d_model,
            "d_head": d_head,
            "n_heads": n_heads,
            "n_layers": n_layers,
            "stats": stats,
            "singular_values": singular_values,
            "mp_predictions": mp_pred,
            "pme_predictions": pme_pred,
            "mp_density_lambda": mp_lam,
            "mp_density_rho": mp_density,
        }

        if args.plot:
            fig = plot_stat_distributions(stats, d_model, d_head)
            plot_path = os.path.join(args.out, f"{model_name}_stats.png")
            fig.savefig(plot_path, dpi=150, bbox_inches='tight')
            plt.close(fig)
            logging.info(f"Saved plot to {plot_path}")

    output_file = os.path.join(args.out, "mp_statistics.npz")
    np.savez_compressed(output_file, **all_data)
    logging.info(f"Saved consolidated data to {output_file}")

    logging.info(f"\nSummary:")
    logging.info(f"Processed {len(all_data)} models")
    logging.info(f"Data file: {output_file}")
    if args.plot:
        logging.info(f"Plots saved to {args.out}/*.png")


if __name__ == "__main__":
    main()