diff --git a/baseline/experiments/mnist_mlp3_tangent_rg/README.md b/baseline/experiments/mnist_mlp3_tangent_rg/README.md index eda9a7cb..06feb01b 100644 --- a/baseline/experiments/mnist_mlp3_tangent_rg/README.md +++ b/baseline/experiments/mnist_mlp3_tangent_rg/README.md @@ -890,3 +890,48 @@ 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. + +For the complete notebook-free state/flow/local-response experiment, run: + +```bash +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: + +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: + 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`. + 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 + `weight_quotient_operators.csv`. +2. **Between-checkpoint RG flow.** Consecutive 10-epoch checkpoints produce + generalized-Gram radial rates, Procrustes-aligned transfer-core rates, + midpoint-ECS top-k Grassmann rates, and relative-polar tilt rates. These are + finite secant/transfer observables and are explicitly never labelled as the + optimizer Jacobian `D beta(W)`. The complete amplitudes and power-law fits + are written to `two_checkpoint_flow_*.csv`. + For the centered log-singular quotient coordinate it also compares the + observed secant `[R(W1)-R(W0)]/delta_s` with the true local map-Jacobian + prediction `D R_W0[W1-W0]/delta_s`, saving relative error and cosine in + `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. + Square FC2 is handled by the same right-singular top-k Grassmann geometry as + wide FC1. + +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. 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 30490dea..0f82581b 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 @@ -33,21 +33,57 @@ "centered_log_singular_radial_pullback": "Centered log-singular radial", "ecs_grassmann_cartan_cover_full_row_shell_pullback": "ECS full-row shell", "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)", + "feshbach_trace_free_log_detx_shell_pullback": "Feshbach trace-free log (detX)", } METHOD_STYLES = { "centered_log_singular_radial_pullback": ("-", "o"), "ecs_grassmann_cartan_cover_full_row_shell_pullback": ("--", "s"), - "ecs_grassmann_cartan_cover_detx_shell_pullback": (":", "^") + "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": ( + (0, (3, 1, 1, 1)), "X" + ), + "feshbach_trace_free_log_detx_shell_pullback": ((0, (1, 1)), "v"), } ECS_METHODS = ( "ecs_grassmann_cartan_cover_full_row_shell_pullback", "ecs_grassmann_cartan_cover_detx_shell_pullback", ) +EXTENDED_ECS_METHODS = ( + "gap_aware_projector_detx_shell_pullback", + "trace_free_log_gram_detx_shell_pullback", + "gram_ridge_resolvent_detx_shell_zratio_0p50_pullback", + "feshbach_trace_free_log_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), + "fc1.weight": ("centered_log_singular_radial_pullback", *ECS_METHODS, *EXTENDED_ECS_METHODS), + "fc2.weight": ("centered_log_singular_radial_pullback", *ECS_METHODS, *EXTENDED_ECS_METHODS), "fc3.weight": ("centered_log_singular_radial_pullback",), } +QUOTIENT_LABELS = { + "midpoint_ecs_control": "Midpoint ECS control", + "gram_ridge": "Gram diagonal subtraction", + "feshbach_downfolding": "Anchor-frozen Feshbach", +} +QUOTIENT_PROFILE_LABELS = { + "midpoint": "no counterterm", + "tau_fraction_0p25": "τ=0.25 λboundary", + "tau_fraction_0p50": "τ=0.50 λboundary", + "tau_fraction_0p75": "τ=0.75 λboundary", + "ridge_ratio_1em2": "ρ=0.01 anchor scale", +} +FLOW_LABELS = { + "two_checkpoint_generalized_gram_radial": "Generalized-Gram radial rate", + "two_checkpoint_aligned_transfer_core": "Aligned transfer-core rate", + "two_checkpoint_ecs_topk_grassmann": "ECS top-k Grassmann rate", + "two_checkpoint_relative_polar_tilt": "Relative-polar tilt rate", + "two_checkpoint_radial_quotient_observed_secant": "Observed radial quotient secant", + "two_checkpoint_radial_jacobian_prediction": "Local radial Jacobian prediction", +} def configure_logging(report_root: Path) -> logging.Logger: @@ -198,6 +234,141 @@ def plot_ecs_quotient_comparison(fits: pd.DataFrame, path: Path) -> Path: return save_figure(fig, path, bottom=0.18, top=0.92) +def plot_single_optimizer_jacobians( + fits: pd.DataFrame, optimizer: str, path: Path +) -> Path: + """Make a readable optimizer-specific view of every local map Jacobian.""" + palette = plt.get_cmap("tab10") + fig, axes = plt.subplots(1, 3, figsize=(17.0, 5.0), sharex=True) + subset_optimizer = fits[fits["optimizer"].astype(str).eq(optimizer)] + method_order = [ + method for method in METHOD_LABELS + if subset_optimizer["method"].astype(str).eq(method).any() + ] + for axis, layer in zip(axes, LAYERS): + subset = subset_optimizer[subset_optimizer["layer"].astype(str).eq(layer)] + for index, method in enumerate(method_order): + curve = subset[subset["method"].astype(str).eq(method)].sort_values("epoch") + if curve.empty: + continue + linestyle, marker = METHOD_STYLES.get(method, ("-", "o")) + axis.plot( + curve["epoch"], curve["alpha"], color=palette(index % 10), + linestyle=linestyle, marker=marker, markersize=3.5, + linewidth=1.6, label=METHOD_LABELS.get(method, method), + ) + axis.axhline(2.0, color="black", linestyle=(0, (1, 2)), linewidth=1.1) + axis.set(title=layer.replace(".weight", ""), xlabel="epoch", ylabel="Jacobian energy alpha") + axis.grid(True, alpha=0.25) + shared_legend(fig, axes, columns=4) + fig.suptitle(f"{OPTIMIZER_LABELS[optimizer]} — single-checkpoint quotient-map Jacobians") + return save_figure(fig, path, bottom=0.18, top=0.93) + + +def plot_weight_quotients_by_optimizer( + quotient_fits: pd.DataFrame, optimizer: str, path: Path +) -> Path: + frame = quotient_fits[ + quotient_fits["optimizer"].astype(str).eq(optimizer) + & quotient_fits["layer"].astype(str).isin(("fc1.weight", "fc2.weight")) + ].copy() + frame = frame[bool_series(frame["fit_ok"])] + profile_colors = { + "midpoint": "#4C78A8", + "tau_fraction_0p25": "#F2CF5B", + "tau_fraction_0p50": "#F58518", + "tau_fraction_0p75": "#B279A2", + "ridge_ratio_1em2": "#54A24B", + } + fig, axes = plt.subplots(2, 2, figsize=(13.5, 8.4), sharex=True) + variants = ("raw", "clip_xmax") + for row, variant in enumerate(variants): + for column, layer in enumerate(("fc1.weight", "fc2.weight")): + axis = axes[row, column] + subset = frame[ + frame["layer"].astype(str).eq(layer) + & frame["fit_variant"].astype(str).eq(variant) + ] + for (method, profile), curve in subset.groupby(["method", "profile_id"]): + curve = curve.sort_values("epoch") + axis.plot( + curve["epoch"], curve["alpha"], + color=profile_colors.get(str(profile), "#777777"), + marker="s" if variant == "clip_xmax" else "o", + markersize=3.5, linewidth=1.7, + label=( + f"{QUOTIENT_LABELS.get(str(method), str(method))}, " + f"{QUOTIENT_PROFILE_LABELS.get(str(profile), str(profile))}" + ), + ) + axis.axhline(2.0, color="black", linestyle=(0, (1, 2)), linewidth=1.1) + axis.set( + title=f"{layer.replace('.weight', '')} — {variant}", + xlabel="epoch", ylabel="WeightWatcher alpha of W′", + ) + axis.grid(True, alpha=0.25) + shared_legend(fig, axes.ravel(), columns=3) + fig.suptitle(f"{OPTIMIZER_LABELS[optimizer]} — transformed weight quotient representatives") + return save_figure(fig, path, bottom=0.14, top=0.94) + + +def plot_checkpoint_flows_by_optimizer( + flow_fits: pd.DataFrame, optimizer: str, path: Path +) -> Path: + primary = flow_fits[ + flow_fits["optimizer"].astype(str).eq(optimizer) + & flow_fits["spectrum_kind"].astype(str).eq("energy_derived_from_amplitude") + & pd.to_numeric(flow_fits["clip_top_k"], errors="coerce").eq(0) + ] + palette = plt.get_cmap("Dark2") + fig, axes = plt.subplots(1, 2, figsize=(13.5, 4.9), sharex=True) + for axis, layer in zip(axes, ("fc1.weight", "fc2.weight")): + subset = primary[primary["layer"].astype(str).eq(layer)] + for index, (method, curve) in enumerate(subset.groupby("method")): + curve = curve.sort_values("epoch_end") + axis.plot( + curve["epoch_end"], curve["alpha"], color=palette(index), + marker="o", markersize=3.5, linewidth=1.7, + label=FLOW_LABELS.get(str(method), str(method)), + ) + axis.axhline(2.0, color="black", linestyle=(0, (1, 2)), linewidth=1.1) + axis.set(title=layer.replace(".weight", ""), xlabel="ending epoch", ylabel="finite-flow energy alpha") + axis.grid(True, alpha=0.25) + shared_legend(fig, axes, columns=2) + fig.suptitle(f"{OPTIMIZER_LABELS[optimizer]} — between-checkpoint RG-flow observables (not Jacobians)") + return save_figure(fig, path, bottom=0.18, top=0.92) + + +def plot_jacobian_transport_by_optimizer( + transport: pd.DataFrame, optimizer: str, path: Path +) -> Path: + """Compare finite quotient flow to the local map-Jacobian prediction.""" + frame = transport[transport["optimizer"].astype(str).eq(optimizer)] + fig, axes = plt.subplots(2, 2, figsize=(13.5, 8.0), sharex=True) + for column, layer in enumerate(("fc1.weight", "fc2.weight")): + curve = frame[frame["layer"].astype(str).eq(layer)].sort_values("epoch_end") + axes[0, column].plot( + curve["epoch_end"], curve["relative_linearization_error"], + color=COLORS[optimizer], marker="o", linewidth=1.8, + ) + axes[1, column].plot( + curve["epoch_end"], curve["cosine_observed_vs_jacobian"], + color=COLORS[optimizer], marker="s", linewidth=1.8, + ) + axes[0, column].set( + title=layer.replace(".weight", ""), + ylabel="relative linearization error", + ) + axes[1, column].set(xlabel="ending epoch", ylabel="observed/JVP cosine") + axes[1, column].axhline(1.0, color="black", linestyle=(0, (1, 2)), linewidth=1.0) + for axis in axes[:, column]: + axis.grid(True, alpha=0.25) + fig.suptitle( + f"{OPTIMIZER_LABELS[optimizer]} — observed radial quotient flow versus local Jacobian" + ) + return save_figure(fig, path, top=0.94) + + def build_method_coverage(primary: pd.DataFrame) -> pd.DataFrame: """Return an explicit expected-versus-observed method inventory.""" rows = [] @@ -406,14 +577,14 @@ def build_html( for name, count in inventory.items() ) html = f""" -Short100 Jacobian report +Short100 quotient and RG-flow report -

Short100 reduced Jacobian analysis

+

Short100 weight quotient, RG-flow, and Jacobian analysis

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. @@ -442,6 +613,59 @@ 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} +

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 +two-sided orthogonal orbit +[W] = {{U W Vᵀ : U∈O(m), V∈O(n)}}. Its invariants are the singular +values. We choose the rectangular-diagonal canonical section and then apply a +declared nuisance-removal hypothesis before running WeightWatcher on the +materialized matrix W′, both raw and with +fix_fingers=clip_xmax. A heavy tail found here is therefore a heavy +tail of that declared quotient representative—not proof that it is the unique +unknown Muon time-orbit quotient.

+ +

Case 2a — flow between checkpoints

+

Two checkpoints identify finite flow/secant observables: generalized-Gram +radial log rates, a Procrustes-aligned rectangular transfer core, top-k ECS +Grassmann principal-angle rates, and relative-polar tilt rates. Their squared +rate amplitudes are fitted for heavy tails. These show how the saved trajectory +moves in quotient coordinates. They are not +Dβ(W), the Jacobian of the optimizer/RG vector field; identifying +that object requires perturbation trajectories or a fully specified training +map with batch and optimizer state.

+

The centered log-singular coordinate additionally supplies a direct local +transport test. For each pair the report compares the observed secant +[R(W₁)−R(W₀)]/Δt with the actual map-Jacobian prediction +D R_W₀[W₁−W₀]/Δt, reporting both spectra, relative error, and cosine. +This tests whether the saved step follows the local quotient-map linearization; +it still does not turn that map derivative into Dβ(W).

+

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 +training-dynamics Jacobian unless the declared map is separately calibrated to +the optimizer step.

+

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.

Data inventory

Analysis-ready tables

Final checkpoint primary fits

{final_table} @@ -482,6 +706,17 @@ 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") + quotient_fits = require_csv( + analysis_root / "weight_quotient_weightwatcher_fits.csv" + ) + quotient_spectra = require_csv(analysis_root / "weight_quotient_spectra.csv") + quotient_operators = require_csv(analysis_root / "weight_quotient_operators.csv") + flow_fits = require_csv(analysis_root / "two_checkpoint_flow_fits.csv") + flow_spectra = require_csv(analysis_root / "two_checkpoint_flow_spectra.csv") + flow_operators = require_csv(analysis_root / "two_checkpoint_flow_operators.csv") + flow_transport = require_csv( + analysis_root / "two_checkpoint_jacobian_transport.csv" + ) primary = fits[ fits["spectrum_kind"].astype(str).eq("energy_derived_from_amplitude") & pd.to_numeric(fits["clip_top_k"], errors="coerce").eq(0) @@ -511,6 +746,13 @@ def main() -> int: table_root / "weightwatcher_raw_and_clip_xmax.csv", table_root / "performance_train_test.csv", table_root / "jacobian_method_coverage.csv", + table_root / "weight_quotient_weightwatcher_fits.csv", + table_root / "weight_quotient_spectra.csv", + table_root / "weight_quotient_operators.csv", + table_root / "two_checkpoint_flow_fits.csv", + table_root / "two_checkpoint_flow_spectra.csv", + table_root / "two_checkpoint_flow_operators.csv", + table_root / "two_checkpoint_jacobian_transport.csv", ] primary.to_csv(tables[0], index=False) fits.to_csv(tables[1], index=False) @@ -519,6 +761,13 @@ def main() -> int: 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) figures = [ plot_jacobian_metric( @@ -543,6 +792,25 @@ def main() -> int: primary, performance, figure_root / "06_alpha_vs_test_accuracy.png" ), ] + for optimizer in OPTIMIZERS: + figures.extend([ + plot_single_optimizer_jacobians( + primary, optimizer, + figure_root / f"optimizer_views/{optimizer}_single_checkpoint_jacobians.png", + ), + plot_weight_quotients_by_optimizer( + quotient_fits, optimizer, + figure_root / f"optimizer_views/{optimizer}_weight_quotients.png", + ), + plot_checkpoint_flows_by_optimizer( + flow_fits, optimizer, + figure_root / f"optimizer_views/{optimizer}_checkpoint_flows.png", + ), + plot_jacobian_transport_by_optimizer( + flow_transport, optimizer, + figure_root / f"optimizer_views/{optimizer}_jacobian_transport.png", + ), + ]) figures.extend(plot_spectral_galleries(spectra, figure_root)) inventory = { "primary Jacobian fit rows": len(primary), @@ -551,6 +819,11 @@ def main() -> int: "operator metadata rows": len(operators), "WeightWatcher control rows": len(weightwatcher), "performance rows": len(performance), + "weight-quotient WeightWatcher rows": len(quotient_fits), + "weight-quotient spectral modes": len(quotient_spectra), + "two-checkpoint flow fit rows": len(flow_fits), + "two-checkpoint flow spectral modes": len(flow_spectra), + "Jacobian transport comparison rows": len(flow_transport), "figures": len(figures), } index = build_html(report_root, figures, tables, primary, coverage, inventory) 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 new file mode 100755 index 00000000..6afff41d --- /dev/null +++ b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_complete_rg_analysis.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash + +# Complete notebook-free analysis: single-checkpoint Jacobians, transformed +# weight quotient representatives, finite between-checkpoint RG flow, and the +# static shareable HTML report. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)" +OUTPUT_ROOT="${RG_MNIST_REDUCED_JACOBIAN_OUTPUT_ROOT:-/private/tmp/rg-mnist-mlp3-short100-jacobians-reduced}" + +export RG_MNIST_REDUCED_JACOBIAN_OUTPUT_ROOT="${OUTPUT_ROOT}" + +bash "${SCRIPT_DIR}/run_short100_jacobians_reduced.sh" "$@" + +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}" + +bash "${SCRIPT_DIR}/build_short100_jacobian_report.sh" 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 de8213a6..1a97222b 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 @@ -46,6 +46,12 @@ "centered_log_singular_radial_pullback", "finite_muon_ns5_pullback", ) +EXTENDED_DETX_METHODS = ( + "gap_aware_projector_detx_shell_pullback", + "trace_free_log_gram_detx_shell_pullback", + "gram_ridge_resolvent_detx_shell_zratio_0p50_pullback", + "feshbach_trace_free_log_detx_shell_pullback", +) class MaxInfoFilter(logging.Filter): @@ -440,6 +446,78 @@ def ecs_fit_amplitudes(record: Any, *, compress_groups: bool) -> tuple[np.ndarra } +def extended_detx_jacobian_spectra( + weight: np.ndarray, + singular_values: np.ndarray, + *, + retained_rank: int, + outer_rank: int, + rcond: float, +) -> 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 + at first order in the checkpoint SVD gauge; that collapse is a scientific + control, not silently interpreted as nontrivial downfolding. + """ + 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 + records = { + "gap_aware_projector_detx_shell_pullback": ( + ecs_jacobians.gap_aware_projector_spectrum( + weight, retained_rank=k, outer_rank=q, rcond=rcond, + precomputed_singular_values=singular_values, + ), + {"jacobian_family": "gap_aware_projector", "retained_rank": k, "outer_rank": q}, + ), + "trace_free_log_gram_detx_shell_pullback": ( + ecs_jacobians.outer_trace_free_log_gram_spectrum( + weight, outer_rank=q, rcond=rcond, + precomputed_singular_values=singular_values, + ), + {"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, + ), + { + "jacobian_family": "feshbach_trace_free_log_effective_core", + "retained_rank": k, "outer_rank": q, + "feshbach_z": feshbach_z, + "first_order_shell_downfolding_active": False, + }, + ), + } + return { + method: ( + np.asarray(record.singular_amplitudes, dtype=float), + record, + metadata, + ) + for method, (record, metadata) in records.items() + } + + def safe_slug(text: str) -> str: return "".join(char if char.isalnum() or char in "-_" else "_" for char in text) @@ -527,7 +605,20 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--minimum-tail", type=int, default=8) parser.add_argument("--ns-steps", type=int, default=5) parser.add_argument("--ns-eps", type=float, default=1e-7) + parser.add_argument( + "--ecs-rcond", type=float, default=1e-9, + help="relative numerical-rank and differentiability tolerance for ECS maps", + ) parser.add_argument("--skip-ecs", action="store_true") + parser.add_argument( + "--extended-ecs-jacobians", + 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" + ), + ) parser.add_argument( "--compress-ecs-groups", action=argparse.BooleanOptionalAction, @@ -646,6 +737,8 @@ def run(args: argparse.Namespace) -> int: "ecs_grassmann_cartan_cover_full_row_shell_pullback", "ecs_grassmann_cartan_cover_detx_shell_pullback", }) + 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: completed += 1 @@ -702,7 +795,9 @@ def same_unit(row: dict[str, Any]) -> bool: 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])) + numerical_rank = int( + np.count_nonzero(singular > args.ecs_rcond * singular[0]) + ) ecs_rank_records = exact_ecs_ranks( identity["seed_dir"], optimizer, seed, int(ref.epoch), int(ref.global_step), layer, numerical_rank, @@ -722,7 +817,7 @@ def same_unit(row: dict[str, Any]) -> bool: ) 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, + weight, retained_rank=k, outer_rank=q, rcond=args.ecs_rcond, precomputed_singular_values=singular, ) amplitudes, compression_metadata = ecs_fit_amplitudes( @@ -739,6 +834,35 @@ def same_unit(row: dict[str, Any]) -> bool: method, k, q, len(amplitudes), int(cover.derivative_rank), args.compress_ecs_groups, ) + if ( + args.extended_ecs_jacobians + and method + == "ecs_grassmann_cartan_cover_detx_shell_pullback" + ): + for extended_method, ( + extended_amplitudes, + extended_record, + extended_metadata, + ) in extended_detx_jacobian_spectra( + weight, + singular, + retained_rank=k, + outer_rank=q, + rcond=args.ecs_rcond, + ).items(): + method_factories[extended_method] = ( + extended_amplitudes, + extended_record, + ) + method_metadata[extended_method] = { + **rank_metadata, + **extended_metadata, + "ecs_shell_variant": "detx_shell", + } + logger.info( + "EXTENDED JACOBIAN method=%s k=%d q=%d n_amplitudes=%d", + extended_method, k, q, len(extended_amplitudes), + ) base = { "optimizer": optimizer, @@ -838,6 +962,7 @@ def same_unit(row: dict[str, Any]) -> bool: not args.skip_ecs and layer in ecs_layers ), "ecs_layers_requested": ",".join(ecs_layers), + "extended_ecs_jacobians": bool(args.extended_ecs_jacobians), "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 3461bbab..5e466060 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 @@ -12,6 +12,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 \ + --extended-ecs-jacobians \ --compress-ecs-groups \ --top-k 0 \ "$@" 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 new file mode 100755 index 00000000..6cbc1de2 --- /dev/null +++ b/baseline/experiments/mnist_mlp3_tangent_rg/scripts/run_short100_quotient_flow_cli.py @@ -0,0 +1,525 @@ +#!/usr/bin/env python3 +"""Command-line weight-quotient and two-checkpoint RG-flow analysis. + +This program deliberately separates two questions that are often conflated: + +1. ``state quotient``: construct an auditable representative of a declared + equivalence class of ``W`` and ask whether its Gram ESD is heavy tailed; +2. ``checkpoint flow``: measure finite radial and angular motion from ``W_t`` + to ``W_{t+1}``. These are secant/transfer observables, not ``D beta(W)``. + +The transformed matrices are actually installed into an MLP3 and analyzed by +WeightWatcher both raw and with ``fix_fingers='clip_xmax'``. Every phase logs +to stdout/stderr, writes resumable CSVs after each unit, and records failures. +""" + +from __future__ import annotations + +import argparse +import copy +from dataclasses import asdict, is_dataclass +from datetime import datetime, timezone +import hashlib +import json +import logging +from pathlib import Path +import sys +import time +import traceback +from types import SimpleNamespace +from typing import Any + +SCRIPT_DIR = Path(__file__).resolve().parent +BASELINE_ROOT = Path(__file__).resolve().parents[3] +for candidate in (SCRIPT_DIR, BASELINE_ROOT): + if str(candidate) not in sys.path: + sys.path.insert(0, str(candidate)) + +import numpy as np +import pandas as pd + +import run_short100_jacobians_cli as jacobian_cli +from rg_baselines.model import MLP3 +from rg_baselines.tangent_rg import single_checkpoint, two_checkpoint, weight_quotients +from rg_baselines.tangent_rg.weightwatcher_fit import analyze_weightwatcher_dual + + +LAYERS = ("fc1.weight", "fc2.weight") +QUOTIENT_PROFILES = ( + ("midpoint_ecs_control", "midpoint", {}), + ("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}), + ( + "feshbach_downfolding", + "ridge_ratio_1em2", + {"regularization_ratio": 1.0e-2}, + ), +) + + +def atomic_frame(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + pd.DataFrame(rows).to_csv(temporary, index=False) + temporary.replace(path) + + +def read_rows(path: Path) -> list[dict[str, Any]]: + if not path.is_file(): + return [] + try: + return pd.read_csv(path).to_dict(orient="records") + except pd.errors.EmptyDataError: + return [] + + +def configure_logging(root: Path, verbose: bool) -> logging.Logger: + root.mkdir(parents=True, exist_ok=True) + logger = logging.getLogger("short100_quotient_flow") + logger.handlers.clear() + logger.setLevel(logging.DEBUG if verbose else logging.INFO) + formatter = logging.Formatter( + "%(asctime)s %(levelname)-8s %(message)s", datefmt="%Y-%m-%d %H:%M:%S" + ) + for stream in (sys.stdout, sys.stderr): + handler = logging.StreamHandler(stream) + handler.setLevel(logging.INFO if stream is sys.stdout else logging.WARNING) + handler.setFormatter(formatter) + logger.addHandler(handler) + file_handler = logging.FileHandler(root / "quotient_flow.log", mode="a") + file_handler.setLevel(logging.DEBUG) + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) + return logger + + +def checkpoint_model(path: Path, fingerprint: str) -> MLP3: + from rg_baselines.tangent_rg.checkpoints import load_analysis_checkpoint + + payload = load_analysis_checkpoint(path, expected_fingerprint=fingerprint) + model = MLP3().to("cpu") + state = payload["model"] + normalized = { + (name.removeprefix("module.")): value for name, value in state.items() + } + model.load_state_dict(normalized, strict=True) + model.eval() + return model + + +def model_matrix(model: MLP3, layer: str) -> np.ndarray: + parameter = dict(model.named_parameters())[layer] + return parameter.detach().cpu().double().numpy() + + +def replace_matrix(model: MLP3, layer: str, value: np.ndarray) -> None: + import torch + + parameter = dict(model.named_parameters())[layer] + candidate = torch.as_tensor(value, dtype=parameter.dtype) + if tuple(candidate.shape) != tuple(parameter.shape): + raise ValueError(f"shape mismatch for {layer}: {candidate.shape} != {parameter.shape}") + with torch.no_grad(): + parameter.copy_(candidate) + + +def midpoint_metadata( + identity: dict[str, Any], optimizer: str, seed: int, ref: Any, layer: str, + numerical_rank: int, +) -> dict[str, Any]: + records = jacobian_cli.exact_ecs_ranks( + identity["seed_dir"], optimizer, seed, int(ref.epoch), + int(ref.global_step), layer, numerical_rank, + ) + if not records: + raise RuntimeError(f"no certified same-checkpoint ECS ranks for {layer} epoch={ref.epoch}") + metadata = dict(records[0][3]) + metadata["ecs_rank"] = weight_quotients.midpoint_ecs_rank( + metadata["k_pl"], metadata["k_tl"], maximum_rank=numerical_rank + ) + return metadata + + +def stable_seed(*parts: Any) -> int: + digest = hashlib.sha256("|".join(map(str, parts)).encode("utf-8")).digest() + return int.from_bytes(digest[:4], "big") + + +def record_metadata(record: Any) -> dict[str, Any]: + raw = asdict(record) if is_dataclass(record) else dict(vars(record)) + return { + key: value for key, value in raw.items() + if not isinstance(value, np.ndarray) + } + + +def run_state_quotients( + *, root: Path, run_root: Path, cache_root: Path, optimizers: tuple[str, ...], + seed: int, epoch_stride: int, logger: logging.Logger, +) -> None: + fit_path = root / "weight_quotient_weightwatcher_fits.csv" + spectrum_path = root / "weight_quotient_spectra.csv" + operator_path = root / "weight_quotient_operators.csv" + 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( + cache_root, optimizer, seed, epoch_stride=epoch_stride, + maximum_checkpoints=None, + ) + anchor_model = checkpoint_model(refs[0].path, identity["fingerprint"]) + for ref in refs: + base_model = checkpoint_model(ref.path, identity["fingerprint"]) + rank_by_layer = {} + for layer in LAYERS: + weight = model_matrix(base_model, layer) + singular = np.linalg.svd(weight, compute_uv=False) + numerical_rank = int(np.count_nonzero(singular > 1.0e-9 * singular[0])) + rank_by_layer[layer] = midpoint_metadata( + identity, optimizer, seed, ref, layer, numerical_rank + ) + for method, profile_id, parameters in QUOTIENT_PROFILES: + key = (optimizer, seed, int(ref.epoch), method, profile_id) + observed = { + (str(row.get("layer")), str(row.get("fit_variant"))) + for row in fits + if ( + str(row.get("optimizer")), int(row.get("seed", -1)), + int(row.get("epoch", -1)), str(row.get("method")), + str(row.get("profile_id")), + ) == 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) + continue + started = time.perf_counter() + logger.info("QUOTIENT START optimizer=%s epoch=%d method=%s profile=%s", optimizer, ref.epoch, method, profile_id) + model = copy.deepcopy(base_model) + results = {} + for layer in LAYERS: + result = weight_quotients.apply_weight_quotient( + method, + model_matrix(base_model, layer), + ecs_rank=int(rank_by_layer[layer]["ecs_rank"]), + anchor_weight=model_matrix(anchor_model, layer), + parameters=parameters, + ) + replace_matrix(model, layer, result.weight) + results[layer] = result + measurement = analyze_weightwatcher_dual( + model, + min_evals=8, + max_fingers=10, + svd_method="accurate", + randomize=True, + analysis_seed=stable_seed(optimizer, seed, ref.epoch, method, profile_id), + primary_variant="clip_xmax", + ) + fits = [ + row for row in fits + if not ( + str(row.get("optimizer")), int(row.get("seed", -1)), + int(row.get("epoch", -1)), str(row.get("method")), + str(row.get("profile_id")), + ) == key + ] + spectra = [ + row for row in spectra + if not ( + str(row.get("optimizer")), int(row.get("seed", -1)), + int(row.get("epoch", -1)), str(row.get("method")), + str(row.get("profile_id")), + ) == key + ] + operators = [ + row for row in operators + if not ( + str(row.get("optimizer")), int(row.get("seed", -1)), + int(row.get("epoch", -1)), str(row.get("method")), + str(row.get("profile_id")), + ) == key + ] + details = measurement.details[ + measurement.details["layer"].astype(str).isin(LAYERS) + ].copy() + details["optimizer"] = optimizer + details["seed"] = seed + details["epoch"] = int(ref.epoch) + details["global_step"] = int(ref.global_step) + details["method"] = method + details["profile_id"] = profile_id + details["analysis_family"] = "weight_state_quotient_representative" + details["equivalence_relation"] = "O(m)_left x O(n)_right orbit plus declared nuisance map" + fits.extend(details.to_dict(orient="records")) + for layer, result in results.items(): + base = { + "optimizer": optimizer, "seed": seed, + "epoch": int(ref.epoch), "global_step": int(ref.global_step), + "layer": layer, "method": method, "profile_id": profile_id, + "ecs_rank": int(result.ecs_rank), + "retained_rank": int(result.retained_rank), + "operator_kind": result.operator_kind, + "map_definition": result.map_definition, + "quotient_parameters": json.dumps(result.parameters, default=str, sort_keys=True), + } + operators.append(base) + spectra.extend({ + **base, + "mode_index_descending": index, + "singular_value": float(value), + "gram_eigenvalue": float(value * value), + } for index, value in enumerate(result.singular_values)) + 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) + + +def ecs_topk_rates(first: np.ndarray, second: np.ndarray, rank: int, delta_s: float) -> np.ndarray: + k = int(rank) + right0 = np.linalg.svd(first, full_matrices=False)[2][:k].T + right1 = np.linalg.svd(second, full_matrices=False)[2][:k].T + cosines = np.linalg.svd(right0.T @ right1, compute_uv=False) + angles = np.arccos(np.clip(cosines, 0.0, 1.0)) + return angles / abs(float(delta_s)) + + +def run_checkpoint_flows( + *, root: Path, run_root: Path, cache_root: Path, optimizers: tuple[str, ...], + seed: int, epoch_stride: int, logger: logging.Logger, +) -> None: + fit_path = root / "two_checkpoint_flow_fits.csv" + spectrum_path = root / "two_checkpoint_flow_spectra.csv" + operator_path = root / "two_checkpoint_flow_operators.csv" + transport_path = root / "two_checkpoint_jacobian_transport.csv" + fits = read_rows(fit_path) + 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( + cache_root, optimizer, seed, epoch_stride=epoch_stride, + maximum_checkpoints=None, + ) + for ref0, ref1 in zip(refs[:-1], refs[1:]): + delta_s = float(ref1.epoch - ref0.epoch) + for layer in LAYERS: + unit = (optimizer, seed, int(ref0.epoch), int(ref1.epoch), layer) + expected = { + "two_checkpoint_generalized_gram_radial", + "two_checkpoint_aligned_transfer_core", + "two_checkpoint_ecs_topk_grassmann", + "two_checkpoint_relative_polar_tilt", + "two_checkpoint_radial_quotient_observed_secant", + "two_checkpoint_radial_jacobian_prediction", + } + observed = { + str(row.get("method")) for row in operators + if ( + str(row.get("optimizer")), int(row.get("seed", -1)), + int(row.get("epoch_start", -1)), int(row.get("epoch_end", -1)), + str(row.get("layer")), + ) == unit + } + if expected.issubset(observed): + logger.info("FLOW SKIP optimizer=%s %d->%d layer=%s", optimizer, ref0.epoch, ref1.epoch, layer) + continue + started = time.perf_counter() + first = jacobian_cli.checkpoint_matrix(ref0.path, identity["fingerprint"], layer) + second = jacobian_cli.checkpoint_matrix(ref1.path, identity["fingerprint"], layer) + singular0 = np.linalg.svd(first, compute_uv=False) + singular1 = np.linalg.svd(second, compute_uv=False) + rank0 = midpoint_metadata(identity, optimizer, seed, ref0, layer, singular0.size)["ecs_rank"] + rank1 = midpoint_metadata(identity, optimizer, seed, ref1, layer, singular1.size)["ecs_rank"] + k = min(int(rank0), int(rank1)) + radial = two_checkpoint.generalized_gram_log_rates(first, second, delta_s) + transfer = two_checkpoint.aligned_rectangular_transfer(first, second, delta_s) + angular = two_checkpoint.relative_polar_angular_flow(first, second, delta_s) + radial_map0 = single_checkpoint.centered_log_singular_map(first) + radial_map1 = single_checkpoint.centered_log_singular_map(second) + radial_jvp = single_checkpoint.centered_log_singular_jvp( + first, second - first + ) + observed_radial = (radial_map1.value - radial_map0.value) / delta_s + predicted_radial = radial_jvp.jvp / delta_s + observed_norm = float(np.linalg.norm(observed_radial)) + predicted_norm = float(np.linalg.norm(predicted_radial)) + residual_norm = float(np.linalg.norm(predicted_radial - observed_radial)) + denominator = max( + observed_norm * predicted_norm, np.finfo(np.float64).tiny + ) + transport = { + "optimizer": optimizer, "seed": seed, "layer": layer, + "epoch_start": int(ref0.epoch), "epoch_end": int(ref1.epoch), + "global_step_start": int(ref0.global_step), + "global_step_end": int(ref1.global_step), + "delta_s": delta_s, + "map": "centered_log_singular_radial_quotient", + "observed_flow_norm": observed_norm, + "jacobian_prediction_norm": predicted_norm, + "linearization_residual_norm": residual_norm, + "relative_linearization_error": residual_norm / max( + observed_norm, np.finfo(np.float64).tiny + ), + "cosine_observed_vs_jacobian": float( + np.dot(observed_radial, predicted_radial) / denominator + ), + "is_actual_map_jvp": True, + "is_optimizer_beta_jacobian": False, + "map_definition": ( + "compare [R(W1)-R(W0)]/delta_s with " + "D R_W0[W1-W0]/delta_s for centered log singular R" + ), + } + cases = ( + ( + "two_checkpoint_generalized_gram_radial", + radial.radial_rate_amplitudes, + radial, + ), + ( + "two_checkpoint_aligned_transfer_core", + transfer.core_rate_amplitudes, + transfer, + ), + ( + "two_checkpoint_ecs_topk_grassmann", + ecs_topk_rates(first, second, k, delta_s), + SimpleNamespace( + operator_kind="two_checkpoint_topk_right_grassmann_geodesic_rates", + map_definition="principal angles between top-k right singular subspaces divided by epoch separation; finite flow, not a Jacobian", + ), + ), + ( + "two_checkpoint_relative_polar_tilt", + angular.tilt_geodesic_rates, + angular, + ), + ( + "two_checkpoint_radial_quotient_observed_secant", + np.abs(observed_radial), + SimpleNamespace( + operator_kind="observed_centered_log_singular_quotient_secant", + map_definition="absolute components of [R(W1)-R(W0)]/delta_s", + ), + ), + ( + "two_checkpoint_radial_jacobian_prediction", + np.abs(predicted_radial), + SimpleNamespace( + operator_kind="centered_log_singular_jacobian_prediction_along_checkpoint_step", + map_definition="absolute components of D R_W0[W1-W0]/delta_s; actual map JVP, not D beta(W)", + ), + ), + ) + base = { + "optimizer": optimizer, "seed": seed, "layer": layer, + "epoch_start": int(ref0.epoch), "epoch_end": int(ref1.epoch), + "global_step_start": int(ref0.global_step), + "global_step_end": int(ref1.global_step), + "epoch": int(ref1.epoch), "delta_s": delta_s, + "ecs_comparison_rank": k, + "analysis_family": "two_checkpoint_finite_rg_flow_not_jacobian", + "is_training_jacobian": False, + } + fits = [row for row in fits if not ( + str(row.get("optimizer")), int(row.get("seed", -1)), + int(row.get("epoch_start", -1)), int(row.get("epoch_end", -1)), + str(row.get("layer")), + ) == unit] + spectra = [row for row in spectra if not ( + str(row.get("optimizer")), int(row.get("seed", -1)), + int(row.get("epoch_start", -1)), int(row.get("epoch_end", -1)), + str(row.get("layer")), + ) == unit] + operators = [row for row in operators if not ( + str(row.get("optimizer")), int(row.get("seed", -1)), + int(row.get("epoch_start", -1)), int(row.get("epoch_end", -1)), + str(row.get("layer")), + ) == unit] + transports = [row for row in transports if not ( + str(row.get("optimizer")), int(row.get("seed", -1)), + int(row.get("epoch_start", -1)), int(row.get("epoch_end", -1)), + str(row.get("layer")), + ) == unit] + transports.append(transport) + for method, raw, record in cases: + amplitudes = np.asarray(raw, dtype=float) + amplitudes = amplitudes[np.isfinite(amplitudes) & (amplitudes > 0)] + metadata = {**base, "method": method} + if amplitudes.size >= 2: + fits.extend(jacobian_cli.fit_spectrum( + amplitudes, record, metadata, (0,), 8 + )) + spectra.extend({ + **metadata, + "mode_index_descending": index, + "flow_rate_amplitude": float(value), + "flow_rate_energy": float(value * value), + } for index, value in enumerate(np.sort(amplitudes)[::-1])) + operators.append({ + **metadata, + "operator_kind": record.operator_kind, + "map_definition": record.map_definition, + **record_metadata(record), + }) + atomic_frame(fit_path, fits) + 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) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-root", type=Path, default=jacobian_cli.DEFAULT_RUN_ROOT) + parser.add_argument("--cache-root", type=Path, default=jacobian_cli.DEFAULT_CACHE_ROOT) + parser.add_argument("--output-root", type=Path, default=Path("/private/tmp/rg-mnist-mlp3-short100-jacobians-reduced")) + parser.add_argument("--optimizers", default="muonclip_rms,adamw") + parser.add_argument("--seed", type=int, default=101) + parser.add_argument("--epoch-stride", type=int, default=10) + parser.add_argument("--skip-state-quotients", action="store_true") + parser.add_argument("--skip-checkpoint-flows", action="store_true") + parser.add_argument("--verbose", action="store_true") + return parser + + +def main() -> int: + args = build_parser().parse_args() + root = args.output_root.expanduser().resolve() + logger = configure_logging(root, args.verbose) + optimizers = jacobian_cli.parse_csv_values(args.optimizers) + started = time.perf_counter() + try: + logger.info("START quotient/flow analysis output_root=%s", root) + if not args.skip_state_quotients: + run_state_quotients( + root=root, run_root=args.run_root.resolve(), + cache_root=args.cache_root.resolve(), optimizers=optimizers, + seed=args.seed, epoch_stride=args.epoch_stride, logger=logger, + ) + if not args.skip_checkpoint_flows: + run_checkpoint_flows( + root=root, run_root=args.run_root.resolve(), + 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) + return 0 + except Exception: + trace = traceback.format_exc() + logger.error("FAILED\n%s", trace) + atomic_frame(root / "quotient_flow_errors.csv", [{ + "failed_at_utc": datetime.now(timezone.utc).isoformat(), + "exception_traceback": trace, + }]) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/baseline/rg_baselines/tangent_rg/ecs_jacobians.py b/baseline/rg_baselines/tangent_rg/ecs_jacobians.py index 76a7a8fa..cccbfef8 100644 --- a/baseline/rg_baselines/tangent_rg/ecs_jacobians.py +++ b/baseline/rg_baselines/tangent_rg/ecs_jacobians.py @@ -3,7 +3,8 @@ Every spectrum in this module is the nonzero singular spectrum of the derivative of a declared map. Squaring those amplitudes gives the nonzero eigenvalues of ``J* J`` in the same energy convention as a WeightWatcher ESD. -The checkpoint-dependent maps freeze their SVD frame at the named checkpoint; +The right-singular maps support wide and square matrices. The checkpoint- +dependent maps freeze their SVD frame at the named checkpoint; that anchoring is part of the map definition and is never hidden. """ @@ -62,8 +63,10 @@ def _matrix(value: ArrayLike, *, name: str = "weight") -> FloatArray: def _wide_matrix(value: ArrayLike, *, name: str = "weight") -> FloatArray: result = _matrix(value, name=name) - if result.shape[0] >= result.shape[1]: - raise ValueError("the ECS row-space Jacobians require a wide matrix") + if result.shape[0] > result.shape[1]: + raise ValueError( + "the right-singular ECS Jacobians require a wide or square matrix" + ) return result diff --git a/baseline/rg_baselines/tangent_rg/single_checkpoint.py b/baseline/rg_baselines/tangent_rg/single_checkpoint.py index fcb8ba66..619e4190 100644 --- a/baseline/rg_baselines/tangent_rg/single_checkpoint.py +++ b/baseline/rg_baselines/tangent_rg/single_checkpoint.py @@ -439,16 +439,16 @@ def _ecs_cover_frames( rcond: float | None, ) -> tuple[FloatArray, FloatArray, FloatArray, FloatArray, float]: rows, columns = matrix.shape - if rows >= columns: + if rows > columns: raise ValueError( - "the row-space ECS Grassmann cover is defined for a wide matrix" + "the right-singular ECS Grassmann cover requires a wide or square matrix" ) k = _integer_rank(retained_rank, name="retained_rank") q = _integer_rank(outer_rank, name="outer_rank") if not (1 <= k < q <= rows): raise ValueError( "ECS cover ranks must satisfy 1 <= retained_rank < outer_rank " - "<= the wide matrix row rank" + "<= the matrix row rank" ) left, singular_values, right_h = np.linalg.svd(matrix, full_matrices=True) largest = float(singular_values[0]) @@ -670,16 +670,16 @@ def ecs_grassmann_cover_analytic_spectrum( matrix = _matrix(weight) rows, columns = matrix.shape - if rows >= columns: + if rows > columns: raise ValueError( - "the row-space ECS Grassmann cover is defined for a wide matrix" + "the right-singular ECS Grassmann cover requires a wide or square matrix" ) k = _integer_rank(retained_rank, name="retained_rank") q = _integer_rank(outer_rank, name="outer_rank") if not (1 <= k < q <= rows): raise ValueError( "ECS cover ranks must satisfy 1 <= retained_rank < outer_rank " - "<= the wide matrix row rank" + "<= the matrix row rank" ) if precomputed_singular_values is None: singular_values = np.linalg.svd(matrix, compute_uv=False) diff --git a/baseline/tests/tangent_rg/test_command_line_jacobians.py b/baseline/tests/tangent_rg/test_command_line_jacobians.py index 4d77dcb0..fc4d6e8d 100644 --- a/baseline/tests/tangent_rg/test_command_line_jacobians.py +++ b/baseline/tests/tangent_rg/test_command_line_jacobians.py @@ -14,6 +14,7 @@ / "scripts" / "run_short100_jacobians_cli.py" ) +EXPERIMENT_SCRIPTS = SCRIPT.parent def load_cli_module(): @@ -46,6 +47,8 @@ def test_cli_duration_and_argument_defaults_are_observable_tmp_paths(): assert args.epoch_stride == 10 assert args.top_k == "0" assert args.ecs_layers == "fc1.weight,fc2.weight" + assert args.ecs_rcond == 1e-9 + assert args.extended_ecs_jacobians is False def test_ecs_group_compression_removes_only_uniform_coordinate_copies(): @@ -60,3 +63,17 @@ def test_ecs_group_compression_removes_only_uniform_coordinate_copies(): assert metadata["ecs_groups_compressed"] is True assert metadata["ecs_uniform_group_multiplicity"] == 4 assert metadata["ecs_expanded_mode_count"] == 12 + + +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() + report = (EXPERIMENT_SCRIPTS / "build_short100_jacobian_report.py").read_text() + assert "analyze_weightwatcher_dual" in source + assert '"gram_ridge"' in source + assert '"feshbach_downfolding"' 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 diff --git a/baseline/tests/tangent_rg/test_ecs_jacobians.py b/baseline/tests/tangent_rg/test_ecs_jacobians.py index adf98651..707f71c9 100644 --- a/baseline/tests/tangent_rg/test_ecs_jacobians.py +++ b/baseline/tests/tangent_rg/test_ecs_jacobians.py @@ -21,6 +21,9 @@ soft_ecs_projector_map, soft_ecs_projector_spectrum, ) +from rg_baselines.tangent_rg.single_checkpoint import ( + ecs_grassmann_cover_analytic_spectrum, +) def _explicit_jacobian(function, base, *, epsilon=1.0e-6): @@ -102,6 +105,21 @@ def test_gap_aware_projector_jvp_and_exact_spectrum(self): ) np.testing.assert_allclose(radial.jvp, 0.0, atol=1.0e-12) + def test_square_fc2_geometry_supports_right_singular_ecs_maps(self): + square = np.diag([5.0, 3.0, 2.0, 1.0]) + cover = ecs_grassmann_cover_analytic_spectrum( + square, retained_rank=2, outer_rank=4, rcond=1.0e-12 + ) + gap = gap_aware_projector_spectrum( + square, retained_rank=2, outer_rank=4, rcond=1.0e-12 + ) + feshbach = feshbach_trace_free_log_spectrum( + square, retained_rank=2, outer_rank=4, z=0.5, rcond=1.0e-12 + ) + self.assertEqual(cover.derivative_rank, 4) + self.assertEqual(gap.derivative_rank, 4) + self.assertEqual(feshbach.derivative_rank, 2) + def test_soft_ecs_projector_jvp_and_exact_spectrum(self): center, temperature = 6.5, 10.0 mapped = lambda candidate: soft_ecs_projector_map( diff --git a/baseline/tests/tangent_rg/test_short100_jacobian_report.py b/baseline/tests/tangent_rg/test_short100_jacobian_report.py index 6b0dfdf3..b888e1e1 100644 --- a/baseline/tests/tangent_rg/test_short100_jacobian_report.py +++ b/baseline/tests/tangent_rg/test_short100_jacobian_report.py @@ -53,6 +53,7 @@ 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 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"])