Backpropagation is the Adjoint Method

25 October 2025 · Download PDF · Download code

Backpropagation is often introduced as something developed within the field of machine learning. However, the reality is that backpropagation is just a special case of the adjoint method where the inversion of the partial derivative $\frac{𝜕 𝑓}{𝜕 𝑢}$ can be done efficiently due to its lower-triangular structure. This note is in two parts. In the first part we review the adjoint method while in the second part we describe how backpropagation is a special case of the adjoint method with a structure that results in scalable computational algorithms.

A key insight is that while the adjoint method is seemingly limited to equality constrained problems, it turns out that all computer programs are trivially equality constrained problems with equalities coming from each assignment¹¹ All lines of code are essentially “ $variable = something$ ” which can be written as “ $variable - something = 0$ ”. In this sense, as shown in a later section, optimizing a neural network is actually an equality constrained optimization problem with constraints coming from each layer of the network. This is opposed to how optimizing a neural network is usually presented as an unconstrained optimization problem.

1. The Adjoint Method

In many engineering applications we are interested in solving equality constrained optimization problems where the optimization variable ( $𝜃$ ) implicitly defines another variable ( $𝑢$ ) that then appears in the objective function²² A famous example of this is Topology Optimization. Written out this means solving optimization problems of the following form

\begin{aligned} min_{𝜃 \in ℝ^{𝑛_{𝜃}}} & 𝐿 (𝑢, 𝜃), & (some objective function) \\ subject to & 𝑓 (𝑢, 𝜃) = 0, & (some equality constraint) \\ 𝑢 \in ℝ^{𝑛_{𝑢}} . \end{aligned}

(1)

Most methods for solving equality constrained optimization problems require the computation of the gradient of the objective function with respect to the optimization variable $𝜃$ . This is where the problem comes in: Naively computing the gradient requires the computation of the sensitivities of the implicitly defined variable $𝑢$ with respect to the optimization variable $𝜃$ $(i.e. \frac{d 𝑢}{d 𝜃} \in ℝ^{𝑛_{𝑢} \times 𝑛_{𝜃}})$ . This results in a computational bottleneck as forming the sensitivities scales as $𝒪︀ (𝑛_{𝑢} 𝑛_{𝜃})$ , meaning that adding a new parameter adds $𝑛_{𝑢}$ additional sensitivities (and similarly adding one more variable would add $𝑛_{𝜃}$ additional sensitivities).

To resolve this computational bottleneck the adjoint method was introduced. The first step in the derivation of the adjoint method is to introduce the Lagrangian of the objective function, i.e.

ℒ︀ (𝑢, 𝜃, 𝜆) = 𝐿 (𝑢, 𝜃) + 𝜆^{⊤} 𝑓 (𝑢, 𝜃) .

(2)

It is easy to see that whenever $𝑢$ satisfies the equality constraint $𝑓 (𝑢, 𝜃) = 0$ , then the Lagrangian is equal to the original objective function. The reason why the introduction of the additional term in the Lagrangian is useful is that $𝜆$ can be set to be anything. In particular, we aim to set it in a way so that we can avoid computing the otherwise computationally expensive sensitivities $\frac{d 𝑢}{d 𝜃}$ . We start by computing the total derivative of the Lagrangian with respect to $𝜃$ ³³ The total derivative is the transpose of the gradient i.e. ${(\frac{d 𝐿}{d 𝜃})}^{⊤} = \nabla_{𝜃} 𝐿$ . That is multiplying the change in $𝜃$ with the total derivative (rather than the transpose of the gradient) gives the change in the function $𝐿$ .

\frac{d ℒ︀}{d 𝜃} = \frac{𝜕 𝐿}{𝜕 𝜃} + \frac{𝜕 𝐿}{𝜕 𝑢} \frac{d 𝑢}{d 𝜃} + 𝜆^{⊤} (\frac{𝜕 𝑓}{𝜕 𝜃} + \frac{𝜕 𝑓}{𝜕 𝑢} \frac{d 𝑢}{d 𝜃}) .

(3)

Now we collect the terms that depend on the undesired $\frac{d 𝑢}{d 𝜃}$ , resulting in

\frac{d ℒ︀}{d 𝜃} = \frac{𝜕 𝐿}{𝜕 𝜃} + 𝜆^{⊤} \frac{𝜕 𝑓}{𝜕 𝜃} + (\frac{𝜕 𝐿}{𝜕 𝑢} + 𝜆^{⊤} \frac{𝜕 𝑓}{𝜕 𝑢}) \frac{d 𝑢}{d 𝜃} .

(4)

As we can choose $𝜆$ freely a natural idea is to set it in a way such that the term in front of the undesired term vanishes. This means that we can choose $𝜆$ as the solution to the equation

\frac{𝜕 𝐿}{𝜕 𝑢} + 𝜆^{⊤} \frac{𝜕 𝑓}{𝜕 𝑢} = 0 \Rightarrow 𝜆^{⊤} = - \frac{𝜕 𝐿}{𝜕 𝑢} {(\frac{𝜕 𝑓}{𝜕 𝑢})}^{- 1} .

(5)

Inserting $𝜆^{⊤}$ back into the equation we find that the gradient of the Lagrangian with respect to $𝜃$ is given by

\frac{d ℒ︀}{d 𝜃} = \frac{𝜕 𝐿}{𝜕 𝜃} \underset{𝜆^{⊤}}{\underset{⏟}{- \frac{𝜕 𝐿}{𝜕 𝑢} {(\frac{𝜕 𝑓}{𝜕 𝑢})}^{- 1}}} \frac{𝜕 𝑓}{𝜕 𝜃} .

(6)

To conclude: The adjoint method is a simple way to avoid the computational bottleneck of computing the sensitivities $\frac{d 𝑢}{d 𝜃}$ by cleverly computing the so-called adjoint variable $𝜆$ in a way that eliminates the need to compute the problematic sensitivities.

Example 1: Linearly constrained problem with structure

Example 1 (Linearly constrained problem with structure) In order to illustrate the adjoint method we consider a simple linearly constrained problem of the form (inspiration from problem 38 in [1] )

\begin{aligned} 𝐿 (𝑢, 𝜃) & = {(𝑐^{⊤} 𝑢 (𝜃))}^{2}, \\ 𝑓 (𝑢, 𝜃) & = 𝐴 (𝜃) 𝑢 - 𝑏 = 0, \end{aligned}

(7)

where $𝐴 (𝜃) \in ℝ^{𝑛 \times 𝑛}$ is a symmetric tridiagonal matrix that depends on the parameters $𝜃 \in ℝ^{2 𝑛 - 1}$ as

𝐴 = [\begin{matrix} 𝜃_{1} & 𝜃_{𝑛 + 1} & 0 \\ 𝜃_{𝑛 + 1} & 𝜃_{2} & ⋱ \\ ⋱ & ⋱ & ⋱ \\ ⋱ & 𝜃_{𝑛 - 1} & 𝜃_{2 𝑛 - 1} \\ 0 & 𝜃_{2 𝑛 - 1} & 𝜃_{𝑛} \end{matrix}] .

(8)

Now in order to compute the gradient of interest we start by computing the adjoint variable $𝜆$ as

𝜆^{⊤} = - \frac{𝜕 𝐿}{𝜕 𝑢} {(\frac{𝜕 𝑓}{𝜕 𝑢})}^{- 1} = - 2 (𝑐^{⊤} 𝑢) 𝑐^{⊤} {𝐴 (𝜃)}^{- 1} .

(9)

Note that in practice we do not form ${𝐴 (𝜃)}^{- 1}$ explicitly but rather compute $𝜆$ by solving the linear system ${𝐴 (𝜃)}^{⊤} 𝜆 = - 2 𝑐 (𝑐^{⊤} 𝑢)$ . Using the adjoint variable we can compute the gradient of $𝐿$ with respect to $𝜃$ as

\frac{d 𝐿}{d 𝜃} = \underset{= 0}{\underset{⏟}{\frac{𝜕 𝐿}{𝜕 𝜃}}} + 𝜆^{⊤} \frac{𝜕 𝑓}{𝜕 𝜃} = 𝜆^{⊤} \frac{𝜕 𝐴}{𝜕 𝜃} 𝑢 .

(10)

In the concrete case of the derivatives of $𝐴$ with respect to $𝜃_{𝑖}$ we have that

\frac{𝜕 𝐴}{𝜕 𝜃_{𝑖}} = [\begin{matrix} 𝛿_{𝑖, 1} & 𝛿_{𝑖, 𝑛 + 1} & 0 \\ 𝛿_{𝑖, 𝑛 + 1} & 𝛿_{𝑖, 2} & ⋱ \\ ⋱ & ⋱ & ⋱ \\ ⋱ & 𝛿_{𝑖, 𝑛 - 1} & 𝛿_{𝑖, 2 𝑛 - 1} \\ 0 & 𝛿_{𝑖, 2 𝑛 - 1} & 𝛿_{𝑖, 𝑛} \end{matrix}] .

(11)

Using this it follows that the $𝑖$ th component of the gradient can be computed as

𝜆^{⊤} \frac{𝜕 𝐴}{𝜕 𝜃_{𝑖}} 𝑢 = {\begin{cases} 𝜆_{𝑖} 𝑢_{𝑖}, & 𝑖 \leq 𝑛, \\ 𝜆_{𝑖 + 1 - 𝑛} 𝑢_{𝑖 - 𝑛} + 𝜆_{𝑖 - 𝑛} 𝑢_{𝑖 + 1 - 𝑛}, & 𝑖 > 𝑛 . \end{cases}

(12)

Below is a Julia implementation of the above example along with a test against ForwardDiff.jl to verify the correctness of the implementation.

julia


1
using LinearAlgebra
2
using ForwardDiff
3
using Random
4
using Test
5
6
Random.seed!(42)
7
n = 1000        # Number of linear constraints
8
c = rand(n)     # Objective vector
9
b = rand(n)     # Linear equality vector
10
θ = rand(2n-1)  # Parameters of the equality constraint
11
12
# Parametrizing linear equality constraint matrix
13
A(θ) = SymTridiagonal(θ[1:n], θ[n+1:end])
14
15
# Objective function
16
f(θ) = (c' * (A(θ) \ b))^2
17
18
# Partial derivatives
19
partial(i,n,λ,u) = i <= n ? λ[i]*u[i] : λ[i+1-n]*u[i-n] + λ[i-n]*u[i+1-n]
20
21
# Gradient computation
22
function ∇f(θ)
23
    n = length(b)
24
    M = A(θ)                # Defining constraint matrix
25
    u = M \ b               # Computing solution
26
    λ = M \ (-2*c*(c'*u))   # Adjoint variables
27
    return [partial(i,n,λ,u) for i = 1:2n-1]
28
end
29
30
# Testing against ForwardDiff
31
@test ForwardDiff.gradient(f,θ) ≈ ∇f(θ)

2. Backpropagation and the Adjoint Method

Section 5.5 of [2] includes an example of how the adjoint method and backpropagation of neural networks are similar. In neural networks we are often interested in minimizing a loss function of the form

𝐿 (𝜃; 𝑢_{0}) = ‖ Φ_{𝑁} (𝜃; 𝑢_{0}) - 𝑦 ‖_{2}^{2} + 𝜇 ℛ︀ (𝜃),

(13)

with $𝜃 \in ℝ^{𝑘}, 𝑢_{0} \in ℝ^{𝑚}, 𝑦 \in ℝ^{𝑛}$ and the notation “ $; 𝑢_{0}$ ” is used in order to highlight that $𝑢_{0}$ is a constant input (i.e. most often the “data”) and $ℛ︀ (𝜃)$ is some regularization function (e.g. $ℛ︀ (𝜃) = ‖ 𝜃 ‖_{2}^{2}$ ). Furthermore $Φ_{𝑁} (𝜃; 𝑢_{0}) : ℝ^{𝑘} \to ℝ^{𝑛}$ is a neural network with $𝑁$ layers that given a set of constant inputs $𝑢_{0}$ maps the parameters $𝜃$ to an output. Now, the above objective function does not include any equality constraints. However, one should realize that an $𝑁$ -layer neural network is nothing more than a composition of $𝑁$ functions

Φ_{𝑁} (𝜃; 𝑢_{0}) = Φ_{𝑁} (Φ_{𝑁 - 1} (\dots (Φ_{1} (𝜃; 𝑢_{0}) \dots, 𝜃; 𝑢_{0}), 𝜃; 𝑢_{0}) .

(14)

Now using the notation that $𝑢_{𝑖}$ is the output of the $𝑖$ th layer of the network we can write the propagation through the network as a large nonlinear system of equations that have to be satisfied

\begin{aligned} 𝑓 (𝑢, 𝜃) & = 𝑢 - Φ (𝑢, 𝜃) \\ = \underset{𝑢}{\underset{⏟}{[\begin{matrix} 𝑢_{1} \\ 𝑢_{2} \\ ⋮ \\ 𝑢_{𝑁} \end{matrix}]}} - \underset{Φ (𝑢, 𝜃)}{\underset{⏟}{[\begin{matrix} Φ_{1} (𝜃; 𝑢_{0}) \\ Φ_{2} (𝑢_{1}, 𝜃; 𝑢_{0}) \\ ⋮ \\ Φ_{𝑁} (𝑢_{1}, \dots, 𝑢_{𝑁 - 1}, 𝜃; 𝑢_{0}) \end{matrix}]}} = [\begin{matrix} 0 \\ 0 \\ ⋮ \\ 0 \end{matrix}] . \end{aligned}

(15)

Using this notation (and removing the explicit dependence on $𝑢_{0}$ ) we see that optimizing a neural network is really nothing more than solving the following equality constrained optimization problem

\begin{aligned} min_{𝜃} & 𝐿 (𝑢, 𝜃) = ‖ 𝑢_{𝑁} - 𝑦 ‖_{2}^{2} + 𝜇 ℛ︀ (𝜃), \\ subject to & 𝑓 (𝑢, 𝜃) = 0 . \end{aligned}

(16)

We can now use the adjoint method to compute the gradient of $𝐿$ with respect to $𝜃$ . As a reminder this means that

\frac{d 𝐿}{d 𝜃} = \frac{𝜕 𝐿}{𝜕 𝜃} \underset{𝜆^{⊤}}{\underset{⏟}{- \frac{𝜕 𝐿}{𝜕 𝑢} {(\frac{𝜕 𝑓}{𝜕 𝑢})}^{- 1}}} \frac{𝜕 𝑓}{𝜕 𝜃},

(17)

where $𝜆$ is referred to as the adjoint variable. Let $𝑛_{𝑖}$ be the dimension of the output $𝑢_{𝑖}$ of layer $𝑖$ , so that the stacked state $𝑢$ has dimension $𝑑 = \sum_{𝑖 = 1}^{𝑁} 𝑛_{𝑖}$ . For simplicity, we assume that the final output $𝑢_{𝑁}$ is scalar, i.e. $𝑛_{𝑁} = 1$ . Then $𝑢_{𝑁} = 𝑒_{𝑑}^{⊤} 𝑢$ , where $𝑒_{𝑑}$ is the last canonical basis vector in $ℝ^{𝑑}$ . In this case we have that

\frac{𝜕 𝐿}{𝜕 𝑢} = 2 (𝑒_{𝑑}^{⊤} 𝑢 - 𝑦) 𝑒_{𝑑}^{⊤} = 2 (𝑢_{𝑁} - 𝑦) 𝑒_{𝑑}^{⊤} = 𝑔_{𝑁}^{⊤} .

(18)

Now what is left to compute is $\frac{𝜕 𝑓}{𝜕 𝑢}$ and $\frac{𝜕 𝑓}{𝜕 𝜃}$ . Starting with $\frac{𝜕 𝑓}{𝜕 𝑢}$ we note that because $𝑓$ only propagates the $𝑢_{𝑖}$ ’s forwards the resulting partial derivative is a lower block triangular matrix of the form

\frac{𝜕 𝑓}{𝜕 𝑢} = [\begin{matrix} 𝐼 & 0 & \dots & 0 \\ - \frac{𝜕 Φ_{2}}{𝜕 𝑢_{1}} & 𝐼 & \dots & 0 \\ ⋮ & ⋮ & ⋱ & ⋮ \\ - \frac{𝜕 Φ_{𝑁}}{𝜕 𝑢_{1}} & \dots & - \frac{𝜕 Φ_{𝑁}}{𝜕 𝑢_{𝑁 - 1}} & 𝐼 \end{matrix}] = 𝐿 .

(19)

Now is a good place to stop and think of a key concept of backpropagation: The resulting computational graph must be a Directed Acyclic Graph (DAG). This is indeed equivalent to stating that $\frac{𝜕 𝑓}{𝜕 𝑢}$ should be a lower triangular matrix. This is an important property as the adjoint method requires us to invert $\frac{𝜕 𝑓}{𝜕 𝑢}$ , which can be done cheaply for a lower triangular matrix. Even more important is that the diagonal blocks are the identity, meaning that forward/backward substitutions can be done without any matrix inversions at all.

What is left is to compute $\frac{𝜕 𝑓}{𝜕 𝜃}$ , which standardly is done as

\frac{𝜕 𝑓}{𝜕 𝜃} = - \underset{𝑀^{⊤}}{\underset{⏟}{[\begin{matrix} \frac{𝜕 Φ_{1}}{𝜕 𝜃} \\ ⋮ \\ \frac{𝜕 Φ_{𝑁}}{𝜕 𝜃} \end{matrix}]}} = - 𝑀^{⊤} .

(20)

Now, in case of the layers not sharing any parameters the matrix $𝑀^{⊤}$ will have a block diagonal structure, i.e.

𝑀^{⊤} = blkdiag (\frac{𝜕 Φ_{1}}{𝜕 𝜃_{1}}, \dots, \frac{𝜕 Φ_{𝑁}}{𝜕 𝜃_{𝑁}}), 𝜃 = [\begin{matrix} 𝜃_{1} \\ ⋮ \\ 𝜃_{𝑁} \end{matrix}],

(21)

where $𝜃_{𝑖}$ are the parameters of layer $𝑖$ .

Putting everything together we find that the gradient of $𝐿$ with respect to $𝜃$ is given by

\begin{aligned} \frac{d ℒ︀}{d 𝜃} & = \frac{𝜕 𝐿}{𝜕 𝜃} - \frac{𝜕 𝐿}{𝜕 𝑢} {(\frac{𝜕 𝑓}{𝜕 𝑢})}^{- 1} \frac{𝜕 𝑓}{𝜕 𝜃} \\ = 𝜇 \frac{𝜕 ℛ︀}{𝜕 𝜃} - (𝑔_{𝑁}^{⊤} 𝐿^{- 1} (- 𝑀^{⊤})) \\ = 𝜇 \frac{𝜕 ℛ︀}{𝜕 𝜃} + 2 (𝑢_{𝑁} - 𝑦) 𝑒_{𝑑}^{⊤} 𝐿^{- 1} 𝑀^{⊤} . \end{aligned}

(22)

As noted previously both $𝐿^{- 1}$ and $𝑀^{⊤}$ are structured, meaning that the above can be computed efficiently.

2.1. How to compute the elements of $𝐿$ and $𝑀^{⊤}$

For a standard linear layer an activation function $𝜎$ is usually applied element wise. In practice this means that we really should look at the rows of $Φ_{𝑖}$ separately. For the $𝑗$ th row the gradients are easily seen to be

\begin{aligned} \nabla_{𝑢_{𝑖 - 1}} 𝜎 (𝑤_{𝑗}^{⊤} 𝑢_{𝑖 - 1} + 𝑏_{𝑗}) & = 𝜎^{'} (𝑤_{𝑗}^{⊤} 𝑢_{𝑖 - 1} + 𝑏_{𝑗}) 𝑤_{𝑗}, & Goes into 𝐿 \\ \nabla_{𝑤_{𝑗}} 𝜎 (𝑤_{𝑗}^{⊤} 𝑢_{𝑖 - 1} + 𝑏_{𝑗}) & = 𝜎^{'} (𝑤_{𝑗}^{⊤} 𝑢_{𝑖 - 1} + 𝑏_{𝑗}) 𝑢_{𝑖 - 1}, & Goes into 𝑀^{⊤} \\ \nabla_{𝑏_{𝑗}} 𝜎 (𝑤_{𝑗}^{⊤} 𝑢_{𝑖 - 1} + 𝑏_{𝑗}) & = 𝜎^{'} (𝑤_{𝑗}^{⊤} 𝑢_{𝑖 - 1} + 𝑏_{𝑗}), & Goes into 𝑀^{⊤} \end{aligned}

(23)

Using the above we see that

\frac{𝜕 Φ_{𝑖}}{𝜕 𝑢_{𝑖 - 1}} = diag (𝜎^{'} (𝑊_{𝑖} 𝑢_{𝑖 - 1} + 𝑏_{𝑖})) 𝑊_{𝑖},

(24)

where $𝜎^{'} (\cdot)$ is applied element wise. Now for $\frac{𝜕 Φ_{𝑖}}{𝜕 𝜃_{𝑖}}$ we have to pick a way of how to vectorize the parameters. Here we choose

𝜃_{𝑖} = [\begin{matrix} vec (𝑊_{𝑖}) \\ 𝑏_{𝑖} \end{matrix}],

(25)

where $vec (𝑊)$ vectorizes by stacking the columns of $𝑊$ . Using the vectorization we can write the sought after sensitivity as

\frac{𝜕 Φ_{𝑖}}{𝜕 𝜃_{𝑖}} = diag (𝜎^{'} (𝑊_{𝑖} 𝑢_{𝑖 - 1} + 𝑏_{𝑖})) ([\begin{matrix} 𝑢_{𝑖 - 1}^{⊤} & 1 \end{matrix}] \otimes 𝐼_{𝑛_{𝑖}}) .

(26)

While the Kronecker product above does result in a sparse matrix, it is even more structured. Using the standard property of Kronecker products, $(𝐴 \otimes 𝐵) 𝑣 = vec (𝐵 mat (𝑣) 𝐴^{⊤})$ , where $mat (𝑣)$ reverses the column-stacking operation $vec$ with dimensions compatible with $𝐴$ and $𝐵$ , its products with vectors can be computed without forming it, e.g. $([\begin{matrix} 𝑢_{𝑖 - 1}^{⊤} & 1 \end{matrix}] \otimes 𝐼_{𝑛_{𝑖}}) 𝜃_{𝑖} = 𝑊_{𝑖} 𝑢_{𝑖 - 1} + 𝑏_{𝑖}$ .

A final remark here is that the diagonal matrices that go into the elements of $𝐿$ and $𝑀^{⊤}$ are the same. That is if we define

\begin{aligned} 𝐷 & = diag (𝜎^{'} (𝑊_{1} 𝑢_{0} + 𝑏_{1}), \dots, 𝜎^{'} (𝑊_{𝑁} 𝑢_{𝑁 - 1} + 𝑏_{𝑁})) \\ = blkdiag (diag (𝜎^{'} (𝑊_{1} 𝑢_{0} + 𝑏_{1})), \dots, diag (𝜎^{'} (𝑊_{𝑁} 𝑢_{𝑁 - 1} + 𝑏_{𝑁}))), \end{aligned}

(27)

then it follows that

\begin{aligned} 𝐿 & = 𝐼 - 𝐷 blkdiag (𝑊_{2}, \dots, 𝑊_{𝑁}, - 1), \\ 𝑀^{⊤} & = 𝐷 blkdiag ([\begin{matrix} 𝑢_{0}^{⊤} & 1 \end{matrix}] \otimes 𝐼_{𝑛_{1}}, \dots, [\begin{matrix} 𝑢_{𝑁 - 1}^{⊤} & 1 \end{matrix}] \otimes 𝐼_{𝑛_{𝑁}}), \end{aligned}

(28)

with the slightly unusual notation that $blkdiag (\cdot, - 1)$ denotes the lower block diagonal matrix. Here its $(𝑖, 𝑖 - 1)$ block is $𝑊_{𝑖}$ for $𝑖 = 2, \dots, 𝑁$ , with all other blocks zero and block row and column sizes $𝑛_{1}, \dots, 𝑛_{𝑁}$ . The weight $𝑊_{1}$ acts on the fixed input $𝑢_{0}$ and therefore does not appear in $𝐿$ .

Below is a Julia implementation of the above backpropagation/adjoint method for a fully connected feedforward neural network with scalar output. In this example we set $𝜇 = 0$ , so the objective contains only the squared error. The implementation uses BlockDiagonalMatrices.jl and Kronecker.jl in order to efficiently represent the involved matrices.

julia


1
using BlockDiagonalMatrices
2
using LinearAlgebra
3
using Kronecker
4
using Random
5
using Test
6
7
Random.seed!(42)
8
9
h(x)  =  exp(-x) # activation function
10
∇h(x) = -exp(-x) # derivative of activation function
11
12
# Forward pass including the derivatives
13
function forward_pass(u0, layers)
14
    u = u0
15
    _, b1 = first(layers)
16
    derivs = typeof(b1)[]
17
    kron_blocks = typeof([u; 1]' ⊗ I(length(b1)))[]
18
    for (W, b) in layers
19
        push!(kron_blocks,[u; 1]' ⊗ I(length(b))) # Lazy Kronecker product
20
        z = W*u + b
21
        u = h.(z)
22
        push!(derivs, ∇h.(z))
23
    end
24
    return kron_blocks, derivs, u
25
end
26
# Block back substitution: solves L^T y = b, i.e. y = L^(-T) b
27
function backsub(derivs, weights, b)
28
    y = Matrix(b)
29
    stop = length(y)
30
    @views for (d,Wt) in zip(reverse(derivs), reverse(transpose.(weights)))
31
        nprev, n = size(Wt)
32
        src = stop-n+1:stop
33
        dst = stop-n-nprev+1:stop-n
34
        tmp = d .* y[src]
35
        # Equivalent to: y[dst] .+= Wt * (d .* y[src])
36
        mul!(y[dst], Wt, tmp, 1, 1)
37
        stop -= n
38
    end
39
    return y
40
end
41
# Helper function that unpacks a vector θ into Ws and bs
42
function unpack_θ(θ, layer_sizes)
43
    Ws, bs = Matrix{eltype(θ)}[], Vector{eltype(θ)}[]
44
    i = 1
45
    for (n_in,n_out) in zip(layer_sizes[1:end-1], layer_sizes[2:end])
46
        j = i + n_out*n_in
47
        push!(Ws, reshape(θ[i:j-1], n_out, n_in))
48
        i = j + n_out
49
        push!(bs, θ[j:i-1])
50
    end
51
    return Ws, bs
52
end
53
# Evaluate the network output using the forward pass
54
function network_output(θ, layer_sizes, u0)
55
    Ws, bs = unpack_θ(θ, layer_sizes)
56
    _, _, uN = forward_pass(u0, zip(Ws,bs))
57
    return only(uN)
58
end
59
# Gradient computation using the adjoint method
60
function ∇L(θ, layer_sizes, u0; y=0.0)
61
    # First unpack the parameter vector
62
    Ws, bs = unpack_θ(θ, layer_sizes)
63
    # Forward pass includes derivative information
64
    kron_blocks, derivs, uN = forward_pass(u0, zip(Ws, bs))
65
    # S has subdiagonal blocks W_2,...,W_N; L = I - D*S and M^T = D*K
66
    D = Diagonal(vcat(derivs...))
67
    K = BlockDiagonal(kron_blocks)
68
    g = zeros(eltype(θ), sum(layer_sizes[2:end]))
69
    g[end] = 2*(only(uN) - y)
70
    # The final step is to evaluate the gradient from the right
71
    grad_adjoint = (backsub(derivs[2:end],Ws[2:end],g')*D)*K
72
    return grad_adjoint'
73
end
74
# Forward difference to test the adjoint gradient implementation
75
function fd(θ, layer_sizes, u0, i; e=1e-6, y=2.0)
76
    uN0 = network_output(θ, layer_sizes, u0)
77
    θ[i] += e
78
    uN1 = network_output(θ, layer_sizes, u0)
79
    θ[i] -= e
80
    return ((uN1 - y)^2 - (uN0 - y)^2)/e
81
end
82
83
# Setting up parameters
84
85
layer_sizes = [50,40,30,20,10,1]
86
N = length(layer_sizes) - 1
87
init(sizes...) = 0.01*randn(sizes...)
88
Ws = [init(layer_sizes[i+1], layer_sizes[i]) for i=1:N]
89
bs = [init(layer_sizes[i+1]) for i = 1:N]
90
u0 = init(layer_sizes[1])
91
layers = zip(Ws, bs)
92
93
# First we compute the forward pass
94
θvec = vcat([[W[:]; b] for (W, b) in layers]...)
95
96
## Testing the gradient
97
y = 3.0 # We aim to have the final output be 3.0
98
grad_adjoint = ∇L(θvec, layer_sizes, u0; y=y)
99
idx = 4000
100
grad_fd = fd(θvec, layer_sizes, u0, idx; e=1e-5, y=y)
101
@test grad_fd ≈ grad_adjoint[idx] atol=1e-6
102
103
# Optimizing to get the desired output y
104
for _ = 1:1000
105
    grad = ∇L(θvec, layer_sizes, u0; y=y)
106
    θvec .-= 0.001 .* grad
107
end
108
@test network_output(θvec, layer_sizes, u0) ≈ y
109
final_gradient = ∇L(θvec, layer_sizes, u0; y=y)
110
@test final_gradient ≈ zeros(length(θvec)) atol=1e-10

Cite this post

bibtex


@misc{paltorp2025adjointmethod,
  author = {Paltorp, Mikkel},
  title = {{Backpropagation is the Adjoint Method}},
  year = {2025},
  month = oct,
  url = {https://mipals.github.io/pdf/2025/10-25-adjoint-method/adjoint-method.pdf},
  note = {Blog post}
}

Bibliography

[1] P. Bright, A. Edelman, and S. G. Johnson, “Matrix Calculus (for Machine Learning and Beyond).” [Online]. Available: https://arxiv.org/abs/2501.14787
[2] A. Edelman, E. Akyürek, and Y. Wang, “Backpropagation through Back Substitution with a Backslash,” SIAM Journal on Matrix Analysis and Applications, vol. 45, no. 1, pp. 429–449, 2024, doi: 10.1137/22M1532871.