[Bug] SVD pullback does not properly take into account the norm of the initial matrix

[leburgel]

When running variational optimization using a cost function which uses MatrixAlgebraKit.svd_trunc in an intermediate step, and whose value is independent of the normalization of the input, the current choice of tolerances for the svd_trunc pullback can give rise to very inaccurate gradients depending on the initial normalization. This came up for example in QuantumKitHub/PEPSKit.jl#276, but the problem persists with the new implementations in MatrixAlgebraKit.jl.

An example would be

using Random
using LinearAlgebra
using MatrixAlgebraKit
using OptimKit
using Zygote

Random.seed!(123)

# some random matrix
O = randn(ComplexF64, 10, 10)

# some cost function which uses svd_trunc
function fg(t)
    e, g = Zygote.withgradient(t) do x
        _, S, _ = svd_trunc(x; trunc = truncrank(10))
        return real(tr(O * (S / norm(S))))
    end
    return e, only(g)
end

# prepare a random input
t0 = randn(ComplexF64, 20, 20)
# but artificially make singular values decay exponentially
U, S, V = svd_full(t0)
S´ = S .* diagm(logrange(1.0e0, 1.0e-12, size(S, 1)))
t = U * S´ * V

Clearly, the value of the cost function is independent of the initial norm of t. Checking the gradient without any prior normalization on t gives a good result,

@show fg(t´)[1]
_, _, dfs1, dfs2 = OptimKit.optimtest(
    fg, t´;
    alpha = LinRange(-1.0e-5, 1.0e-5, 2),
)
@show dfs1
@show dfs2

(fg(t´))[1] = -0.621052432092075
dfs1 = [0.06446034588325489]
dfs2 = [0.06446034587777202]

Running the same for an input with unit norm already gives a significantly less accurate result,

rescaling_factor = norm(t´)
@show fg(t´ / rescaling_factor)[1]
_, _, dfs1, dfs2 = OptimKit.optimtest(
    fg, t´ / rescaling_factor;
    alpha = LinRange(-1.0e-5, 1.0e-5, 2),
)
@show dfs1
@show dfs2

(fg(t´ / rescaling_factor))[1] = -0.6210524320920751
dfs1 = [4.3408483775597695]
dfs2 = [4.909645149555295]

And finally, running the same for an input with a very small norm seems to break things completely,

rescaling_factor = 1.0e3 * norm(t´)
@show fg(t´ / rescaling_factor)[1]
_, _, dfs1, dfs2 = OptimKit.optimtest(
    fg, t´ / rescaling_factor;
    alpha = LinRange(-1.0e-5, 1.0e-5, 2),
)
@show dfs1
@show dfs2

(fg(t´ / rescaling_factor))[1] = -0.6210524320920752
dfs1 = [-2681.1597759100105]
dfs2 = [4.909645149555297e6]

I think the bottom line is that with the default settings, the safe inverse in the SVD pullback uses a tolerance that does not take into account the norm, and therefore simply discards essential information if the norm of the initial tensor is very small. I know that I can fix this by manually setting the tolerance of the pullback to something that scales with the norm of the input, but this is quite annoying to do in practice.

I would really just expect the example here to work out of the box. Is there any reason not to scale (some of) the tolerances in the pullback by the norm of the input of the primal computation?

[Jutho]

You are correct that we should definitely take the norm of the tensor into account. But I am confused (after simply reading your example, not trying anything myself): are you saying that with a tensor that is properly normalized to have norm 1, the current tolerance is already producing bad results?

[leburgel]

For the toy example here that does seem to be the case for some reason, but maybe the example itself is not very good. In more realistic settings normalizing to a norm of 1 before the SVD seems good enough to get an accurate gradient, there I've never seen that a norm of 1 could already be a problem.

[Jutho]

I don't think the example here really works. Note how dfs2 scales exactly as you would expect, it got multiplied with a factor 1e3 between the secon dand third 3, with a rescaling factor that is 1e3 times larger.

So all that you are seeing here is that dfs1, the one computed using finite difference, becomes very inaccurate if you using O(1e-5) changes to a tensor which itself is only order 1e-5 entries to compute the derivative numerically.

[Jutho]

Also, since this test function only depends on the singular values, the ΔU and ΔVᴴ are identically zero and none of the safe_inv stuff really contributes 😄.

[leburgel]

I see, pretty bad example then :)

An actual case would be the example described in QuantumKitHub/PEPSKit.jl#276, which was run on PEPSKit.jl @ commit a609593472ae536048ee0f9327938aae00cb8512, but that one's not very minimal I'm afraid... There I explicitly checked the scaling of the gradient with the initial normalization, and it didn't behave as you would expect when you decrease the norm of the initial PEPS tensor.

[Jutho]

Ok, I guess I will have to look at it. Currently, the tolerance does actually already depend on the size of S, but in this way:

"""
    default_pullback_gaugetol(a)

Default tolerance for deciding to warn if incoming adjoints of a pullback rule
has components that are not gauge-invariant.
"""
function default_pullback_gaugetol(a)
    n = norm(a, Inf)
    return eps(eltype(n))^(3 / 4) * max(n, one(n))
end

So the max(n, one(n)) versus simply using n is something that indeed was commented on before, and we might just want to make this n.

[leburgel]

Just removing the max in default_pullback_gaugetol doesn't seem to solve the problem unfortunately. I have the PEPS example up and running using the MatrixAlgebraKit.jl implementation of svd_trunc and its rrule now, so I'll mess around a bit locally to see if I can find an obvious fix.

[Confusio]

Just removing the max in default_pullback_gaugetol doesn't seem to solve the problem unfortunately. I have the PEPS example up and running using the MatrixAlgebraKit.jl implementation of svd_trunc and its rrule now, so I'll mess around a bit locally to see if I can find an obvious fix.

This issue is related to #7 . You may set

function default_pullback_gaugetol(a)
    n = norm(a, Inf)
    return eps(eltype(n))^(3 / 4) * n
end

Actually, even by this setting, some issue may still take place as in the TensorKit.jl, the norm(a,Inf) will take different values for different symmetry sectors. So I think the way to work around this is just to normalize the matrix (even in TensorKit.jl) and take the current default tol.

[leburgel]

As far as MatrixAlgebraKit.jl is concerned, the issue is actually solved by removing the max from default_pullback_gaugetol as suggested above. There were some additional but very similar issues with the Lorentzian broadening used in PEPSKit.jl (PEPSKit.jl still doesn't use exactly the same SVD pullback as MatrixAlgebraKit.jl yet, there's a broadened inverse used instead of a safe inverse at some point, which I overlooked before). But these can be solved in a similar way (see QuantumKitHub/PEPSKit.jl#285).

This then still leaves the other issue from the above comment regarding the value of norm for a TensorMap with blocks corresponding to symmetry sectors with a non-trivial quantum dimension, meaning (I think) norm(t, Inf) isn't quite representative of the actual numerical entries of the individual blocks that the MatrixAlgebraKit.jl SVD pullback is called on in the end.

However, as far as I can tell, as long as we don't pass an explicit tolerance to the pullback, I think the pullback of an SVD on a TensorMap just calls the MatrixAlgebraKit pullback with a different default tolerance for each block, where that tolerance is really chosen based on the bare numerical entries of the block. So it seems the quantum dimensions could only be a problem when manually setting a tol kwarg based on norm(::TensorMap, Inf).

[lkdvos]

I have to say that I went back and forth a bit about the meaning of these tolerances for the different blocks, and whether relative tolerances have to be scaled with the full norm or the blockwise norm.
I think I've convinced myself that for things that have to do with numerical errors like this, since our implementation doesn't propagate errors between blocks, nor between "different copies" of the same block for non-abelian sectors, it is correct to use norms that are relative to the norm of the block (without quantum dimension), rather than the norm of the full tensor.
I think the current implementation in TensorKit is therefore actually correct, taking relative tolerances with respect to the blocknorm and absolute tolerances as given.

I do think that changing

MatrixAlgebraKit.jl/src/common/defaults.jl

Lines 18 to 21 in c52614b

	function default_pullback_gaugetol(a)
	n = norm(a, Inf)
	return eps(eltype(n))^(3 / 4) * max(n, one(n))
	end

to the suggested solution makes a lot more sense though, I'll open a PR for this.

[edit]
I just remembered why this was there, as the gauge pullbacks are looking at |aUΔU + aVΔV|.
At least the rank_atol and degeneracy_atol are really looking at differences in singular values, so I will just change these values instead.