Update SVD reverse-rule broadening (#194)

- Actually apply broadening in SVD rrule for (quasi) degenerate singular values instead of just using a cutoff
- Add a test for differentiating SVDs with degenerate singular values

---------

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

Author: pbrehmer

src/utility/svd.jl

@@ -438,13 +438,21 @@ function ChainRulesCore.rrule(
     return (U, S, V, info), tsvd!_itersvd_pullback
 end
 
-# scalar inverses with a cutoff tolerance and Lorentzian broadening
-function _safe_inv(x, tol, ε=0)
-    if abs(x) < tol
-        return zero(x)
-    else
-        return iszero(ε) ? inv(x) : _lorentz_broaden(x, ε)
+# scalar inverses with a cutoff tolerance
+_safe_inv(x, tol) = abs(x) < tol ? zero(x) : inv(x)
+
+# compute inverse singular value difference contribution to SVD gradient with broadening ε
+function _broadened_inv_S(S::AbstractVector{T}, tol, ε=0) where {T}
+    F = similar(S, (axes(S, 1), axes(S, 1)))
+    @inbounds for j in axes(F, 2), i in axes(F, 1)
+        F[i, j] = if i == j
+            zero(T)
+        else
+            Δsᵢⱼ = S[j] - S[i]
+            ε > 0 ? _lorentz_broaden(Δsᵢⱼ, ε) : _safe_inv(Δsᵢⱼ, tol)
+        end
     end
+    return F
 end
 
 # Lorentzian broadening for divergent term in SVD rrule, see
@@ -554,9 +562,8 @@ function svd_pullback!(
         @info "`svd` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
     end
 
-    UdΔAV =
-        (aUΔU .+ aVΔV) .* _safe_inv.(Sp' .- Sp, tol, broadening) .+
-        (aUΔU .- aVΔV) .* _safe_inv.(Sp' .+ Sp, tol)
+    inv_S_minus = _broadened_inv_S(Sp, tol, broadening) # possibly divergent/broadened contribution
+    UdΔAV = @. (aUΔU + aVΔV) * inv_S_minus + (aUΔU - aVΔV) * _safe_inv(Sp' .+ Sp, tol)
     if !(ΔS isa ZeroTangent)
         UdΔAV[diagind(UdΔAV)] .+= real.(ΔS)
         # in principle, ΔS is real, but maybe not if coming from an anyonic tensor
@@ -585,16 +592,14 @@ function svd_pullback!(
             VrΔV = fill!(similar(Vd, (r - p, p)), 0)
         end
 
-        X =
-            (1//2) .* (
-                (UrΔU .+ VrΔV) .* _safe_inv.(Sp' .- Sr, tol, broadening) .+
-                (UrΔU .- VrΔV) .* _safe_inv.(Sp' .+ Sr, tol)
-            )
-        Y =
-            (1//2) .* (
-                (UrΔU .+ VrΔV) .* _safe_inv.(Sp' .- Sr, tol, broadening) .-
-                (UrΔU .- VrΔV) .* _safe_inv.(Sp' .+ Sr, tol)
-            )
+        X = @. (1//2) * (
+            (UrΔU + VrΔV) * _safe_inv(Sp' - Sr, tol) +
+            (UrΔU - VrΔV) * _safe_inv(Sp' + Sr, tol)
+        )
+        Y = @. (1//2) * (
+            (UrΔU + VrΔV) * _safe_inv(Sp' - Sr, tol) -
+            (UrΔU - VrΔV) * _safe_inv(Sp' + Sr, tol)
+        )
 
         # ΔA += Ur * X * Vp' + Up * Y' * Vr'
         mul!(ΔA, Ur, X * Vp', 1, 1)

test/utility/svd_wrapper.jl

@@ -21,7 +21,6 @@ rtol = 1e-9
 Random.seed!(123456789)
 r = randn(dtype, ℂ^m, ℂ^n)
 R = randn(space(r))
-broadenings = [10.0^k for k in -16:-4]
 
 full_alg = SVDAdjoint(; rrule_alg=(; alg=:full, broadening=0))
 iter_alg = SVDAdjoint(; fwd_alg=(; alg=:iterative))
@@ -42,13 +41,29 @@ end
     @test g_fullsvd[1] ≈ g_itersvd[1] rtol = rtol
 end
 
-@testset "Truncated SVD with χ=$χ and ε=$ε broadening" for ε in broadenings
-    broadened_alg = @set full_alg.rrule_alg.broadening = ε
-    l_unbroadened, g_unbroadened = withgradient(A -> lossfun(A, full_alg, R, trunc), r)
-    l_broadened, g_broadened = withgradient(A -> lossfun(A, broadened_alg, R, trunc), r)
+@testset "Truncated SVD broadening" begin
+    u, s, v, = tsvd(r)
+    s.data[1:2:m] .= s.data[2:2:m] # make every singular value two-fold degenerate
+    r_degen = u * s * v
 
-    @test l_unbroadened ≈ l_broadened
-    @test 1e1 * norm(g_broadened[1]) * ε > norm(g_unbroadened[1] - g_broadened[1]) > ε
+    no_broadening_no_cutoff_alg = @set full_alg.rrule_alg.broadening = 1e-30
+    small_broadening_alg = @set full_alg.rrule_alg.broadening = 1e-13
+
+    l_only_cutoff, g_only_cutoff = withgradient(
+        A -> lossfun(A, full_alg, R, trunc), r_degen
+    ) # cutoff sets degenerate difference to zero
+    l_no_broadening_no_cutoff, g_no_broadening_no_cutoff = withgradient( # degenerate singular value differences lead to divergent contributions
+        A -> lossfun(A, no_broadening_no_cutoff_alg, R, trunc),
+        r_degen,
+    )
+    l_small_broadening, g_small_broadening = withgradient( # Lorentzian broadening smoothens divergent contributions
+        A -> lossfun(A, small_broadening_alg, R, trunc),
+        r_degen,
+    )
+
+    @test l_only_cutoff ≈ l_no_broadening_no_cutoff ≈ l_small_broadening
+    @test norm(g_no_broadening_no_cutoff[1] - g_small_broadening[1]) > 1e-1 # divergences mess up the gradient
+    @test g_only_cutoff[1] ≈ g_small_broadening[1] rtol = rtol # cutoff and Lorentzian broadening have similar effect
 end
 
 symm_m, symm_n = 18, 24
@@ -80,17 +95,29 @@ symm_R = randn(dtype, space(symm_r))
     @test g_fullsvd_tr[1] ≈ g_itersvd_fb[1] rtol = rtol
 end
 
-@testset "Truncated symmetric SVD with χ=$χ and ε=$ε broadening" for ε in broadenings
-    broadened_alg = @set full_alg.rrule_alg.broadening = ε
-    l_unbroadened, g_unbroadened = withgradient(
-        A -> lossfun(A, full_alg, symm_R, symm_trspace), symm_r
+@testset "Truncated symmetric SVD broadening" begin
+    u, s, v, = tsvd(symm_r)
+    s.data[1:2:m] .= s.data[2:2:m] # make every singular value two-fold degenerate
+    symm_r_degen = u * s * v
+
+    no_broadening_no_cutoff_alg = @set full_alg.rrule_alg.broadening = 1e-30
+    small_broadening_alg = @set full_alg.rrule_alg.broadening = 1e-13
+
+    l_only_cutoff, g_only_cutoff = withgradient(
+        A -> lossfun(A, full_alg, symm_R, symm_trspace), symm_r_degen
+    ) # cutoff sets degenerate difference to zero
+    l_no_broadening_no_cutoff, g_no_broadening_no_cutoff = withgradient( # degenerate singular value differences lead to divergent contributions
+        A -> lossfun(A, no_broadening_no_cutoff_alg, symm_R, symm_trspace),
+        symm_r_degen,
     )
-    l_broadened, g_broadened = withgradient(
-        A -> lossfun(A, broadened_alg, symm_R, symm_trspace), symm_r
+    l_small_broadening, g_small_broadening = withgradient( # Lorentzian broadening smoothens divergent contributions
+        A -> lossfun(A, small_broadening_alg, symm_R, symm_trspace),
+        symm_r_degen,
     )
 
-    @test l_unbroadened ≈ l_broadened
-    @test 1e1 * norm(g_broadened[1]) * ε > norm(g_unbroadened[1] - g_broadened[1]) > ε
+    @test l_only_cutoff ≈ l_no_broadening_no_cutoff ≈ l_small_broadening
+    @test norm(g_no_broadening_no_cutoff[1] - g_small_broadening[1]) > 1e-2 # divergences mess up the gradient
+    @test g_only_cutoff[1] ≈ g_small_broadening[1] rtol = rtol # cutoff and Lorentzian broadening have similar effect
 end
 
 # TODO: Add when IterSVD is implemented for HalfInfiniteEnv