Preconditioning transformations of adjoint systems for evolution equations

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.