From 3f5507efe64e498752cd27812432155002fcd866 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Wed, 3 Jun 2026 07:47:36 -0500 Subject: Add time-varying kernel diagnostic --- scripts/finite_time_kernel_diagnostic.py | 291 +++++++++++++++++++++++++++++++ 1 file changed, 291 insertions(+) create mode 100644 scripts/finite_time_kernel_diagnostic.py (limited to 'scripts/finite_time_kernel_diagnostic.py') diff --git a/scripts/finite_time_kernel_diagnostic.py b/scripts/finite_time_kernel_diagnostic.py new file mode 100644 index 0000000..4827e69 --- /dev/null +++ b/scripts/finite_time_kernel_diagnostic.py @@ -0,0 +1,291 @@ +#!/usr/bin/env python3 +"""Finite-time BP/FA diagnostic with measured time-varying tangent kernels.""" + +from __future__ import annotations + +import argparse +import csv +import json +import sys +from dataclasses import asdict, dataclass +from pathlib import Path + +import numpy as np +import torch + +SCRIPT_DIR = Path(__file__).resolve().parent +if str(SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPT_DIR)) + +import downstream_capacity_sweep as dcs # noqa: E402 +from fa_tangent_kernel_capacity import pseudo_jacobian # noqa: E402 + + +@dataclass(frozen=True) +class DiagnosticConfig: + input_dim: int + output_dim: int + width: int + train_samples: int + test_samples: int + steps: int + lr: float + init_seeds: int + feedback_seeds: int + data_seed: int + feedback_scale: str + device: str + torch_threads: int + outdir: str + + +@dataclass(frozen=True) +class DiagnosticRow: + init_seed: int + feedback_seed: int + steps: int + empirical_bp_loss: float + empirical_fa_loss: float + empirical_gap: float + fixed_bp_loss: float + fixed_fa_loss: float + fixed_gap: float + tv_bp_loss: float + tv_fa_loss: float + tv_gap: float + fixed_gap_error: float + tv_gap_error: float + fixed_bp_error: float + fixed_fa_error: float + tv_bp_error: float + tv_fa_error: float + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Measure finite-time fixedK vs time-varyingK predictions." + ) + parser.add_argument("--input-dim", type=int, default=16) + parser.add_argument("--output-dim", type=int, default=4) + parser.add_argument("--width", type=int, default=64) + parser.add_argument("--train-samples", type=int, default=128) + parser.add_argument("--test-samples", type=int, default=512) + parser.add_argument("--steps", type=int, default=50) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--init-seeds", type=int, default=1) + parser.add_argument("--feedback-seeds", type=int, default=2) + parser.add_argument("--data-seed", type=int, default=123) + parser.add_argument( + "--feedback-scale", + choices=["relu", "fan-in", "unit"], + default="relu", + ) + parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu") + parser.add_argument("--torch-threads", type=int, default=8) + parser.add_argument( + "--outdir", + type=Path, + default=Path("outputs/finite_time_kernel_diagnostic"), + ) + return parser.parse_args() + + +def make_run_config(args: argparse.Namespace) -> dcs.RunConfig: + return dcs.RunConfig( + task="random", + input_dim=args.input_dim, + output_dim=args.output_dim, + teacher_rank=4, + teacher_width=64, + teacher_hidden_layers=2, + normalize_targets=False, + widths=[args.width], + train_samples=args.train_samples, + test_samples=args.test_samples, + probe_samples=8, + steps=args.steps, + lr=args.lr, + optimizer="sgd", + init_seeds=args.init_seeds, + feedback_seeds=args.feedback_seeds, + init_seed_offset=0, + feedback_seed_offset=0, + data_seed=args.data_seed, + noise_std=0.0, + feedback_scale=args.feedback_scale, + capacity_q=0.01, + jacobian_lambda_rel=1e-3, + skip_jacobian=True, + device=args.device, + torch_threads=args.torch_threads, + outdir=str(args.outdir), + plot=False, + ) + + +def loss_from_residual(residual: np.ndarray, samples: int) -> float: + return 0.5 * float(residual @ residual) / samples + + +def sgd_step( + weights: list[torch.Tensor], + x_train: torch.Tensor, + y_train: torch.Tensor, + lr: float, + feedback: list[torch.Tensor] | None, +) -> list[torch.Tensor]: + grads = dcs.gradients(weights, x_train, y_train, feedback) + return [weight - lr * grad for weight, grad in zip(weights, grads)] + + +def kernel_for( + weights: list[torch.Tensor], + x_train: torch.Tensor, + feedback: list[torch.Tensor] | None, +) -> np.ndarray: + j_bp = pseudo_jacobian(weights, x_train, feedback=None).cpu().numpy() + if feedback is None: + return j_bp @ j_bp.T + j_fa = pseudo_jacobian(weights, x_train, feedback=feedback).cpu().numpy() + return j_bp @ j_fa.T + + +def run_one( + config: dcs.RunConfig, + x_train: torch.Tensor, + y_train: torch.Tensor, + init_seed: int, + feedback_seed: int, +) -> DiagnosticRow: + initial_weights = dcs.initialize_weights(config, config.widths[0], init_seed) + feedback = dcs.init_feedback(config, config.widths[0], feedback_seed) + with torch.no_grad(): + r0 = (dcs.predict(initial_weights, x_train) - y_train).reshape(-1).cpu().numpy() + + k_bp0 = kernel_for(initial_weights, x_train, feedback=None) + k_fa0 = kernel_for(initial_weights, x_train, feedback=feedback) + + fixed_bp_residual = r0.copy() + fixed_fa_residual = r0.copy() + tv_bp_residual = r0.copy() + tv_fa_residual = r0.copy() + bp_weights = dcs.clone_weights(initial_weights) + fa_weights = dcs.clone_weights(initial_weights) + scale = config.lr / config.train_samples + + for _step in range(config.steps): + fixed_bp_residual = fixed_bp_residual - scale * (k_bp0 @ fixed_bp_residual) + fixed_fa_residual = fixed_fa_residual - scale * (k_fa0 @ fixed_fa_residual) + + k_bp_t = kernel_for(bp_weights, x_train, feedback=None) + k_fa_t = kernel_for(fa_weights, x_train, feedback=feedback) + tv_bp_residual = tv_bp_residual - scale * (k_bp_t @ tv_bp_residual) + tv_fa_residual = tv_fa_residual - scale * (k_fa_t @ tv_fa_residual) + + bp_weights = sgd_step(bp_weights, x_train, y_train, config.lr, feedback=None) + fa_weights = sgd_step(fa_weights, x_train, y_train, config.lr, feedback=feedback) + + empirical_bp_loss = dcs.mse(bp_weights, x_train, y_train) + empirical_fa_loss = dcs.mse(fa_weights, x_train, y_train) + fixed_bp_loss = loss_from_residual(fixed_bp_residual, config.train_samples) + fixed_fa_loss = loss_from_residual(fixed_fa_residual, config.train_samples) + tv_bp_loss = loss_from_residual(tv_bp_residual, config.train_samples) + tv_fa_loss = loss_from_residual(tv_fa_residual, config.train_samples) + + empirical_gap = empirical_fa_loss - empirical_bp_loss + fixed_gap = fixed_fa_loss - fixed_bp_loss + tv_gap = tv_fa_loss - tv_bp_loss + + return DiagnosticRow( + init_seed=init_seed, + feedback_seed=feedback_seed, + steps=config.steps, + empirical_bp_loss=empirical_bp_loss, + empirical_fa_loss=empirical_fa_loss, + empirical_gap=empirical_gap, + fixed_bp_loss=fixed_bp_loss, + fixed_fa_loss=fixed_fa_loss, + fixed_gap=fixed_gap, + tv_bp_loss=tv_bp_loss, + tv_fa_loss=tv_fa_loss, + tv_gap=tv_gap, + fixed_gap_error=fixed_gap - empirical_gap, + tv_gap_error=tv_gap - empirical_gap, + fixed_bp_error=fixed_bp_loss - empirical_bp_loss, + fixed_fa_error=fixed_fa_loss - empirical_fa_loss, + tv_bp_error=tv_bp_loss - empirical_bp_loss, + tv_fa_error=tv_fa_loss - empirical_fa_loss, + ) + + +def write_rows(path: Path, rows: list[DiagnosticRow]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(DiagnosticRow.__annotations__.keys())) + writer.writeheader() + for row in rows: + writer.writerow(asdict(row)) + + +def main() -> None: + args = parse_args() + if args.torch_threads > 0: + torch.set_num_threads(args.torch_threads) + config = make_run_config(args) + x_train, y_train, _x_test, _y_test, _x_probe, _teacher = dcs.make_data(config) + + rows: list[DiagnosticRow] = [] + total = args.init_seeds * args.feedback_seeds + count = 0 + for init_index in range(args.init_seeds): + init_seed = 10_000 + init_index + for feedback_index in range(args.feedback_seeds): + feedback_seed = 100_000 + init_index * 1000 + feedback_index + count += 1 + print( + f"[{count}/{total}] init={init_seed} feedback={feedback_seed}", + flush=True, + ) + rows.append(run_one(config, x_train, y_train, init_seed, feedback_seed)) + + outdir = Path(args.outdir) + outdir.mkdir(parents=True, exist_ok=True) + write_rows(outdir / "finite_time_kernel_rows.csv", rows) + payload = { + "config": asdict( + DiagnosticConfig( + input_dim=args.input_dim, + output_dim=args.output_dim, + width=args.width, + train_samples=args.train_samples, + test_samples=args.test_samples, + steps=args.steps, + lr=args.lr, + init_seeds=args.init_seeds, + feedback_seeds=args.feedback_seeds, + data_seed=args.data_seed, + feedback_scale=args.feedback_scale, + device=args.device, + torch_threads=args.torch_threads, + outdir=str(outdir), + ) + ), + "rows": [asdict(row) for row in rows], + } + (outdir / "summary.json").write_text(json.dumps(payload, indent=2) + "\n") + + fixed_errors = np.array([row.fixed_gap_error for row in rows]) + tv_errors = np.array([row.tv_gap_error for row in rows]) + print(f"rows: {outdir / 'finite_time_kernel_rows.csv'}") + print( + "fixed gap error mean=" + f"{fixed_errors.mean():.6g}, mae={np.mean(np.abs(fixed_errors)):.6g}" + ) + print( + "time-varying gap error mean=" + f"{tv_errors.mean():.6g}, mae={np.mean(np.abs(tv_errors)):.6g}" + ) + + +if __name__ == "__main__": + main() -- cgit v1.2.3