diff options
| -rw-r--r-- | notes/01_theory_notes.md | 88 | ||||
| -rw-r--r-- | notes/02_experiment_notes.md | 25 | ||||
| -rw-r--r-- | scripts/README.md | 31 | ||||
| -rwxr-xr-x | scripts/minimax_initialization.py | 364 |
4 files changed, 508 insertions, 0 deletions
diff --git a/notes/01_theory_notes.md b/notes/01_theory_notes.md index c0ac8fe..d524000 100644 --- a/notes/01_theory_notes.md +++ b/notes/01_theory_notes.md @@ -215,6 +215,94 @@ Thus: Isotropic initialization attains this bound. +### Proof Sketch + +The proof only needs second moments. For any feedback initialization \(\mu\), define: + +\[ +M_\mu=\mathbb E_\mu[\hat b\hat b^\top]. +\] + +Since \(\|\hat b\|=1\): + +\[ +\operatorname{tr}M_\mu += +\mathbb E_\mu[\operatorname{tr}(\hat b\hat b^\top)] += +\mathbb E_\mu[\|\hat b\|^2] +=1. +\] + +For any fixed target direction \(a\): + +\[ +\mathbb E_\mu[(a^\top \hat b)^2] += +\mathbb E_\mu[a^\top \hat b\hat b^\top a] += +a^\top M_\mu a. +\] + +The worst-case target direction is the minimum-eigenvalue direction of \(M_\mu\): + +\[ +\inf_{\|a\|=1} a^\top M_\mu a += +\lambda_{\min}(M_\mu). +\] + +Because the average eigenvalue is \(1/D\): + +\[ +\lambda_{\min}(M_\mu) +\le +\frac1D. +\] + +Hence: + +\[ +\inf_{\|a\|=1} +\mathbb E_\mu[(a^\top \hat b)^2] +\le +\frac1D +\] + +for every initialization \(\mu\). If \(\mu\) is isotropic, then: + +\[ +M_\mu=\frac1D I, +\] + +so every direction has: + +\[ +a^\top M_\mu a=\frac1D. +\] + +Thus: + +\[ +\sup_\mu +\inf_{\|a\|=1} +\mathbb E_\mu[(a^\top \hat b)^2] += +\frac1D. +\] + +This is a worst-case theorem. If the target direction is itself uniformly random, then: + +\[ +\mathbb E_a[a^\top M_\mu a] += +\frac{\operatorname{tr}M_\mu}{D} += +\frac1D +\] + +for every \(\mu\). Therefore prior-free average-case alignment over uniform targets cannot distinguish initializations by their mean; the distinction is worst-case coverage, tail behavior, and conditioning. + ## Prior-Aware Corollary If target directions have prior covariance: diff --git a/notes/02_experiment_notes.md b/notes/02_experiment_notes.md index ab69cd8..05b108b 100644 --- a/notes/02_experiment_notes.md +++ b/notes/02_experiment_notes.md @@ -259,3 +259,28 @@ Result: - chance-level max total cost: `18.3652` nats at width `128`, layers `16` This cleanly separates the fixed-threshold regime, where total cost scales like \(Ln^2\), from the chance-level regime, where per-layer cost is nearly width-independent. + +## Minimax Initialization Run Log + +Script: + +```bash +python scripts/minimax_initialization.py --dimension 32 --feedback-samples 20000 --target-samples 10000 --seed 11 --subspace-dim 4 --plot +``` + +Result: + +- minimax bound \(1/D\): `0.03125` +- isotropic \(\lambda_{\min}\): `0.029084138` +- rademacher \(\lambda_{\min}\): `0.028946927` +- anisotropic \(\lambda_{\min}\): `0.006149976` +- subspace \(\lambda_{\min}\): `0` +- axis \(\lambda_{\min}\): `0` + +Random-target means remain close to \(1/D\) for all distributions, but worst-case target coverage differs sharply: + +- isotropic and rademacher nearly equalize all target directions; +- anisotropic improves some directions while sacrificing others; +- subspace and axis initializations leave entire orthogonal directions uncovered. + +This empirically illustrates the prior-free minimax theorem: without target or weight prior information, anisotropic feedback cannot improve the worst-case angular bound. diff --git a/scripts/README.md b/scripts/README.md index 3124261..94f9fff 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -53,3 +53,34 @@ The default run compares two regimes: - `chance`: \(q=1/D\), where \(C_{\mathrm{all}}\) grows mostly with \(L\). Outputs are written under `outputs/capacity_scaling/`. + +## Minimax Initialization Bound + +Run: + +```bash +python scripts/minimax_initialization.py --dimension 32 --feedback-samples 20000 --target-samples 10000 --seed 11 --subspace-dim 4 --plot +``` + +This estimates the feedback second-moment matrix: + +\[ +M_\mu=\mathbb E_\mu[\hat b\hat b^\top] +\] + +for several initialization distributions. The worst-case expected squared alignment is: + +\[ +\inf_{\|a\|=1} +\mathbb E_\mu[(a^\top \hat b)^2] += +\lambda_{\min}(M_\mu). +\] + +The prior-free minimax theorem says: + +\[ +\sup_\mu \lambda_{\min}(M_\mu)=\frac1D, +\] + +with equality for isotropic feedback. Outputs are written under `outputs/minimax_initialization/`. diff --git a/scripts/minimax_initialization.py b/scripts/minimax_initialization.py new file mode 100755 index 0000000..ba73b2e --- /dev/null +++ b/scripts/minimax_initialization.py @@ -0,0 +1,364 @@ +#!/usr/bin/env python3 +"""Empirically illustrate the prior-free minimax feedback bound. + +For normalized feedback directions b_hat in S^{D-1}, any initialization +distribution mu induces M_mu = E[b_hat b_hat^T] with trace 1. For any target +direction a, the expected squared alignment is a^T M_mu a, so the worst-case +target receives lambda_min(M_mu) <= 1/D. Isotropic feedback attains 1/D. +""" + +from __future__ import annotations + +import argparse +import json +from dataclasses import asdict, dataclass +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + + +@dataclass(frozen=True) +class RunConfig: + dimension: int + feedback_samples: int + target_samples: int + batch_size: int + seed: int + distributions: list[str] + subspace_dim: int + anisotropy: float + outdir: str + plot: bool + + +@dataclass(frozen=True) +class DistributionSummary: + name: str + trace: float + lambda_min: float + lambda_max: float + minimax_bound: float + isotropic_gap: float + random_target_mean: float + random_target_std: float + random_target_p01: float + random_target_p50: float + random_target_p99: float + min_eigen_target_alignment: float + max_eigen_target_alignment: float + frobenius_distance_to_isotropic: float + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Simulate the prior-free minimax bound for feedback initialization." + ) + parser.add_argument( + "--dimension", + type=int, + default=64, + help="Flattened feedback dimension D.", + ) + parser.add_argument( + "--feedback-samples", + type=int, + default=50_000, + help="Number of feedback directions used to estimate M_mu.", + ) + parser.add_argument( + "--target-samples", + type=int, + default=20_000, + help="Number of random target directions for evaluating a^T M_mu a.", + ) + parser.add_argument( + "--batch-size", + type=int, + default=10_000, + help="Sampling batch size.", + ) + parser.add_argument("--seed", type=int, default=0, help="Random seed.") + parser.add_argument( + "--distribution", + nargs="+", + default=["isotropic", "rademacher", "anisotropic", "subspace", "axis"], + choices=["isotropic", "rademacher", "anisotropic", "subspace", "axis"], + help="Feedback initialization distributions to compare.", + ) + parser.add_argument( + "--subspace-dim", + type=int, + default=8, + help="Active dimension for the subspace distribution.", + ) + parser.add_argument( + "--anisotropy", + type=float, + default=16.0, + help="Variance ratio for the anisotropic Gaussian distribution.", + ) + parser.add_argument( + "--outdir", + type=Path, + default=Path("outputs/minimax_initialization"), + help="Directory for summaries and plots.", + ) + parser.add_argument("--plot", action="store_true", help="Save diagnostic plots.") + return parser.parse_args() + + +def normalize(values: np.ndarray) -> np.ndarray: + norms = np.linalg.norm(values, axis=1, keepdims=True) + return values / np.maximum(norms, np.finfo(values.dtype).tiny) + + +def sample_feedback( + rng: np.random.Generator, + distribution: str, + count: int, + dimension: int, + subspace_dim: int, + anisotropy: float, +) -> np.ndarray: + if distribution == "isotropic": + return normalize(rng.standard_normal((count, dimension))) + + if distribution == "rademacher": + values = rng.choice(np.array([-1.0, 1.0]), size=(count, dimension)) + return values / np.sqrt(dimension) + + if distribution == "axis": + values = np.zeros((count, dimension), dtype=np.float64) + values[:, 0] = rng.choice(np.array([-1.0, 1.0]), size=count) + return values + + if distribution == "subspace": + active = min(subspace_dim, dimension) + values = np.zeros((count, dimension), dtype=np.float64) + values[:, :active] = rng.standard_normal((count, active)) + return normalize(values) + + if distribution == "anisotropic": + variances = np.geomspace(anisotropy, 1.0, num=dimension) + values = rng.standard_normal((count, dimension)) * np.sqrt(variances) + return normalize(values) + + raise ValueError(f"Unknown distribution: {distribution}") + + +def estimate_second_moment( + rng: np.random.Generator, config: RunConfig, distribution: str +) -> np.ndarray: + moment = np.zeros((config.dimension, config.dimension), dtype=np.float64) + remaining = config.feedback_samples + while remaining > 0: + count = min(config.batch_size, remaining) + feedback = sample_feedback( + rng, + distribution, + count, + config.dimension, + config.subspace_dim, + config.anisotropy, + ) + moment += feedback.T @ feedback + remaining -= count + return moment / config.feedback_samples + + +def random_targets( + rng: np.random.Generator, count: int, dimension: int +) -> np.ndarray: + return normalize(rng.standard_normal((count, dimension))) + + +def summarize_distribution( + name: str, + moment: np.ndarray, + rng: np.random.Generator, + target_samples: int, +) -> tuple[DistributionSummary, np.ndarray, np.ndarray]: + dimension = moment.shape[0] + eigenvalues = np.linalg.eigvalsh(moment) + targets = random_targets(rng, target_samples, dimension) + target_alignments = np.einsum("ij,jk,ik->i", targets, moment, targets) + minimax_bound = 1.0 / dimension + isotropic = np.eye(dimension) / dimension + + summary = DistributionSummary( + name=name, + trace=float(np.trace(moment)), + lambda_min=float(eigenvalues[0]), + lambda_max=float(eigenvalues[-1]), + minimax_bound=minimax_bound, + isotropic_gap=float(minimax_bound - eigenvalues[0]), + random_target_mean=float(np.mean(target_alignments)), + random_target_std=float(np.std(target_alignments, ddof=1)), + random_target_p01=float(np.quantile(target_alignments, 0.01)), + random_target_p50=float(np.quantile(target_alignments, 0.50)), + random_target_p99=float(np.quantile(target_alignments, 0.99)), + min_eigen_target_alignment=float(eigenvalues[0]), + max_eigen_target_alignment=float(eigenvalues[-1]), + frobenius_distance_to_isotropic=float(np.linalg.norm(moment - isotropic)), + ) + return summary, eigenvalues, target_alignments + + +def write_outputs( + config: RunConfig, + summaries: list[DistributionSummary], + spectra: dict[str, np.ndarray], + target_alignments: dict[str, np.ndarray], + outdir: Path, +) -> None: + outdir.mkdir(parents=True, exist_ok=True) + payload = { + "config": asdict(config), + "summaries": [asdict(summary) for summary in summaries], + } + (outdir / "summary.json").write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n" + ) + + spectrum_payload = { + name: values.tolist() for name, values in sorted(spectra.items()) + } + (outdir / "eigenvalues.json").write_text( + json.dumps(spectrum_payload, indent=2, sort_keys=True) + "\n" + ) + + target_payload = { + name: values.tolist() for name, values in sorted(target_alignments.items()) + } + (outdir / "random_target_alignments.json").write_text( + json.dumps(target_payload, indent=2, sort_keys=True) + "\n" + ) + + +def save_plots( + summaries: list[DistributionSummary], + spectra: dict[str, np.ndarray], + target_alignments: dict[str, np.ndarray], + outdir: Path, +) -> list[Path]: + outdir.mkdir(parents=True, exist_ok=True) + paths: list[Path] = [] + dimension = len(next(iter(spectra.values()))) + minimax_bound = 1.0 / dimension + + spectrum_path = outdir / "eigenvalue_spectra.png" + plt.figure(figsize=(7, 4.5)) + for name, values in sorted(spectra.items()): + plt.plot(np.arange(1, dimension + 1), np.sort(values), label=name) + plt.axhline(minimax_bound, color="black", linestyle="--", linewidth=1, label="1/D") + plt.xlabel("eigenvalue index") + plt.ylabel("eigenvalue of M_mu") + plt.title("Feedback second-moment spectra") + plt.legend() + plt.tight_layout() + plt.savefig(spectrum_path, dpi=180) + plt.close() + paths.append(spectrum_path) + + worst_case_path = outdir / "worst_case_alignment.png" + names = [summary.name for summary in summaries] + worst = [summary.lambda_min for summary in summaries] + best = [summary.lambda_max for summary in summaries] + x = np.arange(len(names)) + plt.figure(figsize=(7, 4.5)) + plt.bar(x - 0.18, worst, width=0.36, label="worst target") + plt.bar(x + 0.18, best, width=0.36, label="best target") + plt.axhline(minimax_bound, color="black", linestyle="--", linewidth=1, label="1/D") + plt.xticks(x, names, rotation=25, ha="right") + plt.ylabel("expected squared alignment") + plt.title("Worst-case target penalty") + plt.legend() + plt.tight_layout() + plt.savefig(worst_case_path, dpi=180) + plt.close() + paths.append(worst_case_path) + + target_path = outdir / "random_target_alignment_hist.png" + plt.figure(figsize=(7, 4.5)) + for name, values in sorted(target_alignments.items()): + plt.hist(values, bins=70, density=True, histtype="step", linewidth=1.5, label=name) + plt.axvline(minimax_bound, color="black", linestyle="--", linewidth=1, label="1/D") + plt.xlabel("a^T M_mu a for random target a") + plt.ylabel("density") + plt.title("Random target alignment") + plt.legend() + plt.tight_layout() + plt.savefig(target_path, dpi=180) + plt.close() + paths.append(target_path) + + return paths + + +def main() -> None: + args = parse_args() + if args.dimension < 2: + raise ValueError("--dimension must be at least 2.") + if args.feedback_samples < 1: + raise ValueError("--feedback-samples must be positive.") + if args.target_samples < 1: + raise ValueError("--target-samples must be positive.") + if args.batch_size < 1: + raise ValueError("--batch-size must be positive.") + if args.subspace_dim < 1: + raise ValueError("--subspace-dim must be positive.") + if args.anisotropy < 1: + raise ValueError("--anisotropy must be at least 1.") + + config = RunConfig( + dimension=args.dimension, + feedback_samples=args.feedback_samples, + target_samples=args.target_samples, + batch_size=args.batch_size, + seed=args.seed, + distributions=args.distribution, + subspace_dim=args.subspace_dim, + anisotropy=args.anisotropy, + outdir=str(args.outdir), + plot=args.plot, + ) + + rng = np.random.default_rng(args.seed) + summaries: list[DistributionSummary] = [] + spectra: dict[str, np.ndarray] = {} + target_alignments: dict[str, np.ndarray] = {} + + for distribution in args.distribution: + moment = estimate_second_moment(rng, config, distribution) + summary, eigenvalues, alignments = summarize_distribution( + distribution, moment, rng, args.target_samples + ) + summaries.append(summary) + spectra[distribution] = eigenvalues + target_alignments[distribution] = alignments + + write_outputs(config, summaries, spectra, target_alignments, args.outdir) + plot_paths = ( + save_plots(summaries, spectra, target_alignments, args.outdir) + if args.plot + else [] + ) + + print(f"dimension: {args.dimension}") + print(f"minimax bound 1/D: {1.0 / args.dimension:.8g}") + for summary in summaries: + print( + f"{summary.name}: " + f"lambda_min={summary.lambda_min:.8g}, " + f"lambda_max={summary.lambda_max:.8g}, " + f"random_mean={summary.random_target_mean:.8g}, " + f"gap_to_1/D={summary.isotropic_gap:.8g}" + ) + print(f"summary: {args.outdir / 'summary.json'}") + for path in plot_paths: + print(f"plot: {path}") + + +if __name__ == "__main__": + main() |