Preconditioning transformations of adjoint systems for evolution equations

TL;DR

This paper develops a framework for preconditioning adjoint systems in evolution equations by drawing an analogy with preconditioned gradient descent. It identifies two main classes of transformations—one induced by transformations of the state dynamics and another by duality pairings (including state-dependent fiberwise pairings)—and shows these are symplectomorphisms that preserve the adjoint backpropagation property. The authors extend the approach to coupled evolution equations, introducing a scale preconditioner to handle large scale separation, and demonstrate its effectiveness on an inverse problem for radiation diffusion (Marshak-wave) where convergence improves to about 10–15 optimization iterations. The work provides a structured, geometry-informed path to robust, multiscale adjoint-based optimization with potential extensions to state-dependent and spatially multiscale problems.

Abstract

Achieving robust control and optimization in high-fidelity physics simulations is extremely challenging, especially for evolutionary systems whose solutions span vast scales across space, time, and physical variables. In conjunction with gradient-based methods, adjoint systems are widely used in the optimization of systems subject to differential equation constraints. In optimization, gradient-based methods are often transformed using suitable preconditioners to accelerate the convergence of the optimization algorithm. Inspired by preconditioned gradient descent methods, we introduce a framework for the preconditioning of adjoint systems associated to evolution equations, which allows one to reshape the dynamics of the adjoint system. We develop two classes of adjoint preconditioning transformations: those that transform both the state dynamics and the adjoint equation and those that transform only the adjoint equation while leaving the state dynamics invariant. Both classes of transformations have the flexibility to include generally nonlinear state-dependent transformations. Using techniques from symplectic geometry and Hamiltonian mechanics, we further show that these preconditioned adjoint systems preserve the property that the adjoint system backpropagates the derivative of an objective function. We then apply this framework to the setting of coupled evolution equations, where we develop a notion of scale preconditioning of the adjoint equations when the state dynamics exhibit large scale-separation. We demonstrate the proposed scale preconditioning on an inverse problem for the radiation diffusion equations. Naive gradient descent is unstable for any practical gradient descent step size, whereas our proposed scale-preconditioned adjoint descent converges in 10-15 gradient-based optimization iterations, with highly accurate reproduction of the wavefront at the final time.

Figures (9)

• Figure 1: Schematic of the state dynamics for $\mathbf{u}(\cdot)$ evolving on the base space $U$ and backpropagation through the adjoint equation for $\mathbf{p}(\cdot)$ evolving in reverse time on the fibers of $T^*U$ over $\mathbf{u}(\cdot)$.
• Figure 2: Schematic of the mapping arising from the preconditioning of the adjoint system induced by a transformation of the state dynamics. The base space and the fibers are transformed under this map.
• Figure 3: Schematic of the mapping arising from the preconditioning of the adjoint system induced by a duality pairing. Only the fibers are transformed under this map, while the base space is left invariant.
• Figure 4: Snapshots of the evolution of $E$ (left) and $T$ (right) of the thick Marshak wave with a constant equilibrium initial state.
• Figure 5: The observed values $E^*$ compared to $E(t_f)$ produced from the constant equilibrium initial conditions (left) and $T^*$ compared to $T(t_f)$ produced from the constant equilibrium initial conditions (right).

Theorems & Definitions (22)

• Proposition 1.1
• proof
• Proposition 1.2: Proposition 2.2 of TrLe2024
• Proposition 1.3: Proposition 2.3 of TrLe2024
• Definition 2.1: Fréchet Derivative
• Definition 2.2: Duality Pairing
• Example 2.1: Evolution equations with mass matrices
• Proposition 2.1
• Proposition 2.2
• Definition 2.3: Fiberwise duality pairing