Adjoint path-kernel method for backpropagation and data assimilation in unstable diffusions

TL;DR

This paper introduces the adjoint path-kernel method for efficient gradient computation in stochastic differential equations, enabling advanced backpropagation and data assimilation in complex, chaotic systems with high-dimensional parameters.

Contribution

The paper extends backpropagation to non-hyperbolic systems with multiplicative noise and develops a new data assimilation approach for chaotic models with partial observations.

Findings

Effective gradient computation for SDEs with parameter-controlled noise

Successful application to 40D Lorenz 96 system demonstrating backpropagation in unstable systems

Demonstration of stochastic gradient descent for data assimilation in 10D Lorenz 96 with partial observations

Abstract

We derive the adjoint path-kernel method for computing parameter-gradients (linear responses) of SDEs. Its cost is almost independent of the number of parameters, and it works for non-hyperbolic systems with parameter-controlled multiplicative noise. With this new formula, we extend the conventional backpropagation method to settings with gradient explosion, and demonstrate it on the 40-dimensional Lorenz 96 system. Moreover, we consider a difficult version of the 4D-Var data assimilation problem where (1) the deterministic part of the model is chaotic, (2) the loss is a single long-time functional accounting for discrepancies in both the observations and the dynamics, (3) some parameters in the dynamics are unknown, and (4) some coordinates of the states cannot be observed, and cannot be reasonably inferred from other coordinates within a short time. We model the correction term at each time-step separately as a parameterized function of the random state. With our new tool, we can run stochastic gradient descent to find the path and parameters that best match the low-dimensional observation data. We demonstrate this on the 10D Lorenz-96 system with 8D observations.