Optimization helpers

The question this method answers¶

Given a scalar objective $F(x)$, how can we move parameters toward a smaller value while retaining enough diagnostics to decide whether the result is credible? Jaxstro owns small loss, line-search, and stopping mechanics. It does not own an optimizer stack, parameter schedule, or scientific acceptance policy.

Before computation: what should be true?¶

The objective must return one finite floating scalar for parameter arrays of a fixed shape. Scales and units must make the residual definition meaningful. For line search, the proposed direction $p_k$ should be a descent direction: $\nabla F(x_k)^\mathsf{T}p_k<0$. Tolerances should be tied to a requested scientific accuracy, not copied from an unrelated problem.

See Parameters, constraints, and transforms for Jaxstro’s parameter bridge and its cached-derived-leaf boundary.

Define the mathematical objects¶

Let $x\in\mathbb{R}^n$ be a parameter vector, $F:\mathbb{R}^n\to\mathbb{R}$ a scalar objective, and $g_k=\nabla F(x_k)$ its gradient. A search direction $p_k\in\mathbb{R}^n$ specifies the proposed motion, and a positive step length $\alpha_k$ controls its size.

For residuals $r_i(x)$, the half-squared loss is $\rho(r)=r^2/2$. The Huber loss is quadratic near zero and linear in the tails. Its value and first derivative are continuous at $|r|=\delta$, while its second derivative changes there. The pseudo-Huber loss is the smoother approximation

Derive the method¶

Steepest descent chooses $p_k=-g_k$, giving

A step that is too large can increase $F$. Armijo backtracking tests the more general direction $p_k$ against the sufficient-decrease inequality

Starting from $\alpha_0$, the tested sequence is $\alpha_i=\alpha_0\beta^i$ with $0<\beta<1$. The right-hand side decreases below $F(x_k)$ only for a descent direction.

No single stopping statistic proves convergence. Jaxstro reports three. Let $s_0>0$ be the parameter-scale floor used when $\lVert x_k\rVert_2$ is small. The runtime name is scale_floor, with default scale_floor=1e-12:

The helper declares convergence only when all three inequalities pass.

What the algorithm actually does¶

squared_loss, huber_loss, and pseudo_huber_loss act elementwise. objective_summary returns loss, mean loss, RMSE, maximum absolute residual, and element count; supplied weights multiply squared residuals and their sum is the normalizer. A nonpositive normalizer is replaced by one for finite division, but the original zero-weight scientific meaning is not repaired.

armijo_backtracking evaluates a fixed max_steps sequence with lax.scan and records the first accepted candidate in LineSearchResult(step, value, accepted, iterations). If none passes, it returns the last backtracked candidate with accepted=False. The scan count, not the first acceptance, fixes the executed trace.

What JAX differentiates¶

The smooth losses and diagnostics compose with JAX array transforms. Squared and pseudo-Huber losses are smooth on their floating domains. Huber is $C^1$ but not $C^2$ at $|r|=\delta$: its gradient is continuous, while its curvature and Hessian change across the threshold. Armijo’s acceptance predicate selects a discrete branch, so a derivative of the returned step is a derivative of the selected finite program, not an implicit derivative of an optimum.

Using it in Jaxstro¶

import jax
import jax.numpy as jnp

from jaxstro.numerics.optimization import (
    armijo_backtracking,
    convergence_summary,
)


def objective(x):
    return 0.5 * jnp.sum((x - jnp.array([1.0, -2.0])) ** 2)


x = jnp.array([4.0, 1.0])
grad = jax.grad(objective)(x)
search = armijo_backtracking(objective, x, -grad, grad, max_steps=12)
x_new = x + search.step * (-grad)
diagnostics = convergence_summary(
    x_new=x_new,
    x_old=x,
    grad=jax.grad(objective)(x_new),
    loss_new=objective(x_new),
    loss_old=objective(x),
)

The parameter arrays and direction must share a shape. f returns a scalar; max_steps and f are static when the line search is itself JIT-compiled. Diagnostics reduce array inputs to scalar JAX arrays.

How to audit the result¶

First verify the gradient against a central finite difference along at least one declared direction. Confirm $g_k^\mathsf{T}p_k<0$, reproduce every tested Armijo inequality, and retain accepted and iterations. Run from multiple initial conditions when local minima are possible. Refine all three convergence tolerances and check that the scientifically relevant outputs remain stable. For robust losses, finite-difference the Huber gradient from both sides of the threshold to verify first-derivative continuity, and keep curvature or Hessian probes away from $|r|=\delta$. Include outliers large enough to exercise the tails; use pseudo-Huber when a smoother curvature contract is required.

Executable method audits are indexed in Validation methods.

Where the claim stops¶

These helpers do not prove convexity, uniqueness, identifiability, global optimality, or posterior correctness. They do not manage PyTrees, constraints, optimizer state, batching policy, or second-order solves. A converged numerical summary supports only the tested objective and representation.

Connected ideas¶