[Yue-Zhengyuan]

In the Arnoldi fixed-point gradient solver, we should choose the eigenvalue closest to 1, instead of the one with the largest magnitude (which can overshoot due to numerical errors).

This PR fixes global phase differentiation in the gauge-fixed CTMRG iteration.

Previously, the global phase corrections for the corner and edge tensors were computed outside the differentiated iteration and treated as constants.
This left phase gauge modes with eigenvalues of modulus one in the environment Jacobian, which could make the Arnoldi gradient solver select an unwanted eigenvector.

Now we recompute the global phase corrections inside the differentiated iteration, relative to the fixed converged environment.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
...lgorithms/optimization/implicit_differentiation.jl	94.93% <100.00%> (ø)

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[pbrehmer]

Looks good to me in principle! Although I would be curious if you have explicit examples where this overshooting happens and whether this selection criterion actually then helps or whether these cases are pathological anyway. And this shouldn't interfere with the GMRES fallback condition norm(vecs[1][2]) < 1.0e-2 * alg.tol, right?

[Yue-Zhengyuan]

A minimal example with Heisenberg model, D = 3, chi = 20:

Details

using PEPSKit, TensorKit, Zygote, KrylovKit, LinearAlgebra, Random

"""Compare Arnoldi selectors at a verified CTMRG fixed point."""
function main(tol)
    BLAS.set_num_threads(1)
    Random.seed!(1234)
    H = heisenberg_XYZ(InfiniteSquare(1, 1))
    ψ = PEPSKit.peps_normalize(InfinitePEPS(randn, ComplexF64, physicalspace(H)[1], ComplexSpace(3)))
    alg = PEPSKit.CTMRGAlgorithm(; tol, maxiter = 1000, verbosity = 0)
    env, info = leading_boundary(CTMRGEnv(ψ, ComplexSpace(20)), ψ, alg)
    @assert info.converged

    # Linearize one iteration AFTER converging the environment.
    fixed = PEPSKit._set_fixed_truncation(alg)
    nextenv, = PEPSKit.ctmrg_iteration(InfiniteSquareNetwork(ψ), env, fixed)
    signs, cp, ep = PEPSKit.compute_gauge_fix_gauge(nextenv, env, PEPSKit.ScramblingEnvGauge())
    iteration = (A, x) -> PEPSKit.fix_phases(
        first(PEPSKit.ctmrg_iteration(InfiniteSquareNetwork(A), x, fixed)), signs, cp, ep,
    )
    nextenv, back = Zygote.pullback(iteration, ψ, env)
    # verify fixed-point equation
    fixedpoint_error = norm(nextenv - env) / norm(env)
    @show tol info.convergence_error fixedpoint_error
    @assert fixedpoint_error < 10 * tol

    b = only(Zygote.gradient(x -> PEPSKit.cost_function(ψ, x, H), env))
    b = b isa typeof(env) ? b : CTMRGEnv(b.corners, b.edges)
    Jᵀ = v -> back(v)[2]::typeof(env)
    # The eigenvector (y, 1) at eigenvalue one solves (I - Jᵀ)y = b.
    action = X -> (PEPSKit.VI.add!!(Jᵀ(X[1]), b, X[2]), X[2])
    for which in (:LM, EigSorter(λ -> abs(λ - 1)))
        λ, vectors, info = KrylovKit.realeigsolve(action, (b, one(PEPSKit.scalartype(b))), 1, which, Arnoldi(; tol = 1.0e-9))
        y = PEPSKit.VI.scale(vectors[1][1], inv(vectors[1][2]))
        relative_residual = norm(y - Jᵀ(y) - b) / norm(b)
        @show which λ[1] info.numops relative_residual
    end
    return
end
# Compare a converged environment with a more tightly converged one.
main(1.0e-10)
main(1.0e-12)

CTMRG tolerance	Fixed-point error	:LM residual	:LM $\lambda$	Nearest-one residual	Nearest-one $\lambda$
$10^{-10}$	$5.5\times10^{-11}$	$63.8$	1.000000537687291	$2.4\times10^{-9}$	0.999999999981909
$10^{-12}$	$2.0\times10^{-12}$	$1.4\times10^{-11}$	1.000000000000303	$1.1\times10^{-11}$	0.999999999999919

Here :LM/Nearest-one residual is the error in solving the equation (1 - Jᵀ)y = b.

GMRES fallback is not affected.

[lkdvos]

Can you explain a bit more why this is the case?
It looks a bit fishy that we are expecting something with eigenvalue 1 and getting spurious larger values.
IIRC this is happening in 1/(1-x), so this isn't even a converging series if x > 1 (which is what this is relating to, I think?). I definitely might be wrong here, I have not looked at it super closely.
Additionally, Krylov methods are good at getting these extremal eigenvalues and not so much at interior ones.
Maybe @Jutho can also comment on this one, because it looks like just selecting the value closest to one seems like a band-aid fix for something a bit more fundamental?

[Yue-Zhengyuan]

It appears that the extra larger-than-1 modes are related to incompletely removed gauge freedom in the CTMRG environment. I'll further investigate.

[leburgel]

I'm not sure if it's related, but I've recently run into trouble with non-converging gradient solves due to the limited accuracy of our environment gauge fixing routine. This can lead to elementwise convergence of the fixed-point iteration we differentiate that is significantly worse than the actual residual of the forward pass, which sometimes gives a very unstable gradient linear problem at the fixed point. It sounds like this could cause the issue here as well. I have a local implementation of more accurate robust gauge fixing routines that seems to solve this on my problem cases, but I haven't gotten around to self reviewing that and filing it here.

We've been running into spuriously failing Arnoldi gradient solves for a long time now, but I never realized this was the effect at play. This is a great catch! I'll move the gauge fixing routines up in priority, and check if they help with the reproducer you posted.

[leburgel]

It turns out this has nothing to do with gauge fixing accuracy, increasing the accuracy of the gauge fixing actually makes the issue even more pronounced. It turns out that the Jacobian here has six eigenvalues of modulus 1 to within 1e-11. They come from the global phase of the corners and edges, which are treated as constants independent of the actual environment in the gauge-fixed iteration that we differentiate. I introduced this behavior in #364.

The systematic fix is to actually differentiate the global phase fixing, which is the right thing to do since the global phases directly depend on the environment. I confirmed that defining the gauge-fixed iteration as

    env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed)
    signs, = compute_gauge_fix_gauge(env_conv, env, alg_gauge)
    # fix global phases relative to the fixed point inside the iteration
    function gauge_fixed_iteration(A, x)
        return fix_global_phases(
            fix_relative_phases(ctmrg_iteration(InfiniteSquareNetwork(A), x, alg_fixed)[1], signs),
            env,
        )
    end

before differentiating removes the other eigenvalues of magnitude 1. So I think this is the fix we need, and the :LM solve can stay.

[Yue-Zhengyuan]

@leburgel Do you mind directly apply the fix here? I need some time to comprehend...

[leburgel]

I pushed the fix here. The gauge-fixed iteration now recomputes the global phases relative to the converged environment, so they are actually differentiated like they should have been in the first place in #364 . I restored the :LM target, since this now gives exactly the same result as the closest-to-1 targeting on the reproducer:

CTMRG tolerance	Fixed-point error	:LM residual	:LM $\lambda$	Nearest-one residual	Nearest-one $\lambda$
$10^{-10}$	$5.5\times10^{-11}$	$1.7\times10^{-14}$	0.9999999999999993	$1.7\times10^{-14}$	0.9999999999999993
$10^{-12}$	$2.0\times10^{-12}$	$9.2\times10^{-15}$	0.9999999999999996	$9.2\times10^{-15}$	0.9999999999999996

[Yue-Zhengyuan]

After tests pass I think we can directly merge.

[Yue-Zhengyuan]

Should switch to another title, though...