Add BP-calibrated transition theory

2 files changed, 133 insertions, 9 deletions

diff --git a/notes/07_minimal_transition_theory.md b/notes/07_minimal_transition_theory.md
new file mode 100644
index 0000000..ced7e61
--- /dev/null
+++ b/notes/07_minimal_transition_theory.md
@@ -0,0 +1,104 @@
+# Minimal Transition Theory
+
+The cleanest correction is not to add a separate FA burden coefficient first.
+Start by calibrating the task dimension from BP alone.
+
+## Problem
+
+The hard theory used:
+
+\[
+d_{\mathrm{task}}=Nn_L=512.
+\]
+
+With this value, the smooth theory peak is near:
+
+\[
+M_{\mathrm{FA}}\approx -273,
+\]
+
+while 256 empirical FA trajectories peak at:
+
+\[
+M_{\mathrm{FA}}=-318.
+\]
+
+## BP-Only Calibration
+
+Fit the BP train curve with:
+
+\[
+L_{\mathrm{BP}}(P)
+\approx
+s\cdot \operatorname{softplus}_{\tau}(d_{\mathrm{task}}^{\mathrm{eff}}-P).
+\]
+
+The BP-only fit gives:
+
+\[
+d_{\mathrm{task}}^{\mathrm{eff}}\approx 437,
+\qquad
+\tau\approx 50.
+\]
+
+So BP is stronger than the naive parameter count suggests:
+
+\[
+d_{\mathrm{task}}^{\mathrm{eff}}
+\approx
+0.85\,Nn_L.
+\]
+
+## Minimal FA Prediction
+
+Use the same hard FA burden:
+
+\[
+K_{\mathrm{FA}}=\sum_l(D_l-1),
+\]
+
+but replace the task dimension:
+
+\[
+M_{\mathrm{FA}}^{\mathrm{min}}
+=
+P-K_{\mathrm{FA}}-d_{\mathrm{task}}^{\mathrm{eff}}.
+\]
+
+Equivalently, when plotted on the original hard-margin axis
+
+\[
+P-K_{\mathrm{FA}}-Nn_L,
+\]
+
+the theory shifts left by:
+
+\[
+Nn_L-d_{\mathrm{task}}^{\mathrm{eff}}
+\approx
+75.
+\]
+
+This explains the observed left shift without introducing a fitted FA burden
+coefficient.
+
+Generated figure:
+
+`outputs/downstream_capacity_random_main_fast/capacity_transition_bp_calibrated_p2_256.png`
+
+## Recommendation
+
+Use this as the main downstream theory:
+
+1. Matrix/capacity theory remains exact and uses the hard/log-volume definitions.
+2. Downstream performance uses a BP-calibrated effective task dimension.
+3. FA capacity loss is then tested with no extra FA-specific fitting parameter.
+
+Only add an effective FA burden coefficient
+
+\[
+K_{\mathrm{FA}}^{\mathrm{eff}}=\alpha K_{\mathrm{FA}}
+\]
+
+as an appendix ablation if needed. It improves fit but makes the theory look
+more ad hoc.
diff --git a/scripts/plot_capacity_transition_overlay.py b/scripts/plot_capacity_transition_overlay.py
index cc53012..28028c8 100644
--- a/scripts/plot_capacity_transition_overlay.py
+++ b/scripts/plot_capacity_transition_overlay.py
@@ -51,6 +51,12 @@ def parse_args() -> argparse.Namespace:
         default=70.0,
         help="Soft capacity-transition width measured in scalar output constraints.",
     )
+    parser.add_argument(
+        "--effective-task-dim",
+        type=float,
+        default=None,
+        help="Use this task dimension in the smooth theory while keeping the hard-margin x-axis unchanged.",
+    )
     parser.add_argument("--xmin", type=float, default=-430.0)
     parser.add_argument("--xmax", type=float, default=600.0)
     parser.add_argument("--outname", default="capacity_transition_theory_vs_trajectories.png")
@@ -96,7 +102,8 @@ def soft_overlap_gap_counts(
     width: np.ndarray,
     input_dim: float,
     output_dim: float,
-    task_dim: float,
+    hard_task_dim: float,
+    effective_task_dim: float,
     samples: int,
     transition_width: float,
     rng: np.random.Generator,
@@ -105,7 +112,7 @@ def soft_overlap_gap_counts(
     constraint_rank = np.maximum(width * width - 1.0, 0.0) + np.maximum(
         width * output_dim - 1.0, 0.0
     )
-    hard_fa_margin = parameter_count - constraint_rank - task_dim
+    hard_fa_margin = parameter_count - constraint_rank - hard_task_dim
 
     raw_gaps = np.zeros((samples, width.size), dtype=np.float64)
 
@@ -114,8 +121,8 @@ def soft_overlap_gap_counts(
         return transition_width * np.logaddexp(0.0, scaled)
 
     for index, (p_count, k_rank) in enumerate(zip(parameter_count, constraint_rank)):
-        bp_missing = softplus_hinge(task_dim - p_count)
-        task_subspace_dim = min(task_dim, p_count)
+        bp_missing = softplus_hinge(effective_task_dim - p_count)
+        task_subspace_dim = min(effective_task_dim, p_count)
 
         if task_subspace_dim >= p_count or k_rank <= 0:
             overlap = np.full(samples, min(k_rank, p_count), dtype=np.float64)
@@ -134,7 +141,7 @@ def soft_overlap_gap_counts(
 
         finite_width_noise = rng.normal(0.0, transition_width, size=samples)
         fa_missing = softplus_hinge(
-            task_dim - (parameter_count[index] - overlap) + finite_width_noise
+            effective_task_dim - (parameter_count[index] - overlap) + finite_width_noise
         )
         raw_gaps[:, index] = np.maximum(fa_missing - bp_missing, 0.0)
 
@@ -167,6 +174,7 @@ def plot_transition(
     extra_run_dirs: list[Path],
     samples: int,
     transition_width: float,
+    effective_task_dim_arg: float | None,
     outname: str,
     xmin: float,
     xmax: float,
@@ -175,14 +183,24 @@ def plot_transition(
     runs = read_runs(run_dir / "runs.csv", 0)
     for source_id, extra_run_dir in enumerate(extra_run_dirs, start=1):
         runs.extend(read_runs(extra_run_dir / "runs.csv", source_id))
-    input_dim, output_dim, task_dim = infer_dims(run_dir)
+    input_dim, output_dim, hard_task_dim = infer_dims(run_dir)
+    effective_task_dim = (
+        hard_task_dim if effective_task_dim_arg is None else effective_task_dim_arg
+    )
     rng = np.random.default_rng(0)
 
     widths = np.array([row.width for row in summary], dtype=np.float64)
     dense_width = np.linspace(widths.min(), widths.max(), 240)
 
     _, raw_mean_at_widths, _, _ = soft_overlap_gap_counts(
-        widths, input_dim, output_dim, task_dim, samples, transition_width, rng
+        widths,
+        input_dim,
+        output_dim,
+        hard_task_dim,
+        effective_task_dim,
+        samples,
+        transition_width,
+        rng,
     )
     scale = calibrate_scale(summary, raw_mean_at_widths)
 
@@ -190,7 +208,8 @@ def plot_transition(
         dense_width,
         input_dim,
         output_dim,
-        task_dim,
+        hard_task_dim,
+        effective_task_dim,
         samples,
         transition_width,
         np.random.default_rng(1),
@@ -261,7 +280,7 @@ def plot_transition(
     axes[0].set_ylabel("FA train MSE - BP train MSE")
     axes[0].legend(loc="upper right")
     fig.suptitle(
-        f"Random-label memorization transition, scale={scale:.4g}, transition width={transition_width:g}",
+        f"Random-label memorization transition, scale={scale:.4g}, transition width={transition_width:g}, d_eff={effective_task_dim:g}",
         y=1.02,
     )
     fig.tight_layout()
@@ -279,6 +298,7 @@ def main() -> None:
         args.extra_run_dir,
         args.samples,
         args.transition_width,
+        args.effective_task_dim,
         args.outname,
         args.xmin,
         args.xmax,