L-Stable Trajectory-Derivative Optimizer
Implementation & benchmark of arXiv:2608.03256 — An L-Stable Sequential Two-Stage Fourth-Order Method with ADER Trajectory Derivatives for Stiff Transport--Relaxation Systems
Source paper: An L-Stable Sequential Two-Stage Fourth-Order Method with ADER Trajectory Derivatives for Stiff Transport--Relaxation Systems arXiv:2608.03256 ⓘ · analyzed Aug 31, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace an explicit gradient step by an implicit correction using the trajectory derivative \(Dg(\theta)g(\theta)=H(\theta)g(\theta)\), where \(g=\nabla f\) and \(H=\nabla^2 f\). The update should strongly damp high-curvature or stiff modes while preserving fourth-order matching of the local linearized dynamics. Start with a self-contained fourth-order L-stable rational prototype, then compare it with the paper's exact two-stage coefficients after recovering those coefficients from the full paper.
Formulas
Mathematical statement
The paper defines the trajectory derivative as \(\mathcal{G}^{\rm tr}(U)=D\mathcal{L}(U)\,\mathcal{L}(U)\). For optimization, set the state to parameters \(\theta\), vector field \(\mathcal{L}(\theta)=-g(\theta)\), gradient \(g(\theta)=\nabla f(\theta)\), and Hessian \(H(\theta)=\nabla^2 f(\theta)\). Then \(\mathcal{G}^{\rm tr}(\theta)=D(-g)(-g)=H(\theta)g(\theta)\). The paper reports that choosing \(C_q=5/183\) cancels the leading inverse-power term in the stiff amplification factor, changing deep-stiff decay from \(O(|z|^{-1})\) to \(O(|z|^{-2})\), where \(z=h\lambda\), \(h\) is the step size, and \(\lambda\) is an eigenvalue of the linearized vector field. As a self-contained fourth-order L-stable prototype, use the [1/3] rational amplification function \(R(z)=(1+z/4)/(1-3z/4+z^2/4-z^3/24)\), which satisfies \(R(z)=e^z+O(z^5)\) near zero and \(R(z)=O(z^{-2})\) as \(|z|\to\infty\). For gradient flow \(\dot\theta=-g(\theta)\), approximate the local linear dynamics using Hessian-vector products; the denominator becomes the operator \(P_h=I+(3/4)hH+(1/4)h^2H^2+(1/24)h^3H^3\), because the stability argument uses \(z=-h\lambda_H\) for a positive Hessian eigenvalue \(\lambda_H\).
Implementation notes
Integrate this method at the parameter-update line of a PyTorch or JAX training loop, initially for full-batch or large-batch fine-tuning where Hessian-vector products are affordable. At iteration k, compute the minibatch loss \(f(\theta_k)\), gradient \(g_k=\nabla f(\theta_k)\), and the Hessian-vector-product oracle \(\operatorname{HVP}(v)=\nabla_\theta(g_k^\top v)\). Use the local gradient-flow linearization \(\dot\theta=-g_k-H_k(\theta-\theta_k)\). Form the right-hand side corresponding to the [1/3] rational propagator, \(b_k=\theta_k-hg_k+(h^2/4)H_kg_k-(h^3/24)H_k^2g_k\), and solve \(P_h\Delta=b_k-\theta_k\), with \(P_h=I+(3h/4)H_k+(h^2/4)H_k^2+(h^3/24)H_k^3\). Apply \(\theta_{k+1}=\theta_k+\Delta\), optionally applying decoupled AdamW-style weight decay after the solve. Do not materialize H: compute \(HVP(g_k)\), \(HVP(HVP(g_k))\), and repeated HVPs inside the matrix-vector product required by conjugate gradients or MINRES. Terminate the implicit solve when the relative residual is below \(10^{-3}\) or after 10 iterations. Estimate the largest curvature with 3-5 power iterations, \(v\leftarrow Hv/\|Hv\|\), and choose h so \(h\widehat\lambda_{\max}\) lies in a sweep such as \(\{0.1,0.3,1,3,10\}\); reduce h if the solve residual grows. The mathematical quantities are the trajectory derivative, repeated HVPs, and rational stability function. The empirical quantities are the curvature estimate, CG residual, and minibatch gradient noise. First test a 2-layer MLP and ResNet-18 on MNIST or CIFAR-10, comparing SGD, AdamW, implicit Euler, and this update at equal HVP-adjusted FLOPs. Measure training loss versus optimizer steps, loss versus wall-clock time, gradient-norm spikes, and robustness to deliberately oversized h. A useful success signal is stable descent at stepsizes that make SGD diverge, together with lower loss at equal HVP-adjusted compute. As a cheap ablation, replace the fourth-order operator by the second-order L-stable prototype \(R(z)=1/(1-z+z^2/2)\) and test whether the higher-order version improves accuracy rather than merely stability.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.