summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorYurenHao0426 <Blackhao0426@gmail.com>2026-06-05 15:44:19 -0500
committerYurenHao0426 <Blackhao0426@gmail.com>2026-06-05 15:44:19 -0500
commit2df52d3dd7b35053d4a357c3814d7cba6db602ce (patch)
tree96c08971d94cd931eb2727e3479d0b63f7319c5a
parent4dee090530588792b0e0bce9a36981415e4a99a8 (diff)
Derive actual FA initial operator moment
-rw-r--r--notes/18_contribution_roadmap.md34
-rw-r--r--notes/28_contribution_status_proof_vs_fit.md38
-rw-r--r--notes/29_actual_fa_initial_operator_moment_theorem.md327
-rw-r--r--scripts/actual_fa_initial_operator_moments.py293
4 files changed, 686 insertions, 6 deletions
diff --git a/notes/18_contribution_roadmap.md b/notes/18_contribution_roadmap.md
index be4d775..aa26caf 100644
--- a/notes/18_contribution_roadmap.md
+++ b/notes/18_contribution_roadmap.md
@@ -93,6 +93,33 @@ conditioning loss, or effective-dimension loss.
This is the soft capacity-erosion contribution.
+There is now a cleaner actual-FA initialization theorem that should be used
+before the random-subspace null model. For fixed forward weights and residuals,
+with independent zero-mean feedback matrices:
+
+```text
+E_B[speed_FA | W, r] = ||g_output^BP||^2
+```
+
+where:
+
+```text
+speed_FA = sum_l <g_l^BP, g_l^FA>
+speed_BP = sum_l ||g_l^BP||^2
+```
+
+Therefore:
+
+```text
+E_B[e_0 | W, r]
+ = 1 - ||g_output^BP||^2 / sum_l ||g_l^BP||^2
+```
+
+So the true initial FA burden is the hidden-layer BP speed share, not the raw
+hard feedback-matrix constraint count. The hard-k random-subspace model remains
+a useful null model and capacity sanity check, but it should not be presented
+as the exact FA gap theorem.
+
## 3. Prior-Free Minimax Initialization Bound
Let `b` be the normalized feedback direction in `D` dimensions and `a` be the
@@ -169,11 +196,16 @@ Experiments should validate three levels:
```
2. capacity and soft erosion:
FA/BP gap distribution changes smoothly with hard FA capacity margin;
-3. tangent-operator trajectory distributions:
+3. actual-FA initial operator moments:
+ ```text
+ E_B[e_0 | W,r] = hidden BP speed share
+ ```
+4. tangent-operator trajectory distributions:
predicted gap distributions overlap empirical trajectory distributions.
Current strong evidence:
+- initial actual-FA operator moment validation over widths and feedback seeds;
- 256-trajectory finite-time overlap plot for `d=2,w=64,T=50,s=20`;
- stress grid over depth, width, and horizon;
- systematic low-bias analysis showing residual error is kernel-path curvature,
diff --git a/notes/28_contribution_status_proof_vs_fit.md b/notes/28_contribution_status_proof_vs_fit.md
index 40c1118..5507de7 100644
--- a/notes/28_contribution_status_proof_vs_fit.md
+++ b/notes/28_contribution_status_proof_vs_fit.md
@@ -17,7 +17,7 @@ Our current project should avoid calling every component a theorem.
| contribution | current status | fitted? | safe main-paper phrasing |
|---|---|---:|---|
| 1. capacity formalization | definition + exact random-geometry theorem | no | exact under isotropic feedback assumptions |
-| 2. scaling and soft erosion | exact scaling + exact random-subspace null model; real FA extension not yet proved | no, but partly only a null model | capacity cost scales exactly; random-subspace erosion gives a baseline, while real FA needs `k_eff` |
+| 2. scaling and soft erosion | exact scaling + exact random-subspace null model + actual-FA initial moment theorem | no, but finite-time FA still needs dynamics | capacity cost scales exactly; actual FA has a no-fit `e_0` theorem, while `t>0` needs operator dynamics |
| 3. prior-free minimax bound | exact theorem | no | minimax optimality for prior-free angular alignment |
| 4. tangent-operator estimator | first-order tangent derivation + conditional predictor | no scalar fit, but uses early observed operators | estimator for finite-time gap, not architecture-only theorem |
| 5. empirical distribution validation | experiment | no | validates which theoretical distributions and estimators match trajectories |
@@ -132,6 +132,31 @@ hard-k random-subspace formula exactly predicts real FA/BP gap
This is false in our current experiments.
+We now have a stronger actual-FA initialization theorem:
+
+```text
+speed_BP = sum_l ||g_l^BP||^2
+speed_FA = sum_l <g_l^BP, g_l^FA>
+e_0 = 1 - speed_FA / speed_BP
+```
+
+For fixed forward weights and residuals, with independent zero-mean feedback:
+
+```text
+E_B[speed_FA | W,r] = ||g_output^BP||^2
+```
+
+Therefore:
+
+```text
+E_B[e_0 | W,r]
+ = 1 - ||g_output^BP||^2 / sum_l ||g_l^BP||^2
+```
+
+This is actual FA, not a random-subspace proxy. It says the initial effective
+burden is the hidden-layer BP speed share. It does not yet solve the full
+finite-time problem because training makes `W_t` depend on `B`.
+
## 3. Prior-Free Minimax Initialization Bound
This is proof-level.
@@ -305,9 +330,12 @@ The paper can honestly claim:
The current unresolved theoretical gap:
```text
-derive k_eff or e_t distribution for actual FA directly from architecture and
-feedback initialization, rather than measuring early tangent operators.
+derive the t>0 distribution of k_eff(t) or e_t for actual FA directly from
+architecture and feedback initialization, rather than measuring early tangent
+operators.
```
-Until that is solved, contribution 2 should not overpromise exact FA gap
-prediction, and contribution 4 should be the quantitative gap tool.
+The t=0 conditional mean is now solved by the actual-FA initial moment theorem.
+Until the t>0 dynamics are solved, contribution 2 should not overpromise exact
+finite-time FA gap prediction, and contribution 4 should remain the
+quantitative finite-time gap tool.
diff --git a/notes/29_actual_fa_initial_operator_moment_theorem.md b/notes/29_actual_fa_initial_operator_moment_theorem.md
new file mode 100644
index 0000000..5814f06
--- /dev/null
+++ b/notes/29_actual_fa_initial_operator_moment_theorem.md
@@ -0,0 +1,327 @@
+# Actual FA Initial Operator Moment Theorem
+
+This note records the first successful attempt to replace the hard-k
+random-subspace proxy with a theorem about actual FA.
+
+## Problem
+
+The hard random-subspace null model says:
+
+```text
+E[e] = k/P
+```
+
+but experiments showed it overpredicts real FA gaps. The reason is that actual
+FA does not delete a hard `k`-dimensional subspace. It changes the backward
+operator.
+
+We need a quantity closer to real FA.
+
+## Setup
+
+Squared loss with residual vector `r`.
+
+For each parameter block/layer `l`:
+
+```text
+g_l^BP = BP gradient block
+g_l^FA = FA pseudo-gradient block
+```
+
+The BP one-step learning speed is:
+
+```text
+speed_BP = sum_l ||g_l^BP||^2
+```
+
+The FA first-order loss decrease uses the mixed inner product:
+
+```text
+speed_FA = sum_l <g_l^BP, g_l^FA>
+```
+
+Equivalently:
+
+```text
+speed_BP = r^T K_BP r
+speed_FA = r^T K_FA r
+```
+
+where:
+
+```text
+K_BP = J J^T
+K_FA = J J_tilde^T
+```
+
+The local operator erosion is:
+
+```text
+e_0(B) = 1 - speed_FA / speed_BP
+```
+
+## Theorem: Exact Conditional Mean at Initialization
+
+Assume:
+
+```text
+1. forward weights W are fixed;
+2. residual r is fixed;
+3. feedback matrices B are independent of W and r;
+4. each feedback matrix has zero-mean entries or, more generally, zero
+ first moment in every hidden backward direction;
+5. the output layer is trained with the true output gradient, as in standard FA.
+```
+
+Then, at initialization:
+
+```text
+E_B[speed_FA | W, r] = ||g_output^BP||^2
+```
+
+Therefore:
+
+```text
+E_B[ speed_FA / speed_BP | W, r ]
+ = ||g_output^BP||^2 / sum_l ||g_l^BP||^2
+```
+
+and:
+
+```text
+E_B[e_0 | W, r]
+ = 1 - ||g_output^BP||^2 / sum_l ||g_l^BP||^2
+```
+
+So the initial effective FA burden is not raw hard constraint count. It is the
+fraction of BP learning speed carried by hidden-layer gradient blocks:
+
+```text
+e_0_mean = hidden BP speed share
+```
+
+If we want to express it as an effective dimension:
+
+```text
+k_eff,0 = P * e_0_mean
+```
+
+but the more natural object is the dimensionless erosion `e_0_mean`.
+
+## Proof Sketch
+
+The output-layer FA gradient equals the BP gradient:
+
+```text
+g_output^FA = g_output^BP
+```
+
+so its mixed contribution is:
+
+```text
+<g_output^BP, g_output^FA> = ||g_output^BP||^2
+```
+
+Every hidden FA pseudo-gradient contains at least one zero-mean independent
+feedback matrix. Conditional on fixed forward weights and residuals:
+
+```text
+E_B[g_l^FA | W, r] = 0
+```
+
+for hidden layers `l`.
+
+Thus:
+
+```text
+E_B[<g_l^BP, g_l^FA> | W, r] = 0
+```
+
+for every hidden layer.
+
+Adding the layer contributions gives the result.
+
+## Relation to Song, Xu, and Lafferty
+
+For a two-layer scalar-output network, Song/Xu/Lafferty decompose FA dynamics
+as:
+
+```text
+K_FA = G + H_FA
+```
+
+where `G` is the top-layer kernel and `H_FA` is the feedback-dependent hidden
+term. This theorem is the residual-direction version of that decomposition at
+initialization:
+
+```text
+E_B[H_FA | W, r] = 0
+```
+
+so:
+
+```text
+E_B[K_FA | W] = K_output
+```
+
+not:
+
+```text
+E_B[K_FA | W] = (1 - k/P) K_BP
+```
+
+This explains why the hard-k random-subspace theory was too pessimistic.
+
+## Distribution, Not Just Mean
+
+For one-hidden-layer scalar-output FA with Gaussian feedback, the hidden mixed
+speed is exactly Gaussian conditional on `W` and `r`.
+
+It can be written as:
+
+```text
+speed_FA = output_speed + Z
+```
+
+where:
+
+```text
+E_B[Z | W, r] = 0
+Var_B[Z | W, r] = feedback_variance * ||C(W,r)||^2
+```
+
+for an explicit coefficient tensor `C(W,r)` determined by activations, gates,
+residuals, and BP hidden gradients.
+
+Therefore:
+
+```text
+e_0(B) = 1 - (output_speed + Z) / speed_BP
+```
+
+has an explicit conditional Gaussian distribution in the one-hidden-layer
+Gaussian-feedback case.
+
+For deeper FA, hidden terms contain products of feedback matrices. The exact
+mean above still holds, but the finite-width distribution is no longer exactly
+Gaussian. A CLT/Wick expansion should give a computable approximation.
+
+## Infinite-Width Interpretation
+
+At infinite width, layerwise BP speed shares should converge to deterministic
+NTK/NNGP quantities:
+
+```text
+output_share = r^T K_output r / r^T K_BP r
+hidden_share = 1 - output_share
+```
+
+Thus an architecture-level prediction may be obtained from layerwise NTK
+recursions:
+
+```text
+E_B[e_0] -> hidden_share_NTK
+```
+
+This gives a principled route to an architecture/data-level prediction without
+post-hoc scaling.
+
+## Time-Varying Extension
+
+At later time:
+
+```text
+speed_FA,t = output_speed_t + hidden_alignment_gain_t
+```
+
+where:
+
+```text
+hidden_alignment_gain_t = sum_hidden <g_l^BP(t), g_l^FA(t)>
+```
+
+At initialization:
+
+```text
+E_B[hidden_alignment_gain_0 | W,r] = 0
+```
+
+During training, forward weights become dependent on feedback matrices, so this
+term can become positive. This is actual feedback alignment.
+
+Therefore:
+
+```text
+e_t = 1 - (output_speed_t + hidden_alignment_gain_t) / speed_BP,t
+```
+
+The current tangent-operator estimator works because it measures early movement
+of this alignment gain. A future theorem should bound or approximate:
+
+```text
+hidden_alignment_gain_t
+```
+
+from the early dynamics, possibly using the perturbative `H_FA` analysis style
+from Song/Xu/Lafferty.
+
+## Empirical Validation
+
+Script:
+
+```text
+scripts/actual_fa_initial_operator_moments.py
+```
+
+Run:
+
+```text
+python scripts/actual_fa_initial_operator_moments.py \
+ --widths 16 24 32 48 64 96 \
+ --init-seeds 6 \
+ --feedback-samples 512 \
+ --torch-threads 8 \
+ --outdir outputs/actual_fa_initial_operator_moments
+```
+
+Outputs:
+
+```text
+outputs/actual_fa_initial_operator_moments/initial_operator_moment_rows.csv
+outputs/actual_fa_initial_operator_moments/predicted_vs_empirical_initial_erosion_mean.png
+outputs/actual_fa_initial_operator_moments/initial_erosion_mean_error_by_width.png
+outputs/actual_fa_initial_operator_moments/initial_erosion_histograms.png
+```
+
+Representative results:
+
+| width | init | theory mean erosion | empirical mean erosion | error |
+|---:|---:|---:|---:|---:|
+| 16 | 0 | 0.451931 | 0.444348 | -0.00758 |
+| 24 | 0 | 0.312016 | 0.312286 | 0.00027 |
+| 32 | 0 | 0.236058 | 0.239714 | 0.00366 |
+| 48 | 0 | 0.261471 | 0.259858 | -0.00161 |
+| 64 | 0 | 0.209125 | 0.211261 | 0.00214 |
+| 96 | 0 | 0.259765 | 0.259147 | -0.00062 |
+
+Across all tested widths and initializations, the no-fit mean prediction is
+accurate to Monte Carlo error.
+
+## Consequence for Contributions
+
+This improves contribution 2:
+
+```text
+hard-k random-subspace erosion = clean null model
+actual FA initial erosion = hidden BP speed share theorem
+```
+
+Contribution 4 still handles finite-time gap:
+
+```text
+initial theorem gives e_0
+operator estimator tracks e_t and converts it into finite-time gap
+```
+
+This is a stronger and cleaner story than trying to force raw capacity margin
+to predict the full trajectory gap.
diff --git a/scripts/actual_fa_initial_operator_moments.py b/scripts/actual_fa_initial_operator_moments.py
new file mode 100644
index 0000000..6ee4900
--- /dev/null
+++ b/scripts/actual_fa_initial_operator_moments.py
@@ -0,0 +1,293 @@
+#!/usr/bin/env python3
+"""Validate initial actual-FA operator moment predictions.
+
+The theorem tested here is conditional on a fixed forward initialization and
+training residual. With independent zero-mean feedback matrices, the hidden
+FA pseudo-gradient has zero conditional mean. Therefore the expected one-step
+FA/BP speed ratio equals the BP output-layer gradient-energy share.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import sys
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+if str(SCRIPT_DIR) not in sys.path:
+ sys.path.insert(0, str(SCRIPT_DIR))
+
+import downstream_capacity_sweep as dcs # noqa: E402
+
+
+@dataclass(frozen=True)
+class MomentRow:
+ width: int
+ init_seed: int
+ feedback_samples: int
+ bp_speed: float
+ output_speed: float
+ predicted_mean_ratio: float
+ empirical_mean_ratio: float
+ empirical_std_ratio: float
+ predicted_mean_erosion: float
+ empirical_mean_erosion: float
+ empirical_std_erosion: float
+ mean_error: float
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description="Actual FA initial operator moment validation."
+ )
+ parser.add_argument("--input-dim", type=int, default=16)
+ parser.add_argument("--output-dim", type=int, default=4)
+ parser.add_argument("--widths", type=int, nargs="+", default=[16, 24, 32, 48, 64, 96])
+ parser.add_argument("--train-samples", type=int, default=128)
+ parser.add_argument("--test-samples", type=int, default=512)
+ parser.add_argument("--init-seeds", type=int, default=8)
+ parser.add_argument("--feedback-samples", type=int, default=512)
+ parser.add_argument("--data-seed", type=int, default=123)
+ parser.add_argument(
+ "--feedback-scale",
+ choices=["relu", "fan-in", "unit"],
+ default="relu",
+ )
+ parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
+ parser.add_argument("--torch-threads", type=int, default=8)
+ parser.add_argument(
+ "--outdir",
+ type=Path,
+ default=Path("outputs/actual_fa_initial_operator_moments"),
+ )
+ return parser.parse_args()
+
+
+def make_config(args: argparse.Namespace) -> dcs.RunConfig:
+ return dcs.RunConfig(
+ task="random",
+ input_dim=args.input_dim,
+ output_dim=args.output_dim,
+ teacher_rank=4,
+ teacher_width=64,
+ teacher_hidden_layers=2,
+ normalize_targets=False,
+ widths=args.widths,
+ train_samples=args.train_samples,
+ test_samples=args.test_samples,
+ probe_samples=8,
+ steps=1,
+ lr=1e-3,
+ optimizer="sgd",
+ init_seeds=args.init_seeds,
+ feedback_seeds=args.feedback_samples,
+ init_seed_offset=0,
+ feedback_seed_offset=0,
+ data_seed=args.data_seed,
+ noise_std=0.0,
+ feedback_scale=args.feedback_scale,
+ capacity_q=0.01,
+ jacobian_lambda_rel=1e-3,
+ skip_jacobian=True,
+ device=args.device,
+ torch_threads=args.torch_threads,
+ outdir=str(args.outdir),
+ plot=False,
+ )
+
+
+def flatten(grads: list[torch.Tensor]) -> torch.Tensor:
+ return torch.cat([grad.reshape(-1) for grad in grads])
+
+
+def squared_norm(grads: list[torch.Tensor]) -> float:
+ return float(sum(torch.sum(grad * grad).cpu() for grad in grads))
+
+
+def inner_product(left: list[torch.Tensor], right: list[torch.Tensor]) -> float:
+ return float(sum(torch.sum(a * b).cpu() for a, b in zip(left, right)))
+
+
+def run_one(
+ config: dcs.RunConfig,
+ x: torch.Tensor,
+ y: torch.Tensor,
+ width: int,
+ init_seed: int,
+ feedback_samples: int,
+) -> tuple[MomentRow, np.ndarray]:
+ weights = dcs.initialize_weights(config, width, init_seed)
+ bp_grads = dcs.gradients(weights, x, y, feedback=None)
+ bp_speed = squared_norm(bp_grads)
+ output_speed = squared_norm([bp_grads[-1]])
+ predicted_mean_ratio = output_speed / bp_speed
+ predicted_mean_erosion = 1.0 - predicted_mean_ratio
+
+ ratios: list[float] = []
+ for feedback_seed in range(feedback_samples):
+ feedback = dcs.init_feedback(config, width, feedback_seed)
+ fa_grads = dcs.gradients(weights, x, y, feedback=feedback)
+ ratios.append(inner_product(bp_grads, fa_grads) / bp_speed)
+
+ ratio_array = np.array(ratios, dtype=np.float64)
+ erosion_array = 1.0 - ratio_array
+ row = MomentRow(
+ width=width,
+ init_seed=init_seed,
+ feedback_samples=feedback_samples,
+ bp_speed=bp_speed,
+ output_speed=output_speed,
+ predicted_mean_ratio=predicted_mean_ratio,
+ empirical_mean_ratio=float(np.mean(ratio_array)),
+ empirical_std_ratio=float(np.std(ratio_array, ddof=1)),
+ predicted_mean_erosion=predicted_mean_erosion,
+ empirical_mean_erosion=float(np.mean(erosion_array)),
+ empirical_std_erosion=float(np.std(erosion_array, ddof=1)),
+ mean_error=float(np.mean(erosion_array) - predicted_mean_erosion),
+ )
+ return row, erosion_array
+
+
+def write_rows(path: Path, rows: list[MomentRow]) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ with path.open("w", newline="") as handle:
+ writer = csv.DictWriter(handle, fieldnames=list(MomentRow.__annotations__.keys()))
+ writer.writeheader()
+ for row in rows:
+ writer.writerow(asdict(row))
+
+
+def plot_summary(rows: list[MomentRow], outdir: Path) -> list[Path]:
+ outdir.mkdir(parents=True, exist_ok=True)
+ paths: list[Path] = []
+
+ scatter = outdir / "predicted_vs_empirical_initial_erosion_mean.png"
+ fig, ax = plt.subplots(figsize=(6.0, 5.2))
+ widths = sorted({row.width for row in rows})
+ for width in widths:
+ subset = [row for row in rows if row.width == width]
+ ax.scatter(
+ [row.predicted_mean_erosion for row in subset],
+ [row.empirical_mean_erosion for row in subset],
+ s=32,
+ alpha=0.65,
+ label=f"w={width}",
+ )
+ values = [
+ value
+ for row in rows
+ for value in (row.predicted_mean_erosion, row.empirical_mean_erosion)
+ ]
+ lo, hi = min(values), max(values)
+ pad = 0.03 * (hi - lo + 1e-12)
+ ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], color="black", linewidth=1)
+ ax.set_xlabel("theory mean erosion: 1 - output BP speed share")
+ ax.set_ylabel("empirical mean erosion over feedback seeds")
+ ax.set_title("Actual FA initial operator moment")
+ ax.legend(fontsize=8, ncols=2)
+ fig.tight_layout()
+ fig.savefig(scatter, dpi=180)
+ plt.close(fig)
+ paths.append(scatter)
+
+ error = outdir / "initial_erosion_mean_error_by_width.png"
+ fig, ax = plt.subplots(figsize=(7.0, 4.6))
+ positions = []
+ means = []
+ ses = []
+ for width in widths:
+ errors = np.array([row.mean_error for row in rows if row.width == width])
+ positions.append(width)
+ means.append(float(np.mean(errors)))
+ ses.append(float(np.std(errors, ddof=1) / np.sqrt(len(errors))) if len(errors) > 1 else 0.0)
+ ax.errorbar(positions, means, yerr=ses, marker="o", linewidth=1.8, capsize=4)
+ ax.axhline(0.0, color="black", linewidth=1)
+ ax.set_xlabel("width")
+ ax.set_ylabel("empirical mean - theory mean")
+ ax.set_title("No-fit moment error")
+ fig.tight_layout()
+ fig.savefig(error, dpi=180)
+ plt.close(fig)
+ paths.append(error)
+
+ return paths
+
+
+def plot_histograms(samples_by_key: dict[tuple[int, int], np.ndarray], rows: list[MomentRow], outdir: Path) -> Path:
+ selected: list[MomentRow] = []
+ for width in sorted({row.width for row in rows}):
+ width_rows = [row for row in rows if row.width == width]
+ selected.append(width_rows[0])
+ selected = selected[:6]
+
+ cols = 3
+ rows_n = int(np.ceil(len(selected) / cols))
+ fig, axes = plt.subplots(rows_n, cols, figsize=(12.0, 3.3 * rows_n), squeeze=False)
+ for ax in axes.ravel():
+ ax.axis("off")
+
+ for ax, row in zip(axes.ravel(), selected):
+ ax.axis("on")
+ values = samples_by_key[(row.width, row.init_seed)]
+ ax.hist(values, bins=36, density=True, color="tab:orange", alpha=0.55)
+ ax.axvline(row.predicted_mean_erosion, color="tab:blue", linewidth=2, label="theory mean")
+ ax.axvline(row.empirical_mean_erosion, color="tab:orange", linewidth=2, linestyle="--", label="empirical mean")
+ ax.set_title(f"width={row.width}, init={row.init_seed}")
+ ax.set_xlabel("one-step erosion")
+ ax.set_ylabel("density")
+ ax.legend(fontsize=8)
+
+ fig.suptitle("Actual FA initial erosion distributions over random feedback", y=1.02)
+ fig.tight_layout()
+ path = outdir / "initial_erosion_histograms.png"
+ fig.savefig(path, dpi=180, bbox_inches="tight")
+ plt.close(fig)
+ return path
+
+
+def main() -> None:
+ args = parse_args()
+ torch.set_num_threads(args.torch_threads)
+ config = make_config(args)
+ x_train, y_train, *_ = dcs.make_data(config)
+
+ all_rows: list[MomentRow] = []
+ samples_by_key: dict[tuple[int, int], np.ndarray] = {}
+ for width in args.widths:
+ for init_seed in range(args.init_seeds):
+ row, samples = run_one(
+ config,
+ x_train,
+ y_train,
+ width,
+ init_seed,
+ args.feedback_samples,
+ )
+ all_rows.append(row)
+ samples_by_key[(width, init_seed)] = samples
+ print(
+ f"width={width} init={init_seed}: "
+ f"theory={row.predicted_mean_erosion:.6f} "
+ f"empirical={row.empirical_mean_erosion:.6f} "
+ f"error={row.mean_error:.3g}"
+ )
+
+ args.outdir.mkdir(parents=True, exist_ok=True)
+ csv_path = args.outdir / "initial_operator_moment_rows.csv"
+ write_rows(csv_path, all_rows)
+ paths = plot_summary(all_rows, args.outdir)
+ paths.append(plot_histograms(samples_by_key, all_rows, args.outdir))
+
+ print(f"rows: {csv_path}")
+ for path in paths:
+ print(f"plot: {path}")
+
+
+if __name__ == "__main__":
+ main()