Cholesky numerical stability: Forward transform (#357)

* Numerical improvements to LKJCholesky forward transform

* Update implementation of Cholesky forward transform in rrule

* Fix FD Jacobian test

* Fix `_link_lkj_chol...` rrule test

* Fix ChainRules test (properly this time)

* Bump patch

Author: penelopeysm

Project.toml

@@ -1,6 +1,6 @@
 name = "Bijectors"
 uuid = "76274a88-744f-5084-9051-94815aaf08c4"
-version = "0.15.7"
+version = "0.15.8"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

src/bijectors/corr.jl

@@ -293,15 +293,15 @@ which is the above implementation.
 function _link_chol_lkj(W::AbstractMatrix)
     K = LinearAlgebra.checksquare(W)
 
-    y = similar(W) # z is also UpperTriangular. 
+    y = similar(W) # W is upper triangular.
     # Some zero filling can be avoided. Though diagnoal is still needed to be filled with zero.
 
     @inbounds for j in 1:K
-        remainder_sq = one(eltype(W))
-        for i in 1:(j - 1)
+        remainder_sq = W[j, j]^2
+        for i in (j - 1):-1:1
             z = W[i, j] / sqrt(remainder_sq)
-            y[i, j] = atanh(z)
-            remainder_sq -= W[i, j]^2
+            y[i, j] = asinh(z)
+            remainder_sq += W[i, j]^2
         end
         for i in j:K
             y[i, j] = 0
@@ -317,17 +317,18 @@ function _link_chol_lkj_from_upper(W::AbstractMatrix)
 
     y = similar(W, N)
 
-    idx = 1
+    starting_idx = 1
     @inbounds for j in 2:K
-        y[idx] = atanh(W[1, j])
-        idx += 1
-        remainder_sq = 1 - W[1, j]^2
-        for i in 2:(j - 1)
+        y[starting_idx] = atanh(W[1, j])
+        starting_idx += 1
+        remainder_sq = W[j, j]^2
+        for i in (j - 1):-1:2
+            idx = starting_idx + i - 2
             z = W[i, j] / sqrt(remainder_sq)
-            y[idx] = atanh(z)
-            remainder_sq -= W[i, j]^2
-            idx += 1
+            y[idx] = asinh(z)
+            remainder_sq += W[i, j]^2
         end
+        starting_idx += length((j - 1):-1:2)
     end
 
     return y

src/chainrules.jl

@@ -161,21 +161,23 @@ function ChainRulesCore.rrule(::typeof(_link_chol_lkj_from_upper), W::AbstractMa
     N = ((K - 1) * K) ÷ 2
 
     z = zeros(eltype(W), N)
-    tmp_vec = similar(z)
+    remainders = similar(z)
 
-    idx = 1
+    starting_idx = 1
     @inbounds for j in 2:K
-        z[idx] = atanh(W[1, j])
-        tmp = sqrt(1 - W[1, j]^2)
-        tmp_vec[idx] = tmp
-        idx += 1
-        for i in 2:(j - 1)
-            p = W[i, j] / tmp
-            tmp *= sqrt(1 - p^2)
-            tmp_vec[idx] = tmp
-            z[idx] = atanh(p)
-            idx += 1
+        z[starting_idx] = atanh(W[1, j])
+        remainder_sq = W[j, j]^2
+        starting_idx += 1
+        for i in (j - 1):-1:2
+            idx = starting_idx + i - 2
+            remainder = sqrt(remainder_sq)
+            remainders[idx] = remainder
+            zt = W[i, j] / remainder
+            z[idx] = asinh(zt)
+            remainder_sq += W[i, j]^2
         end
+        remainders[starting_idx - 1] = sqrt(remainder_sq)
+        starting_idx += length((j - 1):-1:2)
     end
 
     function pullback_link_chol_lkj_from_upper(Δz_thunked)
@@ -190,7 +192,7 @@ function ChainRulesCore.rrule(::typeof(_link_chol_lkj_from_upper), W::AbstractMa
             ΔW[j, j] = 0
             Δtmp = zero(eltype(Δz))
             for i in (j - 1):-1:2
-                tmp = tmp_vec[idx_up_to_prev_column + i - 1]
+                tmp = remainders[idx_up_to_prev_column + i - 1]
                 p = W[i, j] / tmp
                 ftmp = sqrt(1 - p^2)
                 d_ftmp_p = -p / ftmp
@@ -216,21 +218,23 @@ function ChainRulesCore.rrule(::typeof(_link_chol_lkj_from_lower), W::AbstractMa
     N = ((K - 1) * K) ÷ 2
 
     z = zeros(eltype(W), N)
-    tmp_vec = similar(z)
+    remainders = similar(z)
 
-    idx = 1
+    starting_idx = 1
     @inbounds for i in 2:K
-        z[idx] = atanh(W[i, 1])
-        tmp = sqrt(1 - W[i, 1]^2)
-        tmp_vec[idx] = tmp
-        idx += 1
-        for j in 2:(i - 1)
-            p = W[i, j] / tmp
-            tmp *= sqrt(1 - p^2)
-            tmp_vec[idx] = tmp
-            z[idx] = atanh(p)
-            idx += 1
+        z[starting_idx] = atanh(W[i, 1])
+        remainder_sq = W[i, i]^2
+        starting_idx += 1
+        for j in (i - 1):-1:2
+            idx = starting_idx + j - 2
+            remainder = sqrt(remainder_sq)
+            remainders[idx] = remainder
+            zt = W[i, j] / remainder
+            z[idx] = asinh(zt)
+            remainder_sq += W[i, j]^2
         end
+        remainders[starting_idx - 1] = sqrt(remainder_sq)
+        starting_idx += length((i - 1):-1:2)
     end
 
     function pullback_link_chol_lkj_from_lower(Δz_thunked)
@@ -245,7 +249,7 @@ function ChainRulesCore.rrule(::typeof(_link_chol_lkj_from_lower), W::AbstractMa
             ΔW[i, i] = 0
             Δtmp = zero(eltype(Δz))
             for j in (i - 1):-1:2
-                tmp = tmp_vec[idx_up_to_prev_row + j - 1]
+                tmp = remainders[idx_up_to_prev_row + j - 1]
                 p = W[i, j] / tmp
                 ftmp = sqrt(1 - p^2)
                 d_ftmp_p = -p / ftmp

test/ad/chainrules.jl

@@ -1,3 +1,5 @@
+using Random: Xoshiro
+using LinearAlgebra
 using ChainRulesTestUtils: ChainRulesCore
 
 # HACK: This is a workaround to test `Bijectors._inv_link_chol_lkj` which produces an
@@ -77,30 +79,104 @@ end
     test_rrule(Bijectors._transform_inverse_ordered, b(rand(5, 2)))
 
     # LKJ and LKJCholesky bijector
-    dist = LKJCholesky(3, 4)
     # Run multiple tests because we're working with `undef` entries, and so we
     # want to make sure that we hit cases where the `undef` entries have different values.
     # It's also just useful to test numerical stability for different realizations of `dist`.
+
+    # NOTE(penelopeysm): https://github.com/TuringLang/Bijectors.jl/pull/357
+    # changed the implementation of _link_chol_lkj... to improve its numerical stability.
+    # The new implementation relies on the fact that the `LKJCholesky` distribution
+    # yields samples for which each column is a unit vector. Naively using FiniteDifferences
+    # to calculate a JVP (as ChainRulesTestUtils.test_rrule does) does not work, because
+    # FD does not know about this constraint.
+    # To solve this, we run the FD part of the test with the inputs projected onto a
+    # subspace that has that constraint encoded. We have to then recover the original
+    # output by un-projecting.
+    # The PR linked above has a more detailed explanation.
     for i in 1:30
-        x = rand(dist)
-        test_rrule(
-            Bijectors._link_chol_lkj_from_upper,
-            x.U;
-            testset_name="_link_chol_lkj_from_upper on $(typeof(x)) [$i]",
-        )
-        test_rrule(
-            Bijectors._link_chol_lkj_from_lower,
-            x.L;
-            testset_name="_link_chol_lkj_from_lower on $(typeof(x)) [$i]",
-        )
+        dist = LKJCholesky(3, 4)
+        rng = Xoshiro(i)
+        spl = rand(rng, dist)
 
-        b = bijector(dist)
-        y = b(x)
+        @testset "_inv_link_chol_lkj" begin
+            # This one doesn't need the fancy projection bits, so we can just
+            # use test_rrule as usual.
+            x = spl
+            b = bijector(dist)
+            y = b(x)
+            test_rrule(
+                _inv_link_chol_lkj_wrapper,
+                y;
+                testset_name="_inv_link_chol_lkj on $(typeof(x)) [$i]",
+            )
+        end
 
-        test_rrule(
-            _inv_link_chol_lkj_wrapper,
-            y;
-            testset_name="_inv_link_chol_lkj on $(typeof(x)) [$i]",
-        )
+        # Set up a random tangent.
+        ybar = rand(rng, 3) * 10
+        fdm = FiniteDifferences.central_fdm(5, 1)
+
+        # Functions to convert input to/from free parameters
+        to_free_params(x::UpperTriangular) = [x[1, 2], x[1, 3], x[2, 3]]
+        to_free_params(x::LowerTriangular) = [x[2, 1], x[3, 1], x[3, 2]]
+        function from_x_free(x_free::AbstractVector, uplo::Symbol)
+            x = UpperTriangular(zeros(eltype(x_free), 3, 3))
+            x[1, 1] = 1
+            x[1, 2] = x_free[1]
+            x[1, 3] = x_free[2]
+            x[2, 2] = sqrt(1 - x_free[1]^2)
+            x[2, 3] = x_free[3]
+            x[3, 3] = sqrt(1 - x_free[2]^2 - x_free[3]^2)
+            return uplo == :U ? x : transpose(x)
+        end
+        # Function to reconvert the adjoint back into a triangular matrix
+        function fd_xbar_to_cr_xbar(fd_xbar::AbstractVector, uplo::Symbol)
+            x = UpperTriangular(zeros(eltype(fd_xbar), 3, 3))
+            x[1, 2] = fd_xbar[1]
+            x[1, 3] = fd_xbar[2]
+            x[2, 3] = fd_xbar[3]
+            return uplo == :U ? x : transpose(x)
+        end
+
+        @testset "_link_chol_lkj_from_upper" begin
+            f = Bijectors._link_chol_lkj_from_upper
+            x = spl.U
+
+            # test primal is accurate
+            y = f(x)
+            cr_y, cr_pullback = ChainRulesCore.rrule(f, x)
+            @test isapprox(y, cr_y)
+
+            # test that the primal still works when going via free parameters
+            f_via_free(x_free::AbstractVector) = f(from_x_free(x_free, :U))
+            x_free = to_free_params(x)
+            y_via_free = f_via_free(x_free)
+            @test isapprox(y, y_via_free)
+
+            # test pullback
+            cr_xbar = cr_pullback(ybar)[2]
+            fd_xbar = FiniteDifferences.j′vp(fdm, f_via_free, ybar, x_free)[1]
+            @test isapprox(cr_xbar, fd_xbar_to_cr_xbar(fd_xbar, :U))
+        end
+
+        @testset "_link_chol_lkj_from_lower" begin
+            f = Bijectors._link_chol_lkj_from_lower
+            x = spl.L
+
+            # test primal is accurate
+            y = f(x)
+            cr_y, cr_pullback = ChainRulesCore.rrule(f, x)
+            @test isapprox(y, cr_y)
+
+            # test that the primal still works when going via free parameters
+            f_via_free(x_free::AbstractVector) = f(from_x_free(x_free, :L))
+            x_free = to_free_params(x)
+            y_via_free = f_via_free(x_free)
+            @test isapprox(y, y_via_free)
+
+            # test pullback
+            cr_xbar = cr_pullback(ybar)[2]
+            fd_xbar = FiniteDifferences.j′vp(fdm, f_via_free, ybar, x_free)[1]
+            @test isapprox(cr_xbar, fd_xbar_to_cr_xbar(fd_xbar, :L))
+        end
     end
 end

test/transform.jl

@@ -215,19 +215,42 @@ end
 end
 
 @testset "LKJCholesky" begin
+    # Convert Cholesky factor to its free parameters, i.e. its off-diagonal elements
+    function chol_3by3_to_free_params(x::Cholesky)
+        if x.uplo == :U
+            return [x.U[1, 2], x.U[1, 3], x.U[2, 3]]
+        else
+            return [x.L[2, 1], x.L[3, 1], x.L[3, 2]]
+        end
+    end
+
+    # Reconstruct Cholesky factor from its free parameters
+    # Note that x[i, i] is always positive so we don't need to worry about the sign
+    function free_params_to_chol_3by3(free_params::AbstractVector, uplo::Symbol)
+        x = UpperTriangular(zeros(eltype(free_params), 3, 3))
+        x[1, 1] = 1
+        x[1, 2] = free_params[1]
+        x[1, 3] = free_params[2]
+        x[2, 2] = sqrt(1 - free_params[1]^2)
+        x[2, 3] = free_params[3]
+        x[3, 3] = sqrt(1 - free_params[2]^2 - free_params[3]^2)
+        if uplo == :U
+            return Cholesky(x)
+        else
+            return Cholesky(transpose(x))
+        end
+    end
+
     @testset "uplo: $uplo" for uplo in [:L, :U]
         dist = LKJCholesky(3, 1, uplo)
         single_sample_tests(dist)
         x = rand(dist)
-        J = ForwardDiff.jacobian(z -> link(dist, Cholesky(z, x.uplo, x.info)), x.UL)
-        # Remove columns of Jacobian that are all zero (i.e. those
-        # corresponding to entries above the diagonal for uplo = :U, or below
-        # the diagonal for uplo = :L). This slightly unscientific approach
-        # based on filter() is needed to handle both ForwardDiff 0.10 and 1 as
-        # the exact indices will differ for the two versions; see
-        # https://github.com/JuliaDiff/ForwardDiff.jl/issues/738.
-        inds = filter(i -> !all(iszero, J[:, i]), 1:size(J, 2))
-        J = J[:, inds]
+        # Here, we need to pass ForwardDiff only the free parameters of the
+        # Cholesky factor so that we get a square Jacobian matrix
+        free_params = chol_3by3_to_free_params(x)
+        J = ForwardDiff.jacobian(
+            z -> link(dist, free_params_to_chol_3by3(z, uplo)), free_params
+        )
         logpdf_turing = logpdf_with_trans(dist, x, true)
         @test logpdf(dist, x) - _logabsdet(J) ≈ logpdf_turing
     end