diff --git a/docs/src/assets/pepskit.bib b/docs/src/assets/pepskit.bib index 15949e7c3..00c0c5341 100644 --- a/docs/src/assets/pepskit.bib +++ b/docs/src/assets/pepskit.bib @@ -157,3 +157,18 @@ @misc{zhang_accelerating_2025 primaryClass={cond-mat.str-el}, url={https://arxiv.org/abs/2505.00494}, } + +@article{zhang_accelerating_2026, + title = {Accelerating two-dimensional tensor network optimization by preconditioning}, + author = {Zhang, Xing-Yu and Yang, Qi and Corboz, Philippe and Haegeman, Jutho and Tang, Wei}, + journal = {Phys. Rev. B}, + volume = {113}, + issue = {12}, + pages = {125111}, + numpages = {8}, + year = {2026}, + month = {Mar}, + publisher = {American Physical Society}, + doi = {10.1103/h396-yc28}, + url = {https://link.aps.org/doi/10.1103/h396-yc28} +} diff --git a/src/Defaults.jl b/src/Defaults.jl index 13b51e12e..babafba88 100644 --- a/src/Defaults.jl +++ b/src/Defaults.jl @@ -94,6 +94,20 @@ Module containing default algorithm parameter values and arguments. * `gradient_tol_max=$(Defaults.gradient_tol_max)` : Maximal gradient algorithm tolerance used by `gradient_dynamic_tols`. * `gradient_tol_factor=$(Defaults.gradient_tol_factor)` : Tolerance scaling factor relative to the boundary algorithm's tolerance, used by `gradient_dynamic_tols` (e.g. `10` makes the gradient tolerance ~10x looser than the boundary tolerance). +## Preconditioning + +* `precondition_alg=:$(Defaults.precondition_alg)` : Algorithm variant used for preconditioning the PEPS gradient. + - `:LocalPreconditioner` : Precondition using the leading (local) term of the PEPS metric, see [`PEPSKit.LocalPreconditioner`](@ref). +* `precondition_tol=$(Defaults.precondition_tol)` : Convergence tolerance for the linear problem in the preconditioning step. +* `precondition_maxiter=$(Defaults.precondition_maxiter)` : Maximal number of iterations for the linear problem in the preconditioning step. +* `precondition_verbosity=$(Defaults.precondition_verbosity)` : Preconditioning output information verbosity. +* `precondition_krylovdim=$(Defaults.precondition_krylovdim)` : Krylov dimension for the linear problem in the preconditioning step. +* `precondition_regularization=$(Defaults.precondition_regularization)` : Prefactor setting the regularization strength of the local linear problem, see [`PEPSKit.LocalPreconditioner`](@ref). +* `precondition_dynamic_tols=$(Defaults.precondition_dynamic_tols)` : If `true`, wrap the preconditioner algorithm in a `MPSKit.DynamicTol` that rescales its tolerance based on the current PEPS optimization gradient norm, see [`PEPSKit.PEPSOptimize`](@ref). +* `precondition_tol_min=$(Defaults.precondition_tol_min)` : Minimal preconditioner tolerance used by `precondition_dynamic_tols`. +* `precondition_tol_max=$(Defaults.precondition_tol_max)` : Maximal preconditioner tolerance used by `precondition_dynamic_tols`. +* `precondition_tol_factor=$(Defaults.precondition_tol_factor)` : Tolerance scaling factor used by `precondition_dynamic_tols`. + ## Optimization * `reuse_env=$(Defaults.reuse_env)` : If `true`, the current optimization step is initialized on the previous environment, otherwise a random environment is used. @@ -168,6 +182,18 @@ const gradient_tol_min = 1.0e-10 const gradient_tol_max = 1.0e-1 const gradient_tol_factor = 1.0e1 +# Preconditioning +const precondition_alg = :LocalPreconditioner +const precondition_tol = 1.0e-6 +const precondition_maxiter = 1 +const precondition_verbosity = -1 +const precondition_krylovdim = 30 +const precondition_regularization = 100.0 +const precondition_dynamic_tols = true +const precondition_tol_min = 1.0e-12 +const precondition_tol_max = 1.0e-4 +const precondition_tol_factor = 1.0e-2 + # Optimization const reuse_env = true const optimizer_tol = 1.0e-4 diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 25be6afca..4a6956069 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -143,6 +143,7 @@ include("algorithms/correlator_adapters.jl") include("algorithms/correlators.jl") include("algorithms/optimization/fixed_point_differentiation.jl") +include("algorithms/optimization/preconditioning.jl") include("algorithms/optimization/peps_optimization.jl") include("algorithms/select_algorithm.jl") diff --git a/src/algorithms/optimization/peps_optimization.jl b/src/algorithms/optimization/peps_optimization.jl index ab9222515..4823c7447 100644 --- a/src/algorithms/optimization/peps_optimization.jl +++ b/src/algorithms/optimization/peps_optimization.jl @@ -17,37 +17,43 @@ For a full description, see [`fixedpoint`](@ref). The supported keywords are: * `boundary_alg::Union{NamedTuple,<:CTMRGAlgorithm,...}` * `gradient_alg::Union{NamedTuple,Nothing,<:GradientAlgorithm}` * `optimizer_alg::Union{NamedTuple,<:OptimKit.OptimizationAlgorithm}` +* `precondition_alg::Union{NamedTuple,Nothing,<:PreconditionAlgorithm}` * `reuse_env::Bool=$(Defaults.reuse_env)` * `symmetrization::Union{Nothing,SymmetrizationStyle}=nothing` """ -struct PEPSOptimize{B, G} +struct PEPSOptimize{B, G, P} boundary_alg::B gradient_alg::G optimizer_alg::OptimKit.OptimizationAlgorithm + precondition_alg::P reuse_env::Bool symmetrization::Union{Nothing, SymmetrizationStyle} function PEPSOptimize( # Inner constructor to prohibit illegal setting combinations - boundary_alg::B, gradient_alg::G, optimizer_alg, + boundary_alg::B, gradient_alg::G, optimizer_alg, precondition_alg::P, reuse_env, symmetrization, - ) where {B, G} + ) where {B, G, P} _check_algorithm_combination( parent_alg(boundary_alg), parent_alg(gradient_alg), symmetrization ) - return new{B, G}(boundary_alg, gradient_alg, optimizer_alg, reuse_env, symmetrization) + return new{B, G, P}( + boundary_alg, gradient_alg, optimizer_alg, precondition_alg, + reuse_env, symmetrization, + ) end end function PEPSOptimize(; - boundary_alg = (;), gradient_alg = (;), optimizer_alg = (;), + boundary_alg = (;), gradient_alg = (;), optimizer_alg = (;), precondition_alg = (;), reuse_env = Defaults.reuse_env, symmetrization = nothing, ) boundary_algorithm = _alg_or_nt(CTMRGAlgorithm, boundary_alg) gradient_algorithm = _alg_or_nt(GradientAlgorithm, gradient_alg) optimizer_algorithm = _alg_or_nt(OptimKit.OptimizationAlgorithm, optimizer_alg) + precondition_algorithm = _alg_or_nt(PreconditionAlgorithm, precondition_alg) return PEPSOptimize( - boundary_algorithm, gradient_algorithm, optimizer_algorithm, + boundary_algorithm, gradient_algorithm, optimizer_algorithm, precondition_algorithm, reuse_env, symmetrization, ) end @@ -137,18 +143,34 @@ keyword arguments are: - `:FixedPointGradient` : Compute the gradient via fixed-point differentiation, see [`FixedPointGradient`](@ref) * `solver_alg::Union{Algorithm,NamedTuple}`: Solver algorithm for computing the implicit gradient; see [`FixedPointGradient`](@ref) for supported algorithms. +### Preconditioner algorithm + +Supply preconditioner parameters via `precondition_alg::Union{NamedTuple,Nothing,<:PreconditionAlgorithm}` +using either a `NamedTuple` of keyword arguments, `nothing`, or a `PreconditionAlgorithm` +struct directly. By default, the gradient is preconditioned with the local PEPS metric, see +[`LocalPreconditioner`](@ref); pass `nothing` to disable preconditioning and optimize using +the raw Euclidean gradient. The supported `NamedTuple` keyword arguments are: + +* `alg::Symbol=:$(Defaults.precondition_alg)` : Preconditioner algorithm variant, can be one of the following: + - `:LocalPreconditioner` : Precondition using the leading (local) term of the PEPS metric, see [`LocalPreconditioner`](@ref) +* `tol::Real=$(Defaults.precondition_tol)` : Convergence tolerance of the local linear problem. +* `maxiter::Int=$(Defaults.precondition_maxiter)` : Maximal number of iterations of the local linear problem. +* `verbosity::Int` : Preconditioner output verbosity, ≤0 by default to disable too verbose printing. Should only be >0 for debug purposes. +* `krylovdim::Int=$(Defaults.precondition_krylovdim)` : Krylov dimension of the local linear problem. +* `regularization::Real=$(Defaults.precondition_regularization)` : Prefactor setting the regularization strength of the local linear problem. + ### Dynamic tolerances -The boundary and gradient algorithms each additionally accept the keyword arguments below, -which wrap the corresponding algorithm in an `MPSKit.DynamicTols.DynamicTol` that +The boundary, gradient and preconditioner algorithms each additionally accept the keyword +arguments below, which wrap the corresponding algorithm in an `MPSKit.DynamicTols.DynamicTol` that rescales its tolerance over the course of the optimization. This allows the intermediate problems to be solved only as accurately as the current optimization step requires, which -can significantly reduce the total runtime. The boundary tolerance is scaled relative to the -current gradient norm, while the gradient tolerance is in turn scaled relative to the -effective boundary tolerance. These settings are only available within a +can significantly reduce the total runtime. The boundary and preconditioner tolerances are +scaled relative to the current gradient norm, while the gradient tolerance is in turn scaled +relative to the effective boundary tolerance. These settings are only available within a variational optimization, not for standalone [`leading_boundary`](@ref) calls. -* `dynamic_tols::Bool` : Enable dynamic tolerance scaling for this algorithm. Defaults to `$(Defaults.ctmrg_dynamic_tols)` and `$(Defaults.gradient_dynamic_tols)` for the boundary and gradient algorithm respectively. +* `dynamic_tols::Bool` : Enable dynamic tolerance scaling for this algorithm. Defaults to `$(Defaults.ctmrg_dynamic_tols)`, `$(Defaults.gradient_dynamic_tols)` and `$(Defaults.precondition_dynamic_tols)` for the boundary, gradient and preconditioner algorithm respectively. * `tol_min::Real` : Lower clamp on the dynamically scaled tolerance. * `tol_max::Real` : Upper clamp on the dynamically scaled tolerance. * `tol_factor::Real` : Prefactor of the dynamically scaled tolerance. @@ -227,10 +249,18 @@ function fixedpoint( # normalize the initial guess peps₀ = peps_normalize(peps₀) + # initialize the preconditioner + function precondition(x, g) + precondition_alg = updatetol( + alg.precondition_alg, tracked_finalizer.tol_state.iter, tracked_finalizer.tol_state.gradnorm + ) + return peps_precondition(x, g, tracked_finalizer.tol_state, precondition_alg) + end + # optimize operator cost function (peps_final, env_final), cost_final, ∂cost, numfg, convergence_history = optimize( (peps₀, env₀), alg.optimizer_alg; - retract, inner = real_inner, (transport!) = (peps_transport!), + retract, inner = real_inner, (transport!) = (peps_transport!), precondition, hasconverged, shouldstop, (finalize!) = tracked_finalizer, ) do (peps, env) start_time = time() diff --git a/src/algorithms/optimization/preconditioning.jl b/src/algorithms/optimization/preconditioning.jl new file mode 100644 index 000000000..319d0ff2f --- /dev/null +++ b/src/algorithms/optimization/preconditioning.jl @@ -0,0 +1,125 @@ +abstract type PreconditionAlgorithm end + +const PRECONDITION_ALGORITHM_SYMBOLS = IdDict{Symbol, Type{<:PreconditionAlgorithm}}() + +""" + PreconditionAlgorithm(; kwargs...) + +Keyword argument parser returning the appropriate `PreconditionAlgorithm` algorithm struct. +""" +function PreconditionAlgorithm(; + alg = Defaults.precondition_alg, + tol = Defaults.precondition_tol, + maxiter = Defaults.precondition_maxiter, + verbosity = Defaults.precondition_verbosity, + krylovdim = Defaults.precondition_krylovdim, + regularization = Defaults.precondition_regularization, + ) + # replace symbol with PreconditionAlgorithm alg type + haskey(PRECONDITION_ALGORITHM_SYMBOLS, alg) || + throw(ArgumentError("unknown PreconditionAlgorithm algorithm: $alg")) + alg_type = PRECONDITION_ALGORITHM_SYMBOLS[alg] + + return alg_type(GMRES(; tol, maxiter, verbosity, krylovdim), regularization) +end + +""" +$(TYPEDEF) + +Preconditioner for PEPS ground-state optimization based on the leading term of the PEPS +metric tensor, i.e. the local norm matrix ``N`` obtained by contracting the full CTMRG +environment around a given unit-cell site while leaving the ket and bra PEPS tensor at that +site open. + +Instead of following the raw Euclidean gradient ``g``, the optimizer is fed the solution +``g̃`` of the regularized local linear problem +```math +( N / ⟨ψ|ψ⟩ + δ ) g̃ = g +``` +which is solved separately for every tensor in the unit cell using an iterative solver. +The regularization ``δ`` prevents the (generally singular) metric from being inverted +exactly, and decays during the optimization according to +```math +δ = \\texttt{regularization} × ‖∇f‖^2 / \\texttt{iter}^2 +``` +so that the preconditioner acts conservatively far from the minimum and becomes +increasingly aggressive as the optimization converges. + +## Fields + +$(TYPEDFIELDS) + +## Constructors + + LocalPreconditioner(; kwargs...) + +Construct a local preconditioner algorithm struct based on keyword arguments. The supported +keywords are: + +* `tol::Real=$(Defaults.precondition_tol)` : Convergence tolerance of the local linear problem. +* `maxiter::Int=$(Defaults.precondition_maxiter)` : Maximal number of iterations of the local linear problem. Note that the default performs a single iteration, such that the metric is only inverted approximately. +* `verbosity::Int=$(Defaults.precondition_verbosity)` : Preconditioner output verbosity, ≤0 by default to disable too verbose printing. +* `krylovdim::Int=$(Defaults.precondition_krylovdim)` : Krylov dimension of the local linear problem. +* `regularization::Real=$(Defaults.precondition_regularization)` : Prefactor setting the regularization strength ``δ``, see above. + +Reference: [Phys. Rev. B 113, 125111](@cite zhang_accelerating_2026) +""" +struct LocalPreconditioner{A} <: PreconditionAlgorithm + "solver algorithm used for the local linear problem" + solver_alg::A + + "prefactor setting the regularization strength of the local linear problem" + regularization::Float64 +end +LocalPreconditioner(; kwargs...) = PreconditionAlgorithm(; alg = :LocalPreconditioner, kwargs...) +PRECONDITION_ALGORITHM_SYMBOLS[:LocalPreconditioner] = LocalPreconditioner + +# `LocalPreconditioner` has no top-level `tol` field (it lives on `solver_alg`), so the +# default `MPSKit.DynamicTols._updatetol` (which sets `alg.tol`) doesn't apply +_updatetol(alg::LocalPreconditioner, tol::Real) = @set alg.solver_alg.tol = tol + +""" +$(SIGNATURES) + +Apply the regularized local metric ``N / ⟨ψ|ψ⟩ + δ`` to a single PEPS tensor `g_rc`, where +``N`` is obtained by contracting the CTMRG environment `env` around the unit-cell site +`(r, c)` and `norm_pref` is the corresponding local norm ``⟨ψ|ψ⟩``. +""" +function apply_local_preconditioner(g_rc::PEPSTensor, env::CTMRGEnv, δ, (r, c), norm_pref) + @autoopt @tensor g_rc_prec[d; D_N_below D_E_below D_S_below D_W_below] := + g_rc[d; D_N_above D_E_above D_S_above D_W_above] * + corner(env, NORTHWEST, r - 1, c - 1)[χ_WNW; χ_NNW] * + edge(env, NORTH, r - 1, c)[χ_NNW D_N_above D_N_below; χ_NNE] * + corner(env, NORTHEAST, r - 1, c + 1)[χ_NNE; χ_ENE] * + edge(env, EAST, r, c + 1)[χ_ENE D_E_above D_E_below; χ_ESE] * + corner(env, SOUTHEAST, r + 1, c + 1)[χ_ESE; χ_SSE] * + edge(env, SOUTH, r + 1, c)[χ_SSE D_S_above D_S_below; χ_SSW] * + corner(env, SOUTHWEST, r + 1, c - 1)[χ_SSW; χ_WSW] * + edge(env, WEST, r, c - 1)[χ_WSW D_W_above D_W_below; χ_WNW] + g_rc_prec = twistdual(g_rc_prec, 2:5) + return g_rc_prec / norm_pref + δ * g_rc +end + +""" + peps_precondition(x, g, tol_state::NamedTuple, alg) + +Precondition the PEPS gradient `g` at the point `x = (peps, env)` according to the +preconditioner algorithm `alg`. Passing `alg = nothing` disables preconditioning and returns +`g` unchanged. See [`LocalPreconditioner`](@ref) for details on the local metric +preconditioner, and note that its regularization strength is set based on the current +gradient norm and iteration count tracked in `tol_state`. +""" +peps_precondition(x, g, tol_state::NamedTuple, alg::Nothing) = g + +function peps_precondition(x, g, tol_state::NamedTuple, alg::LocalPreconditioner) + peps, env = x + δ = alg.regularization * tol_state.gradnorm^2 / max(tol_state.iter, 1) + g_prec_unitcell = map(eachcoordinate(g)) do (r, c) + nf = _contract_site((r, c), InfiniteSquareNetwork(peps), env) + g_rc_prec, = linsolve(g[r, c], g[r, c], alg.solver_alg) do g_in + return apply_local_preconditioner(g_in, env, δ, (r, c), nf) + end + return g_rc_prec + end + return InfinitePEPS(g_prec_unitcell) +end diff --git a/src/algorithms/select_algorithm.jl b/src/algorithms/select_algorithm.jl index a571574b7..c121fb3fd 100644 --- a/src/algorithms/select_algorithm.jl +++ b/src/algorithms/select_algorithm.jl @@ -1,6 +1,7 @@ _alg_or_nt(::Type{T}, alg::NamedTuple) where {T} = T(; alg...) _alg_or_nt(::Type{T}, alg::A) where {T, A <: T} = alg _alg_or_nt(::Type{T}, alg::DynamicTol{<:T}) where {T} = alg +_alg_or_nt(::Type, ::Nothing) = nothing _alg_or_nt(T, alg) = throw(ArgumentError("unkown $T: $alg")) """ @@ -62,7 +63,7 @@ function select_algorithm( env₀; tol = Defaults.optimizer_tol, # top-level tolerance verbosity = 3, # top-level verbosity - boundary_alg = (;), gradient_alg = (;), optimizer_alg = (;), + boundary_alg = (;), gradient_alg = (;), optimizer_alg = (;), precondition_alg = (;), symmetrization = nothing, kwargs..., ) # adjust CTMRG tols and verbosity @@ -104,13 +105,31 @@ function select_algorithm( gradient_alg = _dynamic_tol_or_alg(gradient_alg; dynamic_tol_kwargs...) end + # adjust preconditioner verbosity and construct the preconditioner algorithm + if precondition_alg isa NamedTuple + defaults = (; + verbosity = verbosity ≤ 3 ? -1 : 3, + dynamic_tols = Defaults.precondition_dynamic_tols, + tol_min = Defaults.precondition_tol_min, + tol_max = Defaults.precondition_tol_max, + tol_factor = Defaults.precondition_tol_factor, + ) + precondition_kwargs = merge(defaults, precondition_alg) + dynamic_tol_kwargs, precondition_kwargs = _pop_dynamic_tol_kwargs(precondition_kwargs) + precondition_alg = PreconditionAlgorithm(; precondition_kwargs...) + precondition_alg = _dynamic_tol_or_alg(precondition_alg; dynamic_tol_kwargs...) + end + # adjust optimizer tol and verbosity if optimizer_alg isa NamedTuple defaults = (; tol, verbosity) optimizer_alg = merge(defaults, optimizer_alg) end - return PEPSOptimize(; boundary_alg, gradient_alg, optimizer_alg, symmetrization, kwargs...) + return PEPSOptimize(; + boundary_alg, gradient_alg, optimizer_alg, precondition_alg, + symmetrization, kwargs..., + ) end function select_algorithm( diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl index d25b19e7a..a1d31c8ff 100644 --- a/src/states/infinitepeps.jl +++ b/src/states/infinitepeps.jl @@ -142,6 +142,9 @@ function eachcoordinate(A::InfinitePEPS, dirs) return collect(Iterators.product(dirs, axes(A, 1), axes(A, 2))) end +Base.real(A::InfinitePEPS) = InfinitePEPS(real.(unitcell(A))) +Base.complex(A::InfinitePEPS) = InfinitePEPS(complex.(unitcell(A))) + ## Spaces TensorKit.spacetype(::Type{T}) where {T <: InfinitePEPS} = spacetype(eltype(T)) diff --git a/test/examples/heisenberg.jl b/test/examples/heisenberg.jl index 5a06daed9..8441998f6 100644 --- a/test/examples/heisenberg.jl +++ b/test/examples/heisenberg.jl @@ -139,7 +139,7 @@ end # continue with auto differentiation peps_final, env_final, E_final, = fixedpoint( - ham, peps, complex(env); # make environment complex explicitly + ham, complex(peps), complex(env); # make environment complex explicitly optimizer_alg = (; tol = gradtol, maxiter = 25), boundary_alg = (; maxiter = ctmrg_maxiter), gradient_alg = (; solver_alg = (; alg = :GMRES)),