From 31e6bd9b6cd825596d8833fff310b12a01bbdd2d Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Fri, 21 Aug 2026 19:43:34 -0700 Subject: [PATCH] Add auditable Jacobian regularization searches --- .../mnist_mlp3_tangent_rg/README.md | 62 ++- .../scripts/build_short100_jacobian_report.py | 299 +++++++++++++- .../run_short100_complete_rg_analysis.sh | 47 ++- .../scripts/run_short100_jacobians_cli.py | 373 ++++++++++++++++-- .../scripts/run_short100_jacobians_reduced.sh | 7 +- .../scripts/run_short100_quotient_flow_cli.py | 139 ++++++- .../tangent_rg/test_command_line_jacobians.py | 52 +++ .../test_short100_jacobian_report.py | 59 ++- 8 files changed, 956 insertions(+), 82 deletions(-) diff --git a/baseline/experiments/mnist_mlp3_tangent_rg/README.md b/baseline/experiments/mnist_mlp3_tangent_rg/README.md index 06feb01..df2521e 100644 --- a/baseline/experiments/mnist_mlp3_tangent_rg/README.md +++ b/baseline/experiments/mnist_mlp3_tangent_rg/README.md @@ -872,9 +872,13 @@ distance while preventing deterministic coordinate copies from inflating the effective sample size. Outputs default to `/private/tmp/rg-mnist-mlp3-short100-jacobians-reduced`. -The reduced runner also persists every fitted spectral observation to +The reduced runner persists every reported spectral observation to `jacobian_spectra.csv`; one row contains the amplitude, squared Gram eigenvalue, physical observation unit, and represented uniform multiplicity. +For the Tikhonov grid, every candidate fit is kept in +`jacobian_hyperparameter_search.csv`, while only the selected candidate's full +mode spectrum is copied into `jacobian_spectra.csv`. This avoids a many-million +row duplication without discarding the grid-search audit. After the reduced run completes, build the complete static comparison report: ```bash @@ -897,17 +901,25 @@ For the complete notebook-free state/flow/local-response experiment, run: bash baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_complete_rg_analysis.sh ``` -This command executes three scientifically distinct analyses before rebuilding -the static HTML report: +This command prints and times four stages: exact single-checkpoint Jacobians, +transformed-weight quotients, checkpoint flow/transport, and the static HTML +report. Each completed checkpoint or quotient profile is saved atomically, so +rerunning the same command resumes rather than starts over. Use `Ctrl-C` to +interrupt; `Ctrl-Z` suspends the foreground job and is not the stop command. + +The command executes three scientifically distinct analyses before rebuilding +the report: 1. **Weight-state quotient representatives.** On FC1 and FC2, it fixes the midpoint ECS rank from the independently recorded `clip_xmax`/detX audit, chooses the rectangular-diagonal canonical section of the two-sided - `O(m) x O(n)` orbit, and materializes three declared representatives: + `O(m) x O(n)` orbit, and materializes four declared representative families: midpoint truncation, the nonlinear Gram counterterm `lambda -> max(lambda-tau,0)` scanned at - `tau/lambda_boundary in {0.25,0.50,0.75}`, and an - epoch-10-anchor-frozen Feshbach/Schur downfolding with ridge ratio `1e-2`. + `tau/lambda_boundary in {0.10,0.25,0.50,0.75,0.90}`, an + epoch-10-anchor-frozen Feshbach/Schur downfolding with ridge ratios + `{1e-3,1e-2,1e-1}`, and an MP-calibrated optimal Frobenius shrinker with + noise-scale multipliers `{0.75,1.00,1.25}`. Every materialized `W'` is passed through WeightWatcher both raw and with `fix_fingers=clip_xmax`; this phase writes `weight_quotient_weightwatcher_fits.csv`, `weight_quotient_spectra.csv`, and @@ -924,14 +936,48 @@ the static HTML report: `two_checkpoint_jacobian_transport.csv`. 3. **Single-checkpoint Jacobians.** In addition to the centered log-singular radial and exact ECS-cover derivatives, the reduced CLI evaluates the - gap-aware projector, trace-free log Gram, trace-free ridge resolvent, and - Feshbach trace-free log derivatives on the detX shell for FC1 and FC2. + gap-aware projector, trace-free log Gram, a grid-selected Tikhonov + resolvent, Feshbach trace-free log, an optimal-MP hard signal-space + projector, and trace-free log Gram restricted to that MP signal space on + FC1 and FC2. Square FC2 is handled by the same right-singular top-k Grassmann geometry as wide FC1. + The Tikhonov map is + `R_z(W)=Pi_tf[(V_o^T W^T W V_o + z I)^-1]` on the frozen detX space, with + `z/lambda_boundary in {0.03,0.10,0.30,1,3}`. Every fixed-`z` derivative is + exact. The selected curve is the valid PL fit with minimum KS distance, + breaking ties by larger tail span, larger tail count, and then grid order. + Distance of alpha from 2 is never used. Because selection uses the same + spectrum being reported, it is exploratory; all candidates are retained for + a later held-out-checkpoint confirmation. + + The MP signal rank uses the empirical median Gram eigenvalue, the numerical + Marchenko-Pastur median, and the Gavish-Donoho asymptotically optimal + Frobenius hard threshold. That rank is frozen before differentiating both + MP-space maps. The rank-selection discontinuity is not differentiated, and + the report does not assume correlated Muon history is truly iid MP noise. + The single-checkpoint Feshbach map is an intentional collapse control. In the checkpoint's own SVD frame the P-Q Gram coupling is exactly zero, so its shell downfolding contribution vanishes at first order. The state-level Feshbach map avoids that triviality by freezing P/Q from the independent epoch-10 anchor. Neither construction is presented as the unique quotient of an unknown sum of Muon updates; they are falsifiable, fully specified quotient hypotheses. + +The principal outputs under +`/private/tmp/rg-mnist-mlp3-short100-jacobians-reduced` are: + +- `jacobian_powerlaw_fits.csv`: one reported alpha per local-map method and fit + convention; +- `jacobian_hyperparameter_search.csv`: all five Tikhonov candidates, fit + diagnostics, selection rank, and selected flag at every optimizer/layer/epoch; +- `weight_quotient_weightwatcher_fits.csv`: raw and `clip_xmax` WeightWatcher + alphas for every materialized quotient representative; +- `two_checkpoint_flow_fits.csv`: finite between-checkpoint flow alphas; +- `two_checkpoint_jacobian_transport.csv`: observed quotient secant versus the + local Jacobian-vector prediction; +- `status.json`, `quotient_flow_status.json`, `jacobians.log`, and + `quotient_flow.log`: live progress/ETA and full terminal-equivalent logs; +- `report/index.html`: shareable documentation, separate MuonClip-RMS and AdamW + plots, tables, method coverage, and saved-data links. 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 0f82581..af9a517 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 @@ -35,8 +35,10 @@ "ecs_grassmann_cartan_cover_detx_shell_pullback": "ECS detX shell", "gap_aware_projector_detx_shell_pullback": "Gap-aware projector (detX)", "trace_free_log_gram_detx_shell_pullback": "Trace-free log Gram (detX)", - "gram_ridge_resolvent_detx_shell_zratio_0p50_pullback": "Gram ridge resolvent (detX, z=0.5 boundary)", + "tikhonov_resolvent_detx_shell_grid_selected_pullback": "Tikhonov resolvent (best valid KS from grid)", "feshbach_trace_free_log_detx_shell_pullback": "Feshbach trace-free log (detX)", + "optimal_mp_projector_full_shell_pullback": "Optimal-MP signal-space projector", + "optimal_mp_signal_log_gram_pullback": "Trace-free log Gram in optimal-MP space", } METHOD_STYLES = { "centered_log_singular_radial_pullback": ("-", "o"), @@ -44,10 +46,12 @@ "ecs_grassmann_cartan_cover_detx_shell_pullback": (":", "^"), "gap_aware_projector_detx_shell_pullback": ("-.", "D"), "trace_free_log_gram_detx_shell_pullback": ((0, (5, 1)), "P"), - "gram_ridge_resolvent_detx_shell_zratio_0p50_pullback": ( + "tikhonov_resolvent_detx_shell_grid_selected_pullback": ( (0, (3, 1, 1, 1)), "X" ), "feshbach_trace_free_log_detx_shell_pullback": ((0, (1, 1)), "v"), + "optimal_mp_projector_full_shell_pullback": ((0, (5, 2)), "*"), + "optimal_mp_signal_log_gram_pullback": ((0, (2, 1)), "h"), } ECS_METHODS = ( "ecs_grassmann_cartan_cover_full_row_shell_pullback", @@ -56,8 +60,10 @@ EXTENDED_ECS_METHODS = ( "gap_aware_projector_detx_shell_pullback", "trace_free_log_gram_detx_shell_pullback", - "gram_ridge_resolvent_detx_shell_zratio_0p50_pullback", + "tikhonov_resolvent_detx_shell_grid_selected_pullback", "feshbach_trace_free_log_detx_shell_pullback", + "optimal_mp_projector_full_shell_pullback", + "optimal_mp_signal_log_gram_pullback", ) EXPECTED_METHODS_BY_LAYER = { "fc1.weight": ("centered_log_singular_radial_pullback", *ECS_METHODS, *EXTENDED_ECS_METHODS), @@ -68,13 +74,21 @@ "midpoint_ecs_control": "Midpoint ECS control", "gram_ridge": "Gram diagonal subtraction", "feshbach_downfolding": "Anchor-frozen Feshbach", + "calibrated_mp_shrinker": "Calibrated MP shrinker", } QUOTIENT_PROFILE_LABELS = { "midpoint": "no counterterm", + "tau_fraction_0p10": "τ=0.10 λboundary", "tau_fraction_0p25": "τ=0.25 λboundary", "tau_fraction_0p50": "τ=0.50 λboundary", "tau_fraction_0p75": "τ=0.75 λboundary", + "tau_fraction_0p90": "τ=0.90 λboundary", + "ridge_ratio_1em3": "ρ=0.001 anchor scale", "ridge_ratio_1em2": "ρ=0.01 anchor scale", + "ridge_ratio_1em1": "ρ=0.1 anchor scale", + "mp_scale_0p75": "MP scale=0.75", + "mp_scale_1p00": "MP scale=1.00", + "mp_scale_1p25": "MP scale=1.25", } FLOW_LABELS = { "two_checkpoint_generalized_gram_radial": "Generalized-Gram radial rate", @@ -84,6 +98,22 @@ "two_checkpoint_radial_quotient_observed_secant": "Observed radial quotient secant", "two_checkpoint_radial_jacobian_prediction": "Local radial Jacobian prediction", } +ANALYSIS_EPOCHS = tuple(range(10, 101, 10)) +TIKHONOV_Z_RATIOS = (0.03, 0.10, 0.30, 1.00, 3.00) +QUOTIENT_EXPECTED_PROFILES = ( + ("midpoint_ecs_control", "midpoint"), + ("gram_ridge", "tau_fraction_0p10"), + ("gram_ridge", "tau_fraction_0p25"), + ("gram_ridge", "tau_fraction_0p50"), + ("gram_ridge", "tau_fraction_0p75"), + ("gram_ridge", "tau_fraction_0p90"), + ("feshbach_downfolding", "ridge_ratio_1em3"), + ("feshbach_downfolding", "ridge_ratio_1em2"), + ("feshbach_downfolding", "ridge_ratio_1em1"), + ("calibrated_mp_shrinker", "mp_scale_0p75"), + ("calibrated_mp_shrinker", "mp_scale_1p00"), + ("calibrated_mp_shrinker", "mp_scale_1p25"), +) def configure_logging(report_root: Path) -> logging.Logger: @@ -265,6 +295,56 @@ def plot_single_optimizer_jacobians( return save_figure(fig, path, bottom=0.18, top=0.93) +def plot_tikhonov_search_by_optimizer( + search: pd.DataFrame, optimizer: str, path: Path +) -> Path: + """Show every ridge candidate and identify the selected PL fit.""" + + frame = search[ + search["optimizer"].astype(str).eq(optimizer) + & search["layer"].astype(str).isin(("fc1.weight", "fc2.weight")) + ].copy() + frame["selected_bool"] = bool_series(frame["selected"]) + ratios = sorted( + pd.to_numeric(frame["tikhonov_z_boundary_ratio"], errors="coerce") + .dropna().unique() + ) + palette = plt.get_cmap("viridis") + ratio_colors = { + float(ratio): palette(index / max(1, len(ratios) - 1)) + for index, ratio in enumerate(ratios) + } + fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.0), sharex=True) + for axis, layer in zip(axes, ("fc1.weight", "fc2.weight")): + subset = frame[frame["layer"].astype(str).eq(layer)] + for ratio, curve in subset.groupby("tikhonov_z_boundary_ratio"): + curve = curve.sort_values("epoch") + numeric_ratio = float(ratio) + axis.plot( + curve["epoch"], curve["alpha"], + color=ratio_colors[numeric_ratio], linewidth=1.25, alpha=0.78, + label=f"z/λboundary={numeric_ratio:g}", + ) + selected = subset[subset["selected_bool"]].sort_values("epoch") + axis.scatter( + selected["epoch"], selected["alpha"], + color="black", marker="x", s=42, linewidth=1.5, + label="selected: valid minimum KS", + zorder=5, + ) + axis.axhline(2.0, color="black", linestyle=(0, (1, 2)), linewidth=1.0) + axis.set( + title=layer.replace(".weight", ""), xlabel="epoch", + ylabel="Tikhonov Jacobian energy alpha", + ) + axis.grid(True, alpha=0.25) + shared_legend(fig, axes, columns=3) + fig.suptitle( + f"{OPTIMIZER_LABELS[optimizer]} — Tikhonov grid; selection does not target α=2" + ) + return save_figure(fig, path, bottom=0.20, top=0.92) + + def plot_weight_quotients_by_optimizer( quotient_fits: pd.DataFrame, optimizer: str, path: Path ) -> Path: @@ -275,10 +355,17 @@ def plot_weight_quotients_by_optimizer( frame = frame[bool_series(frame["fit_ok"])] profile_colors = { "midpoint": "#4C78A8", + "tau_fraction_0p10": "#ECA82C", "tau_fraction_0p25": "#F2CF5B", "tau_fraction_0p50": "#F58518", "tau_fraction_0p75": "#B279A2", + "tau_fraction_0p90": "#8E6C8A", + "ridge_ratio_1em3": "#86BCB6", "ridge_ratio_1em2": "#54A24B", + "ridge_ratio_1em1": "#2E8B57", + "mp_scale_0p75": "#E45756", + "mp_scale_1p00": "#B44A99", + "mp_scale_1p25": "#7A5195", } fig, axes = plt.subplots(2, 2, figsize=(13.5, 8.4), sharex=True) variants = ("raw", "clip_xmax") @@ -398,6 +485,109 @@ def build_method_coverage(primary: pd.DataFrame) -> pd.DataFrame: return pd.DataFrame(rows) +def build_analysis_contract_coverage( + primary: pd.DataFrame, + tikhonov_search: pd.DataFrame, + quotient_fits: pd.DataFrame, + flow_fits: pd.DataFrame, + transport: pd.DataFrame, +) -> pd.DataFrame: + """Audit every required alpha-producing unit before declaring success.""" + + def audit(name: str, expected: set[tuple], observed_values: list[tuple]) -> dict[str, object]: + observed = set(observed_values) + missing = sorted(expected - observed, key=str) + return { + "analysis_family": name, + "expected_unique_units": len(expected), + "observed_unique_units": len(observed & expected), + "unexpected_unique_units": len(observed - expected), + "duplicate_row_count": max(0, len(observed_values) - len(observed)), + "missing_unit_count": len(missing), + "missing_examples": "; ".join(map(str, missing[:8])), + "coverage_status": "complete" if not missing else "INCOMPLETE", + } + + jacobian_expected = { + (optimizer, layer, epoch, method) + for optimizer in OPTIMIZERS + for layer, methods in EXPECTED_METHODS_BY_LAYER.items() + for epoch in ANALYSIS_EPOCHS + for method in methods + } + jacobian_observed = list(primary[[ + "optimizer", "layer", "epoch", "method" + ]].itertuples(index=False, name=None)) + + tikhonov_expected = { + (optimizer, layer, epoch, ratio) + for optimizer in OPTIMIZERS + for layer in ("fc1.weight", "fc2.weight") + for epoch in ANALYSIS_EPOCHS + for ratio in TIKHONOV_Z_RATIOS + } + tikhonov_observed = [ + (str(row.optimizer), str(row.layer), int(row.epoch), float(row.ratio)) + for row in tikhonov_search.rename( + columns={"tikhonov_z_boundary_ratio": "ratio"} + )[["optimizer", "layer", "epoch", "ratio"]].itertuples(index=False) + ] + selected = tikhonov_search[bool_series(tikhonov_search["selected"])] + selected_expected = { + (optimizer, layer, epoch) + for optimizer in OPTIMIZERS + for layer in ("fc1.weight", "fc2.weight") + for epoch in ANALYSIS_EPOCHS + } + selected_observed = list(selected[[ + "optimizer", "layer", "epoch" + ]].itertuples(index=False, name=None)) + + quotient_expected = { + (optimizer, layer, epoch, method, profile, variant) + for optimizer in OPTIMIZERS + for layer in ("fc1.weight", "fc2.weight") + for epoch in ANALYSIS_EPOCHS + for method, profile in QUOTIENT_EXPECTED_PROFILES + for variant in ("raw", "clip_xmax") + } + quotient_observed = list(quotient_fits[[ + "optimizer", "layer", "epoch", "method", "profile_id", "fit_variant" + ]].itertuples(index=False, name=None)) + + flow_primary = flow_fits[ + flow_fits["spectrum_kind"].astype(str).eq("energy_derived_from_amplitude") + & pd.to_numeric(flow_fits["clip_top_k"], errors="coerce").eq(0) + ] + flow_expected = { + (optimizer, layer, epoch_end, method) + for optimizer in OPTIMIZERS + for layer in ("fc1.weight", "fc2.weight") + for epoch_end in ANALYSIS_EPOCHS[1:] + for method in FLOW_LABELS + } + flow_observed = list(flow_primary[[ + "optimizer", "layer", "epoch_end", "method" + ]].itertuples(index=False, name=None)) + transport_expected = { + (optimizer, layer, epoch_end) + for optimizer in OPTIMIZERS + for layer in ("fc1.weight", "fc2.weight") + for epoch_end in ANALYSIS_EPOCHS[1:] + } + transport_observed = list(transport[[ + "optimizer", "layer", "epoch_end" + ]].itertuples(index=False, name=None)) + return pd.DataFrame([ + audit("single_checkpoint_jacobian_alphas", jacobian_expected, jacobian_observed), + audit("tikhonov_grid_candidates", tikhonov_expected, tikhonov_observed), + audit("tikhonov_selected_curve", selected_expected, selected_observed), + audit("weight_quotient_weightwatcher_alphas", quotient_expected, quotient_observed), + audit("two_checkpoint_flow_alphas", flow_expected, flow_observed), + audit("two_checkpoint_jacobian_transport", transport_expected, transport_observed), + ]) + + 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): @@ -554,6 +744,8 @@ def build_html( primary: pd.DataFrame, coverage: pd.DataFrame, inventory: dict[str, int], + tikhonov_search: pd.DataFrame, + contract_coverage: pd.DataFrame, ) -> Path: final = primary.sort_values("epoch").groupby( ["optimizer", "layer", "method"], as_index=False @@ -562,6 +754,16 @@ def build_html( ["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) + contract_table = contract_coverage.to_html(index=False) + selected_tikhonov = tikhonov_search[ + bool_series(tikhonov_search["selected"]) + ].sort_values("epoch").groupby( + ["optimizer", "layer"], as_index=False + ).tail(1) + selected_tikhonov_table = selected_tikhonov[[ + "optimizer", "epoch", "layer", "tikhonov_z_boundary_ratio", + "alpha", "ks_D", "n_tail", "tail_decades", "fit_ok", + ]].to_html(index=False, float_format=lambda value: f"{value:.5g}") figure_html = "\n".join( f'

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

' f'' @@ -613,6 +815,12 @@ def build_html(

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} +

End-to-end data contract

+

This audit separately requires all local-Jacobian alphas, all five Tikhonov +candidates and exactly one selection per checkpoint, every transformed-weight +WeightWatcher fit, every two-checkpoint flow alpha, and every Jacobian-transport +comparison. The report command fails if any family is incomplete.

+{contract_table}

Scientific questions and what is actually identified

Case 1 — heavy tails on a weight quotient representative

For the state-level experiment, the strictly identifiable quotient is the @@ -629,14 +837,19 @@ def build_html( sector without a counterterm.

  • Gram diagonal subtraction: X=W Wᵀ (or WᵀW on the smaller side) is transformed by -λᵢ↦max(λᵢ−τ,0), scanned at τ/λboundary = 0.25, 0.50, and -0.75. This is nonlinear because of the positive-part threshold; it is not +λᵢ↦max(λᵢ−τ,0), scanned at τ/λboundary = 0.10, 0.25, 0.50, +0.75, and 0.90. This is nonlinear because of the positive-part threshold; it is not a global rescaling and can change the fitted spectral shape.
  • Feshbach downfolding: an independent epoch-10 singular frame freezes P and Q, then -X_eff=A−B(C+ρI)⁻¹Bᵀ with ρ=0.01 times the anchor boundary scale. +X_eff=A−B(C+ρI)⁻¹Bᵀ with ρ/λanchor scanned at 0.001, 0.01, +and 0.1. Freezing the anchor is essential: choosing P from the same X diagonalizes the Gram matrix and makes B=0, reducing the construction to truncation.
    • +
  • Calibrated MP shrinker: estimate a white-noise scale from the +discarded midpoint-ECS bulk, then apply the analytic optimal Frobenius +singular-value shrinker at scale multipliers 0.75, 1.00, and 1.25. This is an +MP-like nuisance model, not a claim that Muon noise is actually iid.

    • Case 2a — flow between checkpoints

    Two checkpoints identify finite flow/secant observables: generalized-Gram @@ -656,16 +869,34 @@ def build_html(

    Case 2b — a Jacobian at one checkpoint

    The single-checkpoint tables contain exact analytic derivatives of explicitly declared weight-only maps: centered log-singular radial, ECS Grassmann/Cartan, -gap-aware projector, trace-free log Gram, ridge-resolvent, and Feshbach -trace-free log. These are genuine Jacobians of those maps at W, but not the +gap-aware projector, trace-free log Gram, a selected Tikhonov resolvent, +Feshbach trace-free log, and two MP-selected signal-space maps. These are +genuine Jacobians of those maps at W, but not the training-dynamics Jacobian unless the declared map is separately calibrated to the optimizer step.

    +

    Tikhonov grid. On the fixed detX outer space, the map is +R_z(W)=Π_tf(V_oᵀWᵀWV_o+zI)⁻¹, with +z/λ_boundary∈{{0.03,0.10,0.30,1,3}}. The derivative is exact for +each fixed z. Selection ranks valid fits by minimum KS D, then larger tail span, +then larger tail count, then grid order. It never uses closeness of α to 2. +All candidates are saved in jacobian_hyperparameter_search.csv; +the main Jacobian table contains the selected curve. This selection is +exploratory and should be validated on held-out checkpoints.

    +

    Optimal MP space. The empirical median Gram eigenvalue and +the Marchenko–Pastur median estimate the noise scale. The Gavish–Donoho +asymptotically optimal Frobenius hard threshold fixes a signal rank. With that +rank frozen, the report differentiates (i) the hard signal-space projector and +(ii) trace-free log Gram restricted to the selected signal space. The +discontinuous rank-selection step itself is not differentiated, and correlated +Muon updates need not satisfy the MP white-noise assumptions.

    The single-checkpoint Feshbach result has a mandatory caveat: in the checkpoint's own SVD frame B=0, so shell-downfolding terms vanish at first order. The report keeps this curve as an explicit collapse/control. Nontrivial state-level Feshbach behavior comes from the independently frozen epoch-10 anchor, while nontrivial first-order Feshbach dynamics would require a frozen frame not diagonalizing the evaluation checkpoint.

    +

    Final-checkpoint selected Tikhonov candidates

    +{selected_tikhonov_table}

    Data inventory

      {inventory_html}

    Analysis-ready tables

      {table_html}

    Final checkpoint primary fits

    {final_table} @@ -706,6 +937,9 @@ def main() -> int: fits = require_csv(analysis_root / "jacobian_powerlaw_fits.csv") spectra = require_csv(analysis_root / "jacobian_spectra.csv") operators = require_csv(analysis_root / "jacobian_operators.csv") + tikhonov_search = require_csv( + analysis_root / "jacobian_hyperparameter_search.csv" + ) quotient_fits = require_csv( analysis_root / "weight_quotient_weightwatcher_fits.csv" ) @@ -738,11 +972,15 @@ def main() -> int: table_root.mkdir(parents=True, exist_ok=True) coverage = build_method_coverage(primary) + contract_coverage = build_analysis_contract_coverage( + primary, tikhonov_search, quotient_fits, flow_fits, flow_transport + ) tables = [ table_root / "jacobian_primary_energy_fits.csv", table_root / "jacobian_all_fits.csv", table_root / "jacobian_spectra_all_modes.csv", table_root / "jacobian_operator_metadata.csv", + table_root / "jacobian_hyperparameter_search.csv", table_root / "weightwatcher_raw_and_clip_xmax.csv", table_root / "performance_train_test.csv", table_root / "jacobian_method_coverage.csv", @@ -753,21 +991,33 @@ def main() -> int: table_root / "two_checkpoint_flow_spectra.csv", table_root / "two_checkpoint_flow_operators.csv", table_root / "two_checkpoint_jacobian_transport.csv", + table_root / "analysis_contract_coverage.csv", ] primary.to_csv(tables[0], index=False) fits.to_csv(tables[1], index=False) spectra.to_csv(tables[2], index=False) 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) - quotient_fits.to_csv(tables[7], index=False) - quotient_spectra.to_csv(tables[8], index=False) - quotient_operators.to_csv(tables[9], index=False) - flow_fits.to_csv(tables[10], index=False) - flow_spectra.to_csv(tables[11], index=False) - flow_operators.to_csv(tables[12], index=False) - flow_transport.to_csv(tables[13], index=False) + tikhonov_search.to_csv(tables[4], index=False) + weightwatcher.to_csv(tables[5], index=False) + performance.to_csv(tables[6], index=False) + coverage.to_csv(tables[7], index=False) + quotient_fits.to_csv(tables[8], index=False) + quotient_spectra.to_csv(tables[9], index=False) + quotient_operators.to_csv(tables[10], index=False) + flow_fits.to_csv(tables[11], index=False) + flow_spectra.to_csv(tables[12], index=False) + flow_operators.to_csv(tables[13], index=False) + flow_transport.to_csv(tables[14], index=False) + contract_coverage.to_csv(tables[15], index=False) + incomplete_contract = contract_coverage[ + contract_coverage["coverage_status"].astype(str).ne("complete") + ] + if not incomplete_contract.empty: + raise RuntimeError( + "analysis data contract is incomplete; inspect " + f"{tables[15]}\n" + + incomplete_contract.to_string(index=False) + ) figures = [ plot_jacobian_metric( @@ -798,6 +1048,10 @@ def main() -> int: primary, optimizer, figure_root / f"optimizer_views/{optimizer}_single_checkpoint_jacobians.png", ), + plot_tikhonov_search_by_optimizer( + tikhonov_search, optimizer, + figure_root / f"optimizer_views/{optimizer}_tikhonov_grid_search.png", + ), plot_weight_quotients_by_optimizer( quotient_fits, optimizer, figure_root / f"optimizer_views/{optimizer}_weight_quotients.png", @@ -817,6 +1071,10 @@ def main() -> int: "all Jacobian fit rows": len(fits), "saved Jacobian spectral modes": len(spectra), "operator metadata rows": len(operators), + "Tikhonov grid candidate rows": len(tikhonov_search), + "complete analysis-contract families": int( + contract_coverage["coverage_status"].astype(str).eq("complete").sum() + ), "WeightWatcher control rows": len(weightwatcher), "performance rows": len(performance), "weight-quotient WeightWatcher rows": len(quotient_fits), @@ -826,7 +1084,10 @@ def main() -> int: "Jacobian transport comparison rows": len(flow_transport), "figures": len(figures), } - index = build_html(report_root, figures, tables, primary, coverage, inventory) + index = build_html( + report_root, figures, tables, primary, coverage, inventory, + tikhonov_search, contract_coverage, + ) (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_complete_rg_analysis.sh b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_complete_rg_analysis.sh index 6afff41..6e91a8d 100755 --- a/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_complete_rg_analysis.sh +++ b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_complete_rg_analysis.sh @@ -11,11 +11,46 @@ OUTPUT_ROOT="${RG_MNIST_REDUCED_JACOBIAN_OUTPUT_ROOT:-/private/tmp/rg-mnist-mlp3 export RG_MNIST_REDUCED_JACOBIAN_OUTPUT_ROOT="${OUTPUT_ROOT}" -bash "${SCRIPT_DIR}/run_short100_jacobians_reduced.sh" "$@" +mkdir -p "${OUTPUT_ROOT}" -python -u "${SCRIPT_DIR}/run_short100_quotient_flow_cli.py" \ - --run-root "${RG_MNIST_TANGENT_ROOT:-/private/tmp/rg-mnist-mlp3-short100-runs}" \ - --cache-root "${RG_MNIST_TANGENT_CHECKPOINT_CACHE_ROOT:-/private/tmp/rg-mnist-mlp3-short100-checkpoints}" \ - --output-root "${OUTPUT_ROOT}" +interrupted() { + printf '\nINTERRUPTED: completed units are already stored atomically.\n' >&2 + printf 'Rerun this same command to resume. Use Ctrl-C, not Ctrl-Z, to stop.\n' >&2 + exit 130 +} +trap interrupted INT TERM -bash "${SCRIPT_DIR}/build_short100_jacobian_report.sh" +stage() { + local label="$1" + shift + printf '\n============================================================\n' + printf 'STAGE %s\n' "${label}" + printf 'STARTED %s\n' "$(date '+%Y-%m-%d %H:%M:%S')" + printf '============================================================\n' + "$@" + printf 'COMPLETED %s at %s\n' "${label}" "$(date '+%Y-%m-%d %H:%M:%S')" +} + +stage "1/4 exact single-checkpoint Jacobians" \ + bash "${SCRIPT_DIR}/run_short100_jacobians_reduced.sh" "$@" + +stage "2/4 transformed-weight quotients" \ + python -u "${SCRIPT_DIR}/run_short100_quotient_flow_cli.py" \ + --run-root "${RG_MNIST_TANGENT_ROOT:-/private/tmp/rg-mnist-mlp3-short100-runs}" \ + --cache-root "${RG_MNIST_TANGENT_CHECKPOINT_CACHE_ROOT:-/private/tmp/rg-mnist-mlp3-short100-checkpoints}" \ + --output-root "${OUTPUT_ROOT}" \ + --epoch-stride 10 \ + --skip-checkpoint-flows + +stage "3/4 between-checkpoint flow and local transport" \ + python -u "${SCRIPT_DIR}/run_short100_quotient_flow_cli.py" \ + --run-root "${RG_MNIST_TANGENT_ROOT:-/private/tmp/rg-mnist-mlp3-short100-runs}" \ + --cache-root "${RG_MNIST_TANGENT_CHECKPOINT_CACHE_ROOT:-/private/tmp/rg-mnist-mlp3-short100-checkpoints}" \ + --output-root "${OUTPUT_ROOT}" \ + --epoch-stride 10 \ + --skip-state-quotients + +stage "4/4 static HTML report" \ + bash "${SCRIPT_DIR}/build_short100_jacobian_report.sh" + +printf '\nCOMPLETE REPORT: %s\n' "${OUTPUT_ROOT}/report/index.html" 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 1a97222..27f6b32 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 @@ -12,6 +12,7 @@ import argparse from dataclasses import asdict, is_dataclass from datetime import datetime, timezone +from functools import lru_cache import json import logging from pathlib import Path @@ -49,9 +50,12 @@ EXTENDED_DETX_METHODS = ( "gap_aware_projector_detx_shell_pullback", "trace_free_log_gram_detx_shell_pullback", - "gram_ridge_resolvent_detx_shell_zratio_0p50_pullback", + "tikhonov_resolvent_detx_shell_grid_selected_pullback", "feshbach_trace_free_log_detx_shell_pullback", + "optimal_mp_projector_full_shell_pullback", + "optimal_mp_signal_log_gram_pullback", ) +TIKHONOV_Z_RATIOS = (0.03, 0.10, 0.30, 1.00, 3.00) class MaxInfoFilter(logging.Filter): @@ -122,6 +126,97 @@ def parse_csv_values(text: str, cast=str) -> tuple[Any, ...]: return values +def fit_ok(value: Any) -> bool: + """Normalize package/string booleans used by resumable CSV tables.""" + + if isinstance(value, (bool, np.bool_)): + return bool(value) + return str(value).strip().lower() in {"1", "true", "yes"} + + +@lru_cache(maxsize=16) +def marchenko_pastur_median(aspect_ratio: float) -> float: + """Numerically return the median of the unit-scale MP law. + + ``aspect_ratio`` is min(m,n)/max(m,n), so it lies in (0,1]. A sine-square + coordinate removes the square-root endpoint behavior, including the hard + edge at zero for a square matrix. This is a deterministic calibration + constant, not a fit to WeightWatcher alpha. + """ + + beta = float(aspect_ratio) + if not np.isfinite(beta) or not 0.0 < beta <= 1.0: + raise ValueError("MP aspect ratio must lie in (0, 1]") + lower = (1.0 - np.sqrt(beta)) ** 2 + upper = (1.0 + np.sqrt(beta)) ** 2 + theta = np.linspace(1.0e-8, np.pi / 2.0 - 1.0e-8, 40001) + sine = np.sin(theta) + cosine = np.cos(theta) + values = lower + (upper - lower) * sine**2 + density = np.sqrt( + np.maximum((upper - values) * (values - lower), 0.0) + ) / (2.0 * np.pi * beta * values) + derivative = 2.0 * (upper - lower) * sine * cosine + integrand = density * derivative + increments = 0.5 * (integrand[1:] + integrand[:-1]) * np.diff(theta) + cdf = np.concatenate(([0.0], np.cumsum(increments))) + cdf /= cdf[-1] + return float(np.interp(0.5, cdf, values)) + + +def optimal_mp_signal_rank( + singular_values: np.ndarray, + matrix_shape: tuple[int, int], + *, + numerical_rank: int, +) -> tuple[int, dict[str, Any]]: + """Select the Gavish--Donoho optimal MP hard-threshold signal space. + + The empirical median Gram eigenvalue estimates the white-noise scale using + the MP median. The asymptotically optimal Frobenius hard threshold then + selects the signal rank. The returned usable rank is clipped only when the + raw decision is 0, 1, or the entire numerical space, because the two exact + Jacobians below require ``2 <= k < rank``. The raw and clipped decisions + are both recorded, so a boundary result is never hidden. + """ + + singular = np.asarray(singular_values, dtype=np.float64) + rank = int(numerical_rank) + if not 3 <= rank <= singular.size: + raise ValueError("MP Jacobian selection requires numerical rank >= 3") + beta = min(matrix_shape) / max(matrix_shape) + mp_median = marchenko_pastur_median(beta) + empirical_median = float(np.median(singular[:rank] ** 2)) + noise_unit = float(np.sqrt(empirical_median / mp_median)) + optimal_known_noise = float( + np.sqrt( + 2.0 * (beta + 1.0) + + 8.0 * beta + / ((beta + 1.0) + np.sqrt(beta**2 + 14.0 * beta + 1.0)) + ) + ) + threshold = optimal_known_noise * noise_unit + raw_rank = int(np.count_nonzero(singular[:rank] > threshold)) + usable_rank = int(np.clip(raw_rank, 2, rank - 1)) + return usable_rank, { + "mp_selection_rule": ( + "Gavish-Donoho asymptotically optimal Frobenius hard threshold; " + "noise scale from empirical median Gram eigenvalue divided by the " + "unit-scale Marchenko-Pastur median" + ), + "mp_aspect_ratio": float(beta), + "mp_unit_scale_median": mp_median, + "mp_empirical_gram_median": empirical_median, + "mp_estimated_noise_singular_unit": noise_unit, + "mp_optimal_known_noise_threshold_multiplier": optimal_known_noise, + "mp_singular_threshold": threshold, + "mp_raw_signal_rank": raw_rank, + "mp_selected_signal_rank": usable_rank, + "mp_rank_clipped_for_jacobian_domain": bool(usable_rank != raw_rank), + "mp_selection_is_frozen_during_differentiation": True, + } + + def load_json(path: Path, description: str) -> dict[str, Any]: if not path.is_file(): raise FileNotFoundError(f"missing {description}: {path}") @@ -456,19 +551,26 @@ def extended_detx_jacobian_spectra( ) -> dict[str, tuple[np.ndarray, Any, dict[str, Any]]]: """Exact additional Jacobians on the independently audited detX shell. - The resolvent is the differentiable ridge/noise-control analogue. The - Feshbach derivative is retained even though its shell term must collapse + The Tikhonov family is handled separately by ``tikhonov_grid_spectra`` so + that all candidates can be fit, audited, and reduced to one selected curve. + The Feshbach derivative is retained even though its shell term must collapse at first order in the checkpoint SVD gauge; that collapse is a scientific - control, not silently interpreted as nontrivial downfolding. + control, not silently interpreted as nontrivial downfolding. Two MP methods + freeze the data-selected optimal hard-threshold rank before differentiating: + the hard signal-space projector and trace-free log Gram inside that space. """ from rg_baselines.tangent_rg import ecs_jacobians k = int(retained_rank) q = int(outer_rank) - boundary_scale = float(singular_values[k - 1] ** 2) - resolvent_z = 0.50 * boundary_scale shell_floor = float(singular_values[q - 1] ** 2) feshbach_z = 0.50 * shell_floor + numerical_rank = int(np.count_nonzero(singular_values > rcond * singular_values[0])) + mp_rank, mp_metadata = optimal_mp_signal_rank( + singular_values, + tuple(weight.shape), + numerical_rank=numerical_rank, + ) records = { "gap_aware_projector_detx_shell_pullback": ( ecs_jacobians.gap_aware_projector_spectrum( @@ -484,18 +586,6 @@ def extended_detx_jacobian_spectra( ), {"jacobian_family": "trace_free_log_gram", "retained_rank": k, "outer_rank": q}, ), - "gram_ridge_resolvent_detx_shell_zratio_0p50_pullback": ( - ecs_jacobians.outer_resolvent_spectrum( - weight, outer_rank=q, z=resolvent_z, trace_free=True, - rcond=rcond, precomputed_singular_values=singular_values, - ), - { - "jacobian_family": "trace_free_gram_ridge_resolvent", - "retained_rank": k, "outer_rank": q, - "resolvent_z": resolvent_z, - "resolvent_z_boundary_ratio": 0.50, - }, - ), "feshbach_trace_free_log_detx_shell_pullback": ( ecs_jacobians.feshbach_trace_free_log_spectrum( weight, retained_rank=k, outer_rank=q, z=feshbach_z, rcond=rcond, @@ -507,6 +597,35 @@ def extended_detx_jacobian_spectra( "first_order_shell_downfolding_active": False, }, ), + "optimal_mp_projector_full_shell_pullback": ( + ecs_jacobians.gap_aware_projector_spectrum( + weight, + retained_rank=mp_rank, + outer_rank=numerical_rank, + rcond=rcond, + precomputed_singular_values=singular_values, + ), + { + "jacobian_family": "optimal_mp_hard_signal_space_projector", + "retained_rank": mp_rank, + "outer_rank": numerical_rank, + **mp_metadata, + }, + ), + "optimal_mp_signal_log_gram_pullback": ( + ecs_jacobians.outer_trace_free_log_gram_spectrum( + weight, + outer_rank=mp_rank, + rcond=rcond, + precomputed_singular_values=singular_values, + ), + { + "jacobian_family": "optimal_mp_signal_space_trace_free_log_gram", + "retained_rank": mp_rank, + "outer_rank": mp_rank, + **mp_metadata, + }, + ), } return { method: ( @@ -518,6 +637,128 @@ def extended_detx_jacobian_spectra( } +def tikhonov_grid_spectra( + weight: np.ndarray, + singular_values: np.ndarray, + *, + retained_rank: int, + outer_rank: int, + rcond: float, + z_ratios: tuple[float, ...] = TIKHONOV_Z_RATIOS, +) -> dict[str, tuple[np.ndarray, Any, dict[str, Any]]]: + """Evaluate an explicit Tikhonov grid on one frozen detX ECS space.""" + + from rg_baselines.tangent_rg import ecs_jacobians + + k = int(retained_rank) + q = int(outer_rank) + boundary_scale = float(singular_values[k - 1] ** 2) + candidates: dict[str, tuple[np.ndarray, Any, dict[str, Any]]] = {} + for index, raw_ratio in enumerate(z_ratios): + ratio = float(raw_ratio) + if not np.isfinite(ratio) or ratio <= 0.0: + raise ValueError("Tikhonov z ratios must be finite and positive") + z = ratio * boundary_scale + record = ecs_jacobians.outer_resolvent_spectrum( + weight, + outer_rank=q, + z=z, + trace_free=True, + rcond=rcond, + precomputed_singular_values=singular_values, + ) + candidate_id = f"tikhonov_zratio_{ratio:.2g}".replace(".", "p") + candidates[candidate_id] = ( + np.asarray(record.singular_amplitudes, dtype=float), + record, + { + "jacobian_family": "trace_free_tikhonov_gram_resolvent", + "retained_rank": k, + "outer_rank": q, + "tikhonov_candidate_index": index, + "tikhonov_z": z, + "tikhonov_z_boundary_ratio": ratio, + "tikhonov_grid_ratios": json.dumps(tuple(map(float, z_ratios))), + "tikhonov_z_is_frozen_during_differentiation": True, + }, + ) + return candidates + + +def select_tikhonov_candidate( + candidates: dict[str, tuple[np.ndarray, Any, dict[str, Any]]], + *, + base_metadata: dict[str, Any], + minimum_tail: int, +) -> tuple[str, tuple[np.ndarray, Any, dict[str, Any]], list[dict[str, Any]]]: + """Select the best PL fit without using distance of alpha from two. + + Ranking is lexicographic: valid fit first, then minimum KS distance, then + maximum fitted tail span, then maximum tail count, and finally the original + grid order. All candidates are returned as audit rows. + """ + + if not candidates: + raise ValueError("Tikhonov search requires at least one candidate") + evaluated: list[tuple[tuple[Any, ...], str, dict[str, Any]]] = [] + for candidate_id, (amplitudes, record, metadata) in candidates.items(): + candidate_fits = fit_spectrum( + amplitudes, + record, + {**base_metadata, **metadata, "candidate_id": candidate_id}, + (0,), + minimum_tail, + ) + primary = next( + row for row in candidate_fits + if str(row["spectrum_kind"]) == "energy_derived_from_amplitude" + and int(row["clip_top_k"]) == 0 + ) + valid = fit_ok(primary.get("fit_ok")) + ks = float(primary["ks_D"]) if pd.notna(primary.get("ks_D")) else np.inf + decades = ( + float(primary["tail_decades"]) + if pd.notna(primary.get("tail_decades")) else -np.inf + ) + n_tail = int(primary["n_tail"]) if pd.notna(primary.get("n_tail")) else 0 + order = int(metadata["tikhonov_candidate_index"]) + score = (not valid, ks, -decades, -n_tail, order) + evaluated.append((score, candidate_id, { + **base_metadata, + **metadata, + "candidate_id": candidate_id, + "alpha": primary.get("alpha"), + "ks_D": primary.get("ks_D"), + "xmin": primary.get("xmin"), + "n_tail": primary.get("n_tail"), + "tail_decades": primary.get("tail_decades"), + "fit_ok": primary.get("fit_ok"), + "selection_rule": ( + "fit_ok first; minimum KS D; maximum tail_decades; maximum " + "n_tail; original grid order; alpha proximity to 2 is never used" + ), + })) + evaluated.sort(key=lambda item: item[0]) + selected_id = evaluated[0][1] + rows = [] + for rank, (_, candidate_id, row) in enumerate(evaluated, start=1): + rows.append({ + **row, + "selection_rank": rank, + "selected": candidate_id == selected_id, + }) + selected_amplitudes, selected_record, selected_metadata = candidates[selected_id] + return selected_id, ( + selected_amplitudes, + selected_record, + { + **selected_metadata, + "tikhonov_selected_candidate_id": selected_id, + "tikhonov_selection_rule": rows[0]["selection_rule"], + }, + ), rows + + def safe_slug(text: str) -> str: return "".join(char if char.isalnum() or char in "-_" else "_" for char in text) @@ -615,8 +856,9 @@ def build_parser() -> argparse.ArgumentParser: action=argparse.BooleanOptionalAction, default=False, help=( - "add gap-aware, trace-free log-Gram, ridge-resolvent, and " - "Feshbach exact Jacobians on each requested layer's detX shell" + "add gap-aware, trace-free log-Gram, selected Tikhonov-resolvent, " + "Feshbach, and MP-selected exact Jacobians on each requested " + "layer's detX shell" ), ) parser.add_argument( @@ -640,6 +882,7 @@ def run(args: argparse.Namespace) -> int: spectrum_data_path = output_root / "jacobian_spectra.csv" error_path = output_root / "errors.csv" completion_path = output_root / "completed_checkpoints.csv" + hyperparameter_path = output_root / "jacobian_hyperparameter_search.csv" optimizers = parse_csv_values(args.optimizers) seeds = parse_csv_values(args.seeds, int) layers = parse_csv_values(args.layers) @@ -665,6 +908,7 @@ def run(args: argparse.Namespace) -> int: spectrum_rows: list[dict[str, Any]] = [] errors: list[dict[str, Any]] = [] completion_rows: list[dict[str, Any]] = [] + hyperparameter_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(): @@ -675,6 +919,8 @@ def run(args: argparse.Namespace) -> int: 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") + if args.resume and hyperparameter_path.is_file(): + hyperparameter_rows = pd.read_csv(hyperparameter_path).to_dict(orient="records") identities: dict[tuple[str, int], dict[str, Any]] = {} work: list[tuple[str, int, str, Any]] = [] @@ -724,7 +970,7 @@ def run(args: argparse.Namespace) -> int: ) observed_methods = { str(row.get("method")) - for row in spectrum_rows + for row in operator_rows if ( str(row.get("optimizer")), int(row.get("seed", -1)), str(row.get("layer")), int(row.get("epoch", -1)), @@ -739,8 +985,8 @@ def run(args: argparse.Namespace) -> int: }) if args.extended_ecs_jacobians: expected_methods.update(EXTENDED_DETX_METHODS) - spectrum_data_available = expected_methods.issubset(observed_methods) - if args.resume and completed_before and spectrum_data_available: + method_data_available = expected_methods.issubset(observed_methods) + if args.resume and completed_before and method_data_available: completed += 1 logger.info("SKIP completed %d/%d %s", completed, total, unit_key) continue @@ -762,6 +1008,9 @@ def same_unit(row: dict[str, Any]) -> bool: spectrum_rows = [row for row in spectrum_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)] + hyperparameter_rows = [ + row for row in hyperparameter_rows if not same_unit(row) + ] unit_started = time.perf_counter() logger.info( @@ -778,6 +1027,15 @@ def same_unit(row: dict[str, Any]) -> bool: tuple(weight.shape), time.perf_counter() - load_started, ref.path, ) singular = np.linalg.svd(weight, compute_uv=False) + base = { + "optimizer": optimizer, + "seed": int(seed), + "layer": layer, + "epoch": int(ref.epoch), + "global_step": int(ref.global_step), + "protocol_fingerprint": identity["fingerprint"], + "checkpoint_path": str(ref.path), + } method_factories: dict[str, tuple[np.ndarray, Any]] = {} method_metadata: dict[str, dict[str, Any]] = {} for method in methods: @@ -863,16 +1121,62 @@ def same_unit(row: dict[str, Any]) -> bool: "EXTENDED JACOBIAN method=%s k=%d q=%d n_amplitudes=%d", extended_method, k, q, len(extended_amplitudes), ) - - base = { - "optimizer": optimizer, - "seed": int(seed), - "layer": layer, - "epoch": int(ref.epoch), - "global_step": int(ref.global_step), - "protocol_fingerprint": identity["fingerprint"], - "checkpoint_path": str(ref.path), - } + grid = tikhonov_grid_spectra( + weight, + singular, + retained_rank=k, + outer_rank=q, + rcond=args.ecs_rcond, + ) + selected_id, selected, search_rows = select_tikhonov_candidate( + grid, + base_metadata={ + **base, + **rank_metadata, + "method": "tikhonov_resolvent_detx_shell_grid_search", + "ecs_shell_variant": "detx_shell", + }, + minimum_tail=args.minimum_tail, + ) + selected_amplitudes, selected_record, selected_metadata = selected + selected_method = ( + "tikhonov_resolvent_detx_shell_grid_selected_pullback" + ) + method_factories[selected_method] = ( + selected_amplitudes, + selected_record, + ) + method_metadata[selected_method] = { + **rank_metadata, + **selected_metadata, + "ecs_shell_variant": "detx_shell", + } + hyperparameter_rows.extend(search_rows) + for candidate in sorted( + search_rows, + key=lambda row: int(row["tikhonov_candidate_index"]), + ): + logger.info( + "TIKHONOV GRID candidate=%s z_ratio=%.3g " + "alpha=%s D=%s fit_ok=%s selected=%s", + candidate["candidate_id"], + float(candidate["tikhonov_z_boundary_ratio"]), + ( + f"{float(candidate['alpha']):.4f}" + if pd.notna(candidate["alpha"]) else "nan" + ), + ( + f"{float(candidate['ks_D']):.4f}" + if pd.notna(candidate["ks_D"]) else "nan" + ), + candidate["fit_ok"], + candidate["selected"], + ) + logger.info( + "TIKHONOV SELECTED candidate=%s rule=min_valid_KS " + "(alpha target not used)", + selected_id, + ) 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)) @@ -936,6 +1240,7 @@ def same_unit(row: dict[str, Any]) -> bool: atomic_csv(fit_path, fit_rows) atomic_csv(operator_path, operator_rows) atomic_csv(spectrum_data_path, spectrum_rows) + atomic_csv(hyperparameter_path, hyperparameter_rows) spectrum_path = ( output_root / "plots" / "spectra" / optimizer / safe_slug(layer) / f"epoch_{int(ref.epoch):05d}.png" @@ -999,6 +1304,8 @@ def same_unit(row: dict[str, Any]) -> bool: atomic_csv(operator_path, operator_rows) if spectrum_rows: atomic_csv(spectrum_data_path, spectrum_rows) + if hyperparameter_rows: + atomic_csv(hyperparameter_path, hyperparameter_rows) atomic_json(status_path, { "state": "error", **error_row, "completed_checkpoint_count": completed, 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 5e46606..47f3052 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,8 +1,9 @@ #!/usr/bin/env bash -# Fast scientific preset: the scale-quotiented radial Jacobian on every layer -# 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. +# Fast scientific preset: the scale-quotiented radial Jacobian on every layer; +# both ECS/Grassmann covers on FC1/FC2; and the gap, Tikhonov-grid, Feshbach, +# and optimal-MP local Jacobians on the audited detX shell. ECS deterministic +# shell copies are compressed to physical retained-core groups before the PL fit. set -euo pipefail diff --git a/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_quotient_flow_cli.py b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_quotient_flow_cli.py index 6cbc1de..55da420 100755 --- a/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_quotient_flow_cli.py +++ b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_quotient_flow_cli.py @@ -47,14 +47,41 @@ LAYERS = ("fc1.weight", "fc2.weight") QUOTIENT_PROFILES = ( ("midpoint_ecs_control", "midpoint", {}), + ("gram_ridge", "tau_fraction_0p10", {"tau_fraction": 0.10}), ("gram_ridge", "tau_fraction_0p25", {"tau_fraction": 0.25}), ("gram_ridge", "tau_fraction_0p50", {"tau_fraction": 0.50}), ("gram_ridge", "tau_fraction_0p75", {"tau_fraction": 0.75}), + ("gram_ridge", "tau_fraction_0p90", {"tau_fraction": 0.90}), + ( + "feshbach_downfolding", + "ridge_ratio_1em3", + {"regularization_ratio": 1.0e-3}, + ), ( "feshbach_downfolding", "ridge_ratio_1em2", {"regularization_ratio": 1.0e-2}, ), + ( + "feshbach_downfolding", + "ridge_ratio_1em1", + {"regularization_ratio": 1.0e-1}, + ), + ( + "calibrated_mp_shrinker", + "mp_scale_0p75", + {"noise_scale_multiplier": 0.75}, + ), + ( + "calibrated_mp_shrinker", + "mp_scale_1p00", + {"noise_scale_multiplier": 1.00}, + ), + ( + "calibrated_mp_shrinker", + "mp_scale_1p25", + {"noise_scale_multiplier": 1.25}, + ), ) @@ -65,6 +92,16 @@ def atomic_frame(path: Path, rows: list[dict[str, Any]]) -> None: temporary.replace(path) +def atomic_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text( + json.dumps(payload, indent=2, sort_keys=True, default=str) + "\n", + encoding="utf-8", + ) + temporary.replace(path) + + def read_rows(path: Path) -> list[dict[str, Any]]: if not path.is_file(): return [] @@ -164,13 +201,20 @@ def run_state_quotients( fits = read_rows(fit_path) spectra = read_rows(spectrum_path) operators = read_rows(operator_path) - - for optimizer in optimizers: - identity = jacobian_cli.resolve_run_identity(run_root, optimizer, seed) - refs = jacobian_cli.selected_checkpoint_refs( + refs_by_optimizer = { + optimizer: jacobian_cli.selected_checkpoint_refs( cache_root, optimizer, seed, epoch_stride=epoch_stride, maximum_checkpoints=None, ) + for optimizer in optimizers + } + total = sum(len(refs) * len(QUOTIENT_PROFILES) for refs in refs_by_optimizer.values()) + progress = 0 + phase_started = time.perf_counter() + + for optimizer in optimizers: + identity = jacobian_cli.resolve_run_identity(run_root, optimizer, seed) + refs = refs_by_optimizer[optimizer] anchor_model = checkpoint_model(refs[0].path, identity["fingerprint"]) for ref in refs: base_model = checkpoint_model(ref.path, identity["fingerprint"]) @@ -194,7 +238,11 @@ def run_state_quotients( ) == key } if observed == {(layer, variant) for layer in LAYERS for variant in ("raw", "clip_xmax")}: - logger.info("QUOTIENT SKIP optimizer=%s epoch=%d method=%s", optimizer, ref.epoch, method) + progress += 1 + logger.info( + "QUOTIENT SKIP %d/%d optimizer=%s epoch=%d method=%s", + progress, total, optimizer, ref.epoch, method, + ) continue started = time.perf_counter() logger.info("QUOTIENT START optimizer=%s epoch=%d method=%s profile=%s", optimizer, ref.epoch, method, profile_id) @@ -276,7 +324,28 @@ def run_state_quotients( atomic_frame(fit_path, fits) atomic_frame(spectrum_path, spectra) atomic_frame(operator_path, operators) - logger.info("QUOTIENT DONE optimizer=%s epoch=%d method=%s seconds=%.2f", optimizer, ref.epoch, method, time.perf_counter() - started) + progress += 1 + elapsed = time.perf_counter() - phase_started + eta = elapsed / progress * (total - progress) + logger.info( + "QUOTIENT DONE %d/%d optimizer=%s epoch=%d method=%s " + "seconds=%.2f elapsed=%s ETA=%s", + progress, total, optimizer, ref.epoch, method, + time.perf_counter() - started, + jacobian_cli.format_duration(elapsed), + jacobian_cli.format_duration(eta), + ) + atomic_json(root / "quotient_flow_status.json", { + "state": "running_state_quotients", + "completed_unit_count": progress, + "total_unit_count": total, + "percent_complete": 100.0 * progress / max(1, total), + "last_optimizer": optimizer, + "last_epoch": int(ref.epoch), + "last_method": method, + "elapsed_seconds": elapsed, + "eta_seconds": eta, + }) def ecs_topk_rates(first: np.ndarray, second: np.ndarray, rank: int, delta_s: float) -> np.ndarray: @@ -300,12 +369,19 @@ def run_checkpoint_flows( spectra = read_rows(spectrum_path) operators = read_rows(operator_path) transports = read_rows(transport_path) - for optimizer in optimizers: - identity = jacobian_cli.resolve_run_identity(run_root, optimizer, seed) - refs = jacobian_cli.selected_checkpoint_refs( + refs_by_optimizer = { + optimizer: jacobian_cli.selected_checkpoint_refs( cache_root, optimizer, seed, epoch_stride=epoch_stride, maximum_checkpoints=None, ) + for optimizer in optimizers + } + total = sum(max(0, len(refs) - 1) * len(LAYERS) for refs in refs_by_optimizer.values()) + progress = 0 + phase_started = time.perf_counter() + for optimizer in optimizers: + identity = jacobian_cli.resolve_run_identity(run_root, optimizer, seed) + refs = refs_by_optimizer[optimizer] for ref0, ref1 in zip(refs[:-1], refs[1:]): delta_s = float(ref1.epoch - ref0.epoch) for layer in LAYERS: @@ -327,7 +403,11 @@ def run_checkpoint_flows( ) == unit } if expected.issubset(observed): - logger.info("FLOW SKIP optimizer=%s %d->%d layer=%s", optimizer, ref0.epoch, ref1.epoch, layer) + progress += 1 + logger.info( + "FLOW SKIP %d/%d optimizer=%s %d->%d layer=%s", + progress, total, optimizer, ref0.epoch, ref1.epoch, layer, + ) continue started = time.perf_counter() first = jacobian_cli.checkpoint_matrix(ref0.path, identity["fingerprint"], layer) @@ -472,7 +552,29 @@ def run_checkpoint_flows( atomic_frame(spectrum_path, spectra) atomic_frame(operator_path, operators) atomic_frame(transport_path, transports) - logger.info("FLOW DONE optimizer=%s %d->%d layer=%s seconds=%.2f", optimizer, ref0.epoch, ref1.epoch, layer, time.perf_counter() - started) + progress += 1 + elapsed = time.perf_counter() - phase_started + eta = elapsed / progress * (total - progress) + logger.info( + "FLOW DONE %d/%d optimizer=%s %d->%d layer=%s " + "seconds=%.2f elapsed=%s ETA=%s", + progress, total, optimizer, ref0.epoch, ref1.epoch, layer, + time.perf_counter() - started, + jacobian_cli.format_duration(elapsed), + jacobian_cli.format_duration(eta), + ) + atomic_json(root / "quotient_flow_status.json", { + "state": "running_checkpoint_flows", + "completed_unit_count": progress, + "total_unit_count": total, + "percent_complete": 100.0 * progress / max(1, total), + "last_optimizer": optimizer, + "last_epoch_start": int(ref0.epoch), + "last_epoch_end": int(ref1.epoch), + "last_layer": layer, + "elapsed_seconds": elapsed, + "eta_seconds": eta, + }) def build_parser() -> argparse.ArgumentParser: @@ -509,7 +611,15 @@ def main() -> int: cache_root=args.cache_root.resolve(), optimizers=optimizers, seed=args.seed, epoch_stride=args.epoch_stride, logger=logger, ) - logger.info("COMPLETE seconds=%.2f", time.perf_counter() - started) + elapsed = time.perf_counter() - started + atomic_json(root / "quotient_flow_status.json", { + "state": "complete", + "state_quotients_ran": not args.skip_state_quotients, + "checkpoint_flows_ran": not args.skip_checkpoint_flows, + "elapsed_seconds": elapsed, + "completed_at_utc": datetime.now(timezone.utc).isoformat(), + }) + logger.info("COMPLETE seconds=%.2f", elapsed) return 0 except Exception: trace = traceback.format_exc() @@ -518,6 +628,11 @@ def main() -> int: "failed_at_utc": datetime.now(timezone.utc).isoformat(), "exception_traceback": trace, }]) + atomic_json(root / "quotient_flow_status.json", { + "state": "error", + "failed_at_utc": datetime.now(timezone.utc).isoformat(), + "exception_traceback": trace, + }) return 1 diff --git a/baseline/tests/tangent_rg/test_command_line_jacobians.py b/baseline/tests/tangent_rg/test_command_line_jacobians.py index fc4d6e8..fd40d29 100644 --- a/baseline/tests/tangent_rg/test_command_line_jacobians.py +++ b/baseline/tests/tangent_rg/test_command_line_jacobians.py @@ -65,6 +65,55 @@ def test_ecs_group_compression_removes_only_uniform_coordinate_copies(): assert metadata["ecs_expanded_mode_count"] == 12 +def test_mp_median_and_optimal_signal_rank_are_auditable(): + module = load_cli_module() + median = module.marchenko_pastur_median(1.0) + assert 0.60 < median < 0.75 + singular = np.array([8.0, 5.0, 1.2, 1.1, 1.0, 0.9]) + rank, metadata = module.optimal_mp_signal_rank( + singular, (6, 6), numerical_rank=6 + ) + assert 2 <= rank < 6 + assert metadata["mp_raw_signal_rank"] >= 0 + assert metadata["mp_selected_signal_rank"] == rank + assert metadata["mp_selection_is_frozen_during_differentiation"] is True + + +def test_tikhonov_grid_selection_uses_ks_not_alpha_target(monkeypatch): + module = load_cli_module() + + def fake_fit(_amplitudes, _record, metadata, _top_k, _minimum_tail): + ratio = float(metadata["tikhonov_z_boundary_ratio"]) + return [{ + "spectrum_kind": "energy_derived_from_amplitude", + "clip_top_k": 0, + "alpha": 2.0 if ratio == 0.1 else 7.0, + "ks_D": 0.20 if ratio == 0.1 else 0.05, + "xmin": 1.0, + "n_tail": 20, + "tail_decades": 1.0, + "fit_ok": True, + }] + + monkeypatch.setattr(module, "fit_spectrum", fake_fit) + candidates = { + "near_two": ( + np.array([1.0, 2.0]), SimpleNamespace(), + {"tikhonov_z_boundary_ratio": 0.1, "tikhonov_candidate_index": 0}, + ), + "better_ks": ( + np.array([1.0, 2.0]), SimpleNamespace(), + {"tikhonov_z_boundary_ratio": 1.0, "tikhonov_candidate_index": 1}, + ), + } + selected_id, _, rows = module.select_tikhonov_candidate( + candidates, base_metadata={}, minimum_tail=2 + ) + assert selected_id == "better_ks" + assert next(row for row in rows if row["selected"])["alpha"] == 7.0 + assert "alpha proximity to 2 is never used" in rows[0]["selection_rule"] + + def test_complete_cli_separates_state_flow_and_local_jacobian_claims(): source = (EXPERIMENT_SCRIPTS / "run_short100_quotient_flow_cli.py").read_text() wrapper = (EXPERIMENT_SCRIPTS / "run_short100_complete_rg_analysis.sh").read_text() @@ -72,8 +121,11 @@ def test_complete_cli_separates_state_flow_and_local_jacobian_claims(): assert "analyze_weightwatcher_dual" in source assert '"gram_ridge"' in source assert '"feshbach_downfolding"' in source + assert '"calibrated_mp_shrinker"' in source assert '"is_training_jacobian": False' in source assert "run_short100_quotient_flow_cli.py" in wrapper assert "Case 1 — heavy tails on a weight quotient representative" in report assert "Case 2a — flow between checkpoints" in report assert "Case 2b — a Jacobian at one checkpoint" in report + assert "Tikhonov grid" in report + assert "Optimal MP space" in report diff --git a/baseline/tests/tangent_rg/test_short100_jacobian_report.py b/baseline/tests/tangent_rg/test_short100_jacobian_report.py index b888e1e..4af1876 100644 --- a/baseline/tests/tangent_rg/test_short100_jacobian_report.py +++ b/baseline/tests/tangent_rg/test_short100_jacobian_report.py @@ -53,8 +53,65 @@ def test_method_coverage_requires_ecs_on_fc1_and_fc2_but_not_fc3(): }) coverage = module.build_method_coverage(pd.DataFrame(rows)) assert coverage["coverage_status"].eq("complete").all() - assert len(coverage) == 30 + assert len(coverage) == sum( + len(methods) for methods in module.EXPECTED_METHODS_BY_LAYER.values() + ) * len(module.OPTIMIZERS) 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"} + + +def test_end_to_end_contract_counts_every_alpha_family(): + module = load_report_module() + primary = pd.DataFrame([ + {"optimizer": optimizer, "layer": layer, "epoch": epoch, "method": method} + for optimizer in module.OPTIMIZERS + for layer, methods in module.EXPECTED_METHODS_BY_LAYER.items() + for epoch in module.ANALYSIS_EPOCHS + for method in methods + ]) + search = pd.DataFrame([ + { + "optimizer": optimizer, "layer": layer, "epoch": epoch, + "tikhonov_z_boundary_ratio": ratio, + "selected": ratio == module.TIKHONOV_Z_RATIOS[0], + } + for optimizer in module.OPTIMIZERS + for layer in ("fc1.weight", "fc2.weight") + for epoch in module.ANALYSIS_EPOCHS + for ratio in module.TIKHONOV_Z_RATIOS + ]) + quotient = pd.DataFrame([ + { + "optimizer": optimizer, "layer": layer, "epoch": epoch, + "method": method, "profile_id": profile, "fit_variant": variant, + } + for optimizer in module.OPTIMIZERS + for layer in ("fc1.weight", "fc2.weight") + for epoch in module.ANALYSIS_EPOCHS + for method, profile in module.QUOTIENT_EXPECTED_PROFILES + for variant in ("raw", "clip_xmax") + ]) + flow = pd.DataFrame([ + { + "optimizer": optimizer, "layer": layer, "epoch_end": epoch, + "method": method, "spectrum_kind": "energy_derived_from_amplitude", + "clip_top_k": 0, + } + for optimizer in module.OPTIMIZERS + for layer in ("fc1.weight", "fc2.weight") + for epoch in module.ANALYSIS_EPOCHS[1:] + for method in module.FLOW_LABELS + ]) + transport = pd.DataFrame([ + {"optimizer": optimizer, "layer": layer, "epoch_end": epoch} + for optimizer in module.OPTIMIZERS + for layer in ("fc1.weight", "fc2.weight") + for epoch in module.ANALYSIS_EPOCHS[1:] + ]) + coverage = module.build_analysis_contract_coverage( + primary, search, quotient, flow, transport + ) + assert coverage["coverage_status"].eq("complete").all() + assert coverage["missing_unit_count"].eq(0).all()