diff options
Diffstat (limited to 'scripts')
| -rw-r--r-- | scripts/README.md | 31 | ||||
| -rwxr-xr-x | scripts/capacity_empirical_validation.py | 505 |
2 files changed, 536 insertions, 0 deletions
diff --git a/scripts/README.md b/scripts/README.md index 37dc68b..e30b309 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -54,6 +54,37 @@ The default run compares two regimes: Outputs are written under `outputs/capacity_scaling/`. +## Large Empirical Capacity Validation + +Run: + +```bash +python scripts/capacity_empirical_validation.py --dimensions 64 128 256 512 1024 2048 4096 --samples 100000 --batch-size 2048 --seed 123 --plot +``` + +This uses rotational invariance to fix the target direction and sample random +feedback directions: + +\[ +Q=\frac{z_1^2}{\|z\|^2}, +\qquad +z\sim \mathcal N(0,I_D). +\] + +It validates: + +- beta-law distribution calibration across dimensions; +- empirical tail probabilities for \(q=c/D\) and fixed \(q\); +- empirical capacity cost \(C(q)=-\log P(Q\ge q)\); +- multilayer product scaling for all-alignment events. + +Outputs are written under `outputs/capacity_empirical_validation/`: + +- `distribution_summary.csv` +- `capacity_tails.csv` +- `multilayer_capacity.csv` +- diagnostic plots when `--plot` is set. + ## Minimax Initialization Bound Run: diff --git a/scripts/capacity_empirical_validation.py b/scripts/capacity_empirical_validation.py new file mode 100755 index 0000000..a6c33dc --- /dev/null +++ b/scripts/capacity_empirical_validation.py @@ -0,0 +1,505 @@ +#!/usr/bin/env python3 +"""Large empirical validation of FA alignment capacity costs. + +By rotational invariance, for an isotropic target A and feedback B in R^D, +cos^2(A, B) has the same law as the squared first coordinate of a random +Gaussian direction: + + Q = z_1^2 / ||z||^2, z ~ N(0, I_D). + +This script uses that fixed-target sampler to validate the beta law, empirical +tail probabilities, log-volume capacity costs, and multilayer product scaling. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +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: + dimensions: list[int] + samples: int + batch_size: int + seed: int + chance_cs: list[float] + fixed_qs: list[float] + multilayer_cs: list[float] + multilayer_counts: list[int] + outdir: str + plot: bool + + +@dataclass(frozen=True) +class DistributionRow: + dimension: int + samples: int + empirical_mean: float + theoretical_mean: float + empirical_var: float + theoretical_var: float + ks_statistic: float + ks_pvalue: float + q50_empirical: float + q50_theoretical: float + q90_empirical: float + q90_theoretical: float + q99_empirical: float + q99_theoretical: float + + +@dataclass(frozen=True) +class CapacityTailRow: + dimension: int + samples: int + threshold_mode: str + threshold_label: str + q_threshold: float + hits: int + empirical_tail: float + empirical_tail_smoothed: float + theoretical_tail: float + expected_hits: float + empirical_cost_nats: float + empirical_cost_smoothed_nats: float + theoretical_cost_nats: float + cost_error_smoothed_nats: float + relative_tail_error: float + + +@dataclass(frozen=True) +class MultilayerRow: + dimension: int + samples: int + grouped_samples: int + layers: int + c: float + q_threshold: float + hits: int + empirical_all_tail: float + empirical_all_tail_smoothed: float + theoretical_all_tail: float + expected_hits: float + empirical_total_cost_smoothed_nats: float + theoretical_total_cost_nats: float + cost_error_smoothed_nats: float + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Large empirical validation for FA capacity costs." + ) + parser.add_argument( + "--dimensions", + type=int, + nargs="+", + default=[64, 128, 256, 512, 1024, 2048, 4096], + help="Flattened matrix dimensions D to validate.", + ) + parser.add_argument( + "--samples", + type=int, + default=100_000, + help="Monte Carlo samples per dimension.", + ) + parser.add_argument( + "--batch-size", + type=int, + default=2048, + help="Gaussian direction samples per batch.", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--chance-cs", + type=float, + nargs="+", + default=[0.5, 1.0, 2.0, 5.0, 10.0], + help="Chance-level thresholds q=c/D.", + ) + parser.add_argument( + "--fixed-qs", + type=float, + nargs="*", + default=[0.001, 0.002, 0.005, 0.01], + help="Fixed thresholds q.", + ) + parser.add_argument( + "--multilayer-cs", + type=float, + nargs="+", + default=[0.5, 1.0], + help="Chance-level c values for multilayer all-event validation.", + ) + parser.add_argument( + "--multilayer-counts", + type=int, + nargs="+", + default=[1, 2, 4, 8], + help="Layer counts for multilayer all-event validation.", + ) + parser.add_argument( + "--outdir", + type=Path, + default=Path("outputs/capacity_empirical_validation"), + ) + parser.add_argument("--plot", action="store_true") + return parser.parse_args() + + +def sample_q_fixed_target( + rng: np.random.Generator, dimension: int, samples: int, batch_size: int +) -> np.ndarray: + values = np.empty(samples, dtype=np.float64) + offset = 0 + while offset < samples: + count = min(batch_size, samples - offset) + z = rng.standard_normal((count, dimension)) + numerator = z[:, 0] * z[:, 0] + denominator = np.einsum("ij,ij->i", z, z) + values[offset : offset + count] = numerator / denominator + offset += count + return values + + +def beta_dist(dimension: int) -> stats.rv_continuous: + return stats.beta(0.5, (dimension - 1) / 2) + + +def distribution_summary(q_values: np.ndarray, dimension: int) -> DistributionRow: + dist = beta_dist(dimension) + ks = stats.kstest(q_values, dist.cdf) + return DistributionRow( + dimension=dimension, + samples=len(q_values), + empirical_mean=float(np.mean(q_values)), + theoretical_mean=float(dist.mean()), + empirical_var=float(np.var(q_values, ddof=1)), + theoretical_var=float(dist.var()), + ks_statistic=float(ks.statistic), + ks_pvalue=float(ks.pvalue), + q50_empirical=float(np.quantile(q_values, 0.50)), + q50_theoretical=float(dist.ppf(0.50)), + q90_empirical=float(np.quantile(q_values, 0.90)), + q90_theoretical=float(dist.ppf(0.90)), + q99_empirical=float(np.quantile(q_values, 0.99)), + q99_theoretical=float(dist.ppf(0.99)), + ) + + +def finite_cost(tail: float) -> float: + if tail <= 0: + return math.inf + return -math.log(tail) + + +def capacity_tail_row( + q_values: np.ndarray, + dimension: int, + threshold_mode: str, + threshold_label: str, + q_threshold: float, +) -> CapacityTailRow: + samples = len(q_values) + dist = beta_dist(dimension) + hits = int(np.count_nonzero(q_values >= q_threshold)) + empirical_tail = hits / samples + empirical_tail_smoothed = (hits + 0.5) / (samples + 1.0) + theoretical_tail = float(dist.sf(q_threshold)) + expected_hits = samples * theoretical_tail + theoretical_cost = finite_cost(theoretical_tail) + smoothed_cost = finite_cost(empirical_tail_smoothed) + relative_error = ( + abs(empirical_tail - theoretical_tail) / theoretical_tail + if theoretical_tail > 0 + else math.nan + ) + return CapacityTailRow( + dimension=dimension, + samples=samples, + threshold_mode=threshold_mode, + threshold_label=threshold_label, + q_threshold=q_threshold, + hits=hits, + empirical_tail=empirical_tail, + empirical_tail_smoothed=empirical_tail_smoothed, + theoretical_tail=theoretical_tail, + expected_hits=expected_hits, + empirical_cost_nats=finite_cost(empirical_tail), + empirical_cost_smoothed_nats=smoothed_cost, + theoretical_cost_nats=theoretical_cost, + cost_error_smoothed_nats=smoothed_cost - theoretical_cost, + relative_tail_error=relative_error, + ) + + +def multilayer_row( + q_values: np.ndarray, dimension: int, layers: int, c: float +) -> MultilayerRow: + samples = len(q_values) + grouped_samples = samples // layers + usable = grouped_samples * layers + q_threshold = c / dimension + grouped = q_values[:usable].reshape(grouped_samples, layers) + hits = int(np.count_nonzero(np.all(grouped >= q_threshold, axis=1))) + empirical_tail = hits / grouped_samples + empirical_tail_smoothed = (hits + 0.5) / (grouped_samples + 1.0) + per_layer_tail = float(beta_dist(dimension).sf(q_threshold)) + theoretical_tail = per_layer_tail**layers + expected_hits = grouped_samples * theoretical_tail + smoothed_cost = finite_cost(empirical_tail_smoothed) + theoretical_cost = finite_cost(theoretical_tail) + return MultilayerRow( + dimension=dimension, + samples=samples, + grouped_samples=grouped_samples, + layers=layers, + c=c, + q_threshold=q_threshold, + hits=hits, + empirical_all_tail=empirical_tail, + empirical_all_tail_smoothed=empirical_tail_smoothed, + theoretical_all_tail=theoretical_tail, + expected_hits=expected_hits, + empirical_total_cost_smoothed_nats=smoothed_cost, + theoretical_total_cost_nats=theoretical_cost, + cost_error_smoothed_nats=smoothed_cost - theoretical_cost, + ) + + +def validate_args(config: RunConfig) -> None: + if any(d < 2 for d in config.dimensions): + raise ValueError("All dimensions must be at least 2.") + if config.samples < 1: + raise ValueError("--samples must be positive.") + if config.batch_size < 1: + raise ValueError("--batch-size must be positive.") + if any(c < 0 for c in config.chance_cs): + raise ValueError("--chance-cs must be non-negative.") + if any(c < 0 for c in config.multilayer_cs): + raise ValueError("--multilayer-cs must be non-negative.") + if any(q < 0 or q > 1 for q in config.fixed_qs): + raise ValueError("--fixed-qs must lie in [0, 1].") + if any(layers < 1 for layers in config.multilayer_counts): + raise ValueError("--multilayer-counts must be positive.") + + +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, + distribution_rows: list[DistributionRow], + capacity_rows: list[CapacityTailRow], + multilayer_rows: list[MultilayerRow], + outdir: Path, +) -> None: + outdir.mkdir(parents=True, exist_ok=True) + write_csv(outdir / "distribution_summary.csv", distribution_rows) + write_csv(outdir / "capacity_tails.csv", capacity_rows) + write_csv(outdir / "multilayer_capacity.csv", multilayer_rows) + payload = { + "config": asdict(config), + "distribution_summary": [asdict(row) for row in distribution_rows], + "capacity_tails": [asdict(row) for row in capacity_rows], + "multilayer_capacity": [asdict(row) for row in multilayer_rows], + } + (outdir / "summary.json").write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n" + ) + + +def save_plots( + distribution_rows: list[DistributionRow], + capacity_rows: list[CapacityTailRow], + multilayer_rows: list[MultilayerRow], + outdir: Path, +) -> list[Path]: + outdir.mkdir(parents=True, exist_ok=True) + paths: list[Path] = [] + + ks_path = outdir / "ks_by_dimension.png" + plt.figure(figsize=(7, 4.5)) + dims = [row.dimension for row in distribution_rows] + ks_values = [row.ks_statistic for row in distribution_rows] + plt.plot(dims, ks_values, marker="o") + plt.xscale("log", base=2) + plt.xlabel("dimension D") + plt.ylabel("KS statistic") + plt.title("Beta-law distribution calibration") + plt.tight_layout() + plt.savefig(ks_path, dpi=180) + plt.close() + paths.append(ks_path) + + chance_path = outdir / "chance_capacity_empirical_vs_theory.png" + chance_rows = [row for row in capacity_rows if row.threshold_mode == "chance"] + plt.figure(figsize=(7, 4.5)) + for label in sorted({row.threshold_label for row in chance_rows}): + rows = sorted( + [row for row in chance_rows if row.threshold_label == label], + key=lambda row: row.dimension, + ) + plt.plot( + [row.theoretical_cost_nats for row in rows], + [row.empirical_cost_smoothed_nats for row in rows], + marker="o", + linestyle="", + label=label, + ) + max_cost = max( + max(row.theoretical_cost_nats, row.empirical_cost_smoothed_nats) + for row in chance_rows + if np.isfinite(row.empirical_cost_smoothed_nats) + ) + plt.plot([0, max_cost], [0, max_cost], color="black", linewidth=1) + plt.xlabel("theoretical capacity cost (nats)") + plt.ylabel("empirical smoothed cost (nats)") + plt.title("Chance-threshold capacity calibration") + plt.legend() + plt.tight_layout() + plt.savefig(chance_path, dpi=180) + plt.close() + paths.append(chance_path) + + multilayer_path = outdir / "multilayer_capacity_empirical_vs_theory.png" + plt.figure(figsize=(7, 4.5)) + for c in sorted({row.c for row in multilayer_rows}): + rows = [row for row in multilayer_rows if row.c == c] + plt.scatter( + [row.theoretical_total_cost_nats for row in rows], + [row.empirical_total_cost_smoothed_nats for row in rows], + label=f"c={c:g}", + alpha=0.8, + ) + max_cost = max( + max(row.theoretical_total_cost_nats, row.empirical_total_cost_smoothed_nats) + for row in multilayer_rows + if np.isfinite(row.empirical_total_cost_smoothed_nats) + ) + plt.plot([0, max_cost], [0, max_cost], color="black", linewidth=1) + plt.xlabel("theoretical total cost (nats)") + plt.ylabel("empirical smoothed total cost (nats)") + plt.title("Multilayer product-capacity calibration") + plt.legend() + plt.tight_layout() + plt.savefig(multilayer_path, dpi=180) + plt.close() + paths.append(multilayer_path) + + return paths + + +def parse_config(args: argparse.Namespace) -> RunConfig: + return RunConfig( + dimensions=args.dimensions, + samples=args.samples, + batch_size=args.batch_size, + seed=args.seed, + chance_cs=args.chance_cs, + fixed_qs=args.fixed_qs, + multilayer_cs=args.multilayer_cs, + multilayer_counts=args.multilayer_counts, + outdir=str(args.outdir), + plot=args.plot, + ) + + +def main() -> None: + args = parse_args() + config = parse_config(args) + validate_args(config) + + rng = np.random.default_rng(config.seed) + distribution_rows: list[DistributionRow] = [] + capacity_rows: list[CapacityTailRow] = [] + multilayer_rows: list[MultilayerRow] = [] + + for dimension in config.dimensions: + q_values = sample_q_fixed_target( + rng, dimension, config.samples, config.batch_size + ) + dist_row = distribution_summary(q_values, dimension) + distribution_rows.append(dist_row) + + for c in config.chance_cs: + q_threshold = min(c / dimension, 1.0) + capacity_rows.append( + capacity_tail_row( + q_values, + dimension, + "chance", + f"c={c:g}/D", + q_threshold, + ) + ) + + for q in config.fixed_qs: + capacity_rows.append( + capacity_tail_row(q_values, dimension, "fixed", f"q={q:g}", q) + ) + + for c in config.multilayer_cs: + for layers in config.multilayer_counts: + if layers <= len(q_values): + multilayer_rows.append( + multilayer_row(q_values, dimension, layers, c) + ) + + print( + f"D={dimension}: " + f"mean_emp={dist_row.empirical_mean:.6g}, " + f"mean_theory={dist_row.theoretical_mean:.6g}, " + f"KS={dist_row.ks_statistic:.6g}" + ) + + outdir = Path(config.outdir) + write_outputs(config, distribution_rows, capacity_rows, multilayer_rows, outdir) + plot_paths = ( + save_plots(distribution_rows, capacity_rows, multilayer_rows, outdir) + if config.plot + else [] + ) + + reliable_rows = [ + row + for row in capacity_rows + if row.expected_hits >= 20 and row.threshold_mode == "chance" + ] + max_abs_cost_error = max( + (abs(row.cost_error_smoothed_nats) for row in reliable_rows), default=math.nan + ) + max_ks = max(row.ks_statistic for row in distribution_rows) + print(f"dimensions: {config.dimensions}") + print(f"samples_per_dimension: {config.samples}") + print(f"max_ks_statistic: {max_ks:.8g}") + print( + "max_abs_chance_cost_error_nats_expected_hits_ge_20: " + f"{max_abs_cost_error:.8g}" + ) + print(f"distribution_summary: {outdir / 'distribution_summary.csv'}") + print(f"capacity_tails: {outdir / 'capacity_tails.csv'}") + print(f"multilayer_capacity: {outdir / 'multilayer_capacity.csv'}") + for path in plot_paths: + print(f"plot: {path}") + + +if __name__ == "__main__": + main() |