Chebyshev-Stabilized SDIRK Neural ODE

(1) Integrate this at the solver used for a continuous-depth MLP, neural ODE, or state-space model: the tensor being modified is each implicit stage state \(Y_i\), not the neural weights. The neural function \(F_\theta\) is evaluated normally, and the outer training graph can either retain all inner iterations or use a custom backward pass. Start with a detached or recomputed inner solve to control memory. (2) Pseudocode: for each outer state \(y_n\), compute the five explicit offsets \(b_i\) from the displayed SDIRK4 coefficients; initialize \(z_0=b_i\), \(z_{-1}=z_0\), and \(r_0=b_i+h\gamma F_\theta(z_0)-z_0\). Estimate \([a,b]\), construct \(\alpha_k,\beta_k\), then loop for \(k=0,\ldots,K-1\): update \(z_{k+1}=z_k+\alpha_k r_k+\beta_k(z_k-z_{k-1})\), set \(r_{k+1}=b_i+h\gamma F_\theta(z_{k+1})-z_{k+1}\), and stop when \(\|r_{k+1}\|/\|b_i\|<\varepsilon\). Set \(Y_i=z_K\), then continue to the next stage. (3) Estimate the spectral interval empirically: use 3--5 power or Lanczos iterations with Jacobian-vector products to estimate the most negative real Jacobian eigenvalue \(\lambda_{\min}\); set \(a=1\) and \(b=1-h\gamma\lambda_{\min}\), inflated by 10--25% for safety. If eigenvalues are not predominantly real and nonpositive, fall back to damped fixed-point updates or a Newton-Krylov baseline. The mathematical part is the SDIRK4 discretization and Chebyshev polynomial recurrence; the spectral bound and inner iteration count \(K\) are estimated. (4) First experiment: train a small 2--4 layer neural ODE MLP on MNIST or CIFAR-10 with a deliberately stiff residual block, such as \(F_\theta(y)=W_2\tanh(W_1y)\) initialized with a large negative symmetric component. Compare adaptive SDIRK4-Chebyshev, explicit RK4, and standard Newton-solved SDIRK4 at equal function/JVP evaluations. Measure training loss versus wall-clock time, gradient norm variance, failed-step frequency, and final accuracy. Success is a stable outer step size at least 4x larger than RK4, lower loss at equal function evaluations, and agreement with Newton SDIRK4 within the stage residual tolerance.