diff --git a/docs/callbacks.md b/docs/callbacks.md index 8089449..8133475 100644 --- a/docs/callbacks.md +++ b/docs/callbacks.md @@ -44,7 +44,10 @@ and `xtol` cannot fire at step zero. Repeated rejections multiply the damping by `damping_increase` without bound, which can overflow in float32. The constructor's `max_damping` clamps the -damping from above; leave it `None` for uncapped classic behavior. +damping from above; leave it `None` for uncapped classic behavior. Accepted +steps cannot underflow damping to zero: `min_damping=None` uses +`jnp.finfo(residual.dtype).tiny`, while an explicit value selects a larger +absolute floor. Status codes are integer constants: @@ -82,10 +85,10 @@ to second-guess a callback's explicit replacement. ### Resettable Hyperparameters `solve()` populates `lm_state.hyper` with an `LMHyperparams` of traced -per-step values: `damping_decrease`, `damping_increase`, `max_damping`, -`geodesic_acceptance_ratio`, `iterative_tol`, `iterative_atol`, and -`iterative_maxiter`. Because they ride in the lm_state, a -callback can reset any of them mid-solve — exactly like a damping reset: +per-step values: `damping_decrease`, `damping_increase`, `min_damping`, +`max_damping`, `geodesic_acceptance_ratio`, `iterative_tol`, `iterative_atol`, +and `iterative_maxiter`. Because they ride in the lm_state, a callback can reset +any of them mid-solve — exactly like a damping reset: ```python new_hyper = dataclasses.replace( diff --git a/docs/index.md b/docs/index.md index e5f436b..02631ed 100644 --- a/docs/index.md +++ b/docs/index.md @@ -140,7 +140,9 @@ so the natural parameter metric is \(M=K\). The actual nonlinear step is accepted only if the unregularized residual sum of squares decreases. On acceptance, damping is multiplied by `damping_decrease`; -on rejection, it is multiplied by `damping_increase`. +on rejection, it is multiplied by `damping_increase`. `min_damping=None` +clamps it to the smallest positive normal value of the residual dtype, avoiding +backend flush-to-zero behavior without imposing a larger regularization floor. Near the interpolation threshold, small-damping LM becomes metric Gauss-Newton, whose step is the minimum-\(M\)-norm solution of the linearized diff --git a/docs/tuning_guide.md b/docs/tuning_guide.md index ac6982b..b44654e 100644 --- a/docs/tuning_guide.md +++ b/docs/tuning_guide.md @@ -192,6 +192,11 @@ few steps. Try them when you see specific signatures: descent). - Long rejection storms in float32 → set `max_damping` (~`1e6`) so damping cannot overflow. +- Literal damping underflow → the default `min_damping=None` already resolves + to the residual dtype's smallest positive normal value. This representation + floor prevents flush-to-zero but does not keep the damping numerically active + in an ill-conditioned system. A standard scale-aware floor is on the order of + `eps * operator_scale`; pass that larger absolute `min_damping` explicitly. - Accept/reject oscillation → bring `damping_decrease`/`damping_increase` closer to 1 (e.g. 0.7 / 2.0) for smoother adaptation. - All steps accepted but progress is slow → lower `init_damping` or decrease diff --git a/src/nlls_gram/gram_lm.py b/src/nlls_gram/gram_lm.py index 8d89927..1b642f2 100644 --- a/src/nlls_gram/gram_lm.py +++ b/src/nlls_gram/gram_lm.py @@ -64,6 +64,7 @@ class LMHyperparams: damping_decrease: jax.Array damping_increase: jax.Array + min_damping: jax.Array max_damping: jax.Array | None geodesic_acceptance_ratio: jax.Array iterative_tol: jax.Array @@ -71,12 +72,23 @@ class LMHyperparams: iterative_maxiter: jax.Array | None +def _damping_floor(min_damping, dtype): + if dtype is None: + seed = 0.0 if min_damping is None else min_damping + dtype = jnp.asarray(seed).dtype + dtype_floor = jnp.asarray(jnp.finfo(dtype).tiny, dtype=dtype) + if min_damping is None: + return dtype_floor + return jnp.maximum(jnp.asarray(min_damping, dtype=dtype), dtype_floor) + + def _cast_hyper(hyper, dtype): if hyper is None: return None return LMHyperparams( jnp.asarray(hyper.damping_decrease, dtype=dtype), jnp.asarray(hyper.damping_increase, dtype=dtype), + _damping_floor(hyper.min_damping, dtype), None if hyper.max_damping is None else jnp.asarray(hyper.max_damping, dtype=dtype), @@ -984,6 +996,15 @@ class LevenbergMarquardt: Take higher-order derivatives of ``solve`` with a fixed metric. See :class:`MetricFactory`. + ``min_damping`` is the absolute lower bound applied before every linear + solve and after every damping update. ``None`` (the default) resolves to + ``jnp.finfo(residual.dtype).tiny``, the smallest positive normal value, so + damping cannot enter a backend's flush-to-zero range. This is only an + underflow floor. A standard conditioning-oriented minimum is instead on + the order of machine epsilon times a representative scale of ``B'B`` (or + ``B B'``); pass that larger absolute value explicitly when the damping must + remain numerically effective in the linear system. + ``solve(...).x`` has a custom implicit AD rule with respect to ``p``, relinearized at the returned solution. ``ad_solver`` explicitly selects one method from ``{"direct", "svd", "qr", "augmented_qr", "gram_cg", @@ -1104,6 +1125,7 @@ def __init__( init_damping=1e-3, damping_decrease=0.5, damping_increase=4.0, + min_damping=None, max_damping=None, linear_solver="auto", jacobian_mode="auto", @@ -1150,6 +1172,10 @@ def __init__( raise ValueError("damping_decrease must be positive") if damping_increase <= 0: raise ValueError("damping_increase must be positive") + if min_damping is not None and min_damping <= 0: + raise ValueError("min_damping must be positive or None") + if min_damping is not None and min_damping > init_damping: + raise ValueError("min_damping must not exceed init_damping") if max_damping is not None and max_damping < init_damping: raise ValueError("max_damping must be at least init_damping") if iterative_tol < 0: @@ -1365,6 +1391,7 @@ def __init__( self.init_damping = init_damping self.damping_decrease = damping_decrease self.damping_increase = damping_increase + self.min_damping = min_damping self.max_damping = max_damping self.linear_solver = linear_solver self.jacobian_mode = jacobian_mode @@ -1445,6 +1472,7 @@ def __init__( init_damping, damping_decrease, damping_increase, + min_damping, max_damping, linear_solver, jacobian_mode, @@ -1489,6 +1517,7 @@ def hyperparams(self, dtype=None): return LMHyperparams( jnp.asarray(self.damping_decrease, dtype=dtype), jnp.asarray(self.damping_increase, dtype=dtype), + _damping_floor(self.min_damping, dtype), None if self.max_damping is None else jnp.asarray(self.max_damping, dtype=dtype), @@ -1508,7 +1537,10 @@ def init(self, x0, args=None, *, p=None): # populates it for its callbacks. self._check_residual_args(args, p) residual, aux = self._residual_and_aux(x0, args, p) - damping = jnp.asarray(self.init_damping, dtype=residual.dtype) + min_damping = _damping_floor(self.min_damping, residual.dtype) + damping = jnp.maximum( + jnp.asarray(self.init_damping, dtype=residual.dtype), min_damping + ) recycle = self._init_recycle_state(residual) precond, precond_valid = self._init_precond(x0, args, p, aux) metric_state, metric_valid = self._init_metric_state(x0, args, p, aux) @@ -1728,7 +1760,6 @@ def reuse_resid_and_jt(_): ) else: resid, Jt, aux = self._dense_resid_jt_aux(residual_flat, theta) - damping = jnp.asarray(lm_state.damping, dtype=resid.dtype) # Traced hyperparameters from the lm_state when present (resettable by # solve callbacks); the None fallback compiles to the same constants # as reading the constructor values directly. @@ -1739,6 +1770,10 @@ def reuse_resid_and_jt(_): ) damping_decrease = jnp.asarray(hyper.damping_decrease, dtype=resid.dtype) damping_increase = jnp.asarray(hyper.damping_increase, dtype=resid.dtype) + min_damping = _damping_floor(hyper.min_damping, resid.dtype) + damping = jnp.maximum( + jnp.asarray(lm_state.damping, dtype=resid.dtype), min_damping + ) # Trace-time shape resolution of the default solver: shapes are concrete # while tracing, so this is a plain Python branch (never a lax.cond). @@ -2231,8 +2266,12 @@ def accelerated_loss(_): new_damping = damping * damping_factor if hyper.max_damping is not None: new_damping = jnp.minimum( - new_damping, jnp.asarray(hyper.max_damping, dtype=resid.dtype) + new_damping, + jnp.maximum( + jnp.asarray(hyper.max_damping, dtype=resid.dtype), min_damping + ), ) + new_damping = jnp.maximum(new_damping, min_damping) loss = jnp.where(improved, loss_candidate, loss_old) # New recycle state: the velocity solve's harvested basis and the (stop- # gradient'd) dual solutions become warm starts for the next step. Threads diff --git a/tests/test_float64_subprocess.py b/tests/test_float64_subprocess.py index ecde70b..f76a624 100644 --- a/tests/test_float64_subprocess.py +++ b/tests/test_float64_subprocess.py @@ -1552,3 +1552,39 @@ def residual(theta): capture_output=True, text=True, ) + + +def test_float64_default_min_damping_uses_float64_normal_floor(): + script = r""" +import jax +jax.config.update("jax_enable_x64", True) +import jax.numpy as jnp + +from nlls_gram import LevenbergMarquardt, LMState + + +def residual(theta): + return theta + + +solver = LevenbergMarquardt( + residual, + cache_jacobian=False, + geodesic_acceleration=False, +) +state = LMState(jnp.asarray(0.0, dtype=jnp.float64)) +x, state, info = solver.update(jnp.ones(1, dtype=jnp.float64), state) +floor = jnp.asarray(jnp.finfo(jnp.float64).tiny, dtype=jnp.float64) + +assert info.accepted +assert jnp.all(jnp.isfinite(x)) +assert state.damping.dtype == jnp.float64 +assert state.damping == floor, (state.damping, floor) +assert info.damping == floor, (info.damping, floor) +""" + subprocess.run( + [sys.executable, "-c", textwrap.dedent(script)], + check=True, + capture_output=True, + text=True, + ) diff --git a/tests/test_gram_lm.py b/tests/test_gram_lm.py index b96cbca..42b74eb 100644 --- a/tests/test_gram_lm.py +++ b/tests/test_gram_lm.py @@ -234,6 +234,13 @@ def test_damping_update_factors_must_be_positive(): LevenbergMarquardt(residual_fn, damping_increase=0.0) +def test_min_damping_must_be_positive_and_not_exceed_initial_damping(): + with pytest.raises(ValueError, match="min_damping must be positive or None"): + LevenbergMarquardt(residual_fn, min_damping=0.0) + with pytest.raises(ValueError, match="min_damping must not exceed"): + LevenbergMarquardt(residual_fn, init_damping=1e-3, min_damping=1e-2) + + def test_iterative_options_must_be_valid(): with pytest.raises(ValueError, match="iterative_tol must be nonnegative"): LevenbergMarquardt(residual_fn, linear_solver="gram_cg", iterative_tol=-1.0) @@ -361,6 +368,56 @@ def test_max_damping_below_init_damping_raises(): LevenbergMarquardt(residual_fn, init_damping=1e-2, max_damping=1e-3) +@pytest.mark.parametrize("jit", [False, True]) +@pytest.mark.parametrize("linear_solver", ["auto", "lsmr"]) +def test_default_min_damping_repairs_zero_and_prevents_float32_underflow( + jit, linear_solver +): + def residual(theta): + return theta + + solver = LevenbergMarquardt( + residual, + init_damping=1e-3, + linear_solver=linear_solver, + cache_jacobian=False, + geodesic_acceleration=False, + ) + state = LMState(jnp.asarray(0.0, dtype=jnp.float32)) + update = solver.update if not jit else jax.jit(lambda x, s: solver.update(x, s)) + x, state, info = update(jnp.ones(1, dtype=jnp.float32), state) + floor = jnp.asarray(jnp.finfo(jnp.float32).tiny, dtype=jnp.float32) + + assert bool(info.accepted) + assert jnp.all(jnp.isfinite(x)) + assert state.damping == floor + assert info.damping == floor + + +@pytest.mark.parametrize("jit", [False, True]) +def test_explicit_min_damping_is_used_by_step_and_update(jit): + def residual(theta): + return theta + + min_damping = 1e-4 + solver = LevenbergMarquardt( + residual, + init_damping=1e-3, + min_damping=min_damping, + cache_jacobian=False, + geodesic_acceleration=False, + ) + state = LMState(jnp.asarray(0.0, dtype=jnp.float32)) + update = solver.update if not jit else jax.jit(lambda x, s: solver.update(x, s)) + x, state, info = update(jnp.ones(1, dtype=jnp.float32), state) + expected_x = min_damping / (1.0 + min_damping) + + assert bool(info.accepted) + assert jnp.allclose(x, expected_x, rtol=1e-4, atol=1e-7) + assert float(state.damping) == pytest.approx(min_damping) + assert float(info.damping) == pytest.approx(min_damping) + + def test_metric_requirements_per_linear_solver(): with pytest.raises(ValueError, match="metric.solve"): LevenbergMarquardt( @@ -2053,6 +2110,7 @@ def build(**overrides): # Any static-setting change (or a different residual function) is a # different solver, so it cannot silently reuse the wrong compiled loop. assert a != build(init_damping=2e-2) + assert a != build(min_damping=1e-8) assert a != build(geodesic_acceleration=False) assert a != LevenbergMarquardt( lambda theta, args, p: theta - args, init_damping=1e-2, cache_jacobian=False @@ -2216,6 +2274,8 @@ def test_hyperparams_typing_and_solve_population(): assert solver.init({"a": 1.0, "b": 0.0}, (ts, ys)).hyper is None hyper = solver.hyperparams(jnp.float32) assert hyper.damping_decrease.dtype == jnp.float32 + assert hyper.min_damping.dtype == jnp.float32 + assert hyper.min_damping == jnp.finfo(jnp.float32).tiny assert hyper.iterative_maxiter.dtype == jnp.int32 assert hyper.max_damping is None result = solver.solve({"a": 1.0, "b": 0.0}, (ts, ys), max_steps=2)