Add Neighborhood Tensor Update (NTU)#144
Conversation
lkdvos
force-pushed
the
master
branch
2 times, most recently
from
ad4945f to
37ace4c
Compare
Yue-Zhengyuan
changed the title
Yue-Zhengyuan
force-pushed
the
bondenv-ntu
branch
from
9258396 to
ecbbee4
Compare
Yue-Zhengyuan
force-pushed
the
bondenv-ntu
branch
2 times, most recently
from
85d2c44 to
ecbbee4
Compare
Yue-Zhengyuan
force-pushed
the
bondenv-ntu
branch
from
1a231ec to
4c5f078
Compare
Yue-Zhengyuan
force-pushed
the
bondenv-ntu
branch
from
7c9ea13 to
f9d029a
Compare
|
I'll keep this as a draft for a while to not waste resources on CI. |
Yue-Zhengyuan
left a comment
Yue-Zhengyuan
left a comment
There was a problem hiding this comment.
Some comments that may help with reviewing.
| twistdual!(ket, ax) | ||
| h === nothing && continue | ||
| axes, biperm = _permute_to_first(axes, ax) | ||
| ket = permute(h, ((1,), (2,))) * permute(ket, biperm) |
There was a problem hiding this comment.
This is slightly unfortunate that I have to permute h every time, because I decide to put the bra legs in front of the ket legs, and all legs in the codomain (which was the result of using ncon).
There was a problem hiding this comment.
I assume we can leave this for follow-up optimizations?
| a, X = if stype1 == :first | ||
| bond_tensor_first(A; trunc = qrtrunc) | ||
| else | ||
| @assert stype1 == :middle | ||
| bond_tensor_midnext(A; trunc = qrtrunc) | ||
| end | ||
| b, Y = if stype2 == :last | ||
| bond_tensor_last(B; trunc = qrtrunc) | ||
| else | ||
| @assert stype2 == :middle | ||
| bond_tensor_midprev(B; trunc = qrtrunc) | ||
| end |
There was a problem hiding this comment.
I may better explain why the middle sites need different treatment from the first/last sites when factoring out the reduced bond tensor.
Consider the 3-site gate MPO acting on PEPS. Its virtual bond dimension is
|
--- X --- --- a ---
| ↘
The dimension of the leg connecting X and a is a (which is going to be optimized with ALS) is
But if we keep the physical leg in X,
|
--- X --- --- a ---
| ↘
then after applying the gate, the leg connecting X, a has dimension a remains a, thus being less efficient.
The same argument applies to the last site, for which it is also better to transfer the physical leg to the bond tensor b to reduce the dimension of the bond environment.
However, things are different for middle sites.
╱
--- M --
╱ |
↓
--- g ---
↓
After applying the gate, the cluster bond dimension increases to
|
--- b --- --- Y ---
↘ |
The leg connecting b, Y have dimension b has size Y,
|
--- b --- --- Y ---
| ↘
then the leg connecting b, Y has a smaller dimension b also having a smaller size
There is also an important subtlety: for middle sites, the unitary tensors X or Y also contains the effect of the gate, in contrast to the first/last tensors, for which X, Y can be found prior to applying the gate. Thus, it seems that varying the bond tensors alone cannot reach the optimal truncation result except for NN gates.
Finally, a technical question on whether it is good to use symbols to distinguish the different cases.
There was a problem hiding this comment.
I think symbols should be fine, the only other thing I can think of would be to create some @enum but I don't think it would matter in the end.
I think I follow the explanation, it seems reasonable that it is not possible to just get rid of the extra legs in the middle and that this doesn't buy us anything there.
| function bond_tensor_midnext(A::PEPSTensor; kwargs...) | ||
| X, a = left_orth!(permute(A, ((1, 2, 4, 5), (3,)); copy = true); kwargs...) | ||
| X = permute(X, ((1,), (2, 5, 3, 4))) | ||
| a = insertrightunit(a, 1) |
There was a problem hiding this comment.
For middle tensors, the physical leg is not transferred to the reduced bond tensor a. A trivial physical leg is added in order to fit the current bond_truncate code.
| showinfo = (iter == 1) || (iter % check_interval == 0) || (iter == it.nstep) | ||
| !showinfo && continue | ||
| # bond weight change | ||
| Δλ = hasproperty(info0, :wts) ? compare_weights(info.wts, info0.wts) : NaN |
There was a problem hiding this comment.
Here the bond weights are simply the bond_truncate SVD spectra. They are not the same as those produced by BP gauge fixing. For the first iteration, since the initialization does not provide information on bond weights, Δλ is set to NaN.
If we incorporate BP gauge fixing in the future, the bond truncation spectra should be replaced by the diagonalized BP messages. In this case, we can also find Δλ for the first step by gauge fixing the initial state.
| struct NTUState{S <: InfiniteState, N <: Number} | ||
| "number of performed iterations" | ||
| iter::Int | ||
| "evolved time" | ||
| t::N | ||
| "PEPS/PEPO" | ||
| psi::S | ||
| end |
There was a problem hiding this comment.
If we incorporate BP gauge fixing, NTUState should also include the BPEnv. It will also be used to construct bond environments. But for now I just leave it out.
Yue-Zhengyuan
force-pushed
the
bondenv-ntu
branch
from
3a8c969 to
03a160f
Compare
|
@lkdvos Hope to get merged soon so I can proceed from here. |
lkdvos
left a comment
lkdvos
left a comment
There was a problem hiding this comment.
Sorry for taking so long to get to this!
| function _svd_cut!(t::AbstractTensorMap) | ||
| t1, s, t2 = svd_trunc!(t; trunc = truncrank(1)) | ||
| return absorb_s(t1, s, t2) | ||
| end |
There was a problem hiding this comment.
Out of curiosity: do you need the weight to be distributed equally? Otherwise this might be replaced with left_orth!(t; trunc = truncrank(1)) or right_orth!(t; trunc = truncrank(1)), which I think removes some intermediate allocations.
There was a problem hiding this comment.
On a separate note, do we need this to be trivially charged for everything to work? It might reduce quite a bit of the number of legs if we can do something like:
A, B = left_orth!(t; trunc = truncspace(oneunit(spacetype(t))))
return removeunit(A, numind(A)), removeunit(B, 1)There was a problem hiding this comment.
do you need the weight to be distributed equally?
Nice observation! Actually no, since this D = 1 leg is in the end contracted, so how we distribute s is irrelevant.
do we need this to be trivially charged for everything to work?
It seems that _svd_cut is always applied to a positive map, which then (if my feeling is right) ensures that the leading singular value has trivial charge. I'll double-check this.
In YASTN they don't need to worry about it, since they only support Abelian symmetries and allow a tensor to have nonzero total charge.
| neighbors = [(-1, 0), (0, -1), (1, 0), (1, 1), (0, 2), (-1, 1)] | ||
| m = collect_neighbors(state, row, col, neighbors) | ||
| X, Y = _prepare_site_tensor(X), _prepare_site_tensor(Y) | ||
| # southwest half |
There was a problem hiding this comment.
Do you need these to be normalized separately? In principle, by putting the entire benv contraction into a single @tensor call we end up with slightly less intermediate permutations. For example, here the benv_sw and benv_ne will be permuted for the final contraction, and since these are the biggest objects this might actually matter quite a bit.
| vecr[D21 D20 D31 D30 -3 -4 -5 -6] | ||
| normalize!(vecr, Inf) | ||
| # combine left and right part | ||
| @tensor benv[-1 -2; -3 -4] := vecl[1 2 -1 -3 3 4] * vecr[1 2 -2 -4 3 4] |
There was a problem hiding this comment.
Same comment here about a single @tensor call removing intermediate permutations
| twistdual!(ket, ax) | ||
| h === nothing && continue | ||
| axes, biperm = _permute_to_first(axes, ax) | ||
| ket = permute(h, ((1,), (2,))) * permute(ket, biperm) |
There was a problem hiding this comment.
I assume we can leave this for follow-up optimizations?
| a, X = if stype1 == :first | ||
| bond_tensor_first(A; trunc = qrtrunc) | ||
| else | ||
| @assert stype1 == :middle | ||
| bond_tensor_midnext(A; trunc = qrtrunc) | ||
| end | ||
| b, Y = if stype2 == :last | ||
| bond_tensor_last(B; trunc = qrtrunc) | ||
| else | ||
| @assert stype2 == :middle | ||
| bond_tensor_midprev(B; trunc = qrtrunc) | ||
| end |
There was a problem hiding this comment.
I think symbols should be fine, the only other thing I can think of would be to create some @enum but I don't think it would matter in the end.
I think I follow the explanation, it seems reasonable that it is not possible to just get rid of the extra legs in the middle and that this doesn't buy us anything there.
2 participants
(Finally, after being delayed for over a year)
This PR adds the Neighborhood Tensor Update (NTU, arXiv 2107.06635), which works for both iPEPS and iPEPO, and supports nearest neighbor 2-site gates and arbitrary MPO gates (in particular, the next-nearest neighbor 3-site MPO).
One caveat to be emphasized is that after applying an N-site gate (N > 2), the updated bonds are truncated sequentially without going back. In other words, it does not perform an ALS optimization of all site tensors that are updated, since producing the environment surrounding a general cluster of sites is complicated, and the following ALS optimization will be extremely costly.
When truncating a bond after applying a time evolution gate, the NTU uses some neighboring tensors as the bond environment, instead of using only a few bond weights (simple update) of the full iPEPS approximated by CTM (full update), "interpolating" between the two extremes. Currently the NN, NN+ and NNN bond environments (see description in the YASTN package) are implemented.