SAI
← All ICML 2026 orals

CoEvol-NO: State and Coordinate Co-Evolution with an Error-Driven Predictor-Corrector Paradigm for Neural Operator Transformer

Jianqiao Zeng, Ruocheng Wang, Yanzhi Liu, Hao Xiong, Junchi Yan

OralReplication score 2%Paper PDFCode repoOpenReview

CoEvol-NO: State and Coordinate Co-Evolution with an Error-Driven Predictor-Corrector Paradigm for Neural Operator Transformer

SAI replication review · Referee report

Summary

CoEvol-NO proposes a neural-operator Transformer that maintains two persistent representations at every layer — a compact latent state SRM×C\mathbf{S}\in\mathbb{R}^{M\times C} and a coordinate/mesh sequence XRN×C\mathbf{X}\in\mathbb{R}^{N\times C} — and updates them jointly with a Predictor–Corrector (PC) block. The conceptual move is to sit between the two dominant linear-complexity paradigms in operator learning: Coordinate-Evolution methods (Transolver-style), which retain fine geometry but reset their features every layer and thus lack dynamic memory, and State-Evolution methods (LNO-style), which evolve latent tokens in isolation and lose geometric coupling. A cross-attention predictor generates a tentative target S^t\hat{\mathbf{S}}_t, and the corrector refines St1\mathbf{S}_{t-1} via a single-step exact gradient descent on a squared correction loss (with a first-order approximation for X\mathbf{X}); the authors argue that direct-substitution and residual updates are recovered as first-order approximations under squared vs. dot-product loss assumptions. The closed-form gradient in Section 4.5, avoiding any N×NN\times N attention matrix, delivers a genuinely linear-in-NN per-layer cost, and Appendix D contributes a modest stability/conditioning analysis of the exact-gradient map versus the first-order one. The framework is evaluated on five standard PDE benchmarks, two industrial tasks (Shape-Net Car, AirfRANS), OOD/super-resolution generalization, and a cross-architecture transfer study. The mechanism is intuitive and the derivation is clean; execution confirms that the analytical gradient matches autograd to machine precision, substantiating the linear-time argument. However, the code executed here could not reproduce the headline accuracy numbers under the available compute budget (roughly 10% of the published training, with Navier-Stokes additionally downsampled), so the central SOTA magnitudes remain unverified rather than confirmed. The main conceptual limitations are that the corrector's monotone-refinement property is asserted informally rather than derived, the loss-form pairing (squared on S\mathbf{S}, dot-product on X\mathbf{X}) is not motivated from physics, and several quantitative claims are undermined by numerical inconsistencies between prose and tables.

Strengths

  • Conceptual contribution. The Coords-Evol / State-Evol / Co-Evol taxonomy is a genuinely useful lens for organizing linear-complexity operator Transformers, and the co-evolutionary framing is a natural extrapolation that resolves the geometry-vs-memory trade-off cleanly on paper.
  • Closed-form linear-time gradient, verified. The Appendix A.1 derivation yielding the explicit form Γ=(SPV)V\mathbf{\Gamma}=(\mathbf{S}-\mathbf{PV})\mathbf{V}^\top and the row-sum structure keeps the exact PC update strictly O(NMC)O(NMC) per layer. Execution confirmed the analytical implementation is numerically equivalent to autograd (17/17 tests, max difference <3×107<3\times10^{-7}), and the component benchmark corroborates roughly linear scaling — the strongest verified result in the package.
  • Unification narrative. Showing that direct substitution and residual updates are first-order approximations of a common error-driven descent under two loss forms is a nice pedagogical framing, and the exact-vs-approximation ablation gives it empirical bite (the ordering reproduced directionally under re-execution).
  • Breadth of evaluation. Five standard benchmarks, two industrial tasks, OOD splits, super-resolution, cross-architecture transfer, depth/latent scaling, and PCA visualizations together make an unusually thorough empirical package on paper.

Weaknesses

  • Numerical inconsistencies between prose and tables. The Elasticity headline (0.0039 vs Transolver 0.0064) does not match Table 2 (0.0036 / 0.0057); the "exact vs approximation" pair 0.0038 vs 0.0046 is Elasticity but is attributed to Airfoil too, where Table 5 shows 0.0047 vs 0.0057; the SS-momentum "32%-32\%" on Pipe is ~27% by the tabulated values; the Navier-Stokes std quoted as 0.0022 contradicts Table 16's 0.0082. Individually minor, together they erode confidence that the quantitative narrative is checked against its own tables.
  • Monotone refinement is asserted, not derived. The (Informal) condition d(St,S)d(St1,S)d(\mathbf{S}_t,\mathbf{S}^*)\le d(\mathbf{S}_{t-1},\mathbf{S}^*) describes descent toward the predicted target S^t\hat{\mathbf{S}}_t, not toward S\mathbf{S}^*; the assumption under which descending on a surrogate reduces the true distance is never stated. Framing it as a condition each layer "satisfies" overclaims what a single gradient step guarantees.
  • Fixed-point motivation glosses the X\mathbf{X} dependence. The Direct Mapping choice is justified via Sf(S)\mathbf{S}^*\approx f(\mathbf{S}^*), but ff is a cross-attention reading from X\mathbf{X}, so f(S)f(\mathbf{S}^*) is only well-defined for fixed mesh state. Any fixed point is joint in (S,X)(\mathbf{S},\mathbf{X}).
  • Loss pairing looks post-hoc. Running squared loss on S\mathbf{S} and dot-product loss on X\mathbf{X} is exactly what lands on the desired architectural blocks; no principled criterion for the pairing is given beyond "residual has better gradient flow", which sits awkwardly next to the deliberate non-residual exact-gradient branch on S\mathbf{S}.
  • Stability argument is conditional but presented as unconditional. Appendix D shows MEG\mathbf{M}_{\text{EG}} symmetrizes error dynamics, but stability still requires ησmax2<2\eta\sigma_{\max}^2<2; the text reads as if symmetrization removes instability by itself.
  • Empirical over-claims. Table 4 shows the Sequence variant beating Dual on AirfRANS, yet the discussion says Dual "matches" it; Table 5 shows the Approximation variant winning several AirfRANS/ShapeNet cells, yet Section 5.1.2 opens with an unqualified claim of "necessity and superiority"; the OOD AoA "superior ranking" is a 0.9950 vs 0.9949 tie.
  • Interpretive overreach on visualization. Claims that PCA alignment "confirms" an invariant operator and that trajectory persistence shows "insensitivity to specific implementation details" are stronger than a 2D projection of a single batch supports, and the latter conflicts with the quantitative sensitivity in Table 11 and Appendix D.
  • Terminology and notation. Ablation variants are named State/Coord, then Latent/Sequence, then STATES/COORDS; Δ\mathbf{\Delta} denotes two different objects (different sign and dimension) between Section 4.5 and Appendix A.1; V\mathbf{V}/Vt\mathbf{V}_t collide (Value vs velocity buffer) on one page; the complexity enumeration omits forming P\mathbf{P} and PVX\mathbf{PV}_X; and the empty () and the mis-parenthesized OOD Reynolds range get in the way of a careful read.

Reproducibility & code

Veritas executed the shipped code (tasks/pde_benchmarks/run.py with the repo's own configs/pde/*.yaml, plus the repo's analytical-equivalence and scaling unit tests) without reimplementing anything. The code is faithful and ran nearly clean; the observations below distinguish what was verified, what could not be reproduced under compute limits, and what was never exercised.

  • Headline accuracy could not be reproduced (environmental). All five standard-benchmark cells came back 3x–85x higher than Table 2 (Elasticity 0.03345 vs ~0.0037, Navier-Stokes ~6.2 vs 0.0731, Darcy 0.01552 vs ~0.0046). The cause is a documented compute cap — epochs cut ~10x (500→50, and 500→10 for NS) and NS downsampled (64×64→32×32, ntrain 1000→128) to fit a shared GPU — not a code defect; loss curves were still descending at cutoff. No baselines (Transolver, LNO) were trained, so the SOTA rankings could not be established either. The central magnitudes are therefore unverified, not refuted.
  • Gradient correctness independently verified. The analytical exact gradient for S\mathbf{S}/X\mathbf{X}/FFN matches autograd to machine precision (17/17 tests, max diff <3×107<3\times10^{-7}), and the component benchmark shows the analytical path is faster and lower-memory with roughly O(N)O(N) scaling. This confirms the Appendix A.1 derivation. The specific Table 19 per-epoch timings and the 831 ms vs 490 ms Dual-vs-Transolver figures were not reproduced — the shipped benchmark measures analytical-vs-autograd internals only, has no Transolver baseline, and stops at 10610^6 (not 4×1064\times10^6) points.
  • Ablation directions partially held. Exact-gradient beat first-order on both benchmarks tested (Elasticity, NS), and State-Evol was worse than Dual on Elasticity — both consistent with the paper's claims in direction, though magnitudes are far off. But the Navier-Stokes ordering the paper leans on (Sequence worse than Dual) did not reproduce: at reduced scale all variants sit at ~6.18–6.20, so the persistent-state-needed-for-dynamics evidence was not corroborated here.
  • Industrial datasets not provisioned. Neither Shape-Net Car nor AirfRANS was run; the task code exists but was never invoked, no Volume/Surface/coefficient/Spearman/OOD metrics were produced, and the release documents no download procedure. Tables 3 and 7 could not be checked.
  • Protocol divergences. The shipped configs hardcode batch_size=8 for every benchmark, contradicting Table 17's per-task values (Elasticity=1, Airfoil=4, Pipe=8, NS=2, Darcy=4); the prose "AdamW + OneCycleLR" contradicts Table 17 (Cosine for Elasticity; Adam for Shape-Net Car and AirfRANS); and the package ships with no setup.py/pyproject.toml, so import coevol_no fails without a manual PYTHONPATH.
  • Large parts of the ablation suite unexercised. The super-resolution sweep (Table 6), OOD splits (Table 7), multi-seed variance (Table 16), momentum/LayerScale/depth/latent ablations (Tables 9–12, 14–15), and cross-architecture PC transfer (Table 8, PC-Transolver and PC-LNO) were not run. Some code paths and config knobs are present, but driver scripts, cross-resolution/OOD harnesses, and seed-sweep aggregation are missing, so these numbers are unverified.

Recommended Changes

Essential

  • Reconcile the prose numbers with the tables. Fix the Elasticity headline (0.0039 vs 0.0064) against Table 2 (0.0036 / 0.0057), correct the Airfoil "0.0038 vs 0.0046" attribution against Table 5, recompute the Pipe momentum percentage from a single authoritative value (0.0032 vs 0.0030), and change the Navier-Stokes std from 0.0022 to the tabulated 0.0082.
  • Enable independent verification of the headline table. Ship a data-availability section (filenames, expected paths, download URLs) for all eight benchmarks, add baseline runner configs (plain Transolver, LNO) so the SOTA rankings can be regenerated, and make the shipped configs match Table 17 (per-task batch_size, and the correct optimizer/scheduler per benchmark) so a reproducer at full compute lands on the published protocol. Add setup.py/pyproject.toml so the package imports without a manual PYTHONPATH.
  • Provision and document the industrial pipelines. Make Shape-Net Car and AirfRANS runnable end-to-end, including the Volume/Surface relative L2L_2, drag/lift coefficient, and Spearman metrics of Tables 3 and 7 and the OOD Reynolds/AoA splits, or clearly cite the source of any number not regenerated.
  • State the assumptions behind the refinement condition. Make explicit that descent is on the predicted target S^t\hat{\mathbf{S}}_t, give the proxy-quality/contraction assumption needed for d(St,S)d(St1,S)d(\mathbf{S}_t,\mathbf{S}^*)\le d(\mathbf{S}_{t-1},\mathbf{S}^*), and restate the fixed-point condition in the joint (S,X)(\mathbf{S},\mathbf{X}) variables.

Suggested

  • Motivate the loss pairing. Justify (or admit as empirical) squared loss on S\mathbf{S} and dot-product loss on X\mathbf{X} from information-bottleneck or gradient-flow considerations, so the unification story does not read as post-hoc.
  • Qualify the stability argument. State the required step-size regime ησmax2<2\eta\sigma_{\max}^2<2 alongside the symmetrization claim, and reword "higher-order dependencies" to "the exact Jacobian of the predictor."
  • Recalibrate the empirical claims. Acknowledge that Sequence beats Dual on AirfRANS in Table 4, open Section 5.1.2 with the qualified exact-gradient statement since Approximation wins some AirfRANS/ShapeNet cells, and note the OOD AoA ranking is a numerical tie; soften the conclusion's unqualified "training stability" to match the Limitations.
  • Soften the visualization claims. Replace "confirms" with "is consistent with", drop the "insensitivity to implementation details" statement, and support the invariance claim with a quantitative cross-sample similarity metric.
  • Fix notation and presentation. Use distinct symbols for the raw residual SC(S)\mathbf{S}-C(\mathbf{S}) and the projected error (both currently Δ\mathbf{\Delta}), rename the momentum buffer away from V\mathbf{V}, complete the complexity enumeration to include P\mathbf{P} and PVX\mathbf{PV}_X, fill the empty (), add the outer parentheses to the Reynolds range ([2,3][5,6])×106([2,3]\cup[5,6])\times10^6, and settle on one ablation-variant naming.
  • Add the missing reproducibility scripts. Provide a cross-resolution super-resolution eval (Table 6), an OOD-split loader for AirfRANS (Table 7), a multi-seed sweep with mean±\pmstd aggregation (Table 16), configs/wrappers for the momentum/LayerScale/depth/latent and PC-Transolver/PC-LNO ablations, and a PCA-figure generator that hooks per-layer St\mathbf{S}_t from a trained checkpoint.