Batched Backpropagation as a Linear Solve

7 October 2026 · Download PDF · Download code

This post extends Backpropagation as a Linear Solve to batches of inputs. The main change is that the (sparse) matrices of scalars from the previous post become (sparse) matrices whose entries are linear maps.¹¹ More often than not the matrices are sparse—which is key for the efficiency of automatic differentiation. To readers used to the usual presentation of backpropagation as a pass over a computational graph, matrices may seem an odd choice. The two views are the same: the sparsity pattern below the diagonal of a lower triangular matrix $\mathrm{ℒ︀}$ is a Directed Acyclic Graph (DAG).

The post ends with a Julia implementation of a neural network whose gradient is ${\nabla }_{𝜃}\mathrm{𝒥︀}=\mathrm{ℳ︀}\left({\left({\mathrm{ℒ︀}}^{\ast }\right)}^{-1}{\nabla }_{𝑈}\ell \right)+{\nabla }_{𝜃}\ell$, the same formula as before, written as a linear solve in the style of Edelman, Akyürek, and Wang [1]. The inverse is never formed: because $\mathrm{ℒ︀}$ is unit lower triangular, ${\mathrm{ℒ︀}}^{\ast }$ is unit upper triangular, and its inverse can be applied efficiently using back substitution.

1. The batched forward pass

The previous post considered a single input ${𝑢}_0$ and a loss $\ell \left({𝑢}_{𝑃}\right)$ of its outputs. However, the loss function is often based on minimizing a data-fitting term of a minibatch of $𝑚$ input-output pairs (${𝑢}_0^{\left(𝑠\right)},{𝑦}^{\left(𝑠\right)}$). Due to the linearity of a sum the full gradient over the $𝑚$ input-output data could be computed by simply summing the gradient contribution for each example. In practice, it is often more efficient to process the data together in a batch.²² Batching turns matrix–vector products into matrix–matrix products, which use modern hardware far better. Moreover, because the weights are shared by all the data points, they only have to be loaded into memory once for a batch, whereas processing each example separately would load them once per example. The first step is to introduce the batched input and output matrices by stacking the data points as columns, i.e.

The only difference between this and the single input case is that the application of layer $𝑖$ now generates $𝑚$ outputs. Writing the outputs with a capital $𝑈$, as for the input data, layer $𝑖$ gives

We denote the collection of all the states by $𝑈\left(𝜃\right)$ and the collection of all the layer maps by $\mathrm{\Phi }\left(𝑈,𝜃\right)$. One can think of them as vectors whose entries are the states and the layer maps, respectively, so that

However, this does not have to be the case, and the only thing we require is that $𝑈\left(𝜃\right)$ is an element of a vector space $\mathrm{𝒰︀}$ with an inner product. Similarly, the parameters $𝜃$ live in a vector space $\mathrm{\Theta }$. The next section makes this precise.³³ See Wikipedia for what is meant by a vector space.

2. Generalization to vector spaces

It is standard to equip the space that each layer output lives in with the Frobenius inner product. This means that for $𝑍,𝐸\in \mathrm{𝒰︀}$ the inner product is a sum of Frobenius inner products over the layers, i.e.

The construction of the parameter space is slightly more complicated, as parameters can be shared between layers.⁴⁴ For example, in a recurrent neural network. The parameters can still be thought of as a vector whose entries are parameter blocks, that is, $𝜃={\left({𝜃}_{𝑘}\right)}_{𝑘=1}^{𝐾}\in \mathrm{\Theta }$. Layer $𝑖$ uses only some of the blocks, and several layers may use the same block. A linear selection map ${\mathrm{𝒮︀}}_{𝑖}:\mathrm{\Theta }\to {\mathrm{\Theta }}_{{\mathrm{𝒮︀}}_{𝑖}}$ picks out the blocks of layer $𝑖$, so that ${\mathrm{\Phi }}_{𝑖}$ depends on $𝜃$ only through ${𝜃}_{{\mathrm{𝒮︀}}_{𝑖}}={\mathrm{𝒮︀}}_{𝑖}\left(𝜃\right)$. In a plain feed-forward network every layer has its own weight block ${𝑊}_{𝑖,𝑖-1}$ and bias block ${𝑏}_{𝑖}$, and ${\mathrm{𝒮︀}}_{𝑖}\left(𝜃\right)=\left({𝑊}_{𝑖,𝑖-1},{𝑏}_{𝑖}\right)$. Similar to $\mathrm{𝒰︀}$, we equip $\mathrm{\Theta }$ with the Frobenius inner product block by block, so that for $ℎ,𝑔\in \mathrm{\Theta }$ it follows that

Lastly, we also aim to equip the output space $\mathrm{𝒴︀}$ with an inner product. We view elements in $\mathrm{𝒴︀}$ as a list of outputs, $𝐴=\left({𝐴}_1,\dots ,{𝐴}_{𝑟}\right)$, where again the inner product is a sum of Frobenius inner products over the list,

We then define a linear output map $\mathrm{𝒫︀}:\mathrm{𝒰︀}\to \mathrm{𝒴︀}$ such that it maps the full state $𝑈\in \mathrm{𝒰︀}$ to a corresponding output. That is $\mathrm{𝒫︀}\left(𝑈\right)\in \mathrm{𝒴︀}$ and we sometimes simply write ${𝑈}_{\mathrm{𝒫︀}}=\mathrm{𝒫︀}\left(𝑈\right)$.⁵⁵ Usually $\mathrm{𝒫︀}$ picks out the last layer, $\mathrm{𝒫︀}\left(𝑈\right)={𝑈}_{𝑁}$. It can also select several layers, or combine them, for example by averaging.

2.1. The loss

We often describe the loss function as a function of the parameters $𝜃$ only, but it is more natural to view it as a function of the states and the parameters. That is, $\mathrm{𝒥︀}\left(𝜃\right)=\ell \left(𝑈\left(𝜃\right),𝜃\right)$, where usually

where $𝑓$ is a data-fitting term that compares the outputs with the targets $𝑌$, and $\mathrm{ℛ︀}$ is a regularization term. The derivation below only needs $\ell$ and $\mathrm{\Phi }$ to be differentiable, which also makes $\mathrm{𝒥︀}$ differentiable. Only the gradients ${\nabla }_{𝑈}\ell$ and ${\nabla }_{𝜃}\ell$ of $\ell$ enter, as the right-hand side of the adjoint equation and as the last term of the gradient.

2.2. Adjoints

For a linear map $\mathrm{𝒜︀}:\mathrm{𝒳︀}\to \mathrm{𝒵︀}$ between two of these spaces, the adjoint ${\mathrm{𝒜︀}}^{\ast }:\mathrm{𝒵︀}\to \mathrm{𝒳︀}$ is defined by

Moving an operator across an inner product therefore changes the inner product from the operator’s output space to its input space.

2.3. Gradients

The gradient is defined through an inner product. For a differentiable scalar function of $𝜃$, such as the loss $\mathrm{𝒥︀}$, the derivative $𝐷\mathrm{𝒥︀}\left(𝜃\right)\in {\mathrm{\Theta }}^{\ast }$ is a linear map from directions $ℎ\in \mathrm{\Theta }$ to numbers.⁷⁷ ${\mathrm{\Theta }}^{\ast }$ denotes the dual space of $\mathrm{\Theta }$, which is the space of all linear functionals on $\mathrm{\Theta }$. This is a technicality that we revisit in Which Way Is Down? First-Order Optimizers as a Choice of Norm. The Riesz representation theorem⁸⁸ In a finite-dimensional inner product space every linear functional $𝜓$ has the form $𝜓\left(ℎ\right)=⟨𝑔,ℎ{⟩}_{\mathrm{\Theta }}$ for exactly one $𝑔\in \mathrm{\Theta }$, its Riesz representative. says that such a map is an inner product with a fixed element of $\mathrm{\Theta }$. That element is the gradient

So to find a gradient, we write the derivative in the form $⟨\cdot ,ℎ{⟩}_{\mathrm{\Theta }}$ and read off the other argument. The same definition applies on $\mathrm{𝒰︀}$ and $\mathrm{𝒴︀}$, each with its own inner product. Changing the inner product therefore changes the gradient, but not the derivative.

3. The batched gradient

Note that the parameters enter the loss function $\mathrm{𝒥︀}\left(𝜃\right)=\ell \left(𝑈\left(𝜃\right),𝜃\right)$ twice: directly as $𝜃$, and through the states $𝑈\left(𝜃\right)$. Differentiating in a parameter direction $ℎ\in \mathrm{\Theta }$, the chain rule gives one term for each,

In order to find the gradient ${\nabla }_{𝜃}\mathrm{𝒥︀}$, we need to write the first term in the form $⟨\cdot ,ℎ{⟩}_{\mathrm{\Theta }}$ so that we can read off the other argument—the Riesz representative of the derivative.

3.1. The forward relation

To relate ${𝑍}_{ℎ}$ to $ℎ$, differentiate the forward relation $𝑈\left(𝜃\right)=\mathrm{\Phi }\left(𝑈\left(𝜃\right),𝜃\right)$ in the direction $ℎ$. The left side gives ${𝑍}_{ℎ}$. On the right, $𝜃$ again enters both directly and through $𝑈\left(𝜃\right)$, and the chain rule gives

The first term accounts for changes through the states, while the second accounts for the explicit parameter dependence.¹⁰¹⁰ Both partial derivatives are evaluated at the current $\left(𝑈,𝜃\right)$. In ${𝐷}_{𝑈}\mathrm{\Phi }$, $𝜃$ is held fixed; in ${𝐷}_{𝜃}\mathrm{\Phi }$, $𝑈$ is held fixed. The total state dependence enters through ${𝑍}_{ℎ}$. Moving the state term to the left gives the linear equation

where $\mathrm{ℒ︀}$ and ${\mathrm{ℳ︀}}^{\ast }$ are the operator versions of $𝐿$ and ${𝑀}^{\top }$ from the previous post. Solving it, $\mathrm{ℒ︀}\left({𝑍}_{ℎ}\right)={\mathrm{ℳ︀}}^{\ast }\left(ℎ\right)$, gives ${𝑍}_{ℎ}$ for one direction $ℎ$. This is forward-mode differentiation, and the full gradient this way would need one solve per direction in a basis of $\mathrm{\Theta }$, that is, $dim\mathrm{\Theta }$ solves.

3.2. The adjoint equation

Solving (17) gives ${𝑍}_{ℎ}={\mathrm{ℒ︀}}^{-1}\left({\mathrm{ℳ︀}}^{\ast }\left(ℎ\right)\right)$. Substituting this into the first term of (15) and moving first ${\mathrm{ℒ︀}}^{-1}$, which maps $\mathrm{𝒰︀}$ to itself, and then ${\mathrm{ℳ︀}}^{\ast }$, which maps $\mathrm{\Theta }$ to $\mathrm{𝒰︀}$, across the inner product using their respective adjoints gives

where $\mathrm{\Lambda }={\left({\mathrm{ℒ︀}}^{-1}\right)}^{\ast }\left({\nabla }_{𝑈}\ell \right)$. Rather than forming the inverse, we compute $\mathrm{\Lambda }$ by solving the adjoint equation¹¹¹¹ For invertible $\mathrm{ℒ︀}$, the inverse and adjoint commute:$${\left({\mathrm{ℒ︀}}^{-1}\right)}^{\ast }={\left({\mathrm{ℒ︀}}^{\ast }\right)}^{-1}.$$Taking adjoints of ${\mathrm{ℒ︀}}^{-1}\mathrm{ℒ︀}=\mathrm{ℐ︀}$ reverses the order, giving ${\mathrm{ℒ︀}}^{\ast }{\left({\mathrm{ℒ︀}}^{-1}\right)}^{\ast }=\mathrm{ℐ︀}$; the reverse product is also the identity. Moving ${\mathrm{ℒ︀}}^{-1}$ across the inner product therefore gives a solve with ${\mathrm{ℒ︀}}^{\ast }$.

whose right-hand side does not depend on $ℎ$, so a single solve serves every direction. This is reverse-mode differentiation: as shown below, $\mathrm{ℒ︀}$ is lower block triangular, so ${\mathrm{ℒ︀}}^{\ast }$ is upper block triangular, and the solve runs from the last layer back to the first. Collecting the terms,

Reading off the Riesz representative of $𝐷\mathrm{𝒥︀}\left(𝜃\right)$ gives

3.3. The operators as block (sparse) matrices

Viewing the states and the parameters as lists, $\mathrm{ℒ︀}$ and ${\mathrm{ℳ︀}}^{\ast }$ can be written as block matrices whose entries are linear maps. For a state change $𝐸\in \mathrm{𝒰︀}$, layer $𝑖$ can only respond to the states that come strictly before it, so entry $𝑖$ of $\mathrm{ℒ︀}\left(𝐸\right)$¹²¹² Recall that $\mathrm{ℒ︀}:\mathrm{𝒰︀}\to \mathrm{𝒰︀}$ maps state changes to state changes, so $\mathrm{ℒ︀}\left(𝐸\right)$ is again a list with one entry per layer. is

so block $\left(𝑖,𝑗\right)$ of $\mathrm{ℒ︀}$ is nonzero only if layer $𝑖$ reads the state ${𝑈}_{𝑗}$. In the same way, for a parameter direction $ℎ\in \mathrm{\Theta }$, layer $𝑖$ only sees the parameters ${𝜃}_{{\mathrm{𝒮︀}}_{𝑖}}={\mathrm{𝒮︀}}_{𝑖}\left(𝜃\right)$, so by the chain rule entry $𝑖$ of ${\mathrm{ℳ︀}}^{\ast }\left(ℎ\right)$¹³¹³ Recall that ${\mathrm{ℳ︀}}^{\ast }:\mathrm{\Theta }\to \mathrm{𝒰︀}$ maps parameter changes to the state changes they cause directly: it takes a list over the $𝐾$ parameter blocks to a list over the $𝑁$ layers. Its adjoint $\mathrm{ℳ︀}:\mathrm{𝒰︀}\to \mathrm{\Theta }$ goes the other way. is

so block $\left(𝑖,𝑘\right)$ of ${\mathrm{ℳ︀}}^{\ast }$ is nonzero only if ${\mathrm{𝒮︀}}_{𝑖}$ selects the parameter block ${𝜃}_{𝑘}$. The indices of ${\mathrm{ℳ︀}}_{𝑘𝑖}^{\ast }$ are written in the order $𝑘𝑖$ because it is the adjoint of the block ${\mathrm{ℳ︀}}_{𝑘𝑖}$ of $\mathrm{ℳ︀}$, at position $\left(𝑘,𝑖\right)$, as shown below. Written out, with $\mathrm{𝒪︀}$ denoting the zero map,

Layer $𝑖$ only reads earlier states, so $\mathrm{ℒ︀}$ is lower block triangular with identities on the diagonal. In practice a layer reads only a few of the states, and most parameter blocks are used by a single layer, so most blocks are $\mathrm{𝒪︀}$, and $\mathrm{ℒ︀}$ and ${\mathrm{ℳ︀}}^{\ast }$ are sparse. Since the inner products on $\mathrm{𝒰︀}$ and $\mathrm{\Theta }$ are sums over the layers and the parameter blocks, the adjoints follow the same pattern as for matrices: transpose the block pattern and take the adjoint of each block, ${\left({\mathrm{ℒ︀}}^{\ast }\right)}_{𝑗𝑖}={\mathrm{ℒ︀}}_{𝑖𝑗}^{\ast }$.¹⁴¹⁴ This relies on the inner products treating the blocks separately. An inner product that couples the blocks, such as $⟨ℎ,\mathrm{𝒜︀}\left(𝑔\right){⟩}_{\mathrm{\Theta }}$ with the Fisher information $\mathrm{𝒜︀}$ of natural gradient descent, gives an adjoint that also involves $\mathrm{𝒜︀}$. In practice one keeps the separable inner product for the derivation and applies ${\mathrm{𝒜︀}}^{-1}$ to the resulting gradient afterwards. This gives

Since ${\mathrm{ℒ︀}}^{\ast }$ is upper block triangular, the adjoint equation is solved by back substitution from the last layer to the first. Both this solve and the application of $\mathrm{ℳ︀}$ only require the action of the adjoint blocks on the entries ${\mathrm{\Lambda }}_{𝑖}$.

4. A code example

The code is adapted from the Julia notebook Backprop with Backslash accompanying [1], which fixes a block bidiagonal $\mathrm{ℒ︀}$ and a block diagonal $\mathrm{ℳ︀}$. Here both are general sparse matrices of operators, so that a layer can read any earlier state and share parameter blocks with other layers.

We introduce a network of $𝑁=6$ layers, which produce the states ${𝑈}_1,\dots ,{𝑈}_6$ from the input ${𝑈}_0$, and include $𝑚=10$ data points. The layers are those of Example 4. Besides the chain weights ${𝑊}_{𝑖,𝑖-1}$, layers 3, 5, and 6 also read ${𝑈}_1$, ${𝑈}_2$, and ${𝑈}_3$ through the skip weights ${𝑊}_{31}$, ${𝑊}_{52}$, and ${𝑊}_{63}$, and layers 3 and 4 share both their chain weight, ${𝑊}_{32}={𝑊}_{43}$, and their bias, ${𝑏}_3={𝑏}_4$. With 8 weight blocks and 5 bias blocks, $𝐾=13$.¹⁷¹⁷ Each of the 6 layers has a chain weight and a bias, but layers 3 and 4 share theirs, which leaves $6-1=5$ pairs. Each of the 3 skip connections adds a weight block of its own, so $𝐾=2\left(6-1\right)+3=13$. The operator $\mathrm{ℒ︀}$ then has the pattern

Already here one can see that $\mathrm{ℒ︀}$ is sparse, with only 14 nonzero blocks out of 36, and only 8 of the 15 blocks below the diagonal. The deeper the network, the sparser it becomes, as each layer reads only a few of the earlier states.

Column $𝑖$ of $\mathrm{ℳ︀}$ has a block in row $𝑘$ for every parameter block ${𝜃}_{𝑘}$ that ${\mathrm{𝒮︀}}_{𝑖}$ selects.¹⁸¹⁸ We number the weights first, ${𝜃}_1={𝑊}_{10}$, ${𝜃}_2={𝑊}_{21}$, ${𝜃}_3={𝑊}_{32}={𝑊}_{43}$, ${𝜃}_4={𝑊}_{54}$, and ${𝜃}_5={𝑊}_{65}$ along the chain and ${𝜃}_6={𝑊}_{31}$, ${𝜃}_7={𝑊}_{52}$, and ${𝜃}_8={𝑊}_{63}$ for the skip connections, and the biases ${𝜃}_9,\dots ,{𝜃}_{13}$ last, so that, for example, ${\mathrm{𝒮︀}}_3\left(𝜃\right)=\left({𝜃}_3,{𝜃}_6,{𝜃}_{11}\right)$ and ${\mathrm{𝒮︀}}_4\left(𝜃\right)=\left({𝜃}_3,{𝜃}_{11}\right)$. This gives

The rows of the shared weight ${𝜃}_3$ and of the shared bias ${𝜃}_{11}$ each hold two blocks, so these two entries of $\mathrm{ℳ︀}\left(\mathrm{\Lambda }\right)$ sum the contributions of layers 3 and 4.

In the code, each block is an operator that stores the map and its adjoint, and $\mathrm{ℒ︀}$ and ${\mathrm{ℳ︀}}^{\ast }$ are sparse matrices of these operators. The only new methods are the back substitution with ${\mathrm{ℒ︀}}^{\ast }$ and the application of $\mathrm{ℳ︀}$, which both visit only the stored blocks. The script prints the patterns of $\mathrm{ℒ︀}$ and ${\mathrm{ℳ︀}}^{\ast }$ (the transpose of the pattern of $\mathrm{ℳ︀}$ above) and compares the gradient with a central difference along a random direction.

backbackback.jl
1
using LinearAlgebra, Random, SparseArrays
2
Random.seed!(42)
3
4
# A linear map stored together with its adjoint
5
struct Operator
6
    operator
7
    adjoint
8
    sym::String
9
end
10
Base.adjoint(A::Operator) = Operator(A.adjoint, A.operator, "$(A.sym)′")
11
Base.:*(A::Operator, X) = A.operator(X)
12
Base.show(io::IO, A::Operator) = print(io, A.sym)
13
Base.zero(::Type{Operator}) = Operator(zero, zero, "𝒪") # For printing
14
const ℐ = Operator(identity, identity, "ℐ")
15
16
# Adjoints of sparse matrices whose entries are operators.
17
const AdjointOperators = Adjoint{<:Any, <:SparseMatrixCSC{Operator}}
18
19
# Back substitution with ℒ*: Λ_j = g_j - ∑_{i>j} ℒ_ij* Λ_i.
20
function Base.:\(A::AdjointOperators, g::Vector)
21
    ℒ = parent(A) # Get the non-adjoint, but apply adjoint later
22
    rows, blocks = rowvals(ℒ), nonzeros(ℒ) # Column j: blocks[r], rows[r]
23
    Λ = copy(g)
24
    for j in length(g):-1:1, r in nzrange(ℒ, j)
25
        i = rows[r]
26
        if i > j
27
            Λ[j] -= blocks[r]' * Λ[i]
28
        end
29
    end
30
    return Λ
31
end
32
33
# ℳ = (ℳ′)*: entry k sums the contributions of the layers that use θ_k.
34
function Base.:*(A::AdjointOperators, Λ::Vector)
35
    ℳ′ = parent(A) # Get the non-adjoint, but apply adjoint later
36
    rows, blocks = rowvals(ℳ′), nonzeros(ℳ′)
37
    return [sum(blocks[r]' * Λ[rows[r]] for r in nzrange(ℳ′, k))
38
            for k in axes(ℳ′, 2)]
39
end
40
41
# Defining the activation function and its derivative
42
σ(x) = exp(-x)
43
σ′(x) = -exp(-x)
44
45
# Layer i reads the state U_j through the weight W_ij = θ_k for every
46
# j => k in W[i], and adds the bias b_i = θ_b[i]. Layers 3 and 4 share
47
# both their weight θ₃ and their bias θ₁₁.
48
d = [5, 4, 4, 4, 4, 3, 1] # Feature sizes of U₀, …, U₆
49
W = [[0 => 1], [1 => 2], [2 => 3, 1 => 6], [3 => 3], [4 => 4, 2 => 7],
50
     [5 => 5, 3 => 8]]
51
b = [9, 10, 11, 11, 12, 13]
52
N, K = 6, 13 # Layers, parameter blocks
53
54
function network(θ)
55
    U, Δ = Matrix{Float64}[], Matrix{Float64}[]
56
    state(j) = j == 0 ? U₀ : U[j]
57
    for i in 1:N
58
        Z = sum(θ[k] * state(j) for (j, k) in W[i]) .+ θ[b[i]]
59
        push!(U, σ.(Z))
60
        push!(Δ, σ′.(Z))
61
    end
62
    return U, Δ
63
end
64
65
# The sparse matrices ℒ = I - D_U Φ and ℳ′ = D_θ Φ, block by block.
66
function operators(θ, U, Δ)
67
    state(j) = j == 0 ? U₀ : U[j]
68
    ℒ, ℳ′ = spzeros(Operator, N, N), spzeros(Operator, N, K)
69
    for i in 1:N
70
        ℒ[i, i] = ℐ
71
        for (j, k) in W[i]
72
            Wij, Uj = θ[k], state(j)
73
            if j > 0
74
                ℒ[i, j] = Operator(E -> -(Δ[i] .* (Wij * E)),
75
                                   Λ -> -(Wij' * (Δ[i] .* Λ)), "ℒ")
76
            end
77
            ℳ′[i, k] = Operator(h -> Δ[i] .* (h * Uj),
78
                                Λ -> (Δ[i] .* Λ) * Uj', "ℳ′")
79
        end
80
        ℳ′[i, b[i]] = Operator(h -> Δ[i] .* h,
81
                               Λ -> vec(sum(Δ[i] .* Λ, dims=2)), "ℳ′")
82
    end
83
    return ℒ, ℳ′
84
end
85
86
# The data and the parameters.
87
m = 10 # Data points
88
α = 1e-2 # Weight of the regularization α/2 ⟨θ, θ⟩
89
init(dims...) = 0.01 * randn(dims...)
90
U₀, Y = init(d[1], m), init(d[end], m)
91
θ = Vector{Array{Float64}}(undef, K)
92
for i in 1:N
93
    for (j, k) in W[i]
94
        θ[k] = init(d[i+1], d[j+1])
95
    end
96
    θ[b[i]] = init(d[i+1])
97
end
98
99
# The gradient of J(θ) = ‖U_N - Y‖²/(2m) + α/2 ⟨θ, θ⟩.
100
J(θ) = norm(network(θ)[1][N] - Y)^2 / (2m) + α / 2 * dot(θ, θ)
101
U, Δ = network(θ)
102
ℒ, ℳ′ = operators(θ, U, Δ)
103
∇Uℓ = [i == N ? (U[N] - Y) / m : zero(U[i]) for i in 1:N] # 𝒫*(∇f)
104
Λ = ℒ' \ ∇Uℓ
105
∇J = ℳ′' * Λ + α * θ # ℳ(Λ) + αθ
106
display(ℒ)
107
display(ℳ′)
108
109
# Central difference of the loss along a random direction.
110
h = [randn(size(p)...) for p in θ]
111
ε = 1e-6
112
directional = (J(θ + ε * h) - J(θ - ε * h)) / (2ε)
113
versus_finite_difference = abs(directional - dot(∇J, h)) / abs(dot(∇J, h))
114
@show versus_finite_difference

5. Conclusion

Going from a single input to a batch only changed the entries of the matrices, from scalars to linear maps. The overall structure stayed the same: $\mathrm{ℒ︀}$ is a unit lower block triangular matrix whose sparsity pattern is a DAG, and both $\mathrm{ℒ︀}$ and $\mathrm{ℳ︀}$ are most often sparse. The formula for the gradient also kept the same form, ${\nabla }_{𝜃}\mathrm{𝒥︀}=\mathrm{ℳ︀}\left({\left({\mathrm{ℒ︀}}^{\ast }\right)}^{-1}{\nabla }_{𝑈}\ell \right)+{\nabla }_{𝜃}\ell$,¹⁹¹⁹ In the previous post the gradient reads ${\nabla }_{𝜃}\mathrm{𝒥︀}=𝑀{𝐿}^{-\top }{𝑃}^{\top }{\nabla }_{{𝑢}_{𝑃}}\ell +𝜇{\nabla }_{𝜃}\mathrm{ℛ︀}$. The matrices $𝑀$ and $𝐿$ become the operators $\mathrm{ℳ︀}$ and $\mathrm{ℒ︀}$, the transpose becomes the adjoint, ${𝑃}^{\top }{\nabla }_{{𝑢}_{𝑃}}\ell$ becomes ${\nabla }_{𝑈}\ell ={\mathrm{𝒫︀}}^{\ast }\left(\nabla 𝑓\right)$, and the regularization is part of ${\nabla }_{𝜃}\ell$. and can still be computed with a single back substitution and a single sparse product.

@misc{paltorp2026batchedbackpropagation,
  author = {Paltorp, Mikkel},
  title = {{Batched Backpropagation as a Linear Solve}},
  year = {2026},
  month = oct,
  url = {https://mipals.github.io/pdf/2026/10-07-batched-backpropagation/batched-backpropagation.pdf},
  note = {Blog post}
}