Coeus Nonlinear Least-Squares Solver

coeus-optim owns the Levenberg-Marquardt (damped Gauss-Newton) solver used by the nonlinear diffusion models in ADR 0036 decision 2. DKI, NODDI, and IVIM fit their parameters through this solver; RITK adds no local nonlinear optimiser.

Why a separate solver contract?

The gradient-descent optimisers in coeus-optim (SGD, Adam, AdamW, RMSProp, AdaGrad) step on gradients already accumulated into parameters, which suits network training. A least-squares solver instead re-evaluates the model at trial points to decide whether a step is acceptable, and exploits the Gauss-Newton curvature approximation \(J^{\!T}\!J \approx H\) that a bare gradient does not expose. The two are different contracts, not two spellings of one.

LeastSquaresProblem trait

LeastSquaresProblem<T> is the contract between a model and the solver. It requires:

The Jacobian is row-major: entry \((i, j)\) at index \(i \cdot n + j\) is \(\partial r_i / \partial p_j\).

The solver requires \(m \ge n\) — fewer residuals than parameters makes \(J^{\!T}\!J\) singular by construction, and damping would mask that rather than solve it.

Domain errors

ProblemError::Domain signals that the model is undefined at the trial parameters (e.g. a negative diffusivity under a square root). The solver treats this as a rejected step — it increases damping to pull the next trial back toward the last accepted point — rather than a failure. ProblemError::Evaluation is for unrecoverable failures.

Levenberg-Marquardt algorithm

Each iteration solves the damped normal equations:

\[ (J^{\!T}\!J + \lambda \cdot \operatorname{diag}(J^{\!T}\!J)) \delta = -J^{\!T}r \]

Damping is scaled by \(\operatorname{diag}(J^{\!T}\!J)\) rather than the identity (Marquardt’s modification), making the step invariant to rescaling of individual parameters — a diffusion model mixing diffusivities near \(10^{-3}\) with signal amplitudes near \(10^3\) is exactly the badly-scaled case that motivates it.

The algorithm proceeds as follows:

1. Evaluate residuals and Jacobian at the current parameters.
2. Check gradient convergence: \(\|J^{\!T}r\|_\infty \le\) gradient_tolerance.
3. Build \(J^{\!T}\!J\) and solve the damped system via Cholesky decomposition (the matrix is symmetric positive-definite for \(\lambda > 0\)).
4. If Cholesky fails, increase \(\lambda\) and retry (up to a cap).
5. Accept the step when cost decreases; otherwise reject and increase \(\lambda\). On acceptance, decrease \(\lambda\).
6. Check step-tolerance and cost-tolerance convergence.

Usage

use coeus_optim::{
    LeastSquaresProblem, LeastSquaresReport, LevenbergMarquardtConfig,
    ProblemError, levenberg_marquardt,
};

struct MyModel { /* data */ }

impl LeastSquaresProblem<f64> for MyModel {
    fn residual_count(&self) -> usize { /* ... */ }
    fn parameter_count(&self) -> usize { /* ... */ }
    fn residuals(&self, p: &[f64], r: &mut [f64]) -> Result<(), ProblemError> { /* ... */ }
    fn jacobian(&self, p: &[f64], j: &mut [f64]) -> Result<(), ProblemError> { /* ... */ }
}

let problem = MyModel { /* ... */ };
let initial = vec![0.0; problem.parameter_count()];
let config = LevenbergMarquardtConfig::default();
let report: LeastSquaresReport<f64> = levenberg_marquardt(
    &problem, &initial, &config,
)?;

Configuration

LevenbergMarquardtConfig<T> carries seven tuning parameters:

The tolerance defaults are \(\sqrt{\varepsilon}\) for the working scalar type — the standard choice for a first-order criterion in floating point: a residual computed to relative accuracy \(\varepsilon\) cannot certify stationarity below roughly \(\sqrt{\varepsilon}\).

Report and termination

LeastSquaresReport<T> carries the result:

Termination has four variants:

termination.is_converged() is true for all variants except IterationLimit.

Error types

SolverError covers solver-level failures:

ProblemError::Domain is handled internally by the solver (treated as a rejected step) and never propagates to the caller.

First-order optimisers

coeus-optim also provides five gradient-descent optimisers — SGD, Adam, AdamW, RMSProp, AdaGrad — used by the Coeus network training stack. These implement the Optimizer<T, B> trait and are not used by the diffusion pipeline, which requires the residual/curvature contract of LeastSquaresProblem.

The DKI and NODDI models in ritk-diffusion each implement LeastSquaresProblem<f64> and call levenberg_marquardt directly. See the Diffusion Models chapter for the model-specific usage.