From afed1d0f0d0b54dc891e9dab7389e031e5bb67da Mon Sep 17 00:00:00 2001 From: Antovigo Date: Sat, 25 Jul 2026 00:31:51 +0000 Subject: [PATCH 1/2] feat(losses): normalize_at_one option for SmoothL0 imp-min MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add an opt-in `normalize_at_one` flag to SmoothL0ImportanceMinimalityLoss that rescales the Geman–McClure penalty by `(1 + gamma^2)`, so a fully-on component (c=1) contributes exactly 1 regardless of gamma. Without it, phi(1) = 1/(1+gamma^2) grows as gamma anneals down, silently ramping the effective coeff on saturated components across the schedule. Threads the flag through the loss dispatch; defaults to False, so existing runs are unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_011XygndeoX7JQqqnuNpiCh5 --- param_decomp/configs.py | 5 +++++ param_decomp/losses.py | 13 +++++++++--- param_decomp/tests/test_smooth_l0_imp_min.py | 21 ++++++++++++++++++-- 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/param_decomp/configs.py b/param_decomp/configs.py index eea12ab23..3e21a3382 100644 --- a/param_decomp/configs.py +++ b/param_decomp/configs.py @@ -151,11 +151,16 @@ class SmoothL0ImportanceMinimalityLossConfig(LossMetricConfig): `gamma` is the width's full schedule; 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/losses.py b/param_decomp/losses.py index 1734b0596..f0625d207 100644 --- a/param_decomp/losses.py +++ b/param_decomp/losses.py @@ -128,12 +128,17 @@ 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` + rescales by `(1 + gamma^2)` so a fully-on component (`c = 1`) always contributes 1.""" gamma_sq = gamma * gamma - return _imp_min_terms(ci_upper, lambda ci: ci**2 / (ci**2 + gamma_sq), reference_token_count) + scale = (1.0 + gamma_sq) if normalize_at_one else 1.0 + return _imp_min_terms( + ci_upper, lambda ci: scale * ci**2 / (ci**2 + gamma_sq), reference_token_count + ) def annealed_imp_min_param( @@ -162,4 +167,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/tests/test_smooth_l0_imp_min.py b/param_decomp/tests/test_smooth_l0_imp_min.py index a4fd62c8b..6a59ed0b2 100644 --- a/param_decomp/tests/test_smooth_l0_imp_min.py +++ b/param_decomp/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)) From 48a78e18ac888ebb4b571245f639c3e305a8554b Mon Sep 17 00:00:00 2001 From: Antovigo Date: Wed, 29 Jul 2026 18:22:44 +0000 Subject: [PATCH 2/2] refactor(losses): pick the full smooth-L0 formula, drop the scale var Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_016QbtkPzfQyKcS5L4tE8LAJ --- param_decomp/losses.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/param_decomp/losses.py b/param_decomp/losses.py index f0625d207..d572bedec 100644 --- a/param_decomp/losses.py +++ b/param_decomp/losses.py @@ -133,12 +133,15 @@ def smooth_l0_importance_minimality_terms( """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`. `normalize_at_one` - rescales by `(1 + gamma^2)` so a fully-on component (`c = 1`) always contributes 1.""" + 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 - scale = (1.0 + gamma_sq) if normalize_at_one else 1.0 - return _imp_min_terms( - ci_upper, lambda ci: scale * 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(