PEPSKit v0.1.0

[lkdvos]

This contains discussions about working towards a first release of PEPSKit:

Currently, the list of features that we would like to have included consists of:

v0.1.0 included

General improvements:

• unified framework for CTM and boundary MPS
• support for fermionic symmetries
• testing framework
• documentation

CTM improvements:

• CTM gauge-fixing
• CTM fixed-point differentiation

AD improvements:

• rrule for svdsolve
• rrule for planar contractions

Future versions

• 2D partition functions
• 3D partition functions
• Environment managers

[pbrehmer]

I would work on improving various things around the CTMRG part of PEPSKit. Here's a preliminary list of things which I think are realistic for a first release:

• Proper differentiation of the CTM gradient: Currently, we differentiate through all CTMRG iterations which blows up the memory. Instead, one should evaluate the gradient based on the derivative of a single CTM iteration. Then one can compute the gradient in three ways. One can compute the geometric sum explicitly, one can iterate the linear problem manually or one can solve the linear problem using an iterative solver. I can implement all three possibilities and add an option to choose among them in the CTMRG algorithm struct.

• Gauge fixing of the CTM tensors: This ensures element-wise convergence of the environment tensor after a CTM iteration, and it is anyway needed for the efficient evaluation of CTM gradients. Here, I would first implement the easiest approach, as shown in Eq. (A6) of Anna & Bram's pre-print.

• Correct adjoint for the truncated SVD: See Eqs. (G1-6) in the pre-print. Since I'm anyway using the iterative SVD solver from KrylovKit, the best option would be to directly implement the adjoint in KrylovKit itself. However, one might need something like Lorentzian broadening on the divergent singular value differences $F_{ij} = (s_j^2 - s_i^2)^{-1}$ in the adjoint, which should probably stay separate from KrylovKit. So possibly one needs a custom SVD adjoint in PEPSKit anyway?

• Small technical improvements: I'm thinking of things like reusing previous environments in a CTMRG run, better output messages, choosing different SVD solvers in the CTMRG algorithm struct, etc.

Let me know what you think and if you have any input on this. I'm not quite sure yet how to split this up nicely into different PRs. I will probably start working on the differentiation and gauge fixing first, while hopefully not breaking the Heisenberg example.

[lkdvos]

I think even the Lorentzian broadening is not a bad idea even in KrylovKit. I don't think we currently have a clean way of hooking into keyword arguments such as tolerances for the backwards pass, which appears in the linear solvers and eigensolvers as well.

[pbrehmer]

Some additional thoughts:

• I would try to keep all contractions between the PEPS and the CTM environment in separate functions that dispatch on the PEPSTensor type. That way one could easily simulate also classical partition functions: just replace the PEPS by a MPO/4-legged tensor network operator and adjust the contraction functions (e.g. northwest_corner) to that type with the corresponding environment (edges with three legs instead of four). Most of the CTMRG code would stay exactly the same.

• Probably there should be an optimizer algorithm struct that is separate from the CTMRG algorithm struct. Things like the gradient evaluation method, reusing environments or not, keyword arguments that are passed to the optimizer and so on, could go in there.

[lkdvos]

Looks great to me!
At first glance, the following things come to my mind:

• Boundary MPS methods that leverage MPSKit for computing environments
• Generalizing CTMRG and boundary methods for the 3 obvious cases: 2d stat-mech with 3-legged edges, 2d quantum with 4-legged edges, and 3d stat-mech with 5-legged edges.
• Proper testing framework
• Proper documentation setup

In particular for many of these things, it would indeed be super nice if we can leverage multiple dispatch to re-use large parts of the code. Nevertheless, I would think that the different leg cases is something that can very easily be added in a future version, so focusing on a single case might not be a bad idea. I quite like the idea of having 2-dimensional partition functions in there asap to make testing a whole lot easier.

[pbrehmer]

Also sounds great!

Nevertheless, I would think that the different leg cases is something that can very easily be added in a future version, so focusing on a single case might not be a bad idea. I quite like the idea of having 2-dimensional partition functions in there asap to make testing a whole lot easier.

I anyway coded up an interface for 2D partition functions (for the square lattice Ising model) so this should be relatively easy to get in the first release. But other than that I agree, we should focus on the 2D quantum with 4-legged edges case.

[lkdvos]

Maybe also a worthy addition:

• Fermionic symmetries

[pbrehmer]

While I was looking through the code yesterday and doing some refactoring, some questions came to my mind:

• What's going on with the @diffset macro? What is it actually used for and why do we need it?

• Should we use PeriodicArrays from MPSKit? For now it does seem easier to just use the helper functions _prev(i, total) = mod1(i - 1, total) and _next(i, total) = mod1(i + 1, total). (This is what I've been doing so far.)

• I would suggest removing peps_above and peps_below from the CTMRGEnv struct. This would make the struct more generally applicable (e.g. also to the stat-mech cases). Also, many functions seem to have a separate PEPS argument anyway (this could again be easily generalized to some physical tensor, accounting for the stat-mech cases). What do you think about this?

[lkdvos]

• @diffset is something @maartenvd wrote to circumvent not being able to perform in-place assignments. You could think about is as replacing setindex! with setindex!! everywhere, which will do an in-place assignment when not performing AD, but create a new array when AD is happening. In particular for systems with unit-cells this becomes relevant, as you typically keep an array-like structure of environment tensors, which you are not allowed to change in-place because of the AD restrictions. (Note that it also does some things for .+=, .-= etc.)
• Originally we used PeriodicArray from MPSKit, which is definitely convenient as it replaces a lot of the _next and _prev calls, but at that point we could not get the AD to behave, so we went back to basics and replaced that with regular Array. I think that the getindex and setindex! rules for PeriodicArray might be silently implemented somewhere in Zygote, making them wrong. I haven't been able to find the time to figure that out, and would definitely rather use PeriodicArray than calling the mod1 everywhere manually, but our experience is definitely that when you are trying to get AD to run, it's better to start as basic as possible, because it is very hard to point at what is going wrong afterwards 😁
• I am 100% aboard with changing the CTMRGEnv like that. This is similar to something that is happening in MPSKit, where the environments know about the operator and state. The idea here is that the struct is really an environment manager, and if you ask it for the environments for a given state/operator, it checks to see whether that state/operator matches its internal one, and if so just returns the stored environments, otherwise recomputes them. While this is not a bad idea, it's far from perfect, as the matching is done by comparing pointers (addresses of the state tensors), and thus would not work if you change the tensors in-place, as then the pointer does not change. In any case, I think that the responsability of checking if the environments need to be recalculated should not be given to the CTMRGEnv struct, and if we really want this behaviour we could even consider writing an actual EnvironmentManager to wrap things in afterwards. (I would even be in favour of doing this in MPSKit as well, but it's quite a big refactor 🥲 )

[pbrehmer]

• That does make a lot of sense and seems very useful. I also encountered similar cases previously, where I ran into trouble with updating environment tensors during AD (in order to update the environment during optimization to reuse previous ones). I was just a bit confused since, IIRC, you implemented this behavior in TensorOperations already but this is only relevant for tensor contractions, I guess.

• Alright, then I will stick to normal arrays for now and just use _prev and _next. Maybe switching to PeriodicArray at a later stage (when the AD rules are figured out) does not seem to hard to do :-)

• The idea of an environment manager is very cool and I also find it natural somehow. Is there maybe a computationally cheap physical check to see whether an environment matches with a given state? (As opposed to just comparing pointers.) But I do see that this would entail a big refactoring effort! So maybe it's best to roll with separate physical state and CTMRGEnv structs for now and possibly add an EnvironmentManager afterwards.

[lkdvos]

What I can think of, in order of most-automated/expensive to least-automated/expensive:

1. In principle we could work with a checksum-type system, and store a hash of the data, which can be used afterwards to check if anything changed. I am not sure how expensive this is, as it would really need to hash all the data, i.e. every element of every tensor. This is what I would consider the equivalent of just storing a copy of the data, and checking if anything changes.

• (+) this is very robust, and automatically should handle even in-place changes of the tensors
• (+) this is not that memory expensive, as a single checksum is enough
• (-) this is potentially quite expensive, as the checksum has to be recomputed every time to check if the data has changed

2. We just use a checksum of the pointers to the tensors (or just an array of pointers) and compare these. This is a lot cheaper, as we just compare a very low amount of pointers, but is not robust against in-place changes of the tensors. In that case, it also requires a mechanism to manually invalidate the environments when such a change does happen, but has the added benefit that if you were to for example do a gauge transformation on the state, you would not necessarily have to invalidate the environments as the state is still the same, but the computer does not necessarily know this.

• (+) a lot cheaper computationally
• (+) more control over when environments are invalidated
• (-) less robust against in-place changes
• (-) possible hard to track bugs from assumptions on when environments are being invalidated and when they are not.

3. Manual management, where we just manually call recompute if and only if we want to do that.

• (+) cheapest option
• (+) most control
• (-) least automated

In my opinion, the way to go is to implement option 3, get it to work, and afterwards experiment with writing a wrapper around that according to either option 1 or 2. This has the added benefit of keeping the code structured a bit better, in the sense that we can keep the environment manager part separate from the environment computing part. Of course, I might be missing some options, so I am definitely open to revising this opinion 😉

[maartenvd]

I'm a bit on the fence for diffset. It might introduce subtle bugs, as setindex! modifies memory in-place, but @diffset will map this to something that does not. It's really there to be able to write "envs[i] = tensor" without zygote imploding. I've now been working with python's ad infrastructure and it's both way simpler (essentially trackedarrays) and far far superior. I also don't think it would be that difficult to write a simple trackedtensormap type, and in the long run it really might be worth the added effort. As far as bundling environments with their state goes - the idea was also that environments are always calculated for a particular state, so why not store a pointer to that state? It should have no cost (a single pointer) and it's now easy to avoid bugs where you forgot to update the environments.

…

On Sun, 18 Feb 2024, 14:07 Paul Brehmer, ***@***.***> wrote: - That does make a lot of sense and seems very useful. I also encountered similar cases previously, where I ran into trouble with updating environment tensors during AD (in order to update the environment during optimization to reuse previous ones). I was just a bit confused since, IIRC, you implemented this behavior in TensorOperations already but this is only relevant for tensor contractions, I guess. - Alright, then I will stick to normal arrays for now and just use _prev and _next. Maybe switching to PeriodicArray at a later stage (when the AD rules are figured out) does not seem to hard to do :-) - The idea of an environment manager is very cool and I also find it natural somehow. Is there maybe a computationally cheap physical check to see whether an environment matches with a given state? (As opposed to just comparing pointers.) But I do see that this would entail a big refactoring effort! So maybe it's best to roll with separate physical state and CTMRGEnv structs for now and possibly add an EnvironmentManager afterwards. — Reply to this email directly, view it on GitHub <#12 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJKVCRJJYK5QIVD7L7MPVTYUH4HZAVCNFSM6AAAAABDIHV7WWVHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM4DKMBXGE2DK> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[pbrehmer]

As far as bundling environments with their state goes - the idea was also that environments are always calculated for a particular state, so why not store a pointer to that state? It should have no cost (a single pointer) and it's now easy to avoid bugs where you forgot to update the environments.

I do agree that it would make sense to store a state pointer. However, we would need a type for general physical states that accounts for all the cases that one would use a CTMRG environment for, i.e. a combined type of peps_above and peps_below, a 4-legged MPO for 2D stat-mech and so on. In that case the CTMRGEnv struct should be fairly generic and widely applicable.

[lkdvos]

I don't think there are any issues with just having a parametric CTMRGEnv struct, storing an arbitrary state. The biggest thing is just that both ways of handling "forgetting to update environments" have their issues. On the one hand, storing a pointer to the state avoids "forgetting to update environment" issues, but adds "forgetting to invalidate environment" issues. On the other hand, not storing a pointer to the state does the opposite.
My biggest gripe with the former is that Julia really does promote in-place assignments, and thus most of the time, you'll have to then invalidate the environment anyways. I would say bugs where you forgot to update the environments are easier to track/comprehend than the ones where you forgot to tell the environment that it should be recomputing, but maybe that's just my personal opinion.

In any case, I do really feel like these things should be separated, in the sense that the environment should handle how to compute environments, and a separate environment manager should try and handle when to do so. In that sense, I'd rather first implement an environment, and only then add the manager layer on top. This sounds like a more organised way of structuring the code, which might be easier to maintain, and then also the manager can be re-used for different types of environments, as a similar discussion can be held for boundary MPS environments.

[lkdvos]

Considering @diffset, I do agree that it is both a bit unclear about what is going on, and might introduce very subtle bugs. On the other hand, working with things like map is often quite tedious, and also will not work in-place when you are not using AD. It's a bit disturbing how hard it is to attempt to both have efficient forwards and backwards code, so maybe in PEPSKit we really just should not be attempting to do any in-place assignments at first, as AD will spoil this anyways

[pbrehmer]

In any case, I do really feel like these things should be separated, in the sense that the environment should handle how to compute environments, and a separate environment manager should try and handle when to do so. In that sense, I'd rather first implement an environment, and only then add the manager layer on top.

Personally, I also feel that way and that's the way I have now started refactoring the CTMRG part (of course without an environment manager). I will probably push some code and open a PR tomorrow, so then we can discuss some designs more concretely.

[pbrehmer]

Regarding the rrule for svdsolve: For me personally it would probably be easier/faster to first implement a rrule for the truncated SVD here in PEPSKit, before getting into KrylovKit. (I would wrap the svdsolve in some function and define a rrule on that.) Since I don't have a proper overview over KrylovKit and its super generic interface yet, this would probably take me a bit. Implementing it here directly should be a matter of a few hours (refactoring the code I already have anyway).

Maybe once we have it here, it should be manageable to transfer it to KrylovKit with a little bit of help :) So should I first open a PR here in PEPSKit?

[lkdvos]

I think it's definitely fine to have it here for now, I think we want to transition to package extensions for krylovkit anyways, so it should probably go there afterwards. I'm back in the office next week, I'll have a look at setting that up