A Note on Adjoint Linear Algebra

TL;DR

This paper introduces a novel proof for adjoint linear systems based on Algorithmic Differentiation, extending to higher-order systems and providing a new perspective on adjoint operations in linear algebra.

Contribution

It presents a new proof method for adjoint linear systems using Algorithmic Differentiation principles, applicable to various matrix operations and higher-order systems.

Findings

New proof for adjoint systems based on Algorithmic Differentiation

Extension to higher-order adjoint linear systems

Alternative proof for matrix-matrix and vector products

Abstract

A new proof for adjoint systems of linear equations is presented. The argument is built on the principles of Algorithmic Differentiation. Application to scalar multiplication sets the base line. Generalization yields adjoint inner vector, matrix-vector, and matrix-matrix products leading to an alternative proof for first- as well as higher-order adjoint linear systems.

Equations

$$\left(\begin{array}{rl}{\mathbf{b}}_{(1)}& =A^{-T}\cdot {\mathbf{x}}_{(1)}\\ A_{(1)}& =-{\mathbf{b}}_{(1)}\cdot {\mathbf{x}}^T\end{array}\right).$$

$$⟨\nabla F\cdot {\mathbf{x}}^{(1)},{\mathbf{y}}_{(1)}⟩=⟨{\mathbf{x}}^{(1)},\nabla F^{\ast }\cdot {\mathbf{y}}_{(1)}⟩$$

$$\begin{array}{rl}{a}_{(1)}& =x\cdot y_{(1)}\\ x_{(1)}& =a\cdot y_{(1)}\end{array}$$

$$y^{(1)}=⟨\left(\begin{array}{c}{a}^{(1)}\\ x^{(1)}\end{array}\right),\left(\begin{array}{c}x\\ a\end{array}\right)⟩$$

$$⟨y^{(1)},y_{(1)}⟩$$

$$\left(\begin{array}{c}{a}_{(1)}\\ x_{(1)}\end{array}\right)=\left(\begin{array}{c}x\\ a\end{array}\right)\cdot y_{(1)}$$

$$y=⟨\mathbf{a},\mathbf{x}⟩\equiv {\mathbf{a}}^T\cdot \mathbf{x}=\sum _{i=0}^{n-1}{a}_i\cdot x_i$$

$$\begin{array}{rl}{\mathbf{a}}_{(1)}& =\mathbf{x}\cdot y_{(1)}\\ {\mathbf{x}}_{(1)}& =\mathbf{a}\cdot y_{(1)}\end{array}$$

$$y^{(1)}$$

$$=\sum _{i=0}^{n-1}{x}_i\cdot a_i^{(1)}+\sum _{i=0}^{n-1}{x}_i^{(1)}\cdot a_i={\mathbf{x}}^T\cdot {\mathbf{a}}^{(1)}+{\mathbf{a}}^T\cdot {\mathbf{x}}^{(1)}=({\mathbf{x}}^T{\mathbf{a}}^T)\cdot \left(\begin{array}{c}{\mathbf{a}}^{(1)}\\ {\mathbf{x}}^{(1)}\end{array}\right).$$

$$y_{(1)}\cdot y^{(1)}=({\mathbf{a}}_{(1)}^T{\mathbf{x}}_{(1)}^T)\cdot \left(\begin{array}{c}{\mathbf{a}}^{(1)}\\ {\mathbf{x}}^{(1)}\end{array}\right)=y_{(1)}\cdot ({\mathbf{x}}^T{\mathbf{a}}^T)\cdot \left(\begin{array}{c}{\mathbf{a}}^{(1)}\\ {\mathbf{x}}^{(1)}\end{array}\right)$$

$$A^{(1)}=({\mathbf{a}}_i^{(1)}{)}_{i=0,\dots ,m-1}=(a_{i,j}^{(1)}{)}_{i=0,\dots ,m-1}^{j=0,\dots ,n-1}$$

$$A_{(1)}=({\mathbf{a}}_{(1)i}{)}_{i=0,\dots ,m-1}=(a_{(1)i,j}{)}_{i=0,\dots ,m-1}^{j=0,\dots ,n-1},$$

$$\mathbf{y}=A\cdot \mathbf{x}\equiv {\left({\mathbf{a}}_i\cdot \mathbf{x}\right)}_{i=0,\dots ,m-1}$$

$$\begin{array}{rl}{\mathbf{x}}_{(1)}& =A^T\cdot {\mathbf{y}}_{(1)}\\ A_{(1)}& ={\mathbf{y}}_{(1)}\cdot {\mathbf{x}}^T\end{array}$$

$${\mathbf{y}}^{(1)}$$

$$={\left({\mathbf{x}}^T\cdot {{\mathbf{a}}_i^{(1)}}^T\right)}_{i=0,\dots ,m-1}+{\left({\mathbf{a}}_i\cdot {\mathbf{x}}^{(1)}\right)}_{i=0,\dots ,m-1}$$

$$={\left({\mathbf{a}}_i^{(1)}\cdot \mathbf{x}\right)}_{i=0,\dots ,m-1}+{\left({\mathbf{a}}_i\cdot {\mathbf{x}}^{(1)}\right)}_{i=0,\dots ,m-1}$$

$$={\left({\mathbf{a}}_i^{(1)}\right)}_{i=0,\dots ,m-1}\cdot \mathbf{x}+{\left({\mathbf{a}}_i\right)}_{i=0,\dots ,m-1}\cdot {\mathbf{x}}^{(1)}=A^{(1)}\cdot \mathbf{x}+A\cdot {\mathbf{x}}^{(1)}.$$

$$⟨{\mathbf{y}}_{(1)},{\mathbf{y}}^{(1)}⟩$$

$$={\left({\mathbf{a}}_{(1)i}^T\right)}_{i=0,\dots ,m-1}^T\cdot {\left({\mathbf{a}}_i^{(1)T}\right)}_{i=0,\dots ,m-1}+{\mathbf{x}}_{(1)}^T\cdot {\mathbf{x}}^{(1)}$$

$$={\mathbf{y}}_{(1)}^T\cdot \left(A^{(1)}\cdot \mathbf{x}+A\cdot {\mathbf{x}}^{(1)}\right)={\mathbf{y}}_{(1)}^T\cdot A^{(1)}\cdot \mathbf{x}+{\mathbf{y}}_{(1)}^T\cdot A\cdot {\mathbf{x}}^{(1)}$$

$$={\left({\left(y_{(1)i}\cdot {\mathbf{a}}_i^{(1)T}\right)}_{i=0,\dots ,m-1}\right)}^T\cdot {\left(\mathbf{x}\right)}_{i=0,\dots ,m-1}+{\mathbf{y}}_{(1)}^T\cdot A\cdot {\mathbf{x}}^{(1)}$$

$$=\underset{={\left(({\mathbf{a}}_{(1)i}^T{)}_{i=0,\dots ,n-1}\right)}^T}{\underbrace{{\left({\left(y_{(1)i}\cdot \mathbf{x}\right)}_{i=0,\dots ,m-1}\right)}^T}}\cdot {\left({\mathbf{a}}_i^{(1)T}\right)}_{i=0,\dots ,m-1}+\underset{={\mathbf{x}}_{(1)}^T}{\underbrace{{\mathbf{y}}_{(1)}^T\cdot A}}\cdot {\mathbf{x}}^{(1)},$$

$$\begin{array}{rl}{A}_{(1)}& =Y_{(1)}\cdot X^T\\ X_{(1)}& =A^T\cdot Y_{(1)}\end{array}$$

$${\mathbf{y}}^{(1)k}$$

$$Y^{(1)}=A^{(1)}\cdot X+A\cdot X^{(1)}.$$

$$⟨{\mathbf{y}}_{(1)}^k,{\mathbf{y}}^{(1)k}⟩$$

$$=\underset{={\left(({\mathbf{a}}_{(1)i}^T{)}_{i=0,\dots ,m-1}\right)}^T}{\underbrace{{\left({\left(y_{(1)i}^k\cdot {\mathbf{x}}^k\right)}_{i=0,\dots ,m-1}\right)}^T}}\cdot {\left({\mathbf{a}}_i^{(1)T}\right)}_{i=0,\dots ,m-1}+\underset{={\mathbf{x}}_{(1)}^{k^T}}{\underbrace{{\mathbf{y}}_{(1)}^{k^T}\cdot A}}\cdot {\mathbf{x}}^{(1)k}$$

$$Y^{(1)}=A\cdot X^{(1)}\cdot B$$

$$X_{(1)}=A^T\cdot Y_{(1)}\cdot B^T.$$

$$Y^{(1)}=Z^{(1)}\cdot B\Rightarrow Z_{(1)}=Y_{(1)}\cdot B^T$$

$$Z^{(1)}=A\cdot X^{(1)}\Rightarrow X_{(1)}=A^T\cdot Z_{(1)}$$

$$Y^{(1)}=\sum _{i=0}^{k-1}{A}_i\cdot X_i^{(1)}\cdot B_i$$

$$X_{i(1)}=A_i^T\cdot Y_{(1)}\cdot B_i^T$$

Full text

{AMS}

15A06, 15A29, 26B05

1 Motivation

Algorithmic Differentiation [3, 5] of numerical programs builds on a set of elemental functions with known partial derivatives with respect to their arguments at the given point of evaluation. The propagation of adjoint derivatives relies on the associativity of the chain rule of differential calculus. Differentiable combinations of elemental functions yield higher-level elementals. Efficient implementation of AD requires the highest possible level of elemental functions.

Basic AD assumes the set of elemental functions to be formed by the arithmetic operators and intrinsic functions built into the given programming language. While its application to linear algebra methods turns out to be straight forward basic AD is certainly not the method of choice from the point of view of computational efficiency. Elementals of the highest possible level should be used. Their derivatives should be formulated as functions of high-level elementals in order to exploit benefits of corresponding optimized implementations.

Following this rationale this note presents a new way to derive adjoint systems of linear equations based on adjoint Basic Linear Algebra Subprograms (BLAS) [4]. It is well known (see [2] and references therein) that for systems $A\cdot{\bf x}={\bf b}$ of $n$ linear equations with invertible $A$ and primal solution ${\bf x}=A^{-1}\cdot{\bf b}$ first-order adjoints $A_{(1)}$ of $A$ (both $\in{{\rm I}\kern-2.5pt{\rm R}}^{n\times n}$ with ${{\rm I}\kern-2.5pt{\rm R}}$ denoting the real numbers) and ${\bf b}_{(1)}$ of ${\bf b}$ (both $\in{{\rm I}\kern-2.5pt{\rm R}}^{n}$) can be evaluated at the primal solution ${\bf x}\in{{\rm I}\kern-2.5pt{\rm R}}^{n}$ as

[TABLE]

The main contribution of this note is an alternative proof for Eqn. (1) that builds naturally on the adjoint BLAS used in the context of state of the art AD. For consistency with related work we follow the notation in [5], that is, $v^{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}$ denotes the value of the first-order directional derivative (or tangent) associated with a variable $v\in{{\rm I}\kern-2.5pt{\rm R}}$ and $v_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}$ denotes the value of its adjoint.

2 Prerequisites

The Jacobian $\nabla F=\nabla F({\bf x})\equiv\frac{dF}{d{\bf x}}({\bf x})\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times n}$ of a differentiable implementation of ${\bf y}=F({\bf x}):{{\rm I}\kern-2.5pt{\rm R}}^{n}\rightarrow{{\rm I}\kern-2.5pt{\rm R}}^{m}$ as a computer program induces a linear mapping ${\bf y}^{(1)}=\nabla F\cdot{\bf x}^{(1)}:{{\rm I}\kern-2.5pt{\rm R}}^{n}\rightarrow{{\rm I}\kern-2.5pt{\rm R}}^{m}$ implementing the tangent of $F.$ The corresponding adjoint operator $\nabla F^{*}=\nabla F^{*}({\bf x})$ is formally defined via the inner vector product identity

[TABLE]

yielding $\nabla F^{*}=\nabla F^{T}$ [1]. In the following all (program) variables are assumed to be alias- and context-free, that is, distinct variables do not overlap in memory and $F$ is assumed to be not embedded in an enclosing computation. We distinguish between active and passive variables. Derivatives of all active outputs of the given program are computed with respect to all active inputs. We are not interested in derivatives of passive outputs nor are we computing derivatives with respect to passive inputs.

3 BLAS Revisited

In its basic form AD builds on known tangents and adjoints of the arithmetic functions and operators built into programming languages. Tangents and adjoints are propagated along the flow of data according to the chain rule of differential calculus. We enumerate entries of vectors ${\bf v}\in{{\rm I}\kern-2.5pt{\rm R}}^{n}$ staring from zero as $v_{0},\ldots,v_{n-1}.$

From the perspective of AD adjoint versions of higher-level BLAS are derived as adjoints of lower-level BLAS. Optimization of the result aims for implementation using the highest possible level of BLAS. For example, adjoint matrix-matrix multiplication (level-3 BLAS) is derived from adjoint matrix-vector multiplication (level-2 BLAS) yielding efficient evaluation as two matrix-matrix products (level-3 BLAS) as shown in Lemma 3.7. Rigorous derivation of this result requires bottom-up investigation of the BLAS hierarchy. We start with basic scalar multiplication (Lemma 3.1) followed by the inner vector (Lemma 3.3) and matrix-vector (Lemma 3.5) products as prerequisites for the matrix-matrix product.

Lemma 3.1**.**

The adjoint of scalar multiplication $y=a\cdot x$ with active $a,x,y\in{{\rm I}\kern-2.5pt{\rm R}}$ is computed as

[TABLE]

*for $y_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}$ yielding $a_{(1)},x_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}.$ *

Proof 3.2**.**

Differentiation of $y=a\cdot x$ with respect to $a$ and $x$ yields the tangent

[TABLE]

for $y^{(1)},a^{(1)},x^{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}.$ Eqn. (2) implies

[TABLE]

yielding

[TABLE]

*and hence Eqn. (3). *

Lemma 3.3**.**

The adjoint of an inner vector product

[TABLE]

with active inputs ${\bf a}\in{{\rm I}\kern-2.5pt{\rm R}}^{n}$ and ${\bf x}\in{{\rm I}\kern-2.5pt{\rm R}}^{n}$ yielding the active output $y\in{{\rm I}\kern-2.5pt{\rm R}}$ is computed as

[TABLE]

*for $y_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}$ yielding ${\bf a}_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}^{n}$ and ${\bf x}_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}^{n}.$ *

Proof 3.4**.**

Differentiation of $y={\bf a}^{T}\cdot{\bf x},$ for ${\bf a}=(a_{i})_{i=0,\ldots,n-1}$ and ${\bf x}=(x_{i})_{i=0,\ldots,n-1}$, with respect to ${\bf a}$ and ${\bf x}$ yields the tangent

[TABLE]

Eqn. (2) implies

[TABLE]

*yielding $({\bf a}_{(1)}^{T}~{}{\bf x}_{(1)}^{T})=y_{(1)}\cdot({\bf x}^{T}~{}{\bf a}^{T})$ and hence Eqn. (4). *

The following derivation of adjoint matrix-vector and matrix-matrix products relies on serialization of matrices. Individual rows of a matrix $A\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times n}$ are denoted as ${\bf a}_{i}\in{{\rm I}\kern-2.5pt{\rm R}}^{1\times n}$ for $i=0,\ldots,m-1;$ columns are denoted as ${\bf a}^{j}\in{{\rm I}\kern-2.5pt{\rm R}}^{m}$ for $i=0,\ldots,n-1.$ (Row) Vectors in ${{\rm I}\kern-2.5pt{\rm R}}^{1\times n}$ are denoted as $\left(v_{j}\right)^{j=0,\ldots,n-1};$ (column) vectors in ${{\rm I}\kern-2.5pt{\rm R}}^{m}$ are denoted as $\left(v_{i}\right)_{i=0,\ldots,m-1};$ Consequently, a row-major serialization of $A$ is given by $\left({\bf a}^{T}_{i}\right)_{i=0,\ldots,m-1}.$ A column-major serialization of $A$ is given by $\left({\bf a}^{j}\right)_{j=0,\ldots,n-1}.$ Tangents and adjoints of the individual entries of $A$ define

[TABLE]

and

[TABLE]

respectively.

Lemma 3.5**.**

The adjoint of a matrix-vector product

[TABLE]

with active inputs $A\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times n}$ and ${\bf x}\in{{\rm I}\kern-2.5pt{\rm R}}^{n}$ yielding the active output ${\bf y}\in{{\rm I}\kern-2.5pt{\rm R}}^{m}$ is computed as

[TABLE]

*for ${\bf y}_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}^{m}$ yielding ${\bf x}_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}^{n}$ and $A_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times n}.$ *

Proof 3.6**.**

Differentiation of ${\bf y}=A\cdot{\bf x},$ where $A=\left({\bf a}_{i}\right)_{i=0,\ldots,m-1},$ ${\bf x}=\left(x_{j}\right)_{j=0,\ldots,n-1}$ and ${\bf y}=\left(y_{i}\right)_{i=0,\ldots,m-1},$ with respect to $A$ and ${\bf x}$ yields the tangent

[TABLE]

Eqn. (2) implies

[TABLE]

*where $\left({\bf x}\right)_{i=0,\ldots,m-1}\in{{\rm I}\kern-2.5pt{\rm R}}^{m\cdot n}$ denotes a concatenation of $m$ copies of ${\bf x}\in{{\rm I}\kern-2.5pt{\rm R}}^{n}$ as a column vector. Eqn. (5) follows immediately. *

Lemma 3.7**.**

The adjoint of a matrix-matrix product $Y=A\cdot X$ with active inputs $A\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times p},$ $X\in{{\rm I}\kern-2.5pt{\rm R}}^{p\times n}$ yielding the active output $Y\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times n}$ is computed as

[TABLE]

*for $Y_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times n}$ yielding $A_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times p}$ and $X_{(1)}\in{{\rm I}\kern-2.5pt{\rm R}}^{p\times n}.$ *

Proof 3.8**.**

Differentiation of $Y=A\cdot X,$ where $A=\left({\bf a}_{i}\right)_{i=0,\ldots,m-1}$, $X=\left({\bf x}^{k}\right)^{k=0,\ldots,p-1}$ and $Y=\left({\bf y}^{k}\right)^{k=0,\ldots,p-1},$ with respect to $A$ and $X$ yields tangents

[TABLE]

for $k=0,\ldots,p-1$ and hence

[TABLE]

Eqn. (2) implies

[TABLE]

*for $k=0,\ldots,p-1$ and hence the Eqn. (6). *

4 Systems of Linear Equations Revisited

Lemmas 4.1 and 4.3 form the basis for the new proof of Eqn. (1).

Lemma 4.1**.**

The tangent

[TABLE]

of $Y=A\cdot X\cdot B$ for active $X\in{{\rm I}\kern-2.5pt{\rm R}}^{n\times q},$ $Y\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times p}$ and passive $A\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times n},$ $B\in{{\rm I}\kern-2.5pt{\rm R}}^{q\times p}$ implies the adjoint

[TABLE]

Proof 4.2**.**

[TABLE]

follows from application of Lemma 3.7 to $Y=Z\cdot B$ with passive $B.$

[TABLE]

*follows from application of Lemma 3.7 to $Z=A\cdot X$ with passive $A.$ Substitution of $Z^{(1)}$ and $Z_{(1)}$ yields Lemma 4.1. *

Lemma 4.3**.**

The tangent

[TABLE]

of $Y=\sum_{i=0}^{k-1}A_{i}\cdot X_{i}\cdot B_{i}=\sum_{i=0}^{k-1}Y_{i}$ with active $X_{i}\in{{\rm I}\kern-2.5pt{\rm R}}^{n_{i}\times q_{i}},$ $Y\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times p}$ and with passive $A_{i}\in{{\rm I}\kern-2.5pt{\rm R}}^{m\times n_{i}},$ $B_{i}\in{{\rm I}\kern-2.5pt{\rm R}}^{q_{i}\times p}$ implies the adjoint

[TABLE]

*for $i=0,\ldots,k-1$ *

Proof 4.4**.**

From

[TABLE]

follows with Lemma 4.1

[TABLE]

*for $i=0,\ldots,k-1.$ Moreover, $Y^{(1)}=\sum_{i=0}^{k-1}Y^{(1)}_{i}$ implies $Y_{i(1)}=Y_{(1)}$ due to identity Jacobians of $Y$ with respect to $Y_{i}$ for $i=0,\ldots,k-1$ and hence Lemma 4.3. *

Theorem 4.5**.**

*Adjoints of systems $A\cdot{\bf x}={\bf b}$ of $n$ linear equations with invertible $A\in{{\rm I}\kern-2.5pt{\rm R}}^{n\times n}$ and right-hand side ${\bf b}\in{{\rm I}\kern-2.5pt{\rm R}}^{n}$ are evaluated at the primal solution ${\bf x}=A^{-1}\cdot{\bf b}\in{{\rm I}\kern-2.5pt{\rm R}}^{n}$ by Eqn. (1). *

Proof 4.6**.**

Differentiation of $A\cdot{\bf x}={\bf b}$ with respect to $A$ and ${\bf b}$ yields the tangent system

[TABLE]

which implies

[TABLE]

with identity $I_{n}\in{{\rm I}\kern-2.5pt{\rm R}}^{n\times n}.$ Lemma 4.3 yields

[TABLE]

*and hence Eqn. (1). *

5 Conclusion

As observed previously by various authors a possibly available factorization of $A$ can be reused both for the tangent ($A\cdot{\bf x}^{(1)}={\bf b}^{(1)}-A^{(1)}\cdot{\bf x}$) and the adjoint ($A^{T}\cdot{\bf b}_{(1)}={\bf x}_{(1)}$) systems. The additional worst case computational cost of $O(n^{3})$ can thus be reduced to $O(n^{2}).$ Higher-order tangents [adjoints] of linear systems amount to repeated solutions of linear systems with the same [transposed] system matrix combined with tangent [adjoint] BLAS.