diff options
Diffstat (limited to 'scripts')
| -rw-r--r-- | scripts/README.md | 30 | ||||
| -rwxr-xr-x | scripts/initialization_distribution_matching.py | 457 |
2 files changed, 487 insertions, 0 deletions
diff --git a/scripts/README.md b/scripts/README.md index efd4208..02b55b0 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -143,6 +143,36 @@ The prior-free minimax theorem says: with equality for isotropic feedback. Outputs are written under `outputs/minimax_initialization/`. +## Initialization Coverage Distribution Matching + +Run: + +```bash +python scripts/initialization_distribution_matching.py --dimension 128 --target-samples 100000 --feedback-samples 100000 --batch-size 8192 --seed 2026 --subspace-dim 8 --anisotropy 64 --plot +``` + +For an initialization distribution \(\mu\), define: + +\[ +M_\mu=\mathbb E[\hat b\hat b^\top]. +\] + +For a random target direction \(a\), theory predicts the coverage distribution: + +\[ +A(a)=a^\top M_\mu a += +\frac{\sum_i \lambda_i G_i}{\sum_i G_i}, +\qquad +G_i\sim\chi^2_1, +\] + +where \(\lambda_i\) are eigenvalues of \(M_\mu\). The script compares this +population-predicted distribution with the distribution induced by an empirical +second-moment estimate from sampled feedback directions. + +Outputs are written under `outputs/initialization_distribution_matching/`. + ## Functional Capacity Overlap Run: diff --git a/scripts/initialization_distribution_matching.py b/scripts/initialization_distribution_matching.py new file mode 100755 index 0000000..f4a14d5 --- /dev/null +++ b/scripts/initialization_distribution_matching.py @@ -0,0 +1,457 @@ +#!/usr/bin/env python3 +"""Match feedback-initialization coverage distributions from M_mu spectra. + +For a feedback initialization distribution mu, define + + M_mu = E[b_hat b_hat^T]. + +For a uniformly random target direction a, the expected squared alignment over +feedback draws is A(a) = a^T M_mu a. If lambda_i are eigenvalues of M_mu, then + + A(a) = sum_i lambda_i G_i / sum_i G_i, G_i ~ chi^2_1. + +This script compares the population-predicted distribution from known spectra +with the empirical distribution obtained by estimating M_mu from sampled +feedback directions. +""" + +from __future__ import annotations + +import argparse +import csv +import json +from dataclasses import asdict, dataclass +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +from scipy import stats + + +@dataclass(frozen=True) +class RunConfig: + dimension: int + target_samples: int + feedback_samples: int + batch_size: int + seed: int + schemes: list[str] + subspace_dim: int + anisotropy: float + outdir: str + plot: bool + + +@dataclass(frozen=True) +class MatchRow: + scheme: str + dimension: int + target_samples: int + feedback_samples: int + population_lambda_min: float + empirical_lambda_min: float + population_lambda_max: float + empirical_lambda_max: float + predicted_mean: float + empirical_mean: float + predicted_std: float + empirical_std: float + predicted_q01: float + empirical_q01: float + predicted_q50: float + empirical_q50: float + predicted_q99: float + empirical_q99: float + ks_2sample_statistic: float + ks_2sample_pvalue: float + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Distribution matching for feedback initialization coverage." + ) + parser.add_argument("--dimension", type=int, default=128) + parser.add_argument("--target-samples", type=int, default=100_000) + parser.add_argument("--feedback-samples", type=int, default=100_000) + parser.add_argument("--batch-size", type=int, default=8192) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--schemes", + nargs="+", + default=["isotropic", "rademacher", "subspace", "axis", "geometric_axis"], + choices=["isotropic", "rademacher", "subspace", "axis", "geometric_axis"], + ) + parser.add_argument("--subspace-dim", type=int, default=8) + parser.add_argument("--anisotropy", type=float, default=64.0) + parser.add_argument( + "--outdir", + type=Path, + default=Path("outputs/initialization_distribution_matching"), + ) + parser.add_argument("--plot", action="store_true") + return parser.parse_args() + + +def validate_config(config: RunConfig) -> None: + if config.dimension < 2: + raise ValueError("--dimension must be at least 2.") + if config.target_samples < 1: + raise ValueError("--target-samples must be positive.") + if config.feedback_samples < 1: + raise ValueError("--feedback-samples must be positive.") + if config.batch_size < 1: + raise ValueError("--batch-size must be positive.") + if config.subspace_dim < 1: + raise ValueError("--subspace-dim must be positive.") + if config.anisotropy < 1: + raise ValueError("--anisotropy must be at least 1.") + + +def population_eigenvalues( + scheme: str, dimension: int, subspace_dim: int, anisotropy: float +) -> np.ndarray: + if scheme in {"isotropic", "rademacher"}: + return np.full(dimension, 1.0 / dimension) + + if scheme == "axis": + values = np.zeros(dimension, dtype=np.float64) + values[-1] = 1.0 + return values + + if scheme == "subspace": + active = min(subspace_dim, dimension) + values = np.zeros(dimension, dtype=np.float64) + values[-active:] = 1.0 / active + return values + + if scheme == "geometric_axis": + weights = np.geomspace(1.0, anisotropy, num=dimension) + return weights / weights.sum() + + raise ValueError(f"Unknown scheme: {scheme}") + + +def sample_feedback( + rng: np.random.Generator, + scheme: str, + count: int, + dimension: int, + eigenvalues: np.ndarray, + subspace_dim: int, +) -> np.ndarray: + if scheme == "isotropic": + values = rng.standard_normal((count, dimension)) + norms = np.linalg.norm(values, axis=1, keepdims=True) + return values / np.maximum(norms, np.finfo(float).tiny) + + if scheme == "rademacher": + values = rng.choice(np.array([-1.0, 1.0]), size=(count, dimension)) + return values / np.sqrt(dimension) + + if scheme == "axis": + values = np.zeros((count, dimension), dtype=np.float64) + values[:, -1] = rng.choice(np.array([-1.0, 1.0]), size=count) + return values + + if scheme == "subspace": + active = min(subspace_dim, dimension) + values = np.zeros((count, dimension), dtype=np.float64) + block = rng.standard_normal((count, active)) + norms = np.linalg.norm(block, axis=1, keepdims=True) + values[:, -active:] = block / np.maximum(norms, np.finfo(float).tiny) + return values + + if scheme == "geometric_axis": + indices = rng.choice(dimension, size=count, p=eigenvalues) + values = np.zeros((count, dimension), dtype=np.float64) + signs = rng.choice(np.array([-1.0, 1.0]), size=count) + values[np.arange(count), indices] = signs + return values + + raise ValueError(f"Unknown scheme: {scheme}") + + +def estimate_moment( + rng: np.random.Generator, + scheme: str, + dimension: int, + eigenvalues: np.ndarray, + subspace_dim: int, + feedback_samples: int, + batch_size: int, +) -> np.ndarray: + moment = np.zeros((dimension, dimension), dtype=np.float64) + offset = 0 + while offset < feedback_samples: + count = min(batch_size, feedback_samples - offset) + feedback = sample_feedback( + rng, scheme, count, dimension, eigenvalues, subspace_dim + ) + moment += feedback.T @ feedback + offset += count + return moment / feedback_samples + + +def target_coverage_from_eigenvalues( + rng: np.random.Generator, eigenvalues: np.ndarray, samples: int, batch_size: int +) -> np.ndarray: + values = np.empty(samples, dtype=np.float64) + dimension = len(eigenvalues) + offset = 0 + while offset < samples: + count = min(batch_size, samples - offset) + g = rng.chisquare(df=1.0, size=(count, dimension)) + values[offset : offset + count] = (g @ eigenvalues) / np.sum(g, axis=1) + offset += count + return values + + +def summarize( + scheme: str, + dimension: int, + target_samples: int, + feedback_samples: int, + population_eigs: np.ndarray, + empirical_eigs: np.ndarray, + predicted: np.ndarray, + empirical: np.ndarray, +) -> MatchRow: + if np.std(predicted) < 1e-14: + ks_statistic = float("nan") + ks_pvalue = float("nan") + else: + ks = stats.ks_2samp(predicted, empirical) + ks_statistic = float(ks.statistic) + ks_pvalue = float(ks.pvalue) + return MatchRow( + scheme=scheme, + dimension=dimension, + target_samples=target_samples, + feedback_samples=feedback_samples, + population_lambda_min=float(np.min(population_eigs)), + empirical_lambda_min=float(np.min(empirical_eigs)), + population_lambda_max=float(np.max(population_eigs)), + empirical_lambda_max=float(np.max(empirical_eigs)), + predicted_mean=float(np.mean(predicted)), + empirical_mean=float(np.mean(empirical)), + predicted_std=float(np.std(predicted, ddof=1)), + empirical_std=float(np.std(empirical, ddof=1)), + predicted_q01=float(np.quantile(predicted, 0.01)), + empirical_q01=float(np.quantile(empirical, 0.01)), + predicted_q50=float(np.quantile(predicted, 0.50)), + empirical_q50=float(np.quantile(empirical, 0.50)), + predicted_q99=float(np.quantile(predicted, 0.99)), + empirical_q99=float(np.quantile(empirical, 0.99)), + ks_2sample_statistic=ks_statistic, + ks_2sample_pvalue=ks_pvalue, + ) + + +def write_csv(path: Path, rows: list[object]) -> None: + if not rows: + return + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="") as handle: + first = asdict(rows[0]) # type: ignore[arg-type] + writer = csv.DictWriter(handle, fieldnames=list(first.keys())) + writer.writeheader() + for row in rows: + writer.writerow(asdict(row)) # type: ignore[arg-type] + + +def write_outputs( + config: RunConfig, + rows: list[MatchRow], + spectra: dict[str, dict[str, list[float]]], + outdir: Path, +) -> None: + outdir.mkdir(parents=True, exist_ok=True) + write_csv(outdir / "distribution_match.csv", rows) + payload = { + "config": asdict(config), + "distribution_match": [asdict(row) for row in rows], + "spectra": spectra, + } + (outdir / "summary.json").write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n" + ) + + +def save_plots( + rows: list[MatchRow], + distributions: dict[str, tuple[np.ndarray, np.ndarray]], + spectra: dict[str, dict[str, list[float]]], + outdir: Path, +) -> list[Path]: + outdir.mkdir(parents=True, exist_ok=True) + paths: list[Path] = [] + + spectrum_path = outdir / "population_vs_empirical_spectra.png" + plt.figure(figsize=(7, 4.5)) + for scheme, values in sorted(spectra.items()): + empirical = np.array(values["empirical"]) + plt.plot(np.sort(empirical), label=scheme) + plt.xlabel("eigenvalue index") + plt.ylabel("empirical eigenvalue of M_mu") + plt.title("Estimated feedback second-moment spectra") + plt.legend() + plt.tight_layout() + plt.savefig(spectrum_path, dpi=180) + plt.close() + paths.append(spectrum_path) + + hist_path = outdir / "coverage_histograms.png" + plt.figure(figsize=(8, 5)) + for scheme, (predicted, empirical) in sorted(distributions.items()): + plt.hist( + empirical, + bins=80, + density=True, + histtype="step", + linewidth=1.5, + label=f"{scheme} empirical", + ) + plt.xlabel("target coverage a^T M_mu a") + plt.ylabel("density") + plt.title("Coverage distributions from estimated M_mu") + plt.legend(fontsize=8) + plt.tight_layout() + plt.savefig(hist_path, dpi=180) + plt.close() + paths.append(hist_path) + + for scheme, (predicted, empirical) in sorted(distributions.items()): + if np.std(predicted) < 1e-14: + continue + qq_path = outdir / f"qq_{scheme}.png" + probs = (np.arange(1, len(predicted) + 1) - 0.5) / len(predicted) + pred_q = np.quantile(predicted, probs) + emp_q = np.quantile(empirical, probs) + max_value = float(max(pred_q[-1], emp_q[-1])) + min_value = float(min(pred_q[0], emp_q[0])) + plt.figure(figsize=(5, 5)) + plt.scatter(pred_q, emp_q, s=4, alpha=0.25) + plt.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1) + plt.xlabel("population-predicted quantile") + plt.ylabel("empirical-estimated quantile") + plt.title(f"Q-Q coverage: {scheme}") + plt.tight_layout() + plt.savefig(qq_path, dpi=180) + plt.close() + paths.append(qq_path) + + worst_path = outdir / "worst_best_coverage.png" + names = [row.scheme for row in rows] + x = np.arange(len(names)) + plt.figure(figsize=(7, 4.5)) + plt.bar( + x - 0.18, + [row.empirical_lambda_min for row in rows], + width=0.36, + label="worst target", + ) + plt.bar( + x + 0.18, + [row.empirical_lambda_max for row in rows], + width=0.36, + label="best target", + ) + plt.axhline(1.0 / rows[0].dimension, color="black", linestyle="--", linewidth=1, label="1/D") + plt.xticks(x, names, rotation=25, ha="right") + plt.ylabel("expected squared alignment") + plt.title("Worst and best target coverage") + plt.legend() + plt.tight_layout() + plt.savefig(worst_path, dpi=180) + plt.close() + paths.append(worst_path) + + return paths + + +def parse_config(args: argparse.Namespace) -> RunConfig: + return RunConfig( + dimension=args.dimension, + target_samples=args.target_samples, + feedback_samples=args.feedback_samples, + batch_size=args.batch_size, + seed=args.seed, + schemes=args.schemes, + subspace_dim=args.subspace_dim, + anisotropy=args.anisotropy, + outdir=str(args.outdir), + plot=args.plot, + ) + + +def main() -> None: + args = parse_args() + config = parse_config(args) + validate_config(config) + rng = np.random.default_rng(config.seed) + + rows: list[MatchRow] = [] + spectra: dict[str, dict[str, list[float]]] = {} + distributions: dict[str, tuple[np.ndarray, np.ndarray]] = {} + + for scheme in config.schemes: + population_eigs = population_eigenvalues( + scheme, config.dimension, config.subspace_dim, config.anisotropy + ) + moment = estimate_moment( + rng, + scheme, + config.dimension, + population_eigs, + config.subspace_dim, + config.feedback_samples, + config.batch_size, + ) + empirical_eigs = np.linalg.eigvalsh(moment) + predicted = target_coverage_from_eigenvalues( + rng, population_eigs, config.target_samples, config.batch_size + ) + empirical = target_coverage_from_eigenvalues( + rng, empirical_eigs, config.target_samples, config.batch_size + ) + row = summarize( + scheme, + config.dimension, + config.target_samples, + config.feedback_samples, + population_eigs, + empirical_eigs, + predicted, + empirical, + ) + rows.append(row) + spectra[scheme] = { + "population": population_eigs.tolist(), + "empirical": empirical_eigs.tolist(), + } + distributions[scheme] = (predicted, empirical) + print( + f"{scheme}: " + f"lambda_min_pop={row.population_lambda_min:.6g}, " + f"lambda_min_emp={row.empirical_lambda_min:.6g}, " + f"KS={row.ks_2sample_statistic:.6g}, " + f"mean_pred={row.predicted_mean:.6g}, " + f"mean_emp={row.empirical_mean:.6g}" + ) + + outdir = Path(config.outdir) + write_outputs(config, rows, spectra, outdir) + plot_paths = save_plots(rows, distributions, spectra, outdir) if config.plot else [] + + finite_ks = [row.ks_2sample_statistic for row in rows if np.isfinite(row.ks_2sample_statistic)] + max_ks = max(finite_ks) if finite_ks else float("nan") + max_mean_error = max(abs(row.predicted_mean - row.empirical_mean) for row in rows) + print(f"max_ks_2sample: {max_ks:.8g}") + print(f"max_abs_mean_error: {max_mean_error:.8g}") + print(f"distribution_match: {outdir / 'distribution_match.csv'}") + for path in plot_paths: + print(f"plot: {path}") + + +if __name__ == "__main__": + main() |