Add BP-calibrated transition theory

1 files changed, 29 insertions, 9 deletions

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,