Fixed-point CTMRG with implicit differentiation#

Module yastn.tn.fpeps.envs.fixed_pt provides yastn.tn.fpeps.envs.fixed_pt.fp_ctmrg(), a CTMRG driver whose output environment is differentiable with respect to the PEPS tensors through the fixed-point condition rather than through the unrolled CTMRG iterations. It requires the torch backend.

Given a converged CTM environment \(C^\ast\) of the state \(A\), one CTMRG sweep \(f\) leaves it invariant, \(C^\ast = f(C^\ast, A)\). The implicit-function theorem then gives the derivative of any observable \(\mathcal{L}(C^\ast, A)\) without storing the forward sweeps,

\[\frac{d\mathcal{L}}{dA} = \frac{\partial \mathcal{L}}{\partial A} + \frac{\partial \mathcal{L}}{\partial C} \Big(1 - \frac{\partial f}{\partial C}\Big)^{-1} \frac{\partial f}{\partial A},\]

where the inverse is expanded as a Neumann series \(\sum_{n \ge 0} (\partial f / \partial C)^n\). This works only if the environment converges element-wise, which CTMRG alone does not guarantee: each sweep is free to change the gauge of the environment bonds (a unitary on each bond and a phase on each tensor). The gauge is therefore fixed explicitly before the fixed-point condition is used.

Forward pass#

yastn.tn.fpeps.envs.fixed_pt.fp_ctmrg() wraps the custom autograd function yastn.tn.fpeps.envs.fixed_pt.FixedPoint, whose forward pass does the following.

Converge. FixedPoint.get_converged_env() runs yastn.tn.fpeps.EnvCTM.ctmrg_() until the corner spectra stop changing (corner_tol) or max_sweeps is reached. An optional stall detector (stuck_block) declares a solve non-convergent when the running minimum of the corner change stops improving across blocks of sweeps. Importantly, max_de – the largest change in environment element moduli between consecutive sweeps – can stay finite while max_dsv, which compares corner spectra, has already reached corner_tol. Both are fields of the CTMRG_out returned by yastn.tn.fpeps.EnvCTM.iterate_(). This is due to residual gauge freedom, which is why the next step is needed.
One more sweep. A single CTM step with the options in ctm_opts_fp (full-rank SVD, so that the step has a well-defined derivative) produces env_new from the converged env_old.
Gauge fixing. find_gauge_multi_sites() finds, for every site and every edge direction, the bond gauge sigma relating the old and new edge tensors T along a full row or column of the unit cell (fast_env_T_gauge_multi_sites(), the leading-eigenvector method of arXiv:2311.11894). The gauges are collected in an EnvGauge, applied to edges and corners by apply_sigma(), and the residual per-tensor U(1) phase is extracted by U1_phase() and applied with apply_U1_(). The transformed env_new then equals env_old element-wise.
Save. The converged environment and the gauge are stored for the backward pass; the environment tensors are returned as one flat torch tensor so that autograd tracks them.

If any step fails, a NoFixedPointError is raised (CTM does not converge, the symmetry sectors of the edge tensors change between sweeps, or no gauge is found).

Backward pass#

FixedPoint.backward() builds the gauge-fixed sweep FixedPoint.fixed_point_iter(), i.e. one CTM step followed by the saved sigma and phase transformations, as a function of the environment data and the PEPS data, and takes its vector-Jacobian products. The gauge transformation part is detached from the computation graph. The Neumann series is summed until the increment falls below corner_tol, or it has not decreased for neumann_patience consecutive steps (the best estimate so far is then kept), or until max_sweeps steps. The series converges only when the spectral radius of \(\partial f/\partial C\) is below one, which is not guaranteed by the forward convergence. The result is the gradient with respect to the raw PEPS tensors in the order of site2index.

Forward options carry over to the backward#

fp_ctmrg takes two option sets. The fixed-point step inherits from the forward one.

ctm_opts_fp is not an independent configuration. It starts as a copy of ctm_opts_fwd and overrides it selectively, with opts_svd merged key by key rather than replaced. For example:

env = fp_ctmrg(env,
               ctm_opts_fwd={'opts_svd': {'D_total': 64, 'tol': 1e-8},
                             'corner_tol': 1e-8,
                             'max_sweeps': 100,
                             'method': '2x2',
                             'use_qr': False},
               ctm_opts_fp={'opts_svd': {'policy': 'fullrank'}})

gives a fixed-point step that keeps the forward’s method, use_qr, max_sweeps and corner_tol, and whose opts_svd is {'D_total': 64, 'tol': 1e-8, 'policy': 'fullrank'} – the forward truncation with the full-rank SVD policy merged in. Only what you name in ctm_opts_fp changes.

That inheritance extends to the backward pass, because FixedPoint reverses the CTM step by hand: the Neumann loop takes its iteration budget from the fixed-point step’s max_sweeps and its gradient tolerance from its corner_tol. Raising the forward sweep budget therefore also allows the backward series more terms, which is usually what you want; setting them apart is done by naming them in ctm_opts_fp.

One backward control has no forward counterpart and so is its own option:

neumann_patience: How many Neumann iterations may pass without the gradient estimate improving before the series is judged non-contracting and stopped, returning the best estimate rather than the diverged tail. The default is 10.

The resolved form of this is yastn.tn.fpeps.envs.FixedPointOpts, which can be passed directly as opts= instead of the two dicts. See CTM options.

Running on several devices: single-node multi-gpu#

For multi-site unit cells, CTM can be parallelized across the unit-cell sites. Such multi-device runs, passing devices=[...] with more than one entry, route the forward convergence, the extra CTM step and the Neumann iterations through the distributed CTM workers of yastn.tn.fpeps.envs._env_ctm_dist_mp_AD. Note that devices=['cpu', 'cpu'] is a legitimate configuration and exercises the full multiprocess path.

C4v-symmetric variant#

yastn.tn.fpeps.envs.fixed_pt_c4v.fp_ctmrg_c4v() is the counterpart for C4v-symmetric single-site iPEPS, used with yastn.tn.fpeps.EnvCTM_c4v. It takes the same two option dicts with the same inheritance.

It is a separate implementation, not a thin wrapper: the gauge-fixing works on the single C and T tensor rather than a full unit cell, and its Neumann loop uses a different stopping heuristic. Do not assume a change to one applies to the other.

Example#

A gradient-based optimization with the CTM workers on several GPUs. The state, the gauge fixing and the main-side autograd graph live on the default_device of the yastn config (the home device); the CTMRG stages of the forward convergence, of the fixed-point step and of the Neumann backward are dispatched to one spawned worker process per entry of devices. Keeping the home device out of devices leaves the \(O(\chi^2)\) backward graph off the worker GPUs. The worker pool is created on the first call and reused by every later fp_ctmrg call with the same devices and config, so the spawn cost is paid once per run; it is shut down at interpreter exit. Because the workers are spawned, the script needs the usual if __name__ == '__main__' guard:

import torch
import yastn
import yastn.tn.fpeps as fpeps
from yastn.tn.fpeps.envs.fixed_pt import fp_ctmrg, NoFixedPointError

def main():
    home = 'cuda:0'                                   # state + autograd graph
    devices = ['cuda:1', 'cuda:2', 'cuda:3']          # CTM worker pool
    config = yastn.make_config(backend='torch', sym='U1', fermionic=True,
                               default_device=home, default_dtype='complex128')

    psi = fpeps.Peps(geometry, tensors=...)           # tensors on `home`
    params = [psi[s]._data.requires_grad_(True) for s in psi.sites()]
    opt = torch.optim.LBFGS(params, lr=1.0, max_iter=1, history_size=10)

    chi = 64
    env_leg = yastn.Leg(config, s=1, t=(0,), D=(chi,))
    ctm_opts_fwd = {'method': '2x2', 'corner_tol': 1e-8, 'max_sweeps': 200,
                    'opts_svd': {'D_total': chi, 'tol': 1e-10,
                                 'eps_multiplet': 1e-8, 'truncate_multiplets': True},
                    'use_qr': False, 'stuck_block': 10, 'verbosity': 0}
    ctm_opts_fp = {'opts_svd': {'policy': 'fullrank'}, 'corner_tol': 1e-8,
                   'max_sweeps': 100, 'neumann_patience': 10}

    env = fpeps.EnvCTM(psi, init='eye', leg=env_leg)

    def closure():
        nonlocal env
        opt.zero_grad()
        try:
            # reuse the previous environment as the starting point
            env = fp_ctmrg(env, ctm_opts_fwd=ctm_opts_fwd,
                           ctm_opts_fp=ctm_opts_fp, devices=devices)
        except NoFixedPointError as e:
            raise                                     # or perturb psi and retry
        loss = energy_per_site(psi, env)              # any observable of env
        loss.backward()                               # IFT gradient into `params`
        return loss

    for step in range(100):
        loss = opt.step(closure)
        torch.cuda.empty_cache()                      # home device; the pool
        print(step, loss.item())                      # frees its own cache

if __name__ == '__main__':
    main()

With devices=[home] (or devices=None) the same script runs serially in the main process without spawning workers.

API#

yastn.tn.fpeps.envs.fixed_pt.fp_ctmrg(env: EnvCTM, ctm_opts_fwd: dict | None = None, ctm_opts_fp: dict | None = None, devices=None, *, opts: FixedPointOpts | None = None) → EnvCTM[source]#

Compute the fixed-point environment for the given state using CTMRG. First, run CTMRG until convergence then find the gauge transformation guaranteeing element-wise convergence of the environment tensors. Enables backward differentiation through the fixed-point iteration to compute the gradients of the environment tensors with respect to the state parameters, via Neumann series expansion of the fixed-point iteration.

Parameters:

env (EnvCTM) – CTM environment.
ctm_opts_fwd (dict | None) – Options for the forward CTMRG convergence. See yastn.tn.fpeps.envs.CTMOpts for the accepted keys.
ctm_opts_fp (dict | None) – Overrides for the gauge-fixing CTM step. Applied on top of ctm_opts_fwd, which it otherwise inherits – including max_sweeps and corner_tol, which the Neumann backward loop then uses as its iteration budget and gradient tolerance.
devices (list[str] | None) – Device list for the CTM step. With one device everything runs serially. With more than one, the forward convergence, the fixed-point CTM step and the Neumann backward all use the AD-aware distributed dispatch on those devices. None falls back to [env.config.default_device].
opts (FixedPointOpts | None) – The two option dicts above, pre-resolved. Takes precedence when given.

Returns:

Environment at the fixed point.

Return type:

EnvCTM

class yastn.tn.fpeps.envs.fixed_pt.FixedPoint(*args: Any, **kwargs: Any)[source]#

static forward(ctx, env: EnvCTM, opts: FixedPointOpts, *state_params)[source]#

Parameters:

env (yastn.tn.fpeps.EnvCTM) – Current environment to converge.
opts (FixedPointOpts) – Resolved options. opts.fwd drives the forward convergence; opts.fp – which inherits from opts.fwd – drives the gauge-fixing step and, through its max_sweeps / corner_tol, the Neumann backward loop.
state_params (Sequence[torch.Tensor]) – Tensors of the underlying Peps state.

Returns:

env_converged (yastn.tn.fpeps.EnvCTM) – The gauge-fixed converged environment.
env_t_meta (Sequence) – Metadata needed to rebuild the environment tensors.
env_slices (Sequence[slice]) – Slices locating each environment tensor within env_1d.
env_1d (torch.Tensor) – The environment data as one flat tensor, the output autograd tracks.

get_converged_env(opts: CTMOpts, devices=None, stuck_block=0, stuck_window=3, stuck_factor=2.0, stuck_min_sweeps=60)[source]#

Run the forward CTMRG loop until the corner spectra stop changing (corner_tol) or max_sweeps is reached.

Early no-fixed-point detection (stuck_block > 0). A CTM solve that will never converge is not distinguishable from a slow one by |delta_C| alone: on marginal states |delta_C| oscillates over a decade sweep-to-sweep, so the raw running minimum is pinned by lucky outliers. What DOES separate them is the running minimum over BLOCKS of sweeps: a converging solve drives it down geometrically, a stuck one leaves it flat (or rising).

The rule: after stuck_min_sweeps, at the end of every block of stuck_block sweeps, the block minimum must have improved by at least a factor stuck_factor relative to stuck_window blocks earlier. If it has not, the solve is declared non-convergent and the loop stops, leaving converged=False – which the caller turns into NoFixedPointError, and the optimizer’s existing handler recovers by perturbing the state. stuck_block=0 disables the check entirely.

Parameters:

env (yastn.tn.fpeps.EnvCTM) – Current environment to converge.
opts (CTMOpts) – Options driving the forward convergence, i.e. the fwd field of the resolved yastn.tn.fpeps.envs.FixedPointOpts. Its corner_tol (or conv_check, when given) is the convergence test applied here, and its max_sweeps bounds the loop.
devices (list[str] | None) – Device list for the CTM step. With one device everything runs serially.
stuck_block (int) – Number of sweeps per block for the stuck-block check. 0 disables the check.
stuck_window (int) – Number of blocks to look back for the stuck-block check.
stuck_factor (float) – Minimum improvement factor for the stuck-block check.
stuck_min_sweeps (int) – Minimum number of sweeps before the stuck-block check is applied.

Returns:

env (yastn.tn.fpeps.EnvCTM) – The converged environment.
converged (bool) – Whether the convergence test was met. False also when the stuck-block check stopped the loop; the caller turns this into NoFixedPointError.
conv_history (list) – Per-sweep CTMRG_out records of the forward loop.
t_ctm (float) – Seconds spent in the CTM sweeps.
t_check (float) – Seconds spent in the convergence test.

yastn.tn.fpeps.envs.fixed_pt_c4v.fp_ctmrg_c4v(env: EnvCTM_c4v, ctm_opts_fwd: dict | None = None, ctm_opts_fp: dict | None = None, *, opts: FixedPointOpts | None = None)[source]#

Parameters:

env (EnvCTM_c4v) – C4v-symmetric CTM environment.
ctm_opts_fwd (dict | None) – Options for the forward CTMRG convergence. See yastn.tn.fpeps.envs.CTMOpts for the accepted keys.
ctm_opts_fp (dict | None) – Overrides for the gauge-fixing CTM step, applied on top of ctm_opts_fwd which it otherwise inherits – including max_sweeps and corner_tol, which the Neumann backward loop then uses as its budget and tolerance.
opts (FixedPointOpts | None) – The two option dicts above, pre-resolved.

Returns:

Environment at the fixed point.

Return type:

EnvCTM_c4v

class yastn.tn.fpeps.envs.fixed_pt.NoFixedPointError(code, message=None)[source]#

class yastn.tn.fpeps.envs.fixed_pt.EnvGauge(geometry)[source]#

Gauge matrices used in fixed-point AD.

classmethod from_dict(d, config=None)[source]#: De-serializes EnvGauge from the dictionary d.

to_dict(level=2)[source]#: Return a dictionary representation of the object.

yastn.tn.fpeps.envs.fixed_pt.find_gauge_multi_sites(env_old, env, verbose=False)[source]#

yastn.tn.fpeps.envs.fixed_pt.fast_env_T_gauge_multi_sites(config, T_olds, T_news)[source]#

yastn.tn.fpeps.envs.fixed_pt.apply_sigma(env, env_gauge)[source]#

yastn.tn.fpeps.envs.fixed_pt.U1_phase(env_old, fixed_env)[source]#

yastn.tn.fpeps.envs.fixed_pt.apply_U1_(env, phases_dict)[source]#