Native TensorFlow trust-region minimizers (tf-trust-exact, tf-trust-ncg, tf-trust-krylov) - #153
Open
davidwalter2 wants to merge 5 commits into
Open
Native TensorFlow trust-region minimizers (tf-trust-exact, tf-trust-ncg, tf-trust-krylov)#153davidwalter2 wants to merge 5 commits into
davidwalter2 wants to merge 5 commits into
Conversation
…vice New --minimizerMethod tf-trust-exact: a TensorFlow port of scipy's trust-exact (More-Sorensen) where the Hessian, the H + lambda*I builds, the Cholesky factorizations and the triangular solves stay on the TF device, so the n x n matrix never round-trips to LAPACK per lambda trial. The python outer loop mirrors scipy's _minimize_trust_region and plugs into the existing callback / early-stopping / restart / result plumbing unchanged, including the preconditioner (applied at the same internal-coordinate boundary as for the scipy methods). Differences from a line-for-line port, all deliberate: - No custom op for the potrf failure index: tf.linalg.cholesky signals a non-PD input by filling the factor with NaNs on every backend, and a failed factorization already proves lambda_current is a valid lower bound, so the safeguarded bracket update converges without the index-accelerated bound (a few extra cheap factorizations at worst). - The Hessian closure runs only on accepted steps; scipy's subproblem constructor consumes Hessian norms and therefore pays a Hessian per proposal. - lambda_new is clamped to >= 0 (the interior-case Newton correction can otherwise push the bracket negative) and the matrix norms are computed as matrix norms -- tf.norm's default axis=None flattens the tensor and its max|H_ij| can sit below |lambda_min|, silently invalidating the lambda_ub bracket. - The rare hard-case refinement downloads the factor once and reuses scipy's host-side smallest-singular-value estimate. Validated against scipy on random PD/indefinite subproblems (worst-case model reduction 96.7% of the exact optimum vs scipy's 96.0%, ~1.6 factorizations per solve), on Rosenbrock and an ill-conditioned quadratic (identical minima to ~1e-8), and through a full Fitter fit on the test tensor against the scipy path. Without a visible GPU a warning points users back to scipy trust-exact: TF's CPU Cholesky kernel is single-threaded Eigen and measured ~7x slower at n=2000. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY
The matrix-free counterpart of tf-trust-exact: the same truncated-CG trust-region subproblem as scipy's trust-ncg (the practical stand-in for trust-krylov, which solves the identical subproblem via GLTR), but the entire CG iteration runs inside a single tf.function while_loop. One solve costs one graph dispatch instead of one python round trip -- x assignment, numpy conversion both ways, forced device sync -- per Hessian-vector product, which is what the scipy callback path pays and what dominates when individual HVPs are fast. Exact improvements over scipy's loop, with identical iterates: - Hz is carried through the CG recurrence, so the boundary and negative-curvature exits price their steps from dot products where scipy spends two extra HVPs; - the model value of the returned step falls out of the same bookkeeping and is handed to the outer loop, which otherwise needs one more HVP per proposal. The preconditioner gains tf_transforms(): graph versions of the block T/T^T applications (gather -> dense matvec -> scatter), because for an HVP subproblem the reparameterisation runs inside the compiled loop, once per CG iteration, where the numpy path cannot. The subproblem also re-pins the fitter's parameter state before each solve, since the outer loop's objective evaluations at proposed points move it in between. Validated: steps agree with scipy's CGSteihaugSubproblem to float precision across PD/indefinite models and radii (same algorithm, so bitwise-deterministic, unlike the nearly-exact solver); Rosenbrock matches scipy trust-ncg; full Fitter fits match the trust-krylov reference with and without preconditioning. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY
The actual trust-krylov algorithm (Gould-Lucidi-Roma-Toint, the method behind trlib/scipy trust-krylov), completing the native set. Where Steihaug-CG stops at the first boundary crossing, GLTR keeps expanding the Krylov subspace and returns the step that is optimal within it: the Lanczos basis tridiagonalizes H, the projected subproblem min gamma0*e1.h + h.T_k h/2, ||h|| <= Delta is solved per iteration by eigendecomposition plus a safeguarded secular solve (hard case included), and the full-space residual comes free as gamma_k |h_k|. Division of labour: HVPs, the Lanczos recurrence and full CGS2 reorthogonalization run on the TF device through one compiled step function against a fixed-shape [kmax, n] basis buffer (zero rows make masking unnecessary, and fixed shapes mean no retracing); the k x k tridiagonal solves run in numpy where they cost microseconds against HVPs costing milliseconds. Boundary-phase residuals use a 10x tighter tolerance than the interior forcing sequence, as in trlib. The Krylov data is radius-independent, so the outer loop's rejected-step path -- re-solving at a shrunken radius -- reuses it and usually costs zero new HVPs, verified by a test counting them. Validated: subproblem reaches >= 99.9% of the exact-optimum model reduction on problems the subspace exhausts (PD and indefinite, all radii) where truncated CG has no such guarantee; explicit hard-case KKT check; Rosenbrock matches scipy; full Fitter fits match the trust-krylov reference with and without preconditioning (32 tests). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY
Fixes the two issues the first GPU run surfaced: - The interior branch estimated the smallest singular value with the host-side Cline et al. recurrence, downloading the full factor per call -- ~1 s per outer iteration at n=4000 over PCIe, erasing the cheap on-device factorizations. Replaced by inverse iteration on L L^T: a handful of O(n^2) triangular solves on device, with only two scalars and one n-vector crossing to the host. Inverse iteration converges fastest exactly in the near-singular regime the hard case lives in; far from singularity it only gives an upper bound on sigma_min, which every use is safe against (lambda_lb just gets looser, and acceptance is guarded, below). - The hard-case acceptance trusted the s_min estimate outright, so an inaccurate estimate could accept a corrected step that raises the model; the outer loop then rejects it, which showed up as a doubled outer iteration count under GPU rounding at n=2000. Both hard-case exits now verify the candidate against the model value (one matvec) and fall through to the normal lambda updates when it fails. Subproblem quality scan is unchanged (worst case 96.6% of the exact optimum over 60 PD/indefinite cases vs scipy's 96.0%, ~1.5 factorizations per solve); new unit test pins the device estimator to the true sigma_min in the near-singular regime. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY
A proposal whose objective overflows to NaN produced a NaN rho, and IEEE comparisons on NaN are all False: the radius was neither shrunk nor the step accepted, so the outer loop spun at a fixed radius until early stopping gave up far from the minimum. First observed on the first preconditioned run at scale, where an internal-coordinate step of norm 1 is an enormous physical step whose exponentials overflow. A non-finite rho now counts as a hard rejection (radius shrinks), with a regression test whose objective is NaN outside a box around the minimum. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Three native TensorFlow implementations of the trust-region minimizers, selectable via
--minimizerMethod, sharing one outer loop (rabbit/minimizer/base.py, mirroring scipy's_minimize_trust_regionand plugging into the existing callback / early-stopping / restart /minimizer_resultplumbing unchanged):tf-trust-exact(exact.py): Moré–Sorensen, withH + λI, the Cholesky factorizations and the triangular solves on the TF device. No custom op for the LAPACKpotrffailure index:tf.linalg.choleskysignals non-PD input by a NaN-filled factor on every backend, and a failed factorization already proves λ is a valid lower bound, so the safeguarded bracket update converges without it. The hard-case machinery stays on device too (inverse-iteration smallest-singular-value estimate) and its acceptance is guarded by an explicit model-value check.tf-trust-ncg(krylov.py): Steihaug–Toint CG with the whole inner loop as a singletf.functionwhile_loop— one graph dispatch per subproblem solve instead of one python round trip (assign, two numpy conversions, forced device sync) per Hessian-vector product.Hzis carried through the CG recurrence, so the boundary/negative-curvature exits and the outer loop's model value cost no extra HVPs (scipy spends 2–3 there).tf-trust-krylov(gltr.py): GLTR (Gould–Lucidi–Roma–Toint, the trust-krylov algorithm). Lanczos with full CGS2 reorthogonalization on device against a fixed-shape basis buffer; the k×k tridiagonal trust-region subproblems (secular solve, hard case included) on the host where they cost microseconds. The Krylov data is radius-independent, so re-solves after rejected outer steps reuse it — usually zero new HVPs.The preconditioner gains
tf_transforms(): graph-side block T/Tᵀ applications, so the reparameterisation runs inside the compiled CG/Lanczos loops (once per HVP, where the numpy path cannot).Validation
tests/test_native_minimizer.py(33 tests):CGSteihaugSubproblemto float precision (same algorithm, deterministic);Fitterfits for all three methods match the trust-krylov reference, with and without--precondition.Benchmarks (GTX 1080 Ti, fp64 at 1/32 rate — the least favorable GPU available)
2D Z in-situ fit (24k bins, 3603 nuisances), identical conditions, identical minimum 11249.6383:
tf-trust-krylov1043 s vs scipytrust-krylov1435 s (93 vs 96 outer iterations — the win is per-iteration subproblem cost, largest in the endgame where scipy's iterations reach ~200 s vs 67–102 s native).Synthetic (analytic-Hessian objective, n=2000/4000):
tf-trust-exact5.7 s / 16.6 s vs scipy trust-exact 7.7 s / 24.7 s, identical minima; raw fp64 Cholesky on device 2–3× faster than 8-core LAPACK. On CPU-only nodes the native paths are slower than scipy (TF's CPU Cholesky kernel is single-threaded Eigen) — a warning is logged and the scipy methods remain the default; nothing changes unless atf-method is selected.Multi-GPU support is intentionally not part of this PR; it will come separately.
🤖 Generated with Claude Code
https://claude.ai/code/session_01VCVDNM26GpDMggr8GDFiwY