Adjoint-Matching Neural Network Surrogates

• Adjoint-matching neural network surrogates are models trained to approximate both the forward outputs and the associated sensitivity (adjoint) information of complex systems.
• They integrate forward simulation with adjoint consistency by using loss functions that jointly minimize prediction and gradient discrepancies, ensuring reliable gradient estimates.
• These surrogates accelerate optimization tasks in data assimilation, neural PDE calibration, and optimal control, achieving significant speedups with minimal accuracy tradeoffs.

Adjoint-matching neural network surrogates are a class of surrogate modeling techniques in which neural networks are trained not only to replicate the outputs of a physical or computational model, but also to accurately reproduce the corresponding adjoint or sensitivity information. These methods have gained prominence in scientific computing, optimal control, inverse problems, data assimilation, and scientific machine learning due to their ability to dramatically accelerate optimization or inference workflows that rely on repeated gradient (adjoint) information. By ensuring consistency between the surrogate’s Jacobian (or its action on key adjoint vectors) and the reference model’s, adjoint-matching surrogates provide reliable, efficient gradient estimates for use in optimization, data assimilation, sampling, and uncertainty quantification.

1. Theoretical Foundations

Adjoint-matching surrogates arise in contexts where the underlying model is a mapping $M:\mathbb{R}^n \to \mathbb{R}^m$ whose derivatives, typically with respect to states or parameters, are central to downstream tasks. Classical adjoint methods enable efficient computation of these gradients, but repeated adjoint solves are often prohibitively expensive. Neural network surrogates can accelerate this process if trained to faithfully approximate both the forward map and its adjoint action.

Formally, given a forward model $M(x)$, its Jacobian $DM(x)$, and a loss $J(x)$, adjoint consistency requires that the surrogate $F_\theta(x)$ satisfies not only

$F_\theta(x) \approx M(x),$

but also

$$\left[\partial_x F_\theta(x)\right]^\top v \approx DM(x)^\top v,$$

for relevant vectors $v$ (entire adjoint matrix or chosen adjoint-vector pairs), where $\partial_x$ denotes differentiation with respect to $x$ (Chennault et al., 2021).

In the context of neural PDEs or optimal control, adjoint-matching encompasses ensuring that neural network parameterizations embedded in PDEs yield correct gradients with respect to parameters or controls, typically via an adjoint PDE. For PDE-constrained optimization, this connection is formalized in the Karush–Kuhn–Tucker (KKT) system, where optimality involves state, adjoint, and control equations (see (Yin et al., 2023, Yuan et al., 21 Dec 2025)).

2. Core Methodologies

Adjoint-matching strategies are instantiated via a variety of techniques, depending on the application and context:

• Adjoint-Augmented Loss Functions: Loss functions incorporate penalties for mismatch not only in function approximation, but also between the surrogate’s Jacobian (or its application to relevant adjoint vectors) and the reference model’s adjoint action. Typical forms include:
  • Full Jacobian loss: $$L_a(\theta) = \|\left[\partial_x F_\theta(x)\right]^\top - DM(x)^\top\|_F^2$$,
  • Adjoint-vector product loss: $$L_{av}(\theta) = \|\left[\partial_x F_\theta(x)\right]^\top v - DM(x)^\top v\|^2$$,
  • where the total loss is a weighted sum $L_{\text{total}} = L_f + \lambda L_a$ (Chennault et al., 2021).
• Direct Adjoint Integration: In parametric optimal control or neural PDEs, neural networks are used to parametrize state, adjoint, and control fields, with training driven by residuals of the KKT conditions or direct solution of forward and adjoint PDEs (see (Yin et al., 2023, Yuan et al., 21 Dec 2025, Riedl et al., 16 Jun 2025)).
• Data Requirements: Full-adjoint matching requires access to adjoint information (e.g., full Jacobian or adjoint matrices from the reference model), which can be costly to generate in high dimensions. Adjoint-vector losses alleviate memory and data requirements at the expense of some performance (Chennault et al., 2021).
• Specialized Architectures and Algorithms: Domain-specific architectures (e.g., equivariant GNNs in molecular sampling (Havens et al., 16 Apr 2025), Fourier neural operators (Sun et al., 22 Aug 2025)) and iterative modular training procedures (e.g., direct–adjoint looping) are employed to ensure effective coupling of forward and adjoint constraints.

3. Representative Applications

Adjoint-matching neural network surrogates have been demonstrated across multiple scientific domains:

Application Area	Methodological Context	Notable Features
Variational Data Assimilation	4D-Var with neural surrogates	Adjoint-matching loss and evaluation of 4D-Var accuracy (Chennault et al., 2021)
Neural PDE Calibration	Neural network embedding in PDE source	Adjoint PDE for gradient computation, convergence to global minimizer (Riedl et al., 16 Jun 2025)
Optimal Control (Parametric/PDE)	All-at-once adjoint-oriented neural networks	KKT system residuals, direct–adjoint looping, adaptive sampling (Yin et al., 2023, Yuan et al., 21 Dec 2025)
Sampling (Diffusion/SOC)	Adjoint Matching for drift parameterization	Stochastic optimal control reformulated as adjoint regression (Havens et al., 16 Apr 2025)
A Posteriori Error Estimation	DWR goal-oriented estimator	NN surrogate solves adjoint problem for error indicators (Roth et al., 2021)
Surrogate Sensitivity	Parametric derivatives for ocean models	Ensemble FNOs for uncertainty and sensitivity (Sun et al., 22 Aug 2025)

In data assimilation, adjoint-matching surrogates trained with full-adjoint or adjoint-vector losses recovered almost all of the accuracy of the reference solver at 10–30× speedup, maintaining descent directions in the 4D-Var optimization (Chennault et al., 2021). For neural PDE calibration, embedding the adjoint via a backward PDE yields an efficient and provably convergent gradient computation framework, with global convergence established even in the infinite-width, infinite-time regime (Riedl et al., 16 Jun 2025).

In parametric PDE-constrained optimal control, all-at-once adjoint-oriented neural networks enforce consistency across state, adjoint, and control nets, delivering accurate parametric solution maps without repeated PDE solves and enabling sub-millisecond inference (Yin et al., 2023). Adaptive variants employing deep adaptive sampling further enable robust modeling in the presence of geometric and solution singularities (Yuan et al., 21 Dec 2025). For high-dimensional error estimation and mesh adaptivity in the DWR estimator, feedforward surrogate networks for the adjoint yield effectivity indices close to 1.0 (Roth et al., 2021).

4. Theoretical Guarantees and Convergence

Several adjoint-matching frameworks provide rigorous convergence or optimality results under suitable regularity and scaling regimes.

• Global Convergence in Neural PDEs: For parabolic PDEs with neural networks embedded in the source term, and under “lazy” parameter scaling ($\beta \in (1/2,1)$ in the network), the infinite-width training dynamics are governed by a Hilbert–Schmidt, nonlocal, non–spectrally gapped kernel operator. Despite nonconvexity, a stopping-time and energy argument establishes weak $L^2$ convergence of the surrogate solution to data, i.e., global minimization, due to positivity of the neural tangent kernel operator (Riedl et al., 16 Jun 2025).
• Optimal Control Fixed-Point Results: In Adjoint Sampling and Adjoint Matching for diffusion SDE control, the reciprocal projection operator ensures monotonic decrease of the stochastic optimal control loss, and the training iterates converge to the unique Schrödinger-bridge drift (Havens et al., 16 Apr 2025).

A key observation across applications is that adjoint-matched surrogates yield reliable descent directions for gradient-based optimization, avoiding pathological stationary points that can arise with naive forward-only fitting (Chennault et al., 2021).

5. Implementation Considerations

Practical deployment of adjoint-matching surrogates requires careful attention to architecture, regularization, data acquisition, and training protocols:

• Activation Functions and Scaling: For neural PDE surrogates, activation functions should be smooth and bounded with Lipschitz derivatives; scaling must ensure that the neural tangent kernel remains in the NTK regime (Riedl et al., 16 Jun 2025).
• Loss Weighting and Data Augmentation: The balance between forward and adjoint terms in the loss must be tuned (e.g., by cross-validation on final task error) (Chennault et al., 2021). Adjoint-vector directions can include random, coordinate, or task-informed vectors (e.g., “Lagrange” directions in data assimilation).
• Regularization and Stopping Criteria: Diminishing step sizes, Tikhonov regularization or penalization, and line search methods improve convergence rates and stability (Riedl et al., 16 Jun 2025, Odot et al., 2023).
• Adaptive Collocation and Sampling: Deep adaptive sampling, based on surrogate residuals, is critical for robust training in high-dimensional or low-regularity settings (Yuan et al., 21 Dec 2025). Generative flow models (e.g., KRnet) can anneal sampling distributions to focus learning on difficult regions.
• Ensemble Methods for UQ: Ensembles of hyperparameter-optimized surrogates provide epistemic uncertainty on both forward and adjoint predictions, calibrated via linearization tests and other surrogate-adjoint error metrics (Sun et al., 22 Aug 2025).

6. Performance Characteristics and Comparative Results

Adjoint-matching surrogates consistently deliver dramatic computational speedups relative to classical adjoint-based approaches, with only minor accuracy tradeoffs:

Context	Speedup	Accuracy Degradation	Remarks
4D-Var Data Assimilation (Lorenz-63)	10–30×	RMSE 0.84 (surrogate) vs. 0.83 (exact)	Adjoint-matched surrogates yield near-optimal analyses (Chennault et al., 2021)
Neural PDE Surrogate Calibration	1 forward + 1 adjoint PDE/step	Global minimization in NTK limit	Global convergence even under nonconvexity (Riedl et al., 16 Jun 2025)
NN-Based Adjoint for Shape Optimization	~2× net iteration speedup; adjoint cost nearly zero	Final drag within 0.2% of full adjoint	Generalizes across wider param ranges (Xu et al., 2020)
Optimal Control Surrogates (AONN)	100–160×	Errors $10^{-3}$–$10^{-2}$ compared to FEM	Direct–adjoint loop; rapid inference (Yin et al., 2023, Yuan et al., 21 Dec 2025)
DWR Error Estimation with NN Adjoint	Lower or comparable to enriched FEM	Effectivity index close to 1	Full adaptivity preserved (Roth et al., 2021)
Molecular Diffusion Sampling (Adjoint Matching)	Orders-of-magnitude fewer E-evals	Matches or outperforms baselines in ESS, $W_2$	Symmetries and periodicity supported (Havens et al., 16 Apr 2025)

In inverse design workflows, it is empirically established that surrogate accuracy provides a hard bound on achievable optimization results—the resimulation error of designs obtained via the neural adjoint (NA) method cannot fall below the surrogate’s forward validation MSE (Fujii et al., 2023).

7. Limitations and Future Directions

Adjoint-matching neural network surrogates are constrained by several factors:

• Adjoint Data Scalability: Full Jacobian or adjoint matrix data is expensive to generate, motivating adjoint-vector or physics-informed alternatives in high-dimensional settings (Chennault et al., 2021, Sun et al., 22 Aug 2025).
• Surrogate Error Floor: The best possible inverse or control performance is fundamentally limited by the surrogate’s validation error (Fujii et al., 2023). Surrogates trained without explicit adjoint matching may yield suboptimal gradients in optimization or control.
• Adjoint-Matching Surrogates for Large-Scale PDEs: Extension to high-dimensional nonlinear systems, turbulent flow, or Earth-system models remains a challenge, with adaptive and ensemble-based techniques being active areas of research (Sun et al., 22 Aug 2025, Yuan et al., 21 Dec 2025).
• Memory and Batch Size Constraints: Large, accurate surrogates may limit candidate pool sizes during inverse optimization, motivating candidate-focusing procedures such as the NeuLag approach (Fujii et al., 2023).

Ongoing research is actively exploring the development of scalable adjoint-matching surrogates, improved training algorithms for data- and computation-limited contexts, and tighter coupling with physical and geometric constraints in scientific machine learning pipelines.