newton_cg#
- t3toolbox.backend.optimizers.newton_cg(problem, x0, max_newton=30, gtol_rel=1e-08, cg_maxiter=200, c_armijo=0.0001, g0norm_newton=None, g0norm_cg=None, cg_forcing_power=0.5, use_jit=False, callback=None)#
def newton_cg( problem: Problem, # the fixed-rank least-squares problem x0: Tangent, # initial cores (U, G) max_newton: int = 30, gtol_rel: float = 1e-8, # stop when ‖g‖ <= gtol_rel * ‖g0‖ (‖g0‖ = g0norm_newton or the initial ‖g‖) cg_maxiter: int = 200, c_armijo: float = 1e-4, g0norm_newton: typ.Optional[float] = None, # ‖g0‖ for the Newton stop; default = initial ‖g‖. Also feeds CG unless g0norm_cg set. g0norm_cg: typ.Optional[float] = None, # ‖g0‖ for the CG forcing term; default = g0norm_newton (else initial ‖g‖) cg_forcing_power: float = 0.5, # η = min(0.5, (‖g‖/‖g0‖)**power); larger => tighter CG, fewer Newton steps use_jit: bool = False, # jit the inner CG (lax.while_loop) when the inputs are jax; else eager callback: typ.Optional[typ.Callable] = None, # callback(NewtonInfo) each iteration (host-side; e.g. a display) ) -> typ.Tuple[Tangent, dict]: # (x_cores, stats)
Inexact Riemannian Newton-CG with an Armijo line search – the manifold workhorse (the gauged
His positive-definite there). Each Newton step builds the local GN model once, solvesH p = −gby CG to an inexact forcing-term tolerance (the inner loop – jit-able via_cg_solve()), then backtracks alongretract(α p). The CG truncates on the gauge-singular corewiseH; the outer line search keeps it robust regardless.use_jitjits only the inner CG (the outer loop, line search, and convergence test stay on the host), auto-converting numpy inputs to jax first (_prepare_jit_inputs()) – so ause_jit=Truecall on numpy data returns a jax-backed result (jax’s default float32 unless x64 is enabled) and raises if jax is not installed.Overriding the reference gradient norm ‖g0‖ (
g0norm_newton/g0norm_cg/cg_forcing_power). Both stopping tests are relative to a reference ‖g0‖: the Newton stop is‖g‖ ≤ gtol_rel·‖g0‖and the CG forcing term isη = min(0.5, (‖g‖/‖g0‖)**cg_forcing_power). By default ‖g0‖ is the initial gradient norm – but in a warm-start continuation loop that norm is misleadingly small (the guess is already near the solution), which over-tightens the Newton stop and slackens CG. Pass a reference reflecting the problem’s true gradient scale (e.g. the initial ‖g‖ from the first continuation stage) to restore the intended behavior. Resolution is a chained fallback:g0norm_newtonsets the Newton reference (and CG inherits it unlessg0norm_cgis also given);g0norm_cgalone sets only the CG reference.cg_forcing_power(default0.5, the conventional Eisenstat-Walker value) tunes CG effort per Newton step: since‖g‖/‖g0‖ < 1near the solution, a larger power (0.75,1.0) tightens CG -> more CG iterations but fewer Newton steps – worth it on the manifold when the retraction is expensive relative to a Hessian-apply. Themin(0.5, …)cap onηis retained regardless.callback, if given, is called with aNewtonInfoeach iteration (including the final converged line) – the hook for a live diagnostic display; it runs host-side (it reads the concrete residual), so it composes withuse_jit(only the inner CG jits) but not with a hypothetical fully-jitted outer loop. Ready-made displays:t3toolbox.backend.optimizer_display.make_newton_display().statsalways carries'history'– one_newton_scalar_record()per iteration.- Parameters:
problem (Problem)
x0 (Tangent)
max_newton (int)
gtol_rel (float)
cg_maxiter (int)
c_armijo (float)
g0norm_newton (t3toolbox.backend.common.typ.Optional[float])
g0norm_cg (t3toolbox.backend.common.typ.Optional[float])
cg_forcing_power (float)
use_jit (bool)
callback (t3toolbox.backend.common.typ.Optional[t3toolbox.backend.common.typ.Callable])
- Return type:
t3toolbox.backend.common.typ.Tuple[Tangent, dict]