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

[Antovigo]

What

Adds an opt-in normalize_at_one: bool = False to SmoothL0ImportanceMinimalityLossConfig. When set, the Geman–McClure per-value penalty is rescaled by (1 + gamma^2):

phi(c) = (1 + gamma^2) * c^2 / (c^2 + gamma^2)

so a fully-on component (c = 1) contributes exactly 1 regardless of gamma.

Why

The unnormalized penalty gives phi(1) = 1/(1 + gamma^2). As gamma anneals down over training (e.g. 1.0 -> 0.1), the contribution of a saturated component climbs from 0.5 toward 1.0. This silently couples the sparsity pressure to the gamma schedule in a way that is hard to reason about. normalize_at_one pins the cost of a fully-active component to 1 for the whole schedule, so coeff means the same thing at every step.

Some tests on Llama-8B’s L18, targeted on addition/subtraction tasks:

• addsub-L18-09 (green): no normalization, coeff = 1e-4
• addsub-L18-09-one-im (red): normalize_at_one, coeff = 5e-5 (so the starting value of the coeff is the same)
  In both case, the gamma parameter is annealed from 1 to 0.01 over the second half of training.

The normalized version behaves very similarly to the non-normalized version:

However, if we look at the importance minimality loss throughout training, the un-normalized version has a clear ramp over the annealing period, reflecting a change in the minimality/reconstruction exchange rate:

Changes

• configs.py: new normalize_at_one field (defaults to False) + docstring.
• losses.py: smooth_l0_importance_minimality_terms takes normalize_at_one; the flag is threaded through the imp_min_terms dispatch. L_p path untouched.
• tests/test_smooth_l0_imp_min.py: new test_normalize_at_one_fixes_saturated_contribution (saturated component → exactly 1; whole lp scales by (1 + gamma^2)); existing call sites updated.

Compatibility

Default False ⇒ existing runs are numerically unchanged.

Testing

• Tested on the targeted decomposition of Llama-8B-L18 on arithmetic tasks
• Tested on the toy model of redundancy, a miniature transformer with RMSnorm
• pytest param_decomp/tests/test_smooth_l0_imp_min.py — 5 passed. ruff + basedpyright clean (pre-commit).

[Antovigo]

Moved upstream: goodfire-ai#994 (same two commits, rebased onto upstream main).