By: Random blog posts about machine learning in Julia
Re-posted from: https://rkube.github.io/jekyll/update/2021/07/06/backpropagating-through-qr.html
One of the most useful decompositions of Linear Algebra is the QR decomposition.
This decomposition is particularly important when we are interested in the
vector space spanned by the columns of a matrix A. Formally, we write the QR
decomposition of \(A \in \mathbb{R}^{m \times n}\), where m ≥ n, as
\[\begin{align}
A = QR =
\left[
\begin{array}{c|c|c|c}
& & & \\
q_1 & q_2 & \cdots & q_n \\
& & &
\end{array}
\right]
%
\left[
\begin{array}{cccc}
r_{1,1} & r_{1,2} & \cdots & r_{1,n} \\
0 & r_{2,2} & \cdots & r_{2,n} \\
0 & 0 & \ddots & \vdots \\
0 & \cdots & & r_{n,n}
\end{array}
\right].
\end{align}\]
Where Q is an orthogonal m-by-n matrix, i.e. the columns of Q are orthogonal
\(q_{i} \cdot q_{j} = \delta_{i,j}\). R is a square n-by-n matrix.
Using the decomposition, we can reconstruct the columns of A as
\[\begin{align}
a_1 & = r_{1,1} q_1 \\
a_2 & = r_{1,2} q_1 + r{2,2} q_2 \\
& \cdots \\
a_n & = \sum_{j=1}^{n} r_{j,n} q_j.
\end{align}\]
Internally, Julia treats QR-factorized matrices through a packed format [1].
This format does not store the matrices Q and R explicitly, but using a packed format.
All matrix algebra and arithmetic is implemented through methods that expect
the packed format as input. And the QR decomposition itself is calling the LAPACK
method
geqrt which
returns just this packed format.
A computational graph where one may want to backpropagate through a QR factorization
may look similar to this one:
While the incoming gradients, \(\bar{f} \partial f / \partial Q\) and
\(\bar{f} \partial f / \partial R\) depend on \(f\), the gradients for the QR
decomposition \(\bar{Q} \partial Q / \partial A\) and \(\bar{R} \partial R / \partial A\) are defined through the QR factorization. While these can be in principle
computed through an automatic differentiation framework, it can be beneficial to
implement a pullback. For one, this saves compilation time before first execution.
An additional benefit is less memory use, as the gradients will be propagated through
fewer functions. A pullback for the QR factorization may thus also aid numerical
stability, as fewer accumulations and propagations are performed.
Formulas for the pullback of the QR factorization are given in [2],
[3] and [4]. Pytorch for example implements the method described in [3], see here.
In this blog post, we implement the pullback for the QR factorization in Julia.
Specifically, we implement a so-called rrule for ChainRules.jl. While you are here, please take a moment and read
the documentation of this package. It is very well written and helped me tremendously
to understand how automatic differentiation works and is implemented n Julia.
Also, if you want to skip ahead, here is my example implementation of the
QR pullback for ChainRules.
Now, let’s take a look at the code and how to implement a custom rrule. The first
thing we need to look at is the definition of the QR struct
struct QR{T,S<:AbstractMatrix{T}} <: Factorization{T}
factors::S
τ::Vector{T}
function QR{T,S}(factors, τ) where {T,S<:AbstractMatrix{T}}
require_one_based_indexing(factors)
new{T,S}(factors, τ)
end
end
There are two fields in the structure, factors and τ. The matrices Q and R are
returned through an accompanying getproperty function:
function getproperty(F::QR, d::Symbol)
m, n = size(F)
if d === :R
return triu!(getfield(F, :factors)[1:min(m,n), 1:n])
elseif d === :Q
return QRPackedQ(getfield(F, :factors), F.τ)
else
getfield(F, d)
end
end
For the QR pullback we first need to implement a pullback for the getproperty
function. The pullback for this function only propagates incoming gradients
backwards. Incoming gradients are described using a Tangent. This struct can only
have fields that are present in the parameter type, in our case
LinearAlgebra.QRCompactWY. This struct has fields factors and τ as discussed
above. Now the incoming gradients would be \(\bar{Q}\) and \(\bar{R}\). Thus
the pullback needs to map between these two. Thus the pullback can be implemented
like this:
function ChainRulesCore.rrule(::typeof(getproperty), F::LinearAlgebra.QRCompactWY, d::Symbol)
function getproperty_qr_pullback(Ȳ)
∂factors = if d === :Q
Ȳ
else
nothing
end
∂T = if d === :R
Ȳ
else
nothing
end
∂F = Tangent{LinearAlgebra.QRCompactWY}(; factors=∂factors, T=∂T)
return (NoTangent(), ∂F)
end
return getproperty(F, d), getproperty_qr_pullback
end
Notice the call signature of the rrule. The first argument is always
::typeof(functionname), where functionname is the name of the function that
we want a custom rrule for. Following this argument are the actual arguments
that one normally passes to functionname.
Finally we need to implement the pullback that actually performs the relevant
calculations. The typical way of implementing custom pullbacks with
ChainRules is write a function
that calculates that returns a tuple containing the result of the forward pass,
as well as the pullback \(\mathcal{B}^{x}_{f}(\bar{y})\). Doing it this way allows
the pullback to re-use results from the forward pass. In other words, by defining
the pullback as a function in the forward pass allows to easily use cached results.
There is more discussion on the reason behind this design choice in the
ChainRules documentation.
Returning to the pullback for the QR factorization, here is a possible implementation:
function ChainRules.rrule(::typeof(qr), A::AbstractMatrix{T}) where {T}
QR = qr(A)
m, n = size(A)
function qr_pullback(Ȳ::Tangent)
function qr_pullback_square_deep(Q̄, R̄, A, Q, R)
M = R̄*R' - Q'*Q̄
# M <- copyltu(M)
M = triu(M) + transpose(triu(M,1))
Ā = (Q̄ + Q * M) / R'
end
Q̄ = Ȳ.factors
R̄ = Ȳ.T
Q = QR.Q
R = QR.R
if m ≥ n
Q̄ = Q̄ isa ChainRules.AbstractZero ? Q̄ : @view Q̄[:, axes(Q, 2)]
Ā = qr_pullback_square_deep(Q̄, R̄, A, Q, R)
else
# partition A = [X | Y]
# X = A[1:m, 1:m]
Y = A[1:m, m + 1:end]
# partition R = [U | V], and we don't need V
U = R[1:m, 1:m]
if R̄ isa ChainRules.AbstractZero
V̄ = zeros(size(Y))
Q̄_prime = zeros(size(Q))
Ū = R̄
else
# partition R̄ = [Ū | V̄]
Ū = R̄[1:m, 1:m]
V̄ = R̄[1:m, m + 1:end]
Q̄_prime = Y * V̄'
end
Q̄_prime = Q̄ isa ChainRules.AbstractZero ? Q̄_prime : Q̄_prime + Q̄
X̄ = qr_pullback_square_deep(Q̄_prime, Ū, A, Q, U)
Ȳ = Q * V̄
# partition Ā = [X̄ | Ȳ]
Ā = [X̄ Ȳ]
end
return (NoTangent(), Ā)
end
return QR, qr_pullback
end
The part of the rrule that calculates the foward pass just calculates the QR
decomposition and returns the result of that. The pullback consumes the
incoming gradients, \(\bar{R}\) and \(\bar{Q}\), which are stored as defined
in the pullback for getproperty. The actual calculation for the case where
the dimensions of the matrix A are equal, m=n, and the case where A is tall and
skinny, m>n are different. But they can re-use some code. That is what the
local function qr_pullback_square_deep does. The rest is mostly matrix slicing
and the occasional matrix product. The implementation is borrows strongly from
pytorch’s autograd implementation.
Finally, I need to verify that my implemenation calculates correct results. For
this, I define two small test functions which each work on exactly one output of the
QR factorization, either the matrix Q or the matrix R. Then I generate some random
input and compare the gradient calculated through reverse-mode AD in Zygote
to the gradient calculated by FiniteDifferences. If they are approximately equal,
I take the result to be correct.
V1 = rand(Float32, (4, 4));
function f1(V) where T
Q, _ = qr(V)
return sum(Q)
end
function f2(V) where T
_, R = qr(V)
return sum(R)
end
res1_V1_ad = Zygote.gradient(f1, V1)
res1_V1_fd = FiniteDifferences.grad(central_fdm(5,1), f1, V1)
@assert res1_V1_ad[1] ≈ res1_V1_fd[1]
[1]
Schreiber et al. A Storage-Efficient WY Representation for Products of Householder Transformations
[2]
HJ Liao, JG Liu et al. Differentiable Programming Tensor Networks
[3]
M. Seeger, A. Hetzel et al. Auto-Differenting Linear Algebra
[4]
Walter and Lehman Walter and Lehmann, 2018, Algorithmic Differentiation of Linear Algebra Functions with Application in Optimum Experimental Design
