# SPDX-License-Identifier: MIT
# Copyright (c) 2025 Lucy Eduardo Arias
"""Newton-Raphson Method for Unconstrained Optimization.
Mathematical Framework
----------------------
The Newton-Raphson method (also called Newton's method) is a second-order
optimization algorithm that uses both gradient and Hessian information to
find stationary points of a function.
Problem Formulation
-------------------
Find x* that minimizes f(x):
x* = argmin_{x ∈ ℝᵖ} f(x)
where f: ℝᵖ → ℝ is twice continuously differentiable.
**Optimality condition**: At a local minimum, ∇f(x*) = 0
Newton's Method Algorithm
-------------------------
Starting from initial guess x⁽⁰⁾, iterate until convergence:
**Step 1: Compute gradient and Hessian**
g⁽ᵗ⁾ = ∇f(x⁽ᵗ⁾) (gradient, p-vector)
H⁽ᵗ⁾ = ∇²f(x⁽ᵗ⁾) (Hessian, p×p matrix)
**Step 2: Solve Newton system**
Compute search direction p⁽ᵗ⁾ by solving:
H⁽ᵗ⁾ p⁽ᵗ⁾ = -g⁽ᵗ⁾
**Step 3: Update parameters**
x⁽ᵗ⁺¹⁾ = x⁽ᵗ⁾ + p⁽ᵗ⁾
**Step 4: Check convergence**
Stop when ||g⁽ᵗ⁾|| < tol or ||p⁽ᵗ⁾|| < tol
Derivation via Taylor Approximation
------------------------------------
The Newton step is derived from the second-order Taylor expansion:
f(x) ≈ f(x⁽ᵗ⁾) + g⁽ᵗ⁾ᵀ(x - x⁽ᵗ⁾) + ½(x - x⁽ᵗ⁾)ᵀ H⁽ᵗ⁾ (x - x⁽ᵗ⁾)
Minimizing this quadratic approximation (setting derivative to zero):
∇[quadratic] = g⁽ᵗ⁾ + H⁽ᵗ⁾(x - x⁽ᵗ⁾) = 0
gives the Newton step:
x - x⁽ᵗ⁾ = -(H⁽ᵗ⁾)⁻¹ g⁽ᵗ⁾
**Interpretation**: At each iteration, Newton's method finds the minimum
of the local quadratic approximation of f.
Convergence Theory
------------------
**Theorem** (Quadratic convergence, Dennis & Schnabel, 1996):
Suppose:
1. f is twice continuously differentiable
2. x* is a local minimum with ∇f(x*) = 0
3. H(x*) is positive definite (strict local minimum)
4. Starting point x⁽⁰⁾ is sufficiently close to x*
Then:
- Newton's method converges to x*
- Convergence is quadratic: ||x⁽ᵗ⁺¹⁾ - x*|| ≤ C ||x⁽ᵗ⁾ - x*||²
**Rate**: Number of correct digits approximately doubles per iteration.
**Comparison**:
- Linear convergence: error × constant (gradient descent)
- Superlinear: error^α, 1 < α < 2 (quasi-Newton)
- Quadratic: error² (Newton's method)
**Global convergence**: Newton's method is NOT globally convergent.
- May diverge from poor starting points
- Requires positive-definite Hessian (not guaranteed away from minimum)
- Often combined with line search or trust regions for globalization
Conditions for Positive Definiteness
-------------------------------------
Newton's method requires H⁽ᵗ⁾ to be positive definite at each iteration.
**When H is positive definite**:
- f is strictly convex (globally or locally)
- Near a strict local minimum
- For GLMs with canonical link: always (Fisher information)
**When H may be indefinite**:
- Saddle points: some eigenvalues negative
- Far from minimum
- Non-convex optimization landscapes
**Remedy**: Modified Newton with regularization:
(H⁽ᵗ⁾ + λI) p⁽ᵗ⁾ = -g⁽ᵗ⁾
where λ > 0 ensures positive definiteness (not implemented here).
Hessian Computation
-------------------
This implementation computes Hessians using three methods:
### 1. Automatic Differentiation (Preferred)
**PyTorch**:
Uses torch.autograd.grad twice (forward-mode AD)
Cost: O(p²) forward passes
**JAX**:
Uses jax.hessian (reverse-over-reverse AD)
Cost: O(p) forward + O(p) backward passes
Most efficient for small to medium p
### 2. Finite Differences (Fallback)
When AD not available, uses central differences:
∂²f/∂xᵢ∂xⱼ ≈ [f(x+eᵢ+eⱼ) - f(x+eᵢ-eⱼ) - f(x-eᵢ+eⱼ) + f(x-eᵢ-eⱼ)] / (4h²)
where eᵢ is the i-th unit vector and h = 10⁻⁴.
**Cost**: O(p²) function evaluations (4 per Hessian entry)
**Accuracy**: O(h²) truncation error, but subject to roundoff for small h
Numerical Stability
-------------------
**Challenges**:
1. **Singular Hessian**: When H is rank-deficient (parameter redundancy)
- Solution fails
- Returns error message
- Suggests using L-BFGS or ridge regularization
2. **Ill-conditioned Hessian**: When condition number κ(H) is large
- Numerical error in solution: O(ε × κ(H))
- Error amplification in gradient
- Common in overparameterized models
3. **Finite-difference errors**: Tradeoff between truncation and roundoff
- Step size h = 10⁻⁴ balances errors
- Can fail for very steep or flat functions
**Improvements** (not implemented):
- Cholesky decomposition with diagonal pivoting
- Condition number monitoring
- Iterative refinement
- Hessian-free methods (conjugate gradient on H·p = -g)
Computational Complexity
------------------------
Per iteration, for p parameters:
**Gradient computation**:
- AD: O(p) backward pass
- Finite differences: O(p) function evals
**Hessian computation**:
- AD (JAX): O(p) gradients = O(p²) total
- AD (PyTorch): O(p²) backward passes
- Finite differences: O(p²) function evals
**System solve** H·p = -g:
- Direct (Cholesky): O(p³)
- Iterative (CG): O(kp²) for k iterations
**Total per iteration**: O(p³) dominated by linear solve
**vs IRLS**: Same O(p³), but Newton computes full Hessian not just X^T W X
**vs L-BFGS**: L-BFGS avoids O(p³) by approximating H⁻¹, cost O(mp) where m ≈ 10
Comparison with Other Methods
------------------------------
**vs Gradient Descent**:
- Newton: Quadratic convergence, expensive per iteration
- GD: Linear convergence, cheap per iteration
- Crossover: Newton better for moderate p, high accuracy needs
**vs Fisher Scoring (IRLS for GLMs)**:
- Newton uses observed Hessian: ∇²ℓ
- Fisher uses expected Hessian: E[∇²ℓ]
- For exponential families: same for canonical link
- Fisher more stable (always positive-definite)
**vs Quasi-Newton (L-BFGS)**:
- Newton: O(p³) per iteration, fewer iterations
- Quasi-Newton: O(p²) per iteration, more iterations
- Quasi-Newton preferred for large p (p > 1000)
**vs Trust Region**:
- Newton: No globalization, may diverge
- Trust region: Guaranteed descent, slower per iteration
- Hybrid approaches combine both
Applications in Aurora-GLM
---------------------------
Newton-Raphson is used for:
1. **Non-canonical GLM links**: When IRLS not equivalent to Fisher scoring
2. **Dispersion parameter estimation**: Optimize profile likelihood
3. **Variance component estimation**: REML in mixed models (via PQL)
4. **General maximum likelihood**: When IRLS doesn't apply
**Not used for**:
- Standard GLM fitting → use IRLS instead (more stable)
- Large-scale problems → use L-BFGS
- Non-smooth objectives → use subgradient methods
Implementation Notes
--------------------
**Multi-backend support**:
Transparently works with NumPy, PyTorch, and JAX arrays through
automatic differentiation when available, falling back to finite
differences.
**Automatic differentiation**:
- Leverages backend AD for exact gradient and Hessian
- No manual derivative implementation required
- Enables rapid prototyping of new models
**Convergence criteria**:
- Gradient norm: ||g|| < tol (first-order optimality)
- Step size: ||p|| < tol (stationarity)
- Both checked for robustness
References
----------
**Core Newton method theory**:
- Dennis, J. E., & Schnabel, R. B. (1996). *Numerical Methods for Unconstrained
Optimization and Nonlinear Equations*. SIAM.
https://doi.org/10.1137/1.9781611971200
- Nocedal, J., & Wright, S. J. (2006). *Numerical Optimization* (2nd ed.).
Springer. Chapter 3: Line Search Methods.
https://doi.org/10.1007/978-0-387-40065-5
**Convergence analysis**:
- Ortega, J. M., & Rheinboldt, W. C. (2000). *Iterative Solution of Nonlinear
Equations in Several Variables*. SIAM.
(Classic convergence rate proofs)
**Automatic differentiation**:
- Griewank, A., & Walther, A. (2008). *Evaluating Derivatives: Principles and
Techniques of Algorithmic Differentiation* (2nd ed.). SIAM.
https://doi.org/10.1137/1.9780898717761
- Baydin, A. G., et al. (2018). \"Automatic differentiation in machine learning:
A survey.\" *Journal of Machine Learning Research*, 18(153), 1-43.
**Numerical linear algebra**:
- Golub, G. H., & Van Loan, C. F. (2013). *Matrix Computations* (4th ed.).
Johns Hopkins University Press. Chapter 4: Linear systems.
**Modified Newton methods**:
- Gill, P. E., Murray, W., & Wright, M. H. (1981). *Practical Optimization*.
Academic Press. Chapter 4: Modifications for indefinite Hessians.
See Also
--------
aurora.core.optimization.irls : IRLS for GLMs (Fisher scoring)
aurora.core.optimization.lbfgs : Quasi-Newton method (Hessian-free)
aurora.models.glm.fitting : GLM fitting algorithms
Notes
-----
For detailed mathematical derivations, see REFERENCES.md in the repository root.
Newton's method is the gold standard for small to moderate-sized optimization
problems where Hessian computation is feasible. Its quadratic convergence rate
makes it highly efficient near the solution, though care must be taken with
initialization and Hessian conditioning.
"""
from __future__ import annotations
from collections.abc import Callable
import numpy as np
from ..types import Array, OptimizationCallback
from .result import OptimizationResult
[docs]
def newton_raphson(
loss_fn: Callable,
init_params: Array,
*,
backend=None,
args: tuple = (),
kwargs: dict | None = None,
max_iter: int = 100,
tol: float = 1e-6,
callback: OptimizationCallback | None = None,
) -> OptimizationResult:
"""Run the Newton-Raphson method with automatic differentiation support.
Minimises ``loss_fn`` using the classical Newton-Raphson algorithm. The
Hessian is computed via automatic differentiation (PyTorch or JAX) when
available, falling back to central finite differences otherwise.
Parameters
----------
loss_fn : callable
Objective function to minimise. Must accept a parameter array as its
first argument and return a scalar.
init_params : array-like, shape (p,)
Initial parameter values.
backend : object, optional
Array backend (NumPy, PyTorch, JAX). Defaults to JAX when available.
args : tuple, default=()
Additional positional arguments forwarded to ``loss_fn``.
kwargs : dict, optional
Additional keyword arguments forwarded to ``loss_fn``.
max_iter : int, default=100
Maximum number of Newton iterations.
tol : float, default=1e-6
Convergence tolerance on gradient norm and step size.
callback : callable, optional
Function called as ``callback(iteration, params, loss_value)`` after
each iteration.
Returns
-------
OptimizationResult
Result containing the solution, final loss, gradient, convergence
status, and diagnostic counts.
Notes
-----
Convergence is quadratic near the solution when the Hessian is positive
definite. The method may diverge from poor starting points; consider
using ``modified_newton`` for robustness.
Raises
------
numpy.linalg.LinAlgError
(Returned, not raised) When the Hessian is singular at the current
iterate, the result is returned with ``success=False``.
See Also
--------
modified_newton : Newton-Raphson with Levenberg-Marquardt regularisation.
irls : IRLS for GLMs (preferred for canonical-link models).
lbfgs : L-BFGS quasi-Newton method for large-scale problems.
"""
if kwargs is None:
kwargs = {}
if backend is None:
from ..backends import get_backend
backend = get_backend("jax")
try:
grad_fn = backend.grad(loss_fn)
except NotImplementedError:
# Backend without autodiff (e.g. NumPy): central finite differences,
# mirroring the Hessian fallback in _compute_hessian.
from ..autodiff.gradient import gradient
grad_fn = gradient(loss_fn)
x = backend.array(init_params)
nfev = 0
njev = 0
nhev = 0
for iteration in range(max_iter):
grad = grad_fn(x, *args, **kwargs)
njev += 1
grad_np = np.asarray(backend.as_numpy(grad), dtype=float)
grad_norm = np.linalg.norm(grad_np)
if grad_norm < tol:
f_val = backend.as_numpy(loss_fn(x, *args, **kwargs))
nfev += 1
return OptimizationResult(
x=backend.as_numpy(x),
fun=float(f_val),
grad=backend.as_numpy(grad),
success=True,
message="Converged: gradient norm below tolerance",
nit=iteration,
nfev=nfev,
njev=njev,
nhev=nhev,
)
hess_np, evals = _compute_hessian(loss_fn, x, backend, args, kwargs)
nfev += evals
nhev += 1
try:
step = np.linalg.solve(hess_np, grad_np)
except np.linalg.LinAlgError as exc: # pragma: no cover - rare singular cases
f_val = backend.as_numpy(loss_fn(x, *args, **kwargs))
nfev += 1
return OptimizationResult(
x=backend.as_numpy(x),
fun=float(f_val),
grad=backend.as_numpy(grad),
success=False,
message=f"Hessian is singular: {exc}",
nit=iteration,
nfev=nfev,
njev=njev,
nhev=nhev,
)
x_np = np.asarray(backend.as_numpy(x), dtype=float)
x_new_np = x_np - step
x = backend.array(x_new_np, dtype=getattr(x, "dtype", None))
f_val = backend.as_numpy(loss_fn(x, *args, **kwargs))
nfev += 1
if callback is not None:
callback(iteration, backend.as_numpy(x), float(f_val))
if np.linalg.norm(step) < tol:
final_grad = backend.as_numpy(grad_fn(x, *args, **kwargs))
njev += 1
return OptimizationResult(
x=backend.as_numpy(x),
fun=float(f_val),
grad=final_grad,
success=True,
message="Converged: step size below tolerance",
nit=iteration + 1,
nfev=nfev,
njev=njev,
nhev=nhev,
)
f_val = backend.as_numpy(loss_fn(x, *args, **kwargs))
nfev += 1
final_grad = backend.as_numpy(grad_fn(x, *args, **kwargs))
njev += 1
return OptimizationResult(
x=backend.as_numpy(x),
fun=float(f_val),
grad=final_grad,
success=False,
message="Maximum iterations reached",
nit=max_iter,
nfev=nfev,
njev=njev,
nhev=nhev,
)
def _compute_hessian(loss_fn, params, backend, args, kwargs):
"""Compute Hessian via autograd when available, otherwise finite differences."""
params_np = np.asarray(backend.as_numpy(params), dtype=float)
dim = params_np.size
hessian = np.zeros((dim, dim), dtype=float)
try:
import torch # type: ignore
if isinstance(params, torch.Tensor):
params_t = params.detach().clone().requires_grad_(True)
result = loss_fn(params_t, *args, **kwargs)
grad = torch.autograd.grad(result, params_t, create_graph=True)[0]
rows = []
for grad_component in grad:
if not grad_component.requires_grad:
rows.append(torch.zeros_like(params_t))
continue
second = torch.autograd.grad(
grad_component,
params_t,
retain_graph=True,
allow_unused=True,
)[0]
if second is None:
second = torch.zeros_like(params_t)
rows.append(second)
hess_tensor = torch.stack(rows)
return hess_tensor.detach().cpu().numpy(), 1
except ImportError: # pragma: no cover - optional dependency missing
pass
try:
import jax
import jax.numpy as jnp
if isinstance(params, jnp.ndarray): # type: ignore[attr-defined]
hess = jax.hessian(lambda p: loss_fn(p, *args, **kwargs))(params)
return np.asarray(hess), 1
except ImportError: # pragma: no cover - optional dependency missing
pass
eps = 1e-4
evaluations = 0
for i in range(dim):
for j in range(i, dim):
e_i = np.zeros(dim)
e_j = np.zeros(dim)
e_i[i] = eps
e_j[j] = eps
f_pp = backend.as_numpy(loss_fn(backend.array(params_np + e_i + e_j), *args, **kwargs))
f_pm = backend.as_numpy(loss_fn(backend.array(params_np + e_i - e_j), *args, **kwargs))
f_mp = backend.as_numpy(loss_fn(backend.array(params_np - e_i + e_j), *args, **kwargs))
f_mm = backend.as_numpy(loss_fn(backend.array(params_np - e_i - e_j), *args, **kwargs))
evaluations += 4
value = float((f_pp - f_pm - f_mp + f_mm) / (4 * eps * eps))
hessian[i, j] = value
hessian[j, i] = value
return hessian, evaluations
[docs]
def modified_newton(
loss_fn: Callable,
init_params: Array,
*,
backend=None,
args: tuple = (),
kwargs: dict | None = None,
max_iter: int = 100,
tol: float = 1e-6,
lambda_init: float = 1e-3,
lambda_factor: float = 10.0,
lambda_max: float = 1e8,
callback: OptimizationCallback | None = None,
) -> OptimizationResult:
"""Modified Newton-Raphson with Levenberg-Marquardt regularization.
Mathematical Framework
----------------------
When the Hessian H is indefinite or ill-conditioned, the standard Newton
step p = -H⁻¹g may not be a descent direction. Modified Newton adds
regularization:
(H + λI) p = -g
where λ > 0 ensures positive definiteness.
Levenberg-Marquardt Strategy
----------------------------
The regularization parameter λ is adjusted adaptively:
1. **Start with small λ**: Allow near-Newton steps when H is well-behaved
2. **Increase λ if step rejected**: When function doesn't decrease sufficiently,
multiply λ by `lambda_factor` to trust gradient more than curvature
3. **Decrease λ if step accepted**: After successful steps, reduce λ
to allow faster convergence via better Hessian approximation
**Interpretation**:
- λ → 0: Pure Newton (quadratic convergence near solution)
- λ → ∞: Steepest descent (slow but robust)
The algorithm smoothly interpolates between these extremes.
Guaranteeing Positive Definiteness
----------------------------------
For symmetric H with eigenvalues λ₁ ≤ λ₂ ≤ ... ≤ λₙ:
- H + λI has eigenvalues λ₁ + λ, λ₂ + λ, ..., λₙ + λ
- If λ > -λ₁ (where λ₁ is the smallest eigenvalue), H + λI is positive definite
In practice, we use the Cholesky factorization attempt:
- If Cholesky succeeds: H + λI is positive definite
- If Cholesky fails: Increase λ and retry
Convergence Properties
----------------------
**Theorem** (Dennis & Schnabel, 1996): Modified Newton with adaptive λ
converges globally for any starting point if:
1. f is bounded below
2. ∇f is Lipschitz continuous
3. f is twice continuously differentiable
**Local rate**: Near the solution where H is positive definite:
- If λ → 0 sufficiently fast: Quadratic convergence
- If λ bounded away from 0: Linear convergence
Parameters
----------
loss_fn : callable
Objective function to minimize
init_params : array
Initial parameter values
backend : object, optional
Array backend (NumPy, PyTorch, JAX)
args : tuple
Additional positional arguments for loss_fn
kwargs : dict, optional
Additional keyword arguments for loss_fn
max_iter : int, default=100
Maximum number of iterations
tol : float, default=1e-6
Convergence tolerance for gradient norm
lambda_init : float, default=1e-3
Initial regularization parameter
lambda_factor : float, default=10.0
Factor to increase/decrease λ
lambda_max : float, default=1e8
Maximum allowed λ (switches to pure gradient descent)
callback : callable, optional
Function called after each iteration
Returns
-------
OptimizationResult
Contains solution, final function value, convergence status
References
----------
[1] Levenberg, K. (1944). "A method for the solution of certain non-linear
problems in least squares." Quarterly of Applied Mathematics, 2(2), 164-168.
[2] Marquardt, D. W. (1963). "An algorithm for least-squares estimation of
nonlinear parameters." Journal of SIAM, 11(2), 431-441.
[3] Dennis, J. E., & Schnabel, R. B. (1996). Numerical Methods for
Unconstrained Optimization and Nonlinear Equations. SIAM.
[4] Nocedal, J., & Wright, S. J. (2006). Numerical Optimization (2nd ed.).
Springer. Chapter 4: Trust-Region Methods.
Examples
--------
>>> def rosenbrock(x):
... return (1 - x[0])**2 + 100*(x[1] - x[0]**2)**2
>>> result = modified_newton(rosenbrock, [-1.0, 1.0])
>>> print(result.x) # Should be close to [1, 1]
"""
if kwargs is None:
kwargs = {}
if backend is None:
from ..backends import get_backend
backend = get_backend("jax")
grad_fn = backend.grad(loss_fn)
x = backend.array(init_params)
x_np = np.asarray(backend.as_numpy(x), dtype=float)
nfev = 0
njev = 0
nhev = 0
# Current function value
f_val = float(backend.as_numpy(loss_fn(x, *args, **kwargs)))
nfev += 1
lam = lambda_init
for iteration in range(max_iter):
# Compute gradient
grad = grad_fn(x, *args, **kwargs)
njev += 1
grad_np = np.asarray(backend.as_numpy(grad), dtype=float)
grad_norm = np.linalg.norm(grad_np)
# Check convergence
if grad_norm < tol:
return OptimizationResult(
x=x_np,
fun=f_val,
grad=grad_np,
success=True,
message="Converged: gradient norm below tolerance",
nit=iteration,
nfev=nfev,
njev=njev,
nhev=nhev,
)
# Compute Hessian
hess_np, evals = _compute_hessian(loss_fn, x, backend, args, kwargs)
nfev += evals
nhev += 1
# Modified Newton with adaptive λ
step_found = False
for _ in range(20): # Max attempts to find good λ
try:
# Form H + λI
H_mod = hess_np + lam * np.eye(len(x_np))
# Attempt Cholesky factorization (tests positive definiteness)
L = np.linalg.cholesky(H_mod)
# Solve (H + λI)p = -g via Cholesky
# L L^T p = -g => L y = -g, then L^T p = y
y = np.linalg.solve(L, -grad_np)
step = np.linalg.solve(L.T, y)
# Evaluate new point
x_new_np = x_np + step
x_new = backend.array(x_new_np, dtype=getattr(x, "dtype", None))
f_new = float(backend.as_numpy(loss_fn(x_new, *args, **kwargs)))
nfev += 1
# Armijo condition: sufficient decrease
# f(x + p) ≤ f(x) + c₁ ∇f(x)ᵀp
c1 = 1e-4
directional_deriv = np.dot(grad_np, step)
if f_new <= f_val + c1 * directional_deriv:
# Accept step
x_np = x_new_np
x = x_new
f_val = f_new
step_found = True
# Decrease λ for next iteration (trust Hessian more)
lam = max(lam / lambda_factor, 1e-10)
break
else:
# Increase λ (trust gradient more)
lam = min(lam * lambda_factor, lambda_max)
except np.linalg.LinAlgError:
# Cholesky failed - matrix not positive definite
# Increase λ and retry
lam = min(lam * lambda_factor, lambda_max)
if not step_found:
# Fall back to gradient descent step
step = -grad_np * (1.0 / (grad_norm + 1e-8))
x_new_np = x_np + step * 0.1 # Small step
x_new = backend.array(x_new_np, dtype=getattr(x, "dtype", None))
f_new = float(backend.as_numpy(loss_fn(x_new, *args, **kwargs)))
nfev += 1
if f_new < f_val:
x_np = x_new_np
x = x_new
f_val = f_new
if callback is not None:
callback(iteration, x_np, f_val)
# Check step size convergence
if step_found and np.linalg.norm(step) < tol:
final_grad = np.asarray(backend.as_numpy(grad_fn(x, *args, **kwargs)), dtype=float)
njev += 1
return OptimizationResult(
x=x_np,
fun=f_val,
grad=final_grad,
success=True,
message="Converged: step size below tolerance",
nit=iteration + 1,
nfev=nfev,
njev=njev,
nhev=nhev,
)
# Maximum iterations reached
final_grad = np.asarray(backend.as_numpy(grad_fn(x, *args, **kwargs)), dtype=float)
njev += 1
return OptimizationResult(
x=x_np,
fun=f_val,
grad=final_grad,
success=False,
message="Maximum iterations reached",
nit=max_iter,
nfev=nfev,
njev=njev,
nhev=nhev,
)
__all__ = ["newton_raphson", "modified_newton"]