diff options
Diffstat (limited to 'scripts/plot_capacity_transition_overlay.py')
| -rw-r--r-- | scripts/plot_capacity_transition_overlay.py | 38 |
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, |