Implicit differentiation of tensor network algorithms

Technical Summary: Implicit Differentiation of Tensor Network Algorithms

Problem Statement
The variational optimization of Projected Entangled-Pair States (PEPS) for ground-state calculations relies heavily on the accurate computation of energy gradients with respect to local variational parameters. The current state-of-the-art approach utilizes Automatic Differentiation (AD) to compute these gradients. However, this method faces three significant challenges:

1. Efficiency: Standard reverse-mode AD requires backpropagating through every iteration of the contraction algorithm (e.g., CTMRG or Boundary MPS). Even with fixed-point formulations that avoid storing the full computational graph, the presence of iterative subroutines (such as sparse eigensolvers or SVDs) creates a nested iterative procedure for gradient computation, leading to high computational costs.
2. Stability: Differentiating the primitive subroutines used in contraction algorithms (e.g., truncated eigenvalue decompositions or SVDs) introduces numerical instabilities. These instabilities often arise from near-degenerate eigenvalues or singular values, where the derivatives of these subroutines diverge. Crucially, these instabilities are artifacts of the differentiation procedure; the physical energy gradient itself remains stable with respect to the underlying variations.
3. Implementation Complexity: Applying AD to PEPS optimization requires every primitive operation in the contraction algorithm to be compatible with the AD engine. Defining custom differentiation rules for non-native subroutines is laborious and limits flexibility.

Methodology
The authors propose a framework based on implicit differentiation to address these challenges. Instead of differentiating through the iterative steps of the contraction algorithm, the method reformulates the converged contraction environment as the root of a single set of algebraic characteristic equations.

The core steps of the methodology are:

1. Algebraic Characterization: The converged environment (e.g., corner and edge tensors in CTMRG, or boundary MPS tensors) is characterized not by the iterative update rule $x = f(x, p)$, but by a set of equations $F(y, p) = 0$. Here, $y$ represents a set of environment variables (potentially including intermediate variables like isometries or singular values) that implicitly depend on the variational parameters $p$.
2. Implicit Differentiation: By applying the implicit function theorem to $F(y, p) = 0$, the gradient computation reduces to solving a single, non-nested system of linear equations:
   $$\tilde{y} \partial_y F = -\tilde{y} \partial_p F$$
   where $\tilde{y}$ is the adjoint variable. This eliminates the need to backpropagate through the iterative subroutines.
3. Stable Parametrization: A key innovation is the choice of the variable set $y$ and the parametrization of the characteristic equations. By exploiting the intrinsic gauge symmetries of the contraction environment, the authors construct parametrizations that explicitly remove the degrees of freedom responsible for numerical instabilities (e.g., variations within degenerate subspaces). For instance, in CTMRG, the isometry variations are restricted to their null spaces, bypassing the divergent terms associated with differentiating truncated eigensolvers.
4. Modular Implementation: The approach decouples the contraction algorithm from the gradient computation. The forward pass remains unchanged (using standard contraction routines), while the reverse pass only requires differentiating the algebraic characteristic equations, which consist of simple tensor contractions.

Key Contributions
The paper provides explicit formulations of characteristic equations for three major PEPS contraction schemes:

• $C_{4v}$ Symmetric CTMRG: The authors derive a set of characteristic equations involving corner tensors, edge tensors, and a parametrized isometry. This formulation removes the need to differentiate the truncated eigenvalue decomposition (eigh) and resolves instabilities associated with degenerate eigenvalues.
• Boundary MPS (VUMPS): A set of characteristic equations is formulated for the top and bottom boundary MPS eigenvectors and mixed environments. This approach avoids the need for gauge-fixing the resulting boundary MPS tensors during gradient computation and eliminates nested linear problems arising from differentiating the dominant eigensolvers.
• Asymmetric CTMRG (No Symmetries): The method is extended to general systems without spatial symmetries. The authors introduce modified corner and edge tensors and custom "square root" variables for singular values to ensure gauge invariance and stability when differentiating truncated SVDs.

Results
The authors benchmarked the implicit differentiation approach against fixed-point differentiation and straightforward reverse-mode AD using the PEPSKit.jl and MPSKit.jl packages:

• Accuracy: The gradients computed via implicit differentiation agree with reference finite-difference values and black-box reverse-mode AD within the tolerance of the linear solver, even for moderately converged environments.
• Performance:
  • For $C_{4v}$ CTMRG and Boundary MPS, the implicit approach consistently outperforms fixed-point differentiation, with speedups increasing as the bond dimension ($D$) and environment dimension ($\chi$) grow.
  • For Asymmetric CTMRG on non-trivial unit cells, the implicit method shows superior asymptotic scaling. While fixed-point differentiation may be faster at low dimensions due to implementation overhead, the implicit method becomes significantly more efficient at high dimensions where sparse matrix decompositions are required.
• Scaling: In many cases, the cost of the gradient computation using implicit differentiation becomes a subleading step compared to the contraction cost itself, whereas fixed-point differentiation often remains a dominant bottleneck.

Significance and Claims
The paper claims that this framework represents a paradigm shift in tensor network programming. By differentiating algebraic optimality conditions rather than the algorithmic subroutines, the method:

• Simplifies Implementation: It removes the requirement for every subroutine in the contraction algorithm to be differentiable, allowing for the use of standard, highly optimized linear algebra routines in the forward pass.
• Enhances Stability: It naturally removes numerical instabilities arising from the differentiation of subroutines by exploiting gauge symmetries to construct stable parameterizations.
• Improves Efficiency: It reduces the computational complexity of gradient evaluation by un-nesting the linear problems, leading to better scaling with system size.

The authors emphasize that while they demonstrated the approach on three specific contraction schemes, the core algorithm (Algorithm 1 and 2) is general and applicable to any optimization problem where the solution can be defined by differentiable optimality conditions. They do not claim to have solved all PEPS challenges but assert that this approach makes stable, gradient-based PEPS optimization significantly more practical and robust.