diff --git a/param_decomp/core/configs.py b/param_decomp/core/configs.py index 02a71ea17..84cc2e536 100644 --- a/param_decomp/core/configs.py +++ b/param_decomp/core/configs.py @@ -151,11 +151,16 @@ class SmoothL0ImportanceMinimalityLossConfig(LossMetricConfig): `gamma` is the width's full schedule (SPEC S9′); annealing it down (e.g. `fn_type=linear, final_val_frac < 1`) sharpens the count. Warmup is refused where the term is built. + + With `normalize_at_one`, `phi` is rescaled by `(1 + gamma^2)` so a fully-on component + (`c = 1`) contributes exactly 1 regardless of `gamma`. Otherwise `phi(1) = 1/(1+gamma^2)` + grows as `gamma` anneals, silently ramping the effective `coeff` on saturated components. """ type: Literal["SmoothL0ImportanceMinimalityLoss"] = "SmoothL0ImportanceMinimalityLoss" gamma: ScheduleConfig frequency: FrequencyMinimalityConfig | None = None + normalize_at_one: bool = False # The two imp-min penalties share the `coeff` + optional `frequency` surface and the diff --git a/param_decomp/core/losses.py b/param_decomp/core/losses.py index 263e647cc..0e731e98f 100644 --- a/param_decomp/core/losses.py +++ b/param_decomp/core/losses.py @@ -128,12 +128,20 @@ def smooth_l0_importance_minimality_terms( ci_upper: dict[str, Float[Array, "*leading _"]], gamma: Float[Array, ""], reference_token_count: int | None, + normalize_at_one: bool, ) -> tuple[Float[Array, ""], Float[Array, ""]]: """Geman–McClure smooth-L0 imp-min terms: per-value penalty `c^2 / (c^2 + gamma^2)`. Flat at the origin (`phi'(0)=0`) and bounded (`|phi'| <= 0.65/gamma`) — no singularity, - no `eps` floor. Approaches the true `L_0` count as `gamma -> 0`.""" + no `eps` floor. Approaches the true `L_0` count as `gamma -> 0`. `normalize_at_one` + switches to `(1 + gamma^2) c^2 / (c^2 + gamma^2)`, so a fully-on component (`c = 1`) + always contributes exactly 1.""" gamma_sq = gamma * gamma - return _imp_min_terms(ci_upper, lambda ci: ci**2 / (ci**2 + gamma_sq), reference_token_count) + per_value_penalty = ( + (lambda ci: (1.0 + gamma_sq) * ci**2 / (ci**2 + gamma_sq)) + if normalize_at_one + else (lambda ci: ci**2 / (ci**2 + gamma_sq)) + ) + return _imp_min_terms(ci_upper, per_value_penalty, reference_token_count) def annealed_imp_min_param( @@ -162,4 +170,6 @@ def imp_min_terms( case ImportanceMinimalityLossConfig(): return importance_minimality_terms(ci_upper, annealed_param, cfg.eps, ref) case SmoothL0ImportanceMinimalityLossConfig(): - return smooth_l0_importance_minimality_terms(ci_upper, annealed_param, ref) + return smooth_l0_importance_minimality_terms( + ci_upper, annealed_param, ref, cfg.normalize_at_one + ) diff --git a/param_decomp/core/tests/test_smooth_l0_imp_min.py b/param_decomp/core/tests/test_smooth_l0_imp_min.py index 0c1d79b84..03e151d18 100644 --- a/param_decomp/core/tests/test_smooth_l0_imp_min.py +++ b/param_decomp/core/tests/test_smooth_l0_imp_min.py @@ -56,7 +56,7 @@ def test_terms_match_manual_per_site_structure(): gamma = 0.1 n_positions = 2 # both sites have 2 rows; a' = B·T reproduces the old `log2(1 + sum)` lp, freq = smooth_l0_importance_minimality_terms( - ci, jnp.asarray(gamma), reference_token_count=n_positions + ci, jnp.asarray(gamma), reference_token_count=n_positions, normalize_at_one=False ) exp_lp = jnp.zeros(()) @@ -88,6 +88,23 @@ def test_anneal_and_dispatch(): ci = {"a": jnp.array([[0.0, 0.5, 1.0], [0.2, 0.0, 0.9]])} param = annealed_imp_min_param(jnp.asarray(float(last)), total, cfg) via_dispatch = imp_min_terms(ci, cfg, param) - direct = smooth_l0_importance_minimality_terms(ci, param, reference_token_count=64) + direct = smooth_l0_importance_minimality_terms( + ci, param, reference_token_count=64, normalize_at_one=False + ) assert jnp.allclose(via_dispatch[0], direct[0]) assert jnp.allclose(via_dispatch[1], direct[1]) + + +def test_normalize_at_one_fixes_saturated_contribution(): + """`normalize_at_one` makes a fully-on component contribute exactly 1 at any gamma, and + scales the whole `lp` by `(1 + gamma^2)` relative to the unnormalized penalty.""" + ci = {"a": jnp.array([[1.0, 1.0]])} # both components fully on + for gamma in (jnp.asarray(1.0), jnp.asarray(0.1)): + lp_norm, _ = smooth_l0_importance_minimality_terms( + ci, gamma, reference_token_count=None, normalize_at_one=True + ) + lp_raw, _ = smooth_l0_importance_minimality_terms( + ci, gamma, reference_token_count=None, normalize_at_one=False + ) + assert jnp.allclose(lp_norm, 2.0) # two components, each exactly 1 + assert jnp.allclose(lp_norm, lp_raw * (1.0 + gamma**2))