Merge pull request #282 from PyAutoLabs/feature/nnls-target-kappa-fix

fix(jax): bump default NNLS target_kappa to 1e-2 for finite backward gradients

Author: Jammy2211

autoarray/config/general.yaml

@@ -7,6 +7,7 @@ inversion:
   no_regularization_add_to_curvature_diag_value : 1.0e-3 # The default value added to the curvature matrix's diagonal when regularization is not applied to a linear object, which prevents inversion's failing due to the matrix being singular.
   use_border_relocator: false          # If True, by default a pixelization's border is used to relocate all pixels outside its border to the border.
   nnls_jacobi_preconditioning: true   # If True (default), the curvature matrix passed to jaxnnls.solve_nnls_primal is Jacobi-preconditioned (D Q D y = D q, x = D y). Fixes NaN backward-pass gradients on ill-conditioned Q and roughly halves forward solve time. Set False to restore the raw unpreconditioned solve.
+  nnls_target_kappa: 1.0e-2           # Central-path relaxation parameter passed to jaxnnls.solve_nnls_primal. Larger values smooth the relaxed-KKT backward pass and prevent NaN gradients on ill-conditioned Q; smaller values tighten the primal solve. 1.0e-2 is the smallest value empirically verified to produce finite gradients across MGE pipelines. jaxnnls's own default (1e-3) is too aggressive for the backward pass.
   reconstruction_vmax_factor: 0.5     # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor.
 numba:
   use_numba: true

autoarray/inversion/inversion/inversion_util.py

@@ -287,6 +287,19 @@ def reconstruction_positive_only_from(
             # explicitly disables preconditioning in the shadowing config.
             use_jacobi = True
 
+        try:
+            target_kappa = conf.instance["general"]["inversion"][
+                "nnls_target_kappa"
+            ]
+        except KeyError:
+            # jaxnnls's hardcoded default (1e-3) produces NaN in the relaxed-KKT
+            # backward pass on ill-conditioned curvature matrices, even after
+            # Jacobi preconditioning. 1e-2 is the smallest value empirically
+            # verified to produce finite gradients across MGE pipelines (see
+            # autolens_workspace_developer/jax_profiling/imaging/mge_gradients.py
+            # _diagnose_kappa).
+            target_kappa = 1.0e-2
+
         if use_jacobi:
             # Ill-conditioned Q makes jaxnnls's relaxed-KKT backward pass
             # produce NaN gradients. Rescale Q so its diagonal is unit:
@@ -297,9 +310,14 @@ def reconstruction_positive_only_from(
             D = 1.0 / d
             Q_pc = (curvature_reg_matrix * D[:, None]) * D[None, :]
             q_pc = data_vector * D
-            return jaxnnls.solve_nnls_primal(Q_pc, q_pc) * D
+            return (
+                jaxnnls.solve_nnls_primal(Q_pc, q_pc, target_kappa=target_kappa)
+                * D
+            )
 
-        return jaxnnls.solve_nnls_primal(curvature_reg_matrix, data_vector)
+        return jaxnnls.solve_nnls_primal(
+            curvature_reg_matrix, data_vector, target_kappa=target_kappa
+        )
 
     try: