# Compiled forward second-order jet residuals

- ID: 1551
- Canonical URL: https://synthcore.org/idea/1551/compiled-forward-second-order-jet-residuals
- API JSON: https://synthcore.org/api/idea/1551.json
- API Markdown: https://synthcore.org/api/idea/1551.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2607.28733](https://arxiv.org/abs/2607.28733)
- Category: training
- Solves: speedup, memory, stability
- ML areas: training, optimizer, inference-speedup
- Math tags: differential-geometry, numerical-analysis, linear-algebra, pde
- Ratings: usefulness 8/10; difficulty 5/10; novelty 6/10

## Idea description

Build a forward-mode second-order jet interpreter for the PINN and evaluate the entire PDE residual in one compiled graph. Each intermediate carries its value, first derivative, and Hessian with respect to the collocation coordinates, avoiding repeated nested reverse-mode autodiff calls for every residual component.

## Mathematical statement

A second-order jet for a scalar function $f:\mathbb R^{d}\to\mathbb R$ at an input is the tuple $J^{2}f=(v,p,H)$ with $v=f(x)$, $p=\nabla_xf\in\mathbb R^d$, and $H=\nabla_x^{2}f\in\mathbb R^{d\times d}$. For a scalar composition $y=\phi(v)$, forward propagation uses $y=\phi(v)$, $p_y=\phi'(v)p$, and $H_y=\phi''(v)pp^{\top}+\phi'(v)H$. For a vector affine layer $y=Wv+c$, the jet is $(y,Wp,WH)$, with the same matrix multiplication applied to every derivative slice. These rules exactly propagate all first and second coordinate derivatives needed by a second-order PDE. The residual uses the pullback metric $g=X^{-2}J^{\top}J$ and the tension field $\tau(u)=\operatorname{tr}_{u^{*}g}\nabla du$; therefore one jet pass supplies $u$, $J$, and the Hessians required to assemble $g$, its inverse, and $\tau$. The computational graph should be traced or compiled once for a fixed batch shape rather than rebuilt at every optimizer step.

## Key formulas

- $$J^{2}f(x)=\bigl(f(x),\nabla f(x),\nabla^{2}f(x)\bigr)=\bigl(v,p,H\bigr).$$
- $$y=\phi(v),\qquad p_y=\phi'(v)p,\qquad H_y=\phi''(v)pp^{\top}+\phi'(v)H.$$
- $$y=Wv+c,\qquad p_y=Wp,\qquad H_y=WH.$$
- $$g_{ab}=X^{-2}(J^{\top}J)_{ab},\qquad \mathcal L_{\mathrm{PDE}}=\frac1N\sum_{i=1}^{N}\|\tau(u(z_i))\|_2^2.$$

## Implementation notes

Modify the residual evaluator, not the optimizer. Represent every scalar activation value by a Jet object containing `val` with shape `[B]`, `grad` with shape `[B,d]`, and `hess` with shape `[B,d,d]`, where $d$ is the number of collocation coordinates. For each linear layer apply `val=W@val+c`, `grad=W@grad`, and `hess=W@hess`; for each elementwise activation apply `val=phi(val)`, `grad=phi_prime(val)*grad`, and `hess=phi_double_prime(val)*outer(grad,grad)+phi_prime(val)*hess$. For batched vector outputs, retain one jet per output channel. Implement common activations with explicit derivatives, preferably tanh or sine; avoid ReLU because its second derivative is undefined. After the final layer, apply the exact-boundary output transform if used, propagate its jets analytically, and assemble $J$, $g=X^{-2}J^    op J$, and the target PDE residual from the jets. Trace this function once with `torch.compile`, XLA, JAX `jit`, or a custom CUDA kernel; keep collocation batch size and tensor shapes static. Validate jet derivatives against finite differences and PyTorch nested autodiff on a tiny network before benchmarking. The first experiment should use a 3-layer width-64 PINN for a Poisson equation and then the hyperbolic minimal-disc residual, comparing nested reverse-mode autodiff, forward-over-reverse, and jets at equal batch size. Measure milliseconds per step, peak memory, and residual-loss trajectory; the claimed success signal is at least a 10x step-time reduction, ideally approaching the paper's reported roughly 40--50x reduction, with numerically indistinguishable residuals and final solutions.

## Verification

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: no_effect
- Verdict: Built a forward second-order Jet interpreter for tanh MLPs, including batched values, gradients, Hessians, and Laplacian residual evaluation. Jet derivatives matched nested PyTorch autodiff to 3.6e-16 maximum error across dimensions 1, 2, 4, and 8, while finite-difference curvature exhibited the predicted second-order scaling with observed slope 1.9997. The eager jet evaluator was 91.4x faster than the tested nested-reverse baseline; torch.compile could not run because Triton does not support the Tesla P40.

### Mechanism check

- Verdict: Built a forward second-order Jet interpreter for tanh MLPs, including batched values, gradients, Hessians, and Laplacian residual evaluation. Jet derivatives matched nested PyTorch autodiff to 3.6e-16 maximum error across dimensions 1, 2, 4, and 8, while finite-difference curvature exhibited the predicted second-order scaling with observed slope 1.9997. The eager jet evaluator was 91.4x faster than the tested nested-reverse baseline; torch.compile could not run because Triton does not support the Tesla P40.
- Confidence: 9/10
- Limitations: Only a small tanh MLP and Laplacian residual were tested. The pullback metric, tension-field hyperbolic minimal-disc residual, exact-boundary transform, optimizer trajectory, memory peak, and full PINN training were not implemented or benchmarked. The compiled-graph speedup was not measured because the available Tesla P40 is unsupported by the installed Triton backend; the reported speedup compares eager jets against an intentionally expensive per-sample nested-Hessian evaluator.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.4699
- Paired wins: 5/8
- Benchmark verdict: no measurable effect

## Artifacts

- [bench_jet_pinn.py](https://synthcore.org/code/390/bench_jet_pinn.py)
- [bench_report.json](https://synthcore.org/code/390/bench_report.json)
- [jet_experiment.py](https://synthcore.org/code/390/jet_experiment.py)
- [report.md](https://synthcore.org/code/390/report.md)
- [report_bench_2026-08-31T110642.md](https://synthcore.org/code/390/report_bench_2026-08-31T110642.md)
- [results.json](https://synthcore.org/code/390/results.json)
- [Download all files as ZIP](https://synthcore.org/download/390)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.