SmoothL0: normalize_at_one option (fix gamma-anneal coeff ramp)

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

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
+            )

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))