Differentiable Programming Tensor Networks

Hai-Jun Liao; Jin-Guo Liu; Lei Wang; Tao Xiang

Source-linked AI summary

Differentiable Programming Tensor Networks

Hai-Jun Liao, Jin-Guo Liu, Lei Wang, Tao Xiang

arXiv:1903.09650v2Mar 22, 2019cond-mat.str-elquant-ph

TL;DR

The paper addresses the difficulty of deriving accurate, scalable gradients for complex tensor-network algorithms, especially in higher-dimensional and infinite systems. It formulates entire tensor-network programs as computation graphs and develops stable tensor-decomposition differentiation plus memory-efficient fixed-point backpropagation. The resulting framework computes higher-order physical observables and supports state-of-the-art variational optimization for infinite tensor-network states.

Problem
Gradient-based optimization and higher-order calculations for complex tensor networks are limited by costly contractions, nonlocal dependencies, unstable decomposition derivatives, and cumbersome manual gradient derivations.
Method
The paper expresses complete tensor-network algorithms as differentiable computation graphs and develops stable differentiation through SVD plus efficient backpropagation through fixed-point iterations.
Results
The framework computes Ising specific heat by second-order differentiation of free energy and achieves the lowest variational energy for the infinite-square-lattice Heisenberg model at D = 4, 5, 6, 7.
Takeaways & Limitations
Automatic differentiation reduces gradient-derivation effort while retaining forward-contraction complexity and extends tensor-network programs to gradient- and Hessian-based optimization and physical observables.

Abstract

from arXiv · show

Differentiable programming is a fresh programming paradigm which composes parameterized algorithmic components and trains them using automatic differentiation (AD). The concept emerges from deep learning but is not only limited to training neural networks. We present theory and practice of programming tensor network algorithms in a fully differentiable way. By formulating the tensor network algorithm as a computation graph, one can compute higher order derivatives of the program accurately and efficiently using AD. We present essential techniques to differentiate through the tensor networks contractions, including stable AD for tensor decomposition and efficient backpropagation through fixed point iterations. As a demonstration, we compute the specific heat of the Ising model directly by taking the second order derivative of the free energy obtained in the tensor renormalization group calculation. Next, we perform gradient based variational optimization of infinite projected entangled pair states for quantum antiferromagnetic Heisenberg model and obtain start-of-the-art variational energy and magnetization with moderate efforts. Differentiable programming removes laborious human efforts in deriving and implementing analytical gradients for tensor network programs, which opens the door to more innovations in tensor network algorithms and applications.

I. INTRODUCTION

The paper frames tensor-network algorithms as fully differentiable programs, addressing the difficulty of optimizing higher-dimensional networks and manually deriving gradients. Automatic differentiation enables efficient derivatives, including higher-order quantities, while supporting gradient-based optimization across general tensor-network structures.

Motivation: Tensor-network optimization in two or higher dimensions is difficult because contractions are costly and efficient general optimization schemes are lacking.The challenge is especially pronounced for infinite translationally invariant systems, where one tensor influences the objective nonlocally.
Motivation: Numerical derivatives have limited accuracy and efficiency, while manually deriving gradients is manageable only for simple tensor-network structures.These limitations have restricted broad adoption of gradient-based optimization for more complex systems.
Contribution: Differentiable programming composes the whole tensor-network program as a computation graph and computes higher-order derivatives efficiently to numerical precision.The approach is presented as general across lattice geometries, Hamiltonians, and contraction schemes.
Contribution: The paper focuses on two-dimensional infinite tensor networks, requiring stable differentiation through SVD and memory-efficient backpropagation through fixed-point iterations.These techniques support variational optimization of tensor-network states.
Automatic differentiation: A computation graph represents data and dependencies between elementary computation steps, allowing reverse-mode differentiation to propagate gradients backward from a scalar output.In general graphs, adjoints accumulate contributions from multiple downstream paths.
Automatic differentiation: Reverse-mode automatic differentiation computes gradients through local vector-Jacobian products in one forward and one backward traversal of the graph.Primitive operations provide backward functions without explicitly storing full Jacobian matrices.

B. Computation graphs of tensor network contractions

Tensor-network programs map input tensors to scalar physical outputs and can be differentiated by expressing their contractions as computation graphs. The paper emphasizes approximate contractions, especially infinite networks whose algorithms combine truncation and variational steps.

Program representation: Tensor-network programs map input tensors representing partition functions or wavefunctions to scalar physical quantities through contraction computations.Automatic differentiation is applied by representing these contractions as a computation graph.
Program representation: The infinite one-dimensional Ising partition function can be evaluated by repeatedly squaring its Boltzmann-weight matrix and tracing the result.Differentiation then backpropagates through matrix trace and multiplication operations.
Approximate contractions: The paper focuses on approximate tensor contractions involving truncated factorizations or variational approximations, particularly for infinite tensor networks.The bulk tensor is the fundamental data structure, and contraction schemes include coarse-graining and corner-transfer-matrix methods.

1. Tensor renormalization group

TRG contracts tensor networks by iteratively factorizing and blocking bulk tensors. Its computation graph consists of SVD, tensor contractions, permutations, and rescaling operations repeated until the thermodynamic limit is reached.

TRG iteration: TRG splits each bulk tensor in two ways using truncated SVD, retaining singular components up to bond dimension χ.The resulting tensors are assembled into a new tensor for the next iteration.
TRG iteration: TRG grows the lattice size exponentially, reaching the thermodynamic limit after a few tens of iterations.Tensor elements are rescaled after each iteration for numerical stability.
TRG computation graph: Figure 2 treats each tensor as a computation-graph node and identifies SVD and tensor contractions as primitive functions.The TRG panel depicts one iteration step within this graph.
TRG complexity: TRG has computational cost O(χ6) and memory cost O(χ4).After unrolling, its computation graph resembles a chain, with differentiable tensor permutations, truncated SVD, and contractions.

2. Corner transfer matrix renormalization group

CTMRG iteratively updates corner and edge tensors that represent the environment of a bulk tensor, using contractions, truncated SVD, and isometric projections until convergence. Differentiating this process requires handling complex dependencies, practical memory costs, and stable backward rules for tensor algebra.

CTMRG iteration: CTMRG seeks converged corner and edge tensors representing the environment degrees of freedom of the bulk tensor.Its computation graph has a more complex topology than the TRG graph.
CTMRG iteration: Each CTMRG iteration contracts the bulk, corner, and edge tensors, performs truncated SVD, and applies the resulting isometry to update corner and edge tensors.The process repeats until convergence.
CTMRG dependencies: The converged environment tensors depend on the bulk tensor in a complicated way.This dependence is a central reason to formulate CTMRG as a differentiable computation graph.
CTMRG complexity: CTMRG has complexity O(d3χ3) and memory cost O(d2χ2), both smaller in cutoff-bond-dimension scaling than TRG.Because CTMRG grows system size linearly, it may require more iterations to converge.
Differentiation challenges: Reverse-mode differentiation requires identifying tensor-network primitives and defining their vector-Jacobian products for backpropagation.Relevant primitives include BLAS operations and LAPACK routines such as eigensolvers, SVD, and QR factorization.
Differentiation challenges: Backward differentiation through eigensolvers and SVD can be unstable near degeneracies, while reverse mode can consume substantial memory.These issues can prevent reaching the bond dimensions of ordinary tensor-network programs.

A. Stable backward through linear algebra operations

The section develops stable reverse-mode derivatives for eigendecomposition and SVD, emphasizing numerical treatment of near-degenerate spectra and truncated singular spaces.

A. Stable backward through linear algebra operations: Matrix derivatives can be integrated modularly into differentiable tensor-network programs by specifying local backward functions.The approach focuses on stable numerical implementations for linear-algebra operations.
A. Stable backward through linear algebra operations: The symmetric eigensolver uses A = UDUT, with D containing eigenvalues and U containing corresponding orthogonal eigenvectors.
A. Stable backward through linear algebra operations: The eigensolver backward rule produces a symmetric gradient and connects, when eigenstate dependence vanishes, to the Hellmann-Feynman theorem.
A. Stable backward through linear algebra operations: The eigenvalue-difference denominator reflects nondegenerate perturbation theory, while Lorentzian broadening with ε = 10^-12 stabilizes calculations at a small gradient error.
A. Stable backward through linear algebra operations: For SVD, A = UDVT, with singular values in diagonal D and dimensions determined by k = min(m, n).
A. Stable backward through linear algebra operations: Degenerate singular values are handled with Lorentzian broadening, and vanishing singular values are generally harmless when the corresponding space is truncated.

3. QR factorization

The QR section specifies forward and backward differentiation rules for both tall and wide matrix factorizations, including a chain-rule treatment for the wide case.

3. QR factorization: QR factorization writes A = QR with QTQ = I and upper-triangular R, and its backward rule depends on matrix dimensions.
3. QR factorization: For m ≥ n, R is n × n, and the backward pass uses the corresponding dimension-specific rule.
3. QR factorization: The wide-case backward computation uses copyltu to symmetrize a matrix and solves the multiplication by R-T through a triangular linear system.
3. QR factorization: For m < n, the decomposition splits A = (X, Y) and R = (U, V), first determining Q and U from X = QU, then calculating V = QTY.

B. Memory efficient reverse mode automatic differentiation with checkpointing function

The section addresses reverse-mode memory growth in deep tensor-network computations with checkpointing and implicit differentiation through fixed-point iterations.

B. Memory efficient reverse mode automatic differentiation with checkpointing function: Reverse-mode differentiation stores forward intermediates, causing memory consumption proportional to computation-graph depth.
B. Memory efficient reverse mode automatic differentiation with checkpointing function: Checkpointing trades computation for memory by storing selected tensors and recomputing intermediate segments during backpropagation.The stated overhead is no more than twice the computational effort.
B. Memory efficient reverse mode automatic differentiation with checkpointing function: Checkpointing can encapsulate large computation-graph regions as customized primitives with local backward rules that rerun their forward computations.
C. Backward through fixed point iteration: Unrolling fixed-point iterations creates large memory use when convergence requires many steps.
C. Backward through fixed point iteration: The implicit function theorem differentiates T* = f(T*, θ) without storing the full iteration history, using single-step vector-Jacobian products and convergent adjoint iteration.
C. Backward through fixed point iteration: The fixed-point method applies to CTMRG, reduces reverse-mode memory, and permits accelerated forward iterations, but requires attention to redundant global gauge freedom.

D. Higher order derivatives

Automatic differentiation can compute higher-order derivatives of tensor-network programs, supporting Hessian-based optimization while avoiding explicit construction of the full Hessian.

D. Higher order derivatives: Applying automatic differentiation again to a gradient computation graph enables arbitrary higher-order derivatives in principle.
D. Higher order derivatives: Higher derivatives support Hessian-based methods such as Newton optimization for tensor-network states.
D. Higher order derivatives: Hessian-vector products avoid explicitly constructing or inverting the full Hessian and suffice for iterative linear-equation solvers.
D. Higher order derivatives: The applications demonstrate the approach for the classical Ising model and quantum Heisenberg model on the infinite square lattice.

A. Higher order derivative of the free energy

The paper differentiates the free energy produced by tensor renormalization to obtain Ising energy density and specific heat, including a second-order derivative without finite-difference error.

A. Higher order derivative of the free energy: Energy density and specific heat are computed as the first and second derivatives of the TRG free energy.The free energy comes from 30 TRG iterations with cutoff bond dimension χ = 30.
A. Higher order derivative of the free energy: Automatic differentiation yields results free from the finite difference error associated with numerical differentiation.The approach computes higher-order derivatives directly by backpropagating through the TRG procedure.
A. Higher order derivative of the free energy: Reverse-mode automatic differentiation is used through the entire TRG procedure, including the SVD, to compute higher-order physical observables.A numerically stable SVD backward function is crucial because the Ising model has degenerate singular values from Z2 symmetry.

B. Gradient based optimization of iPEPS

The paper applies automatic differentiation to optimize infinite projected entangled pair states for the square-lattice antiferromagnetic Heisenberg model. For bond dimensions D = 4–7, it reports the lowest variational energies, while staggered magnetization approaches extrapolated QMC results at larger D.

B. Gradient based optimization of iPEPS: The variational ansatz is an infinite projected entangled pair state whose tensor elements are optimized for the square-lattice antiferromagnetic Heisenberg model.The infinite network is contracted with CTMRG to evaluate the expected energy.
B. Gradient based optimization of iPEPS: The calculations use cutoff bond dimensions χ = 30, 50, 80, 100, 144, 160 for D = 2, 3, ..., 7, respectively.For D > 4, checkpointing or fixed-point iteration keeps memory within a single Nvidia P100 GPU with 12G memory.
B. Gradient based optimization of iPEPS: Automatic differentiation computes gradients through the CTMRG contraction, and L-BFGS uses them to optimize the iPEPS after a few hundred evaluations.The gradient computation takes time comparable to the forward expected-energy evaluation.
B. Gradient based optimization of iPEPS: For D = 4, 5, 6, 7, the calculation reaches the lowest variational energy reported for the infinite square-lattice Heisenberg model.At D = 2, 3, the energy is slightly higher than Ref. [34] because the present ansatz has half as many variational parameters.
B. Gradient based optimization of iPEPS: The optimized state's staggered magnetization approaches the extrapolated QMC result at larger bond dimensions.Figure 4 compares energy relative error and staggered magnetization against QMC and previous variational results.
B. Gradient based optimization of iPEPS: The approach is stated to apply to finite systems, larger unit cells, more complex Hamiltonians, and more efficient contraction schemes.These extensions are presented as promising for obtaining new quantum many-body results.

V. DISCUSSIONS

The discussion emphasizes automatic differentiation as a scalable way to extend tensor-network algorithms while reducing manual gradient work. It distinguishes differentiating an approximated contraction from approximating an analytical derivative and identifies broader applications.

V. DISCUSSIONS: Automatic differentiation has the same computational complexity as the forward tensor-network contraction while reducing gradient-calculation effort.This lets researchers focus on core contraction algorithms rather than manually deriving gradients.
V. DISCUSSIONS: The method computes the numerical-exact gradient of an approximated energy density, unlike approaches that approximate analytically derived gradients after contraction.The paper notes that differentiating the approximation can be advantageous for infinite systems with complicated analytical derivatives.
V. DISCUSSIONS: The demonstrated applications differentiate through complete contraction algorithms to optimize tensor networks and compute physical observables.The same techniques are also applicable to lower-level tasks such as optimal truncation bases and variational tensor-network transformations.
V. DISCUSSIONS: The techniques are also relevant to variational contraction of tensor networks, and the authors expect them to become part of the standard tensor-network toolbox.The paper additionally highlights GPU acceleration from implementing tensor-network programs in deep-learning frameworks.