diff options
| author | YurenHao0426 <Blackhao0426@gmail.com> | 2026-06-05 15:44:19 -0500 |
|---|---|---|
| committer | YurenHao0426 <Blackhao0426@gmail.com> | 2026-06-05 15:44:19 -0500 |
| commit | 2df52d3dd7b35053d4a357c3814d7cba6db602ce (patch) | |
| tree | 96c08971d94cd931eb2727e3479d0b63f7319c5a | |
| parent | 4dee090530588792b0e0bce9a36981415e4a99a8 (diff) | |
Derive actual FA initial operator moment
| -rw-r--r-- | notes/18_contribution_roadmap.md | 34 | ||||
| -rw-r--r-- | notes/28_contribution_status_proof_vs_fit.md | 38 | ||||
| -rw-r--r-- | notes/29_actual_fa_initial_operator_moment_theorem.md | 327 | ||||
| -rw-r--r-- | scripts/actual_fa_initial_operator_moments.py | 293 |
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() |
