Implement QR pullback #306

Kolaru · 2020-11-15T13:13:18Z

Implement the pullback for the QR decomposition, following:

Walter and Lehmann, 2018, Algorithmic Differentiation of Linear Algebra Functions with Application in Optimum Experimental Design

oxinabox · 2020-11-16T12:40:17Z

src/rulesets/LinearAlgebra/factorization.jl

+    if size(F.R, 2) != size(F.R, 1)
+        throw(ArgumentError("Pullback for QR decomposition is only supported for m × n matrices with m >= n"))
+    end


Because I have not found a reference for that case.

Section 3.2 of https://arxiv.org/pdf/2009.10071.pdf gives an rrule for qr for wide matrices (m<n). I haven't tested it, though.

I will look into adding that.

oxinabox · 2020-11-16T12:44:16Z

test/rulesets/LinearAlgebra/factorization.jl

+
+                    # Explicitely convert to Matrix since FiniteDifferences seem to
+                    # be broken for LinearAlgebra.QRCompactWYQ (infinite to_vec
+                    # recursion)


Can an issue openned for this an linked here?

Done here JuliaDiff/FiniteDifferences.jl#114

oxinabox · 2020-11-16T12:48:46Z

test/rulesets/LinearAlgebra/factorization.jl

+
+                F, dX_pullback = rrule(qr, X)
+                for p in [:Q, :R]
+                    Y, dF_pullback = rrule(getproperty, F, p)


does rrule_test not work on this?
If it doesn't an issue needs to be openned on ChainRulesTestUtils.jl

It does not work on getproperty because of the failure of FiniteDifferences.jl mentionned elsewhere on this PR.

It does not work directly on qr either, because the object returned by qr can not be collected. See JuliaDiff/ChainRulesTestUtils.jl#74

oxinabox · 2020-11-16T12:49:20Z

test/rulesets/LinearAlgebra/factorization.jl

@@ -138,4 +138,40 @@ using ChainRules: level2partition, level3partition, chol_blocked_rev, chol_unblo
            @test chol_unblocked_rev(B̄, B, true) ≈ chol_blocked_rev(B̄, B, 10, true)
        end
    end
+    @testset "qr" begin
+        @testset "the thing" begin


Please explain

I just reused the same structure as the test for cholesky. I can give it a more explicit name.

oxinabox · 2020-11-16T12:51:38Z

src/rulesets/LinearAlgebra/factorization.jl

+        C = Composite{T}
+        ∂F = if x === :Q
+            C(Q=Ȳ,)
+        elseif x === :R
+            C(R=Ȳ,)
+        end


Is it worth defining C ?

Suggested change

C = Composite{T}

∂F = if x === :Q

C(Q=Ȳ,)

elseif x === :R

C(R=Ȳ,)

end

∂F = if x === :Q

Composite{T}(Q=Ȳ,)

elseif x === :R

Composite{T}(R=Ȳ,)

end

I did it analog to the svd case. I think I could go even simpler with

∂F = Composite{T}(; x => Ȳ)

Should I port that kind of cleanup to svd and cholesky as well?

oxinabox · 2020-11-16T12:52:10Z

src/rulesets/LinearAlgebra/utils.jl

@@ -27,3 +27,17 @@ function _eyesubx!(X::AbstractMatrix)
 end

 _extract_imag(x) = complex(0, imag(x))
+
+# Lower triangle of X - X' overwrite X if possible


this seems to always overwrite?

src/rulesets/LinearAlgebra/utils.jl

oxinabox · 2020-11-16T12:53:41Z

src/rulesets/LinearAlgebra/factorization.jl

+    if x === :Q
+        # Return thing Q for consistency
+        n = size(F.R, 1)
+        return F.Q[:, 1:n], getproperty_qr_pullback


Is this on to do to the primal value?

This was because the anything else does not contribute, and the tests used the result of this to infer dimensions. However after pondering on it I figured out it is probably better to let the getproperty pullback do just the obvious and have the qr pullback make sure to only take in account what is relevant and make sure the dimensions match.

oxinabox · 2020-11-16T12:53:58Z

src/rulesets/LinearAlgebra/factorization.jl

+        return NO_FIELDS, ∂F, DoesNotExist()
+    end
+    if x === :Q
+        # Return thing Q for consistency


I was very confused

Suggested change

# Return thing Q for consistency

# Return thin Q for consistency

sethaxen · 2020-11-21T05:28:41Z

I wonder if the QR rule implied by Seeger et al in https://arxiv.org/pdf/1710.08717.pdf is more performant than the one in Walter and Lehmann? (they actually define an LQ rule, but the same approach produces a QR rule). The below reimplementation of this PR's qr_rev using this rule seems to outperform qr_rev in a simple benchmark while yielding the same result.

function qr_rev2(QR_::ChainRules.QR_TYPE, Q̄, R̄)
    Q, R = QR_
    Q = Matrix(Q)
    Q̄ = Q̄ isa Zero ? Q̄ : @view Q̄[:, axes(Q, 2)]
    V = R̄*R' - Q'*Q̄
    Ā = (Q̄ + Q * Hermitian(V)) / R'
    return Ā
end

julia> A = randn(4, 4);

julia> F = qr(A);

julia> ΔF = Composite{typeof(F)}(Q = randn(eltype(F.Q), size(Matrix(F.Q))), R = randn(eltype(F.R), size(F.R)));

julia> @btime ChainRules.qr_rev($F, $(ΔF.Q), $(ΔF.R));
  16.374 μs (58 allocations: 7.11 KiB)

julia> @btime qr_rev2($F, $(ΔF.Q), $(ΔF.R));
  4.323 μs (12 allocations: 2.27 KiB)

julia> ChainRules.qr_rev(F, ΔF.Q, ΔF.R) ≈ qr_rev2(F, ΔF.Q, ΔF.R)
true

julia> A = randn(10, 4);

julia> F = qr(A);

julia> ΔF = Composite{typeof(F)}(Q = randn(eltype(F.Q), size(Matrix(F.Q))), R = randn(eltype(F.R), size(F.R)));

julia> @btime ChainRules.qr_rev($F, $(ΔF.Q), $(ΔF.R));
  56.299 μs (134 allocations: 20.86 KiB)

julia> @btime qr_rev2($F, $(ΔF.Q), $(ΔF.R));
  5.112 μs (12 allocations: 3.39 KiB)

julia> ChainRules.qr_rev(F, ΔF.Q, ΔF.R) ≈ qr_rev2(F, ΔF.Q, ΔF.R)
true

Or is the rule you have implemented expected to be more numerically stable?

Kolaru · 2020-11-22T19:32:34Z

Or is the rule you have implemented expected to be more numerically stable?

I have no idea. The two references we are using do not directly compare each other, and I do not know how to determine this myself.

sethaxen · 2021-02-04T21:54:48Z

Or is the rule you have implemented expected to be more numerically stable?

I have no idea. The two references we are using do not directly compare each other, and I do not know how to determine this myself.

The article I linked in #306 (comment) (https://arxiv.org/pdf/2009.10071.pdf), which covers wide and tall matrices as well, also uses the simpler rule from the Seeger et al paper. While I didn't implement their rules for wide and tall matrices, I ended up using a similar approach for the LU decomposition of wide and tall matrices in #354. For these reasons, I'm thinking the Seeger approach is preferable.

Kolaru · 2021-07-22T02:13:43Z

I finally had time to come back to this, and it has been kind of a nightmare, because QR decompositions are represented in a weird way that do not play nicely with the tests and the comparison with FiniteDifferences.

After quite a lot of experimentations, I gave up on trying to make everything work with the default type returned by qr. Instead I define a new custom type ExplicitQR that stores the Q and R matrices explicitely. I use this struct for the tests, and compare the end result with the one from the qr method to ensure the latter is correct as well.

I hope this is sufficient. Otherwise I must admit I am out of idea about what should be done to test qr directly.

I am aware of #469 that has a somewhat different approach. I am not currently sure which is better.

Alos I implemented the algorithm suggested by @sethaxen.

So provided my way of testing is okay, this should be ready.

sethaxen · 2021-07-22T20:24:08Z

Thanks, @Kolaru! I'll review #469 first, then this, then I'll recommend how to proceed.

sethaxen · 2021-10-08T22:37:53Z

What I like about this approach is that it completely sidesteps much of the complexity of the objects returned by the qr methods, as described in #469 (comment).

What I don't like is that the object being returned by the rrule is completely different from the one the user requested, which will cause some code that worked before to suddenly fail. For example, this code is totally valid with the qr return values in the std library:

A = randn(10, 5)
Q, _ = qr(A)
v = randn(5)
w = randn(10)
y = Q*w + Q*v
@assert size(y) == (10,)

This works because AbstractQ objects are treated as thin or full depending on what they're multiplied by, but if Q is dense, then this no longer works.

Kolaru · 2021-10-09T11:19:14Z

I tried the following

using ChainRules: rrule, ExplicitQR
using LinearAlgebra

A = randn(10, 5)
Q, _ = qr(A)
v = randn(5)
w = randn(10)

F, F_pullback = rrule(qr, A)
Q, Q_pullback = rrule(getproperty, F, :Q)
y1, y1_pullback = rrule(*, Q, v)
y2, y2_pullback = rrule(*, Q, w)

ȳ1 = rand(10)
_, Q̄1 = y1_pullback(ȳ1)
_, F̄1 = Q_pullback(Q̄1)
_, Ā1 = F_pullback(F̄1) 

ȳ2 = rand(10)
_, Q̄2 = y2_pullback(ȳ2)
_, F̄2 = Q_pullback(Q̄2)
_, Ā2 = F_pullback(F̄2)

and everything seems to be fine (i.e. nothing error, I haven't tested correctness). Is this the correct way to test your point? Whether this PR or #469 is used, adding proper test for quirks of Q seems unavoidable.

Note that the ExplicitQR object I define is only used in the test (and maybe its definition should be moved to the tests for clarity).

As far as I understand, ChainRules should be able to handle the QR objects properly. The problem I had was to create the tangent objects for the tests. The main issues is that we are taking derivative with respect to fields that are not stored in the object directly[*]. After some tries, I just gave up and sidestepped the whole testing issue, hoping this is still enough to ensure correctness.

[*] The most common error I had was:

ArgumentError: Tangent fields do not match primal fields.
Tangent fields: (:Q,). Primal (LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}}) fields: (:factors, :T)

Using a custom struct with explicit fields Q and R naturally solved it.

oxinabox reviewed Nov 16, 2020

View reviewed changes

src/rulesets/LinearAlgebra/utils.jl Outdated Show resolved Hide resolved

oxinabox reviewed Nov 16, 2020

View reviewed changes

Kolaru mentioned this pull request Nov 16, 2020

Stackoverflow with QR JuliaDiff/FiniteDifferences.jl#114

Closed

Kolaru mentioned this pull request Dec 25, 2020

Implement to_vec for matrix factorizations JuliaDiff/FiniteDifferences.jl#128

Merged

Kolaru marked this pull request as draft June 29, 2021 16:05

Rebase and try to make work

430adea

Kolaru force-pushed the qr_adjoint branch from 771ba9f to 430adea Compare June 29, 2021 16:15

Kolaru added 2 commits July 22, 2021 01:31

Use explicit QR struct for the tests

b985f41

Support for wide matrices

61db5ad

Kolaru marked this pull request as ready for review July 22, 2021 02:13

oxinabox requested a review from sethaxen July 22, 2021 10:28

Fix small bug when Rbar is ZeroTangent

e03fce8

sethaxen mentioned this pull request Jul 21, 2022

rrule for QR-decomposition #651

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement QR pullback #306

Implement QR pullback #306

Kolaru commented Nov 15, 2020

oxinabox Nov 16, 2020

Kolaru Nov 16, 2020

sethaxen Nov 21, 2020

Kolaru Nov 22, 2020

oxinabox Nov 16, 2020

Kolaru Nov 16, 2020

oxinabox Nov 16, 2020 •

edited

Loading

Kolaru Nov 16, 2020

oxinabox Nov 16, 2020

Kolaru Nov 16, 2020

oxinabox Nov 16, 2020

oxinabox Nov 16, 2020

Kolaru Nov 16, 2020

oxinabox Nov 16, 2020

oxinabox Nov 16, 2020

Kolaru Nov 16, 2020

oxinabox Nov 16, 2020

sethaxen commented Nov 21, 2020

Kolaru commented Nov 22, 2020

sethaxen commented Feb 4, 2021

Kolaru commented Jul 22, 2021

sethaxen commented Jul 22, 2021

sethaxen commented Oct 8, 2021

Kolaru commented Oct 9, 2021

	# Return thing Q for consistency
	# Return thin Q for consistency

Implement QR pullback #306

Are you sure you want to change the base?

Implement QR pullback #306

Conversation

Kolaru commented Nov 15, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

oxinabox Nov 16, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sethaxen commented Nov 21, 2020

Kolaru commented Nov 22, 2020

sethaxen commented Feb 4, 2021

Kolaru commented Jul 22, 2021

sethaxen commented Jul 22, 2021

sethaxen commented Oct 8, 2021

Kolaru commented Oct 9, 2021

oxinabox Nov 16, 2020 •

edited

Loading