diff --git a/baseline/experiments/mnist_mlp3_tangent_rg/README.md b/baseline/experiments/mnist_mlp3_tangent_rg/README.md index 8326ab4..f288200 100644 --- a/baseline/experiments/mnist_mlp3_tangent_rg/README.md +++ b/baseline/experiments/mnist_mlp3_tangent_rg/README.md @@ -853,3 +853,18 @@ Jacobians and adds the exact ECS cover variants for `fc1.weight` whenever the same-checkpoint WeightWatcher/trace rank records certify them. The primary PL fit uses no post-hoc top-mode search (`--top-k 0`); explicit sensitivity values can be requested, for example, with `--top-k 0,1,2,3,4,5`. + +For the reduced experiment, use: + +```bash +bash baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_jacobians_reduced.sh +``` + +This separate preset retains the centered log-singular radial Jacobian for all +three layers and both ECS/Grassmann covers for `fc1.weight`. It omits the large +ambient polar, Gram, log-Gram, and finite-NS5 spectra. For ECS, the uniform +`q-k` copies of each `2/sigma_i` core amplitude are represented once. This +preserves the empirical distribution, fitted alpha, selected xmin, and KS +distance while preventing deterministic coordinate copies from inflating the +effective sample size. Outputs default to +`/private/tmp/rg-mnist-mlp3-short100-jacobians-reduced`. diff --git a/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_jacobians_cli.py b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_jacobians_cli.py index 77e7176..962f4b1 100755 --- a/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_jacobians_cli.py +++ b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_jacobians_cli.py @@ -407,6 +407,38 @@ def fit_spectrum( return rows +def ecs_fit_amplitudes(record: Any, *, compress_groups: bool) -> tuple[np.ndarray, dict[str, Any]]: + """Optionally replace uniform ECS coordinate copies by physical groups. + + Each retained-core amplitude ``2/sigma_i`` is repeated ``q-k`` times in + the ambient Jacobian. Uniform repetition leaves the empirical CDF, MLE + alpha, package-selected xmin, and KS distance unchanged. Compression avoids + treating deterministic copies as independent observations, so the reported + uncertainty is based on physical core groups. + """ + + multiplicity = int(record.deterministic_shell_multiplicity) + expanded = np.asarray(record.singular_amplitudes, dtype=float) + if not compress_groups: + return expanded, { + "ecs_fit_observation_unit": "expanded_jacobian_mode", + "ecs_uniform_group_multiplicity": multiplicity, + "ecs_groups_compressed": False, + } + core = 2.0 / np.asarray(record.retained_singular_values, dtype=float) + return core, { + "ecs_fit_observation_unit": "physical_retained_core_amplitude_group", + "ecs_uniform_group_multiplicity": multiplicity, + "ecs_groups_compressed": True, + "ecs_expanded_mode_count": int(expanded.size), + "ecs_physical_group_count": int(core.size), + "ecs_compression_invariance": ( + "uniform replication removed; empirical CDF, alpha, xmin, and KS D " + "are invariant; sigma reflects physical groups" + ), + } + + def safe_slug(text: str) -> str: return "".join(char if char.isalnum() or char in "-_" else "_" for char in text) @@ -475,6 +507,11 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--optimizers", default=",".join(DEFAULT_OPTIMIZERS)) parser.add_argument("--seeds", default="101") parser.add_argument("--layers", default=",".join(DEFAULT_LAYERS)) + parser.add_argument( + "--methods", + default=",".join(BASE_METHODS), + help="comma-separated base Jacobians; ECS covers are controlled separately", + ) parser.add_argument("--epoch-stride", type=int, default=10) parser.add_argument("--maximum-checkpoints", type=int) parser.add_argument("--top-k", default="0", help="PL clipping sensitivities; default 0 only") @@ -482,6 +519,12 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--ns-steps", type=int, default=5) parser.add_argument("--ns-eps", type=float, default=1e-7) parser.add_argument("--skip-ecs", action="store_true") + parser.add_argument( + "--compress-ecs-groups", + action=argparse.BooleanOptionalAction, + default=False, + help="fit one physical ECS amplitude per uniformly repeated shell group", + ) parser.add_argument("--resume", action=argparse.BooleanOptionalAction, default=True) parser.add_argument("--fail-fast", action="store_true") parser.add_argument("--verbose", action="store_true") @@ -495,9 +538,14 @@ def run(args: argparse.Namespace) -> int: fit_path = output_root / "jacobian_powerlaw_fits.csv" operator_path = output_root / "jacobian_operators.csv" error_path = output_root / "errors.csv" + completion_path = output_root / "completed_checkpoints.csv" optimizers = parse_csv_values(args.optimizers) seeds = parse_csv_values(args.seeds, int) layers = parse_csv_values(args.layers) + methods = parse_csv_values(args.methods) + unknown_methods = set(methods) - set(BASE_METHODS) + if unknown_methods: + raise ValueError(f"unknown --methods values: {sorted(unknown_methods)}") top_k_values = parse_csv_values(args.top_k, int) if top_k_values[0] != 0 or any(value < 0 for value in top_k_values): raise ValueError("--top-k must begin with 0 and contain nonnegative integers") @@ -507,12 +555,15 @@ def run(args: argparse.Namespace) -> int: fit_rows: list[dict[str, Any]] = [] operator_rows: list[dict[str, Any]] = [] errors: list[dict[str, Any]] = [] + completion_rows: list[dict[str, Any]] = [] if args.resume and fit_path.is_file(): fit_rows = pd.read_csv(fit_path).to_dict(orient="records") if args.resume and operator_path.is_file(): operator_rows = pd.read_csv(operator_path).to_dict(orient="records") if args.resume and error_path.is_file(): errors = pd.read_csv(error_path).to_dict(orient="records") + if args.resume and completion_path.is_file(): + completion_rows = pd.read_csv(completion_path).to_dict(orient="records") identities: dict[tuple[str, int], dict[str, Any]] = {} work: list[tuple[str, int, str, Any]] = [] @@ -542,18 +593,20 @@ def run(args: argparse.Namespace) -> int: for optimizer, seed, layer, ref in work: unit_key = (optimizer, int(seed), layer, int(ref.epoch), int(ref.global_step)) - already = { - str(row.get("method")) - for row in fit_rows - if ( + completed_before = any( + ( str(row.get("optimizer")), int(row.get("seed", -1)), str(row.get("layer")), int(row.get("epoch", -1)), int(row.get("global_step", -1)), ) == unit_key - and str(row.get("spectrum_kind")) == "amplitude" - and int(row.get("clip_top_k", -1)) == 0 - } - if args.resume and set(BASE_METHODS).issubset(already): + and str(row.get("base_methods_requested", "")) == ",".join(methods) + and str(row.get("ecs_groups_compressed", "")).strip().lower() + == str(bool(args.compress_ecs_groups)).lower() + and str(row.get("ecs_skipped", "")).strip().lower() + == str(bool(args.skip_ecs)).lower() + for row in completion_rows + ) + if args.resume and completed_before: completed += 1 logger.info("SKIP completed %d/%d %s", completed, total, unit_key) continue @@ -573,6 +626,7 @@ def same_unit(row: dict[str, Any]) -> bool: fit_rows = [row for row in fit_rows if not same_unit(row)] operator_rows = [row for row in operator_rows if not same_unit(row)] errors = [row for row in errors if not same_unit(row)] + completion_rows = [row for row in completion_rows if not same_unit(row)] unit_started = time.perf_counter() logger.info( @@ -590,7 +644,8 @@ def same_unit(row: dict[str, Any]) -> bool: ) singular = np.linalg.svd(weight, compute_uv=False) method_factories: dict[str, tuple[np.ndarray, Any]] = {} - for method in BASE_METHODS: + method_metadata: dict[str, dict[str, Any]] = {} + for method in methods: build_started = time.perf_counter() logger.info("JACOBIAN START method=%s", method) method_factories[method] = compute_one_spectrum( @@ -614,8 +669,20 @@ def same_unit(row: dict[str, Any]) -> bool: weight, retained_rank=k, outer_rank=q, rcond=args.ns_eps, precomputed_singular_values=singular, ) - method_factories[method] = (cover.singular_amplitudes, cover) - logger.info("ECS ranks method=%s k=%d q=%d", method, k, q) + amplitudes, compression_metadata = ecs_fit_amplitudes( + cover, compress_groups=args.compress_ecs_groups + ) + method_factories[method] = (amplitudes, cover) + method_metadata[method] = { + **rank_metadata, + **compression_metadata, + } + logger.info( + "ECS ranks method=%s k=%d q=%d fit_amplitudes=%d " + "expanded_modes=%d compressed=%s", + method, k, q, len(amplitudes), int(cover.derivative_rank), + args.compress_ecs_groups, + ) base = { "optimizer": optimizer, @@ -629,13 +696,18 @@ def same_unit(row: dict[str, Any]) -> bool: for method, (amplitudes, record) in method_factories.items(): method_started = time.perf_counter() logger.info("METHOD START method=%s n_amplitudes=%d", method, len(amplitudes)) - metadata = {**base, "method": method} + metadata = {**base, **method_metadata.get(method, {}), "method": method} rows = fit_spectrum( np.asarray(amplitudes, dtype=float), record, metadata, top_k_values, args.minimum_tail, ) fit_rows.extend(rows) - operator_rows.append({**base, "method": method, **record_row(record)}) + operator_rows.append({ + **base, + **method_metadata.get(method, {}), + "method": method, + **record_row(record), + }) spectra[method] = np.asarray(amplitudes, dtype=float) elapsed_method = time.perf_counter() - method_started primary = next( @@ -678,6 +750,15 @@ def same_unit(row: dict[str, Any]) -> bool: ] alpha_path = output_root / "plots" / "alpha_progress" / f"{optimizer}_{safe_slug(layer)}.png" save_alpha_progress(block_rows, alpha_path, f"{optimizer} seed {seed} {layer}") + completion_rows.append({ + **base, + "methods": ",".join(method_factories), + "base_methods_requested": ",".join(methods), + "ecs_groups_compressed": bool(args.compress_ecs_groups), + "ecs_skipped": bool(args.skip_ecs), + "completed_at_utc": utc_now(), + }) + atomic_csv(completion_path, completion_rows) completed += 1 elapsed = time.perf_counter() - started eta = elapsed / completed * (total - completed) if completed else None diff --git a/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_jacobians_reduced.sh b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_jacobians_reduced.sh new file mode 100755 index 0000000..f94b34e --- /dev/null +++ b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_jacobians_reduced.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash + +# Fast scientific preset: the scale-quotiented radial Jacobian on every layer +# plus both ECS/Grassmann covers on fc1. ECS deterministic shell copies are +# compressed to physical retained-core groups before the power-law fit. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)" +export RG_MNIST_JACOBIAN_CLI_OUTPUT_ROOT="${RG_MNIST_REDUCED_JACOBIAN_OUTPUT_ROOT:-/private/tmp/rg-mnist-mlp3-short100-jacobians-reduced}" + +exec bash "${SCRIPT_DIR}/run_short100_jacobians_cli.sh" \ + --methods centered_log_singular_radial_pullback \ + --compress-ecs-groups \ + --top-k 0 \ + "$@" diff --git a/baseline/tests/tangent_rg/test_command_line_jacobians.py b/baseline/tests/tangent_rg/test_command_line_jacobians.py index 8cd2281..f1149e6 100644 --- a/baseline/tests/tangent_rg/test_command_line_jacobians.py +++ b/baseline/tests/tangent_rg/test_command_line_jacobians.py @@ -2,6 +2,7 @@ import importlib.util from pathlib import Path +from types import SimpleNamespace import numpy as np @@ -44,3 +45,17 @@ def test_cli_duration_and_argument_defaults_are_observable_tmp_paths(): assert str(args.output_root).startswith("/private/tmp/") assert args.epoch_stride == 10 assert args.top_k == "0" + + +def test_ecs_group_compression_removes_only_uniform_coordinate_copies(): + module = load_cli_module() + record = SimpleNamespace( + deterministic_shell_multiplicity=4, + retained_singular_values=np.array([2.0, 4.0, 8.0]), + singular_amplitudes=np.repeat(np.array([1.0, 0.5, 0.25]), 4), + ) + compressed, metadata = module.ecs_fit_amplitudes(record, compress_groups=True) + assert np.allclose(compressed, [1.0, 0.5, 0.25]) + assert metadata["ecs_groups_compressed"] is True + assert metadata["ecs_uniform_group_multiplicity"] == 4 + assert metadata["ecs_expanded_mode_count"] == 12