diff --git a/baseline/experiments/mnist_mlp3_tangent_rg/README.md b/baseline/experiments/mnist_mlp3_tangent_rg/README.md index f2d3254..eda9a7c 100644 --- a/baseline/experiments/mnist_mlp3_tangent_rg/README.md +++ b/baseline/experiments/mnist_mlp3_tangent_rg/README.md @@ -849,8 +849,9 @@ bash baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_jacobians_c The default run analyzes seed 101 for MuonClip-RMS and AdamW at epochs 10,20,...,100. It computes all five universally defined analytic weight-only -Jacobians and adds the exact ECS cover variants for `fc1.weight` whenever the -same-checkpoint WeightWatcher/trace rank records certify them. The primary PL +Jacobians and adds the exact ECS cover variants for `fc1.weight` and +`fc2.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`. @@ -861,7 +862,9 @@ bash baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_jacobians_r ``` 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 +three layers and both ECS/Grassmann covers for `fc1.weight` and `fc2.weight`. +`fc3.weight` remains radial-only because its 10-row output geometry is treated +as a small-rank diagnostic. The preset 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 @@ -883,4 +886,7 @@ The report combines MuonClip-RMS and AdamW Jacobian alpha trajectories, KS and tail-support diagnostics, raw and `fix_fingers=clip_xmax` WeightWatcher controls, train/test accuracy and loss, alpha-versus-test-accuracy plots, selected spectral CCDF galleries, analysis-ready CSVs, and a browsable HTML -index. Report generation does not recompute any Jacobian. +index. It includes an expected-versus-observed method coverage audit and a +dedicated FC1/FC2 ECS plot whose marker positions are slightly offset for +visibility when full-row and detX fits coincide; the underlying epoch values +are never changed. Report generation does not recompute any Jacobian. diff --git a/baseline/experiments/mnist_mlp3_tangent_rg/scripts/build_short100_jacobian_report.py b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/build_short100_jacobian_report.py index d7923c5..30490de 100755 --- a/baseline/experiments/mnist_mlp3_tangent_rg/scripts/build_short100_jacobian_report.py +++ b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/build_short100_jacobian_report.py @@ -39,6 +39,15 @@ "ecs_grassmann_cartan_cover_full_row_shell_pullback": ("--", "s"), "ecs_grassmann_cartan_cover_detx_shell_pullback": (":", "^") } +ECS_METHODS = ( + "ecs_grassmann_cartan_cover_full_row_shell_pullback", + "ecs_grassmann_cartan_cover_detx_shell_pullback", +) +EXPECTED_METHODS_BY_LAYER = { + "fc1.weight": ("centered_log_singular_radial_pullback", *ECS_METHODS), + "fc2.weight": ("centered_log_singular_radial_pullback", *ECS_METHODS), + "fc3.weight": ("centered_log_singular_radial_pullback",), +} def configure_logging(report_root: Path) -> logging.Logger: @@ -129,12 +138,95 @@ def plot_jacobian_metric( if reference is not None: axis.axhline(reference, color="black", linestyle=(0, (1, 2)), linewidth=1.2) axis.set(title=layer.replace(".weight", ""), xlabel="epoch", ylabel=ylabel) + available = [ + METHOD_LABELS.get(method, method) + for method in EXPECTED_METHODS_BY_LAYER[layer] + if subset["method"].astype(str).eq(method).any() + ] + axis.text( + 0.01, 0.02, "present: " + ", ".join(available), + transform=axis.transAxes, fontsize=6.5, color="#555555", + ha="left", va="bottom", + ) axis.grid(True, alpha=0.25) shared_legend(fig, axes, columns=3) fig.suptitle(title) return save_figure(fig, path, bottom=0.16, top=0.94) +def plot_ecs_quotient_comparison(fits: pd.DataFrame, path: Path) -> Path: + """Plot ECS-only alpha for FC1/FC2 while exposing coincident fits. + + Lines remain at the exact analysis epochs. Marker locations receive a tiny + horizontal display-only offset, so equal full-row and detX values remain + visible instead of one marker painting over the other. + """ + fig, axes = plt.subplots(1, 2, figsize=(13.2, 4.9), sharex=True) + marker_offsets = {ECS_METHODS[0]: -0.45, ECS_METHODS[1]: 0.45} + for axis, layer in zip(axes, ("fc1.weight", "fc2.weight")): + subset = fits[ + fits["layer"].astype(str).eq(layer) + & fits["method"].astype(str).isin(ECS_METHODS) + ] + for (optimizer, method), curve in subset.groupby(["optimizer", "method"]): + curve = curve.sort_values("epoch") + linestyle, marker = METHOD_STYLES[str(method)] + label = ( + f"{OPTIMIZER_LABELS.get(str(optimizer), str(optimizer))} — " + f"{METHOD_LABELS[str(method)]}" + ) + x = pd.to_numeric(curve["epoch"], errors="coerce").to_numpy(float) + y = pd.to_numeric(curve["alpha"], errors="coerce").to_numpy(float) + axis.plot( + x, y, color=COLORS[str(optimizer)], linestyle=linestyle, + linewidth=1.8, label=label, + ) + axis.scatter( + x + marker_offsets[str(method)], y, color=COLORS[str(optimizer)], + marker=marker, s=28, zorder=4, + ) + axis.axhline(2.0, color="black", linestyle=(0, (1, 2)), linewidth=1.2) + axis.set( + title=layer.replace(".weight", ""), xlabel="epoch", + ylabel="ECS Jacobian energy alpha", + ) + axis.grid(True, alpha=0.25) + shared_legend(fig, axes, columns=2) + fig.suptitle( + "ECS quotient comparison — lines use exact epochs; markers offset ±0.45 for visibility" + ) + return save_figure(fig, path, bottom=0.18, top=0.92) + + +def build_method_coverage(primary: pd.DataFrame) -> pd.DataFrame: + """Return an explicit expected-versus-observed method inventory.""" + rows = [] + for optimizer in OPTIMIZERS: + for layer in LAYERS: + layer_rows = primary[ + primary["optimizer"].astype(str).eq(optimizer) + & primary["layer"].astype(str).eq(layer) + ] + expected = EXPECTED_METHODS_BY_LAYER[layer] + for method in expected: + observed = layer_rows[layer_rows["method"].astype(str).eq(method)] + epochs = sorted( + pd.to_numeric(observed["epoch"], errors="coerce") + .dropna().astype(int).unique() + ) + rows.append({ + "optimizer": optimizer, + "layer": layer, + "method": method, + "method_label": METHOD_LABELS.get(method, method), + "expected_epoch_count": 10, + "observed_epoch_count": len(epochs), + "observed_epochs": ",".join(map(str, epochs)), + "coverage_status": "complete" if len(epochs) == 10 else "INCOMPLETE", + }) + return pd.DataFrame(rows) + + def plot_fit_quality(fits: pd.DataFrame, path: Path) -> Path: fig, axes = plt.subplots(2, 3, figsize=(17.0, 8.2), sharex=True) for column, layer in enumerate(LAYERS): @@ -289,6 +381,7 @@ def build_html( figures: list[Path], tables: list[Path], primary: pd.DataFrame, + coverage: pd.DataFrame, inventory: dict[str, int], ) -> Path: final = primary.sort_values("epoch").groupby( @@ -297,6 +390,7 @@ def build_html( final_table = final[ ["optimizer", "epoch", "layer", "method", "alpha", "ks_D", "n_tail", "tail_decades", "fit_ok"] ].to_html(index=False, float_format=lambda value: f"{value:.5g}") + coverage_table = coverage.to_html(index=False) figure_html = "\n".join( f'

{escape(path.stem.replace("_", " ").title())}

' f'' @@ -320,7 +414,34 @@ def build_html( code{{background:#f3f3f3;padding:2px 4px}} .note{{background:#eef6ff;padding:14px;border-left:4px solid #0072B2}}

Short100 reduced Jacobian analysis

-

MuonClip-RMS versus AdamW, seed 101, epochs 10–100. ECS curves use one physical retained-core amplitude per uniformly repeated shell group. The complete mode-level data, fit tables, WeightWatcher controls, and performance data are linked below.

+

Scope. MuonClip-RMS versus AdamW, seed 101, +epochs 10,20,…,100. The centered log-singular radial Jacobian is evaluated on +FC1, FC2, and FC3. Both exact ECS quotient covers are evaluated on FC1 and FC2. +FC3 is intentionally radial-only because its 10-row output geometry is a +small-rank diagnostic rather than the large-layer ECS comparison.

+

What each curve means

+ +

ECS fits use one physical retained-core amplitude for each uniformly repeated +shell group. The omitted coordinate copies have identical values: compressing +them leaves the empirical spectral shape, α, xmin, and KS distance unchanged, +while avoiding a fictitiously large independent sample size.

+

How to read coincident curves. Full-row and detX ECS fits can be +exactly equal. In the all-method figure one line may cover another. The dedicated +ECS figure keeps lines at the true epochs but offsets full-row and detX markers by +−0.45 and +0.45 epoch solely for visibility. The CSV retains the exact epochs.

+

Coverage audit

+

Every expected method should have ten observations. Any row marked +INCOMPLETE means computation is genuinely missing; visual overlap is +not classified as missing.

{coverage_table}

Data inventory

Analysis-ready tables

Final checkpoint primary fits

{final_table} @@ -365,6 +486,7 @@ def main() -> int: fits["spectrum_kind"].astype(str).eq("energy_derived_from_amplitude") & pd.to_numeric(fits["clip_top_k"], errors="coerce").eq(0) ].copy() + primary = primary[pd.to_numeric(primary["epoch"], errors="coerce").gt(0)] performance_frames = [] weightwatcher_frames = [] for optimizer in OPTIMIZERS: @@ -380,6 +502,7 @@ def main() -> int: performance = performance.sort_values(["optimizer", "epoch"]) table_root.mkdir(parents=True, exist_ok=True) + coverage = build_method_coverage(primary) tables = [ table_root / "jacobian_primary_energy_fits.csv", table_root / "jacobian_all_fits.csv", @@ -387,6 +510,7 @@ def main() -> int: table_root / "jacobian_operator_metadata.csv", table_root / "weightwatcher_raw_and_clip_xmax.csv", table_root / "performance_train_test.csv", + table_root / "jacobian_method_coverage.csv", ] primary.to_csv(tables[0], index=False) fits.to_csv(tables[1], index=False) @@ -394,6 +518,7 @@ def main() -> int: operators.to_csv(tables[3], index=False) weightwatcher.to_csv(tables[4], index=False) performance.to_csv(tables[5], index=False) + coverage.to_csv(tables[6], index=False) figures = [ plot_jacobian_metric( @@ -401,6 +526,9 @@ def main() -> int: "MuonClip-RMS versus AdamW Jacobian alpha", figure_root / "01_jacobian_alpha_comparison.png", reference=2.0, ), + plot_ecs_quotient_comparison( + primary, figure_root / "01b_ecs_fc1_fc2_alpha_comparison.png" + ), plot_fit_quality(primary, figure_root / "02_jacobian_fit_quality.png"), plot_jacobian_metric( primary, "n_tail", "package-selected tail modes", @@ -425,7 +553,7 @@ def main() -> int: "performance rows": len(performance), "figures": len(figures), } - index = build_html(report_root, figures, tables, primary, inventory) + index = build_html(report_root, figures, tables, primary, coverage, inventory) (report_root / "report_manifest.json").write_text( json.dumps({ "analysis_root": str(analysis_root), 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 1dee42e..de8213a 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 @@ -38,6 +38,7 @@ DEFAULT_OUTPUT_ROOT = Path("/private/tmp/rg-mnist-mlp3-short100-jacobians") DEFAULT_OPTIMIZERS = ("muonclip_rms", "adamw") DEFAULT_LAYERS = ("fc1.weight", "fc2.weight", "fc3.weight") +DEFAULT_ECS_LAYERS = ("fc1.weight", "fc2.weight") BASE_METHODS = ( "polar_pullback", "normalized_gram_pullback", @@ -507,6 +508,14 @@ 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( + "--ecs-layers", + default=",".join(DEFAULT_ECS_LAYERS), + help=( + "comma-separated layers receiving the two exact ECS/Grassmann " + "cover analyses; defaults to fc1.weight,fc2.weight" + ), + ) parser.add_argument( "--methods", default=",".join(BASE_METHODS), @@ -543,6 +552,7 @@ def run(args: argparse.Namespace) -> int: optimizers = parse_csv_values(args.optimizers) seeds = parse_csv_values(args.seeds, int) layers = parse_csv_values(args.layers) + ecs_layers = parse_csv_values(args.ecs_layers) methods = parse_csv_values(args.methods) unknown_methods = set(methods) - set(BASE_METHODS) if unknown_methods: @@ -552,6 +562,12 @@ def run(args: argparse.Namespace) -> int: raise ValueError("--top-k must begin with 0 and contain nonnegative integers") if args.epoch_stride < 1: raise ValueError("--epoch-stride must be positive") + unknown_ecs_layers = set(ecs_layers) - set(layers) + if unknown_ecs_layers: + raise ValueError( + "--ecs-layers must be a subset of --layers; unknown values: " + f"{sorted(unknown_ecs_layers)}" + ) fit_rows: list[dict[str, Any]] = [] operator_rows: list[dict[str, Any]] = [] @@ -575,6 +591,11 @@ def run(args: argparse.Namespace) -> int: logger.info("run_root=%s", args.run_root.resolve()) logger.info("cache_root=%s", args.cache_root.resolve()) logger.info("output_root=%s", output_root) + logger.info("base_methods=%s", methods) + logger.info( + "ECS coverage=%s (full-row and detX covers); other layers are radial-only", + ecs_layers if not args.skip_ecs else "disabled", + ) logger.info("Preflight is lightweight: selected checkpoints are validated when loaded") for optimizer in optimizers: for seed in seeds: @@ -610,14 +631,22 @@ def run(args: argparse.Namespace) -> int: == str(bool(args.skip_ecs)).lower() for row in completion_rows ) - spectrum_data_available = any( - ( + observed_methods = { + str(row.get("method")) + for row in spectrum_rows + if ( 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 - for row in spectrum_rows - ) + } + expected_methods = set(methods) + if not args.skip_ecs and layer in ecs_layers: + expected_methods.update({ + "ecs_grassmann_cartan_cover_full_row_shell_pullback", + "ecs_grassmann_cartan_cover_detx_shell_pullback", + }) + spectrum_data_available = expected_methods.issubset(observed_methods) if args.resume and completed_before and spectrum_data_available: completed += 1 logger.info("SKIP completed %d/%d %s", completed, total, unit_key) @@ -670,14 +699,28 @@ def same_unit(row: dict[str, Any]) -> bool: method, time.perf_counter() - build_started, len(method_factories[method][0]), ) - if not args.skip_ecs and layer == "fc1.weight": + if not args.skip_ecs and layer in ecs_layers: from rg_baselines.tangent_rg import single_checkpoint numerical_rank = int(np.count_nonzero(singular > args.ns_eps * singular[0])) - for method, k, q, rank_metadata in exact_ecs_ranks( + ecs_rank_records = exact_ecs_ranks( identity["seed_dir"], optimizer, seed, int(ref.epoch), int(ref.global_step), layer, numerical_rank, - ): + ) + observed_ecs_methods = {record[0] for record in ecs_rank_records} + missing_ecs_methods = set({ + "ecs_grassmann_cartan_cover_full_row_shell_pullback", + "ecs_grassmann_cartan_cover_detx_shell_pullback", + }) - observed_ecs_methods + if missing_ecs_methods: + raise RuntimeError( + f"requested ECS analysis is unavailable for {layer} at " + f"epoch {int(ref.epoch)}; missing {sorted(missing_ecs_methods)}. " + "The exact same-checkpoint clip_xmax WeightWatcher fit and " + "certified trace-log rank audit must both exist and define " + "nonempty full-row and detX quotient shells." + ) + for method, k, q, rank_metadata in ecs_rank_records: cover = single_checkpoint.ecs_grassmann_cover_analytic_spectrum( weight, retained_rank=k, outer_rank=q, rcond=args.ns_eps, precomputed_singular_values=singular, @@ -791,6 +834,10 @@ def same_unit(row: dict[str, Any]) -> bool: "base_methods_requested": ",".join(methods), "ecs_groups_compressed": bool(args.compress_ecs_groups), "ecs_skipped": bool(args.skip_ecs), + "ecs_requested_for_layer": bool( + not args.skip_ecs and layer in ecs_layers + ), + "ecs_layers_requested": ",".join(ecs_layers), "completed_at_utc": utc_now(), }) atomic_csv(completion_path, completion_rows) 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 index f94b34e..3461bba 100755 --- 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 @@ -1,7 +1,7 @@ #!/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 +# plus both ECS/Grassmann covers on fc1 and fc2. ECS deterministic shell copies are # compressed to physical retained-core groups before the power-law fit. set -euo pipefail @@ -11,6 +11,7 @@ export RG_MNIST_JACOBIAN_CLI_OUTPUT_ROOT="${RG_MNIST_REDUCED_JACOBIAN_OUTPUT_ROO exec bash "${SCRIPT_DIR}/run_short100_jacobians_cli.sh" \ --methods centered_log_singular_radial_pullback \ + --ecs-layers fc1.weight,fc2.weight \ --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 f1149e6..4d77dcb 100644 --- a/baseline/tests/tangent_rg/test_command_line_jacobians.py +++ b/baseline/tests/tangent_rg/test_command_line_jacobians.py @@ -45,6 +45,7 @@ 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" + assert args.ecs_layers == "fc1.weight,fc2.weight" def test_ecs_group_compression_removes_only_uniform_coordinate_copies(): diff --git a/baseline/tests/tangent_rg/test_short100_jacobian_report.py b/baseline/tests/tangent_rg/test_short100_jacobian_report.py index af319f6..6b0dfdf 100644 --- a/baseline/tests/tangent_rg/test_short100_jacobian_report.py +++ b/baseline/tests/tangent_rg/test_short100_jacobian_report.py @@ -4,6 +4,7 @@ from pathlib import Path import numpy as np +import pandas as pd SCRIPT = ( @@ -37,3 +38,22 @@ def test_report_defaults_to_reduced_tmp_output(): "/private/tmp/rg-mnist-mlp3-short100-jacobians-reduced" ) assert args.seed == 101 + + +def test_method_coverage_requires_ecs_on_fc1_and_fc2_but_not_fc3(): + module = load_report_module() + rows = [] + for optimizer in module.OPTIMIZERS: + for layer, methods in module.EXPECTED_METHODS_BY_LAYER.items(): + for method in methods: + for epoch in range(10, 101, 10): + rows.append({ + "optimizer": optimizer, "layer": layer, + "method": method, "epoch": epoch, + }) + coverage = module.build_method_coverage(pd.DataFrame(rows)) + assert coverage["coverage_status"].eq("complete").all() + fc2 = coverage[coverage["layer"].eq("fc2.weight")] + fc3 = coverage[coverage["layer"].eq("fc3.weight")] + assert set(fc2["method"]) == set(module.EXPECTED_METHODS_BY_LAYER["fc2.weight"]) + assert set(fc3["method"]) == {"centered_log_singular_radial_pullback"}