feat(solver): add autoregressive surrogate and solver-loop studies

[andrinr]

Outcome

Stacked on #116. This PR evaluates two distinct uses of learned physics models:

study	learned model's role	question	current result
3D XLB surrogate	replaces a fixed 100-step solver map	does forward fidelity preserve inverse behavior?	No. The 16k checkpoint improves forward error but worsens projected-L-BFGS IC recovery; recovery is also optimizer-dependent.
2D solver-in-loop corrector	acts sparsely inside a native solver rollout	when does differentiation through the solver materially help?	In the terminal-credit regime, WIG improves the native rollout for all six one-seed admission cells; the three-seed production matrix is still running.

The 3D model is a full-field neural operator applied autoregressively for 20
macro-steps. Its learned map has a finite-difference-consistent VJP, but it is
neither a kernel-level speedup nor a drop-in inverse for XLB observations. The
study shows that forward error and local Jacobian summaries are insufficient
selection criteria for an inverse task: the learned objective can admit a
different low-residual IC branch.

The opt-in 2D study asks a different question. A corrector acts every four
native steps, two warm-up intervals are detached, and the loss is evaluated
only after the final solver interval. Under this objective NOG is exactly the
native-solver control, while WIG is trainable only through the solver VJP. At
the admission budget, WIG/native rollout improvement ranges from 1.194× to
1.503× across JAX-CFD, INS.jl, PhiFlow, PICT, Warp-NS, and XLB, and all six
suffix finite-difference checks pass (0.036%–1.170% best relative error).
These are admission results, not final production statistics; jobs
1699012–1699029 are running offline on Kander.

Together, the studies separate replacing a solver with a surrogate from
learning a corrector through a solver. They also separate recurrent exposure
from solver-gradient credit: one-step supervision is solver-dependent, while
the terminal-credit construction makes the numerical solver the only path from
a learned correction to the training objective.

Scope and implementation

The Tesseract accepts and returns the complete 16×16×16×3 velocity field for
one fixed periodic task:

quantity	fixed value
domain	[0, 2π]³, fully periodic
resolution	N=16
viscosity	ν=0.01
solver step	dt=0.02
horizon	100 XLB steps
recovery initialization	zero velocity
recovery seeds	0, 1, 2

One width-32, six-mode, six-block 3D FNO advances five XLB steps
(ΔT=0.1). The same weights are reused for 20 autoregressive macro-steps; this
is not an IC-to-final-state regressor. An exact viscous skip and Helmholtz
projection are applied at every macro-step. Drag is exactly zero, matching the
other solvers on this obstacle-free task; there is no learned drag head.

The PR:

• adds the xlb-3d-surrogate Tesseract and its fixed-task exclusions;
• admits it only to the two matching 3D recovery cells;
• packages the shared inference/training model and final checkpoint;
• keeps trajectory generation, training-data construction, and curriculum
  training reproducible beside the Tesseract;
• keeps cluster submission, datasets, intermediate checkpoints, and evaluation
  orchestration outside the repository.

Only the inference API, shared model, and final weights enter the runtime image.

Training and held-out evaluation

The checkpoint uses 16,384 native continuous XLB KBC D3Q27 trajectories, split
12,288/2,048/2,048 for training/validation/test. Each trajectory contains the
IC plus 20 full-field snapshots at five-step intervals. XLB populations evolve
continuously and are not reconstructed from equilibrium between snapshots. The
generated final frame matches the canonical 100-step float32 teacher call with
maximum absolute difference zero.

Training uses rollout curriculum 1 → 2 → 4 → 8 → 12 → 20, followed by
12,000 full-horizon updates. Recovery seeds 0/1/2 are excluded from the
dataset.

• dataset SHA-256:
  4836fba4e6a8524af7a552c5977721118e726afa21db9a9f4d0b612a879a0005
• checkpoint SHA-256:
  1ea04a7333981d1bfb836461d6fd6d89ae12f31c64ea40701c2607f03fb4107f

On the three excluded recovery seeds, the final-field relative L2 error is
4.418%, 4.191%, and 4.317% (4.309% mean); mean field cosine is
0.999076.

The held-out test split has 2,048 trajectories. Relative-error aggregates use
the 1,982 trajectories with IC amplitude ≥0.05: mean final error is 7.219%,
median 5.348%, and p95 17.163%. The other 66 trajectories have amplitude
0<a<0.05; relative error is ill-conditioned as the target norm approaches
zero, so those cases are reported separately with absolute RMS error rather
than folded into the headline relative-error statistic.

A follow-up against the current recurrent-state XLB image reproduces the
forward and JVP headline values exactly, including the 7.219% held-out mean,
4.309% excluded-seed mean, 18.24% low-frequency JVP error, and 26.00%
full-spectrum JVP error.

Primary benchmark: solver-self recovery

The benchmark contract is solver-self recovery: each solver generates its own
final observation and differentiates through itself to recover the IC. Both
optimizer variants start from zero, run for at most 100 iterations, use the
same zoom line search, and evaluate seeds 0/1/2.

The plots retain PhiFlow, Warp-NS, and Exponax as benchmark context, but the
surrogate claim below uses only the matched XLB teacher and its two checkpoints;
it does not depend on ranking solvers with different forward solutions.

optimizer / solver checkpoint	seed 0	seed 1	seed 2	mean
L-BFGS / XLB	4.96%	5.17%	6.12%	5.42%
L-BFGS / surrogate 4k	6.94%	7.81%	9.56%	8.10%
L-BFGS / surrogate 16k	17.31%	15.70%	17.10%	16.70%
projected L-BFGS / XLB	4.70%	5.02%	5.94%	5.22%
projected L-BFGS / surrogate 4k	6.97%	7.85%	9.54%	8.12%
projected L-BFGS / surrogate 16k	17.32%	15.59%	17.19%	16.70%

bfgs_proj follows the paper definition: it projects the gradient before the
L-BFGS update; it does not separately project the quasi-Newton iterate. For
surrogate seed 0, unconstrained/projected final max|∇·u₀| is
1.49e-2/1.43e-2, with maxima 1.52e-2/1.43e-2 over optimization.

The 16k checkpoint reaches an objective around 1e-8 and slightly lower
divergence than the 4k checkpoint while recovering an IC with roughly twice
the error. The learned objective therefore admits a different inverse branch;
low self-target residual is not evidence of correct IC recovery.

Optimizer interaction control

Adam changes the result materially. Its learning rate was selected independently
for each solver from {1e-4, 3e-4, 1e-3, 3e-3, 1e-2, 3e-2} using only mean
self-target MSE on calibration seeds 100/101/102. The selected rate was then
frozen before evaluating seeds 0/1/2. No test IC error was used for selection.

Adam projects both the gradient and every accepted iterate onto the
divergence-free subspace. The iterate projection is necessary because Adam's
coordinatewise moment scaling does not preserve a projected search direction.
Without it, the recovered fields reach max|∇·u₀|≈2.4–3.8 and are not a valid
physical comparison.

solver / checkpoint	projected Adam, 300 updates	projected Adam, 1,000 updates
XLB	40.63%	37.07%
surrogate 4k	46.95%	43.62%
surrogate 16k	33.99%	30.37%
Exponax (context)	5.19%	—
Warp-NS (context)	0.116%	—
PhiFlow (context)	4.38% over two finite trials; 1/3 non-finite	—

At the tested finite budgets, the matched XLB/surrogate ordering reverses: the
16k surrogate has the lowest mean IC error of those three under this
projected-Adam protocol. Absolute Adam recovery nevertheless remains worse
than L-BFGS for every matched checkpoint. No optimizer state is transferred
from network training. The reversal establishes sensitivity to the recovery
optimizer and protocol; it does not isolate whether the cause is Adam's update
geometry, iterate projection, learning-rate or horizon choice, the
parameter-training procedure, or transient convergence to different
low-residual regions.

The contextual solvers make the interaction even clearer, but are not used for
the XLB-surrogate claim because solver-self targets do not establish physical
forward agreement. L-BFGS's zoom-line-search oracle evaluations were not stored
in the earlier artifacts, so optimizer steps in this control are deliberately
not presented as a compute-matched efficiency comparison. The constraint
handling also differs: Adam projects both gradients and iterates, while the
existing bfgs_proj path projects gradients only. Adam rates were calibrated
for 100 updates using observable residual rather than IC error, then reused for
300/1,000 updates; XLB and the 4k surrogate selected the upper grid boundary
3e-2, so their rates are not bracketed.

Why improved forward fidelity does not predict recovery

A matched 4k/16k checkpoint comparison separates conventional held-out
overfitting from differences in inverse geometry:

metric	4,096 trajectories	16,384 trajectories
excluded-seed forward error	7.473%	4.309%
held-out final error (a≥0.05)	9.837%	7.219%
low-frequency JVP error (`	k	≤4`)
full-spectrum JVP error	38.28%	26.00%
L-BFGS self-recovery mean IC error	8.10%	16.70%
projected-L-BFGS self-recovery mean IC error	8.12%	16.70%

Forward-error improvements do not track recovery under the primary L-BFGS
protocol; for these checkpoints, forward-only model selection is insufficient.

At exactly zero, the two checkpoints have nearly identical self-target descent
alignment with the true-IC direction. Immediately away from zero, the 16k
checkpoint turns away:

amplitude on the true-IC ray	4k cosine	16k cosine
0.00	0.847	0.845
0.05	0.825	0.419
0.10	0.807	0.150
0.25	0.776	0.184

The learned correction is amplitude-gated. At the zero cold start its first
derivative is effectively the fixed viscous skip, so additional trajectory
data cannot repair the measured 49.34% radial JVP error at the starting point
without changing the architecture or objective.

Derivative-informed training audit

Two leakage-safe follow-ups test the obvious remedies. The first removes the
amplitude gate while preserving the zero fixed point through the bias-free
network. The second fine-tunes that model with a 0.1 directional-JVP loss
against XLB, using training-base directions only and selecting on validation
field/JVP error.

metric	gated baseline	no gate	no gate + JVP loss
held-out final error (a≥0.05)	7.219%	7.313%	6.859%
low-frequency JVP error	18.24%	16.58%	16.48%
full-spectrum JVP error	26.00%	24.85%	24.64%
zero-start radial JVP error	49.34%	17.73%	20.45%
projected-L-BFGS recovery	16.70%	33.69%	23.63%
projected-Adam recovery (300 updates)	33.99%	55.57%	65.81%

Removing the gate fixes the specific zero-derivative restriction, and explicit
JVP training improves the ordinary forward/JVP metrics further. Neither
ablation beats the gated baseline under the tested projected-L-BFGS or
300-update projected-Adam protocols, so the packaged gated checkpoint is
retained. This is consistent with random training-base JVP directions being a
weak proxy for recovery-path sensitivities, but it does not isolate that as the
cause or show that derivative-informed training generally fails. A stronger
follow-up should sample optimization trajectories and supervise their
task-relevant sensitivities directly.

Task-aware Sobolev training

The stronger follow-up trains the surrogate to reproduce the gradient actually
consumed by recovery. For the final-field MSE L(x, y), it matches projected
directional derivatives
<P ∇x Lsurrogate(x, y), v> against XLB. This is first-derivative,
task-aware Sobolev training; differentiating
that derivative loss with respect to network parameters uses mixed
second-order autodiff ("double backprop"). It does not supervise an explicit
teacher Hessian.

The leakage-safe dataset uses 96 training and 24 validation optimization paths
from the original split, with snapshots at updates
{0, 1, 2, 5, 10, 20, 40, 80} and four deterministic unit divergence-free
directions per state (two |k|≤4, two full-spectrum). Benchmark seeds 0–2 and
Adam calibration seeds 100–102 are excluded. All arms start from the packaged
gated 16k checkpoint. A field-only continuation and Sobolev weights
{1e-1, 1e-2, 1e-3} receive 1,000 matched updates; validation
field relative L2 + task-gradient relative L2 selects λ=1e-3. Only that
selected derivative arm is evaluated on benchmark seeds.

metric	packaged start	field-only continuation	task-Sobolev λ=1e-3
validation task-gradient relative L2	5.305	4.568	4.088
validation task-gradient cosine	0.553	0.571	0.599
held-out final-field error	7.219%	7.219%	7.506%
excluded-seed forward error	4.309%	4.343%	4.528%
low-frequency JVP error	18.24%	17.55%	17.17%
full-spectrum JVP error	26.00%	25.67%	25.26%
projected-L-BFGS self-recovery	16.70%	—	17.05%
projected-Adam self-recovery, LR 1e-2	33.99%	—	34.10%
XLB-target cross-model IC error	37.43%	—	42.19%

Residual-only Adam calibration initially appears to reverse this conclusion:
it selects 1e-2 for the packaged checkpoint and the upper-grid 3e-2 for
task-Sobolev, producing 33.99% and 63.06% IC error. Crossed fixed-rate controls
identify a learning-rate/inverse-branch effect instead: at 1e-2, packaged and
Sobolev obtain 33.99% and 34.10%; at 3e-2, they obtain 63.43% and 63.06%.
The larger rate reaches a slightly smaller surrogate residual while recovering
a much worse IC. The matched-rate result, not the calibration-selected
cross-checkpoint contrast, is the appropriate comparison.

The selected Sobolev checkpoint therefore improves held-out average path
gradient agreement but does not improve self- or XLB-target recovery. This does
not show that task gradients or Sobolev training are generally ineffective.
The present study uses four directional sketches, Adam-path data only, one
fine-tuning seed, and retains the zero-start amplitude gate. It instead shows
that average local sensitivity agreement is insufficient to preserve the
global inverse branch. This is consistent with derivative-informed neural
operator work such as
DINO, while the closest
optimization-trajectory study recommends solver feedback when a learned
optimizer path leaves its training distribution
(Cheng et al.). The packaged checkpoint is
therefore retained.

Cross-model negative control

Self-recovery does not establish that the surrogate can invert an
XLB-generated observation. In the cross-model control, projected L-BFGS
differentiates through the surrogate while targeting an XLB final field:

metric	seed 0	seed 1	seed 2	mean
final IC error	39.22%	40.95%	32.12%	37.43%
best saved IC error	17.52%	18.89%	17.56%	17.99%
surrogate residual	1.62%	1.71%	2.02%	1.78%
XLB re-evaluation residual	14.86%	19.79%	18.65%	17.77%

The surrogate objective decreases while the IC moves onto a
surrogate-specific inverse branch. This checkpoint must not be presented as a
drop-in inverse for XLB observations.

Derivative checks

Finite differences

Central finite differences use seeds 0/1/2, ten shared unit-norm directions per
seed, and a 12-point relative-ε sweep:

objective / best-ε aggregate	XLB	surrogate
paper energy sum(u_T²) median error	6.77e-6	6.63e-3
paper energy mean cosine	≈1.000000	0.999984
recovery MSE at zero median error	1.32e-1	2.56e-4
recovery MSE at zero mean cosine	0.991669	≈1.000000

The surrogate VJP is tightly finite-difference-consistent for the actual
zero-start recovery objective. Its distinct recovery behavior is therefore
not an autodiff implementation error: it accurately differentiates the
learned map.
For XLB, the zero-start direction is close but the best finite-difference
magnitude remains approximate, so this table is not used to claim exact XLB
gradient magnitude at zero.

Local Jacobian diagnostics

These are complementary local audits, not predictors of the global
zero-to-solution optimization basin:

1. Restricted 512D: the complete real divergence-free Fourier input
   subspace through |k|≤4, with the full N=16 output retained.
2. Adapted block-grid 1536D: an orthonormal 8³ block-grid
   lift/restriction around the fixed N=16 recovery map, followed by a dense
   SVD.

The second audit covers every coordinate of the adapted coarse block-grid map.
It is not the full 12,288D production Jacobian and is distinct from the
paper's native N=8 Taylor–Green physics.

audit	metric	XLB	surrogate 4k	surrogate 16k
restricted 512D	mean κ(J)	6.864	7.740	7.142
restricted 512D	mean κ(JᵀJ)	47.17	59.98	51.05
restricted 512D	Frobenius relative error	—	13.50%	8.84%
restricted 512D	Frobenius cosine	—	0.9914	0.9961
adapted block-grid 1536D	raw κ(J)	2.790e7	3.747e10	7.476e10
adapted block-grid 1536D	resolved κ(J)	5.319e3	4.890e3	5.445e3
adapted block-grid 1536D	Frobenius relative error	—	32.98%	32.94%
adapted block-grid 1536D	Frobenius cosine	—	0.9442	0.9443

The similar scalar condition numbers do not imply equal Jacobians. Scaling to
16k improves the restricted low-frequency Jacobian, leaves the adapted
block-grid agreement essentially unchanged, and worsens the unresolved raw
spectral tail. The raw tails fall below the float32 rank tolerance; resolved
condition numbers are reported separately.

Runtime

Two matched RTX 5090 blocks counterbalance solver order and contribute 40 warm
trials per solver:

RTX 5090 kernel	XLB	surrogate	surrogate / XLB
forward	4.794 ms	7.350 ms	1.53×
VJP	15.198 ms	14.771 ms	0.97×

The surrogate executes all 20 shared operator applications in forward and
reverse. It is not a forward speedup; VJP cost is effectively at parity
(0.97× in the counterbalanced result and 1.03× in the earlier independent
measurement). Solver-scoped three-seed harness times are also close: 32.81 s
surrogate versus 32.09 s XLB for L-BFGS, and 31.51 s versus 32.50 s with
projection. Including result-script setup and serialization gives 34.73 s
versus 34.21 s and 33.38 s versus 35.78 s, respectively. These wall times
include RPC, callbacks, optimizer bookkeeping, projection, and line search,
and should not be interpreted as kernel timings.

Autoregressive solver-in-loop study

The opt-in optimization/solver_in_loop_curriculum experiment targets a
regime where the numerical solver is the only route from a learned action to
the training objective. A zero-initialized corrector acts once every four native
steps inside a recurrent free rollout. After two stopped-gradient pushforward
intervals, the loss is evaluated only on the terminal state returned by the
last solver interval; no local corrected-state target is present.

The controls are:

• ONE, immediate one-step supervision from reference states;
• NOG, the identical recurrent terminal objective with every solver response
  detached; with zero local weight it is exactly the native-solver control;
• WIG, the identical recurrent rollout with full credit through the solver.

An analytical test proves that terminal WIG and NOG have equal forward
objectives while only WIG receives a non-zero parameter gradient.

quantity	value
candidate grid / native step	32² / dt=0.02
correction cadence	every 4 native solver steps
production reference	128² pseudo-spectral
convergence audit	gated 256² pseudo-spectral
IC spectrum	k₀=4, σₖ=0.75, amplitude 0.5
train / held-out ICs	32 / 8
warm-up	2 stopped-gradient correction intervals
differentiated curriculum	2 → 4 → 8 → 16 intervals
maximum differentiated horizon	64 native solver steps
free-running evaluation	120 correction intervals / 480 native steps

The result bundle retains per-seed raw FD/AD slopes, horizon-conditioned FD
sweeps, training and native-step accounting, correlations, physics diagnostics,
spectra, full fields, and the reference/solver/ONE/NOG/WIG rollout GIF.

JAX-CFD admission result

Budget-matched 275-update calibrations separate the candidate regimes:

regime	WIG reduction	NOG reduction	WIG/NOG	best valid FD error
dense direct loss (1698666)	0.909×	0.636×	1.429×	0.213%
sparse downstream loss (1698667)	1.351×	1.282×	1.054×	—
terminal downstream loss (1698866)	1.519×	1.000×	1.519×	0.128%

The sparse regime shows the literature distinction clearly: recurrent forward
exposure supplies most of its gain (NOG is 30.95× better than ONE), while the
solver VJP adds 5.4%. The dense control has a larger relative VJP separation but
neither learned arm improves the native rollout. The canonical terminal regime
uses a task for which the solver VJP is genuinely required and gives the best
absolute correction at this calibration budget.

Terminal WIG is harmful at the early horizon-2 and horizon-4 checkpoints, then
crosses 1.063× at horizon 8 and reaches 1.488× at horizon 16. The corrected FD
check freezes the detached warm-up state and perturbs only the differentiated
suffix. Its best relative errors are 0.000%, 0.012%, 0.209%, 0.497%, and 0.128%
at horizons 1, 2, 4, 8, and 16; at horizon 16 the raw slopes are FD −0.428728
and AD −0.429829. The horizon-1 zero/zero result confirms that the omitted
production stage has no trainable terminal credit.

Six-solver admission

The exact canonical four-stage schedule passes its one-seed admission cell for
every solver:

solver	job	WIG / native	ONE / native	best suffix FD
JAX-CFD	1698887	1.494×	0.041×	1.170%
INS.jl	1698888	1.352×	0.027×	0.241%
PhiFlow	1698889	1.250×	non-finite rollout	0.036%
PICT	1698890	1.503×	1.400×	0.269%
Warp-NS	1698891	1.194×	1.148×	0.289%
XLB	1698892	1.372×	0.082×	0.308%

NOG is the exact native-solver control in every cell. All six references pass
the common accuracy gate, all six WIG correctors improve their native rollout,
and all six solver VJPs pass the suffix FD check. ONE is useful for PICT and
Warp-NS, harmful for JAX-CFD, INS.jl, and XLB, and non-finite for PhiFlow. This
is why one-step supervision is not the primary comparison. PICT takes 48:47 for
275 updates versus 27–29 minutes for the other cells, so the 4,000-update draft
default does not fit a fair four-hour allocation. The canonical production
budget is therefore 1,000 updates per seed (100→180→270→450 by horizon).
The three-seed, six-solver offline matrix is running as jobs 1699012–1699029;
production artifacts are not yet present.

All experiments use the same Kander Slurm and Pyxis tooling and solver images as
the parent study. No hosted benchmark run is used.

This design follows the distinction between inference-distribution exposure and
long-term gradients in
Differentiability in Unrolled Training of Neural Physics Simulators
and the downstream solver-response objective used by
Solver-in-the-Loop.

Offline provenance and validation

No hosted benchmark run is required for this draft; benchmark:none remains
attached. All published experiments ran offline through Slurm and
Pyxis/Enroot.

• native trajectories: 1693649; curriculum: 1693872; fine-tune: 1694476
• forward/JVP evaluation: 1696191; 4k control: 1693974
• restricted conditioning: 1696373; adapted block-grid conditioning:
  1696374
• cross-model recovery: 1696375; inverse-path diagnostic: 1696590
• finite differences: XLB 1697019; surrogate 1697017
• packaged timing: 1696546
• 4k L-BFGS/projected recovery: 1691080 / 1691085
• 16k L-BFGS/projected recovery: 1696547 / 1696548
• runtime package build and Pyxis round-trip: 1696556
• rebased full validation: 1697321; rebuilt package round-trip: 1697322
• matched timing blocks: 1697367 and 1697363; held-out/JVP follow-up:
  1697360
• artifact-manifest verification: 1697378; immutable-source focused gate:
  1697379

The runtime image passed its API and Pyxis/Enroot round-trip checks. The
published artifact manifest verifies the dataset, split, checkpoint, source,
figures, JSON results, and dense-Jacobian checksums.

Validation:

• prior 3D surrogate head full suite: 522 passed, 3 skipped
• final immutable-source surrogate gate: 5 passed
• combined-head integration tests: 40 passed
• ruff check, ruff format --check, and git diff --check pass
• combined-head GitHub CI passes on Python 3.12 and 3.14
• focused solver-loop suite: 37 passed

Raw results and checksums:
artifact index,
manifest,
evaluation,
current-XLB evaluation follow-up,
order-balanced XLB timing,
order-balanced surrogate timing,
inverse-path diagnostic,
XLB finite differences,
surrogate finite differences,
restricted conditioning,
4k restricted conditioning,
adapted block-grid conditioning,
4k adapted block-grid conditioning,
dense adapted Jacobians,
surrogate recovery,
4k surrogate recovery,
and XLB-target recovery.

[PasteurBot]

📊 View the full benchmark results

No benchmarks ran for this PR, so there is no status report.

---