Nonlinear System Solve - Pushforward/Jvp rule

The solution of nonlinear systems of equations is crucial in scientific computing, like the integration of nonlinear PDEs (e.g., the Navier-Stokes equations). How can we employ the implicit function theorem to propagate tangent information over a Newton-Raphson solver. Here are the notes: raw.githubusercontent.com/Cey...
Next to the numerical solution of differential equations, you also find nonlinear solvers for a bunch of other applications like nonlinear optimization problems or, more recently, deep equilibrium problems. Being able to propagate tangent information is a crucial component in order to get a differentiable chain of computations to then use forward-mode Automatic Differentiation to obtain derivative estimates.
For implicitly given problems, like the roots of nonlinear systems of equations, there are two approaches: One can either differentiate through the solver (also called unrolling or piggybacking) or differentiate over the solver by employing the implicit function theorem and thereby deriving a custom forward rule. Often times, the latter is also the only way of doing so, as these numerical kernels tend to be implemented either in a non-differentiable fashion or rely on non-differentiable third-party libraries.
Timestamps:
00:00 Nonlinear System Solving as a function
01:00 Applications
01:15 Solution by e.g. Newton Raphson
01:53 Dimensionalities involved
02:37 Task: Forward Propagation of tangent information
03:22 Without unrolling by the forward-mode AD engine
04:41 General Pushforward/Jvp rule
06:36 Total derivative of optimality criterion/zero condition
08:15 Identifying the (full and dense) Jacobian
09:15 Plug Jacobian back into general pushforward/Jvp expression
10:20 Requires solution to a LINEAR system of equations
10:56 Full Pushforward rule
11:49 How about the additional derivatives?
12:15 Finding right-hand side with a Jacobian-vector product
12:36 Solve linear system matrix-free Jacobian-vector product
14:03 Summary
16:00 Outro
github.com/Ceyron/machine-lea...
-------
📝 : Check out the GitHub Repository of the channel, where I upload all the handwritten notes and source-code files (contributions are very welcome): github.com/Ceyron/machine-lea...
📢 : Follow me on LinkedIn or Twitter for updates on the channel and other cool Machine Learning & Simulation stuff: / felix-koehler and / felix_m_koehler
💸 : If you want to support my work on the channel, you can become a Patreon here: / mlsim
🪙: Or you can make a one-time donation via PayPal: www.paypal.com/paypalme/Felix...
-------
⚙️ My Gear:
(Below are affiliate links to Amazon. If you decide to purchase the product or something else on Amazon through this link, I earn a small commission.)
- 🎙️ Microphone: Blue Yeti: amzn.to/3NU7OAs
- ⌨️ Logitech TKL Mechanical Keyboard: amzn.to/3JhEtwp
- 🎨 Gaomon Drawing Tablet (similar to a WACOM Tablet, but cheaper, works flawlessly under Linux): amzn.to/37katmf
- 🔌 Laptop Charger: amzn.to/3ja0imP
- 💻 My Laptop (generally I like the Dell XPS series): amzn.to/38xrABL
- 📱 My Phone: Fairphone 4 (I love the sustainability and repairability aspect of it): amzn.to/3Jr4ZmV
If I had to purchase these items again, I would probably change the following:
- 🎙️ Rode NT: amzn.to/3NUIGtw
- 💻 Framework Laptop (I do not get a commission here, but I love the vision of Framework. It will definitely be my next Ultrabook): frame.work
As an Amazon Associate I earn from qualifying purchases.
-------

[@MachineLearningSimulation]

Some recommended pointers:
Steven Johnson's notes on adjoint methods (also briefly covers forward propagation: math.mit.edu/~stevenj/18.336/adjoint.pdf
The paper on "Efficient and Modular Implicit Differentiation": arxiv.org/abs/2105.15183
This Juliacon talk for a Julia package that builds on the aforementioned paper: ua-cam.com/video/TkVDcujVNJ4/v-deo.html