Automatic Differentiation - reverse accumulation of linear system solve

Question

I am studying the reverse mode of automatic differentiation.

The reverse mode of automatic differentiation allows the efficient computation of a the derivative of a single dependent variable $y$ with respect to as many independent variables $x_i$ as you want. One assigns to each intermediate variable $v$ an adjoint variable $\bar{v}$ which is the derivative of a chosen dependent variable with respect to the subexpression $$\bar{v} \rightarrow \bar{v} = \frac{\partial y}{\partial v}$$

so the assignment $$y = v_1 \sin(v_2)$$ corresponds to the adjoint assignments

\begin{eqnarray} \bar{v_1} &=& \bar{y} \sin(v_2) \\ \bar{v_2} &=& \bar{y} v_1 \cos(v_2) \end{eqnarray}

where $\bar{y} = 1$ .

I am interested in the situation where you solve a linear system in the program:

$$Ax = b$$

where $y$ might be another function of $x$: $$y(x)$$

According to this tutorial of the CoDiPack software, the corresponding adjoint statements are

\begin{eqnarray} \bar{A} &=& - \lambda x^T \\ \bar{b} &=& \lambda \\ \end{eqnarray} where $\lambda$ is the the solution of the adjoint equation $$A^T \lambda = \bar{x}$$

I found the same algorithm in several other documents, for example here in section 7, Iteration and equation solving.

It is not clear to me how to arrive at these statements. I think the derivation must be similar as in the case where one wants to optimize $w(x)$ where $x$ is subject to the constraint $$Ax = b$$ See for example this document.

Answer 1

Starting with $b=Ax$ the forward differentiation gives $\dot b=\dot A x+A\dot x$, or, with $\dot y = \dot A x$:

$$\dot x = -A^{-1} \dot y + A^{-1} \dot b$$

Let's write this in terms of the local transformation Jacobian for the tangent linear variables $\left( \dot A, \dot y, \dot b, \dot x\right)$:

$$ \left[ \begin{matrix} \dot A \\ \dot y\\ \dot b\\ \dot x\\ \end{matrix} \right] = \left[ \begin{matrix} 1 & 0 & 0 & 0 \\ x & 0 & 0 & 0 \\ 0 & 0 & 1 & 0 \\ 0 & -A^{-1} & A^{-1} & 0 \\ \end{matrix} \right] \left[ \begin{matrix} \dot A \\ \dot y\\ \dot b\\ \dot x\\ \end{matrix} \right] $$

The adjoint operator is then the transpose matrix acting on the (local) adjoint variables $\left( \bar A, \bar y, \bar b, \bar x\right)$:

$$ \left[ \begin{matrix} \bar A \\ \bar y\\ \bar b\\ \bar x\\ \end{matrix} \right] = \left[ \begin{matrix} 1 & x^T & 0 & 0 \\ 0 & 0 & 0 & -A^{-T} \\ 0 & 0 & 1 & A^{-T} \\ 0 & 0 & 0 & 0 \\ \end{matrix} \right] \left[ \begin{matrix} \bar A \\ \bar y\\ \bar b\\ \bar x\\ \end{matrix} \right], $$

or:

$$ \begin{align} \bar y &= -A^{-T} \bar x \\ \bar A &\mathrel{+}= x^T \bar y \\ \bar b &\mathrel{+}= A^{-T} \bar x \\ \bar x &= 0. \end{align} $$

Reference:

Giering and Kaminski, Recipes for Adjoint Code Construction, ACM TOMS, 1998

Answer 2

The identity connecting forward and reverse mode for $x=\text{solve}(A,b)$ is $$\bar x^⊤\dot x=\bar A(\dot A)+\bar b^⊤\dot b.$$ Where $\bar A$ is a linear functional on matrices which may be realized as the $trace(\bar A^⊤\dot A)$ which is related to the (Frobenius) scalar product for matrices $\langle X,Y\rangle=trace(X^⊤Y)=trace(XY^⊤)$, which is kind-of the euclidean scalar product as it is the sum over the element-wise product of $X$ and $Y$.

With $b=Ax$ the forward differentiation gives $\dot b=\dot A x+A\dot x$, inserting you get $$ \bar x^⊤\dot x=\bar A(\dot A)+\bar b^⊤(\dot A x+A\dot x). \tag{1} $$ As that has to be valid for all choices of $\dot A, \dot x$, one gets \begin{align} 0&=trace(\bar A^⊤\dot A)+ \bar b^⊤\dot A x \\ &=trace(\bar A^⊤\dot A)+ trace(\bar b^⊤\dot A x) \\ &=trace(\bar A^⊤\dot A)+ trace(x\bar b^⊤\dot A) \\ &=trace((\bar A+\bar bx^⊤)^⊤\dot A), \\\text{and }~\bar b A\dot x&=\bar x\dot x \tag{2} \end{align} which implies $$ \bar b^⊤A=\bar x^⊤\iff \bar b=\text{solve}(A^⊤,\bar x),\tag3 \\ \bar A = -\bar bx^⊤. \\ $$