Extracting diagonal of $J^TJ$ via automatic differentiation like techniques

124 Views Asked by At

First of all please, let me know if this question is more suited for scicomp.stackexchange.com or or.stackexchange.com, and sorry for my English and math skills.

Pretty often i do some numerical optimization stuff with Gauss-Newton and automatic differentiation. It's pretty handy because Gauss-Newton requires to solve following system: $$ J^TJ\Delta x = -J^Tr $$ where $r(x)$ is residuals function, $J$ is jacobian of $r$, $\Delta x$ is step. This system can be solved via Conjugate gradient method which requires $Av$ like products. Automatic differentiation allow me to compute $Jv$ and $J^Tv$ via forward and reverse mode of graph based automatic differentiation, so i can easily combine it with CG.

Unfortunately CG also needs preconditioner for fast convergence. Obvious choice for preconditioner is diagonal of $J^TJ$.

My question is - how can i quickly compute diagonal of $J^TJ$ if i know expression graph for $r$ and have implementation of $J^Tv$ and $Jv$ for each elementary function? Without symbolic differentiation, if it possible.

1

There are 1 best solutions below

4
On

A recent paper:

Siskind, Jeffrey Mark. "Automatic Differentiation: Inverse Accumulation Mode." (2019).

should solve your problem. You can efficiently apply the inverse Jacobian and its transpose to the right side.